diff --git a/tardis/spectrum/__init__.py b/tardis/spectrum/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tardis/spectrum/base.py b/tardis/spectrum/base.py new file mode 100644 index 00000000000..5cdb0eeca53 --- /dev/null +++ b/tardis/spectrum/base.py @@ -0,0 +1,177 @@ +import warnings + +import numpy as np +from astropy import units as u + +from tardis.spectrum.formal_integral import IntegrationError +from tardis.spectrum.spectrum import TARDISSpectrum +from tardis.util.base import ( + quantity_linspace, +) + +class SpectrumSolver: + def __init__(self, transport_state, spectrum_frequency): + self.transport_state = transport_state + self.spectrum_frequency = spectrum_frequency + self._montecarlo_virtual_luminosity = u.Quantity( + np.zeros_like(self.spectrum_frequency.value), "erg / s" + ) + self._integrator = None + self.integrator_settings = None + self._spectrum_integrated = None + + @property + def spectrum(self): + return TARDISSpectrum( + self.spectrum_frequency, self.montecarlo_emitted_luminosity + ) + + @property + def spectrum_reabsorbed(self): + return TARDISSpectrum( + self.spectrum_frequency, self.montecarlo_reabsorbed_luminosity + ) + + @property + def spectrum_virtual(self): + if np.all(self.montecarlo_virtual_luminosity == 0): + warnings.warn( + "MontecarloTransport.spectrum_virtual" + "is zero. Please run the montecarlo simulation with" + "no_of_virtual_packets > 0", + UserWarning, + ) + + return TARDISSpectrum( + self.spectrum_frequency, self.montecarlo_virtual_luminosity + ) + + @property + def spectrum_integrated(self): + if self._spectrum_integrated is None: + # This was changed from unpacking to specific attributes as compute + # is not used in calculate_spectrum + try: + self._spectrum_integrated = self.integrator.calculate_spectrum( + self.spectrum_frequency[:-1], + points=self.integrator_settings.points, + interpolate_shells=self.integrator_settings.interpolate_shells, + ) + except IntegrationError: + # if integration is impossible or fails, return an empty spectrum + warnings.warn( + "The FormalIntegrator is not yet implemented for the full " + "relativity mode or continuum processes. " + "Please run with config option enable_full_relativity: " + "False and continuum_processes_enabled: False " + "This RETURNS AN EMPTY SPECTRUM!", + UserWarning, + ) + return TARDISSpectrum( + np.array([np.nan, np.nan]) * u.Hz, + np.array([np.nan]) * u.erg / u.s, + ) + return self._spectrum_integrated + + @property + def integrator(self): + if self._integrator is None: + warnings.warn( + "MontecarloTransport.integrator: " + "The FormalIntegrator is not yet available." + "Please run the montecarlo simulation at least once.", + UserWarning, + ) + if self.enable_full_relativity: + raise NotImplementedError( + "The FormalIntegrator is not yet implemented for the full " + "relativity mode. " + "Please run with config option enable_full_relativity: " + "False." + ) + return self._integrator + + @property + def montecarlo_reabsorbed_luminosity(self): + return u.Quantity( + np.histogram( + self.transport_state.reabsorbed_packet_nu, + weights=self.reabsorbed_packet_luminosity, + bins=self.spectrum_frequency, + )[0], + "erg / s", + ) + + @property + def montecarlo_emitted_luminosity(self): + return u.Quantity( + np.histogram( + self.transport_state.emitted_packet_nu, + weights=self.emitted_packet_luminosity, + bins=self.spectrum_frequency, + )[0], + "erg / s", + ) + + @property + def montecarlo_virtual_luminosity(self): + return ( + self._montecarlo_virtual_luminosity[:-1] + / self.transport_state.time_of_simulation.value + ) + + def calculate_emitted_luminosity( + self, luminosity_nu_start, luminosity_nu_end + ): + """ + Calculate emitted luminosity. + + Parameters + ---------- + luminosity_nu_start : astropy.units.Quantity + luminosity_nu_end : astropy.units.Quantity + + Returns + ------- + astropy.units.Quantity + """ + luminosity_wavelength_filter = ( + self.transport_state.emitted_packet_nu > luminosity_nu_start + ) & (self.transport_state.emitted_packet_nu < luminosity_nu_end) + + return self.emitted_packet_luminosity[ + luminosity_wavelength_filter + ].sum() + + def calculate_reabsorbed_luminosity( + self, luminosity_nu_start, luminosity_nu_end + ): + """ + Calculate reabsorbed luminosity. + + Parameters + ---------- + luminosity_nu_start : astropy.units.Quantity + luminosity_nu_end : astropy.units.Quantity + + Returns + ------- + astropy.units.Quantity + """ + luminosity_wavelength_filter = ( + self.reabsorbed_packet_nu > luminosity_nu_start + ) & (self.reabsorbed_packet_nu < luminosity_nu_end) + + return self.reabsorbed_packet_luminosity[ + luminosity_wavelength_filter + ].sum() + + @classmethod + def from_config(cls, config): + spectrum_frequency = quantity_linspace( + config.spectrum.stop.to("Hz", u.spectral()), + config.spectrum.start.to("Hz", u.spectral()), + num=config.spectrum.num + 1, + ) + + return cls(transport_state=None, spectrum_frequency=spectrum_frequency) diff --git a/tardis/spectrum/spectrum.py b/tardis/spectrum/spectrum.py index 2479206bfe3..2caa10fe269 100644 --- a/tardis/spectrum/spectrum.py +++ b/tardis/spectrum/spectrum.py @@ -1,6 +1,8 @@ import warnings + import numpy as np from astropy import units as u + from tardis.io.util import HDFWriterMixin @@ -37,9 +39,9 @@ def __init__(self, _frequency, luminosity): self._frequency = _frequency.to("Hz", u.spectral()) self.luminosity = luminosity.to("erg / s") - l_nu_unit = u.def_unit("erg\ s^{-1}\ Hz^{-1}", u.Unit("erg/(s Hz)")) + l_nu_unit = u.def_unit(r"erg s^{-1} Hz^{-1}", u.Unit("erg/(s Hz)")) l_lambda_unit = u.def_unit( - "erg\ s^{-1}\ \\AA^{-1}", u.Unit("erg/(s AA)") + r"erg s^{-1} \AA^{-1}", u.Unit("erg/(s AA)") ) self.frequency = self._frequency[:-1] diff --git a/tardis/transport/montecarlo/base.py b/tardis/transport/montecarlo/base.py index 7c508427e6c..7a10e09f8e1 100644 --- a/tardis/transport/montecarlo/base.py +++ b/tardis/transport/montecarlo/base.py @@ -121,7 +121,6 @@ def initialize_transport_state( transport_state = MonteCarloTransportState( packet_collection, estimators, - spectrum_frequency=self.spectrum_frequency, geometry_state=geometry_state, opacity_state=opacity_state, ) @@ -180,7 +179,7 @@ def run( transport_state.opacity_state, self.montecarlo_configuration, transport_state.radfield_mc_estimators, - transport_state.spectrum_frequency.value, + self.spectrum_frequency.value, number_of_vpackets, iteration=iteration, show_progress_bars=show_progress_bars, diff --git a/tardis/transport/montecarlo/montecarlo_transport_state.py b/tardis/transport/montecarlo/montecarlo_transport_state.py index 2c58fe60bee..21a457828f3 100644 --- a/tardis/transport/montecarlo/montecarlo_transport_state.py +++ b/tardis/transport/montecarlo/montecarlo_transport_state.py @@ -59,7 +59,6 @@ def __init__( self, packet_collection, radfield_mc_estimators, - spectrum_frequency, geometry_state, opacity_state, rpacket_tracker=None, @@ -67,13 +66,6 @@ def __init__( ): self.packet_collection = packet_collection self.radfield_mc_estimators = radfield_mc_estimators - self.spectrum_frequency = spectrum_frequency - self._montecarlo_virtual_luminosity = u.Quantity( - np.zeros_like(self.spectrum_frequency.value), "erg / s" - ) - self._integrator = None - self.integrator_settings = None - self._spectrum_integrated = None self.enable_full_relativity = False self.enable_continuum_processes = False self.geometry_state = geometry_state @@ -166,152 +158,6 @@ def emitted_packet_luminosity(self): def reabsorbed_packet_luminosity(self): return -self.packet_luminosity[~self.emitted_packet_mask] - @property - def montecarlo_reabsorbed_luminosity(self): - return u.Quantity( - np.histogram( - self.reabsorbed_packet_nu, - weights=self.reabsorbed_packet_luminosity, - bins=self.spectrum_frequency, - )[0], - "erg / s", - ) - - @property - def montecarlo_emitted_luminosity(self): - return u.Quantity( - np.histogram( - self.emitted_packet_nu, - weights=self.emitted_packet_luminosity, - bins=self.spectrum_frequency, - )[0], - "erg / s", - ) - - @property - def montecarlo_virtual_luminosity(self): - return ( - self._montecarlo_virtual_luminosity[:-1] - / self.packet_collection.time_of_simulation - ) - - @property - def spectrum(self): - return TARDISSpectrum( - self.spectrum_frequency, self.montecarlo_emitted_luminosity - ) - - @property - def spectrum_reabsorbed(self): - return TARDISSpectrum( - self.spectrum_frequency, self.montecarlo_reabsorbed_luminosity - ) - - @property - def spectrum_virtual(self): - if np.all(self.montecarlo_virtual_luminosity == 0): - warnings.warn( - "MontecarloTransport.spectrum_virtual" - "is zero. Please run the montecarlo simulation with" - "no_of_virtual_packets > 0", - UserWarning, - ) - - return TARDISSpectrum( - self.spectrum_frequency, self.montecarlo_virtual_luminosity - ) - - @property - def spectrum_integrated(self): - if self._spectrum_integrated is None: - # This was changed from unpacking to specific attributes as compute - # is not used in calculate_spectrum - try: - self._spectrum_integrated = self.integrator.calculate_spectrum( - self.spectrum_frequency[:-1], - points=self.integrator_settings.points, - interpolate_shells=self.integrator_settings.interpolate_shells, - ) - except IntegrationError: - # if integration is impossible or fails, return an empty spectrum - warnings.warn( - "The FormalIntegrator is not yet implemented for the full " - "relativity mode or continuum processes. " - "Please run with config option enable_full_relativity: " - "False and continuum_processes_enabled: False " - "This RETURNS AN EMPTY SPECTRUM!", - UserWarning, - ) - return TARDISSpectrum( - np.array([np.nan, np.nan]) * u.Hz, - np.array([np.nan]) * u.erg / u.s, - ) - return self._spectrum_integrated - - @property - def integrator(self): - if self._integrator is None: - warnings.warn( - "MontecarloTransport.integrator: " - "The FormalIntegrator is not yet available." - "Please run the montecarlo simulation at least once.", - UserWarning, - ) - if self.enable_full_relativity: - raise NotImplementedError( - "The FormalIntegrator is not yet implemented for the full " - "relativity mode. " - "Please run with config option enable_full_relativity: " - "False." - ) - return self._integrator - - def calculate_emitted_luminosity( - self, luminosity_nu_start, luminosity_nu_end - ): - """ - Calculate emitted luminosity. - - Parameters - ---------- - luminosity_nu_start : astropy.units.Quantity - luminosity_nu_end : astropy.units.Quantity - - Returns - ------- - astropy.units.Quantity - """ - luminosity_wavelength_filter = ( - self.emitted_packet_nu > luminosity_nu_start - ) & (self.emitted_packet_nu < luminosity_nu_end) - - return self.emitted_packet_luminosity[ - luminosity_wavelength_filter - ].sum() - - def calculate_reabsorbed_luminosity( - self, luminosity_nu_start, luminosity_nu_end - ): - """ - Calculate reabsorbed luminosity. - - Parameters - ---------- - luminosity_nu_start : astropy.units.Quantity - luminosity_nu_end : astropy.units.Quantity - - Returns - ------- - astropy.units.Quantity - """ - luminosity_wavelength_filter = ( - self.reabsorbed_packet_nu > luminosity_nu_start - ) & (self.reabsorbed_packet_nu < luminosity_nu_end) - - return self.reabsorbed_packet_luminosity[ - luminosity_wavelength_filter - ].sum() - @property def virt_packet_nus(self): try: