Skip to content

Commit

Permalink
Removed spectrum elements from the transport state
Browse files Browse the repository at this point in the history
  • Loading branch information
andrewfullard committed Jun 28, 2024
1 parent 1ec39a2 commit 14dea10
Show file tree
Hide file tree
Showing 5 changed files with 182 additions and 158 deletions.
Empty file added tardis/spectrum/__init__.py
Empty file.
177 changes: 177 additions & 0 deletions tardis/spectrum/base.py
Original file line number Diff line number Diff line change
@@ -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)
6 changes: 4 additions & 2 deletions tardis/spectrum/spectrum.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
import warnings

import numpy as np
from astropy import units as u

from tardis.io.util import HDFWriterMixin


Expand Down Expand Up @@ -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]
Expand Down
3 changes: 1 addition & 2 deletions tardis/transport/montecarlo/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
)
Expand Down Expand Up @@ -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,
Expand Down
154 changes: 0 additions & 154 deletions tardis/transport/montecarlo/montecarlo_transport_state.py
Original file line number Diff line number Diff line change
Expand Up @@ -59,21 +59,13 @@ def __init__(
self,
packet_collection,
radfield_mc_estimators,
spectrum_frequency,
geometry_state,
opacity_state,
rpacket_tracker=None,
vpacket_tracker=None,
):
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
Expand Down Expand Up @@ -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:
Expand Down

0 comments on commit 14dea10

Please sign in to comment.