Refactor spectrum

tardis/montecarlo/base.py

@@ -12,10 +12,10 @@
 from tardis.io.util import to_hdf
 
 import numpy as np
-import pandas as pd
 
 logger = logging.getLogger(__name__)
 
+
 class MontecarloRunner(object):
     """
     This class is designed as an interface between the Python part and the
@@ -31,7 +31,7 @@ class MontecarloRunner(object):
 
     def __init__(self, seed, spectrum_frequency, virtual_spectrum_range,
                  sigma_thomson, enable_reflective_inner_boundary,
-                 inner_boundary_albedo, line_interaction_type, distance=None):
+                 inner_boundary_albedo, line_interaction_type):
 
         self.seed = seed
         self.packet_source = packet_source.BlackBodySimpleSource(seed)
@@ -41,10 +41,6 @@ def __init__(self, seed, spectrum_frequency, virtual_spectrum_range,
         self.enable_reflective_inner_boundary = enable_reflective_inner_boundary
         self.inner_boundary_albedo = inner_boundary_albedo
         self.line_interaction_type = line_interaction_type
-        self.spectrum = TARDISSpectrum(spectrum_frequency, distance)
-        self.spectrum_virtual = TARDISSpectrum(spectrum_frequency, distance)
-        self.spectrum_reabsorbed = TARDISSpectrum(spectrum_frequency, distance)
-
 
     def _initialize_estimator_arrays(self, no_of_shells, tau_sobolev_shape):
         """
@@ -56,13 +52,12 @@ def _initialize_estimator_arrays(self, no_of_shells, tau_sobolev_shape):
         model: ~Radial1DModel
         """
 
-        #Estimators
+        # Estimators
         self.j_estimator = np.zeros(no_of_shells, dtype=np.float64)
         self.nu_bar_estimator = np.zeros(no_of_shells, dtype=np.float64)
         self.j_blue_estimator = np.zeros(tau_sobolev_shape)
         self.Edotlu_estimator = np.zeros(tau_sobolev_shape)
 
-
     def _initialize_geometry_arrays(self, model):
         """
         Generate the cgs like geometry arrays for the montecarlo part
@@ -94,30 +89,32 @@ def _initialize_packets(self, T, no_of_packets, no_of_virtual_packets=None):
         self.last_interaction_type = -1 * np.ones(no_of_packets, dtype=np.int64)
         self.last_interaction_in_nu = np.zeros(no_of_packets, dtype=np.float64)
 
+        self._montecarlo_virtual_luminosity = u.Quantity(
+                np.zeros_like(self.spectrum_frequency.value),
+                'erg / s'
+                )
 
-        self.legacy_montecarlo_virtual_luminosity = np.zeros_like(
-            self.spectrum_frequency.value)
-
-    def legacy_update_spectrum(self, no_of_virtual_packets):
-        montecarlo_reabsorbed_luminosity = np.histogram(
-            self.reabsorbed_packet_nu,
-            weights=self.reabsorbed_packet_luminosity,
-            bins=self.spectrum_frequency.value)[0] * u.erg / u.s
-
-        montecarlo_emitted_luminosity = np.histogram(
-            self.emitted_packet_nu,
-            weights=self.emitted_packet_luminosity,
-            bins=self.spectrum_frequency.value)[0] * u.erg / u.s
+    @property
+    def spectrum(self):
+        return TARDISSpectrum(
+                self.spectrum_frequency,
+                self.montecarlo_emitted_luminosity)
 
-        self.spectrum.update_luminosity(montecarlo_emitted_luminosity)
-        self.spectrum_reabsorbed.update_luminosity(
-            montecarlo_reabsorbed_luminosity)
+    @property
+    def spectrum_reabsorbed(self):
+        return TARDISSpectrum(
+                self.spectrum_frequency,
+                self.montecarlo_reabsorbed_luminosity)
 
-        if no_of_virtual_packets > 0:
-            self.montecarlo_virtual_luminosity = (
-                self.legacy_montecarlo_virtual_luminosity *
-                1 * u.erg / self.time_of_simulation)[:-1]
-            self.spectrum_virtual.update_luminosity(
+    @property
+    def spectrum_virtual(self):
+        # if no_of_virtual_packets > 0:
+        #     montecarlo_virtual_luminosity = (
+        #             self.montecarlo_virtual_luminosity *
+        #             1 * u.erg / self.time_of_simulation)[:-1]
+
+        return TARDISSpectrum(
+                self.spectrum_frequency,
                 self.montecarlo_virtual_luminosity)
 
     def run(self, model, plasma, no_of_packets, no_of_virtual_packets=0, nthreads=1,last_run=False):
@@ -229,18 +226,37 @@ def emitted_packet_nu(self):
     def reabsorbed_packet_nu(self):
         return self.output_nu[~self.emitted_packet_mask]
 
+    @property
+    def emitted_packet_luminosity(self):
+        return self.packet_luminosity[self.emitted_packet_mask]
+
     @property
     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.value)[0],
+                'erg / s'
+                )
 
     @property
-    def emitted_packet_luminosity(self):
-        return self.packet_luminosity[self.emitted_packet_mask]
+    def montecarlo_emitted_luminosity(self):
+        return u.Quantity(
+                np.histogram(
+                    self.emitted_packet_nu,
+                    weights=self.emitted_packet_luminosity,
+                    bins=self.spectrum_frequency.value)[0],
+                'erg / s'
+                )
 
     @property
-    def reabsorbed_packet_luminosity(self):
-        return -self.packet_luminosity[~self.emitted_packet_mask]
+    def montecarlo_virtual_luminosity(self):
+        return self._montecarlo_virtual_luminosity[:-1] / self.time_of_simulation.value
 
     def calculate_emitted_luminosity(self, luminosity_nu_start,
                                      luminosity_nu_end):
@@ -374,5 +390,5 @@ def from_config(cls, config):
                    sigma_thomson=sigma_thomson,
                    enable_reflective_inner_boundary=config.montecarlo.enable_reflective_inner_boundary,
                    inner_boundary_albedo=config.montecarlo.inner_boundary_albedo,
-                   line_interaction_type=config.plasma.line_interaction_type,
-                   distance=config.supernova.get('distance', None))
+                   line_interaction_type=config.plasma.line_interaction_type
+                   )

tardis/montecarlo/montecarlo.pyx

@@ -229,7 +229,7 @@ cdef initialize_storage_model(model, plasma, runner, storage_model_t *storage):
     storage.spectrum_delta_nu = runner.spectrum_frequency.to('Hz').value[1] - runner.spectrum_frequency.to('Hz').value[0]
 
     storage.spectrum_virt_nu = <double*> PyArray_DATA(
-        runner.legacy_montecarlo_virtual_luminosity)
+        runner._montecarlo_virtual_luminosity.value)
 
     storage.sigma_thomson = runner.sigma_thomson.cgs.value
     storage.inverse_sigma_thomson = 1.0 / storage.sigma_thomson

tardis/montecarlo/spectrum.py

@@ -1,99 +1,121 @@
-import os
 import numpy as np
-from astropy import constants, units as u
 
-from tardis.io.util import to_hdf
+from astropy import units as u
+
+
+def require_field(name):
+    def _require_field(f):
+        def wrapper(self, *args, **kwargs):
+            if not getattr(self, name, None):
+                raise AttributeError(
+                        '{} is required as attribute of'
+                        '{} to calculate {}'.format(
+                            name,
+                            self.__class__.__name__,
+                            f.__name__
+                            )
+                        )
+            else:
+                return f(self, *args, **kwargs)
+        return wrapper
+    return _require_field
 
 
 class TARDISSpectrum(object):
     """
-    TARDIS Spectrum object
-    """
+    TARDISSpectrum(_frequency, luminosity)
 
-    def __init__(self, frequency, distance=None):
-        self._frequency = frequency
-        self.wavelength = self.frequency.to('angstrom', u.spectral())
+    _frequency: astropy.units.Quantity with unit 'Hz' or a length
+        These are bin edges of frequency or wavelenght bins for the spectrum.
 
-        self.distance = distance
+    luminosity: astropy.units.Quantity with unit Energy per second
+        The luminosity in each bin of the spectrum.
 
+    After manually adding a distance attribute, the properties 'flux_nu' and
+    'flux_lambda' become available
+    """
 
+    def __init__(self, _frequency, luminosity):
+        if not _frequency.shape[0] == luminosity.shape[0] + 1:
+            raise ValueError(
+                    "shape of '_frequency' and 'luminosity' are not compatible"
+                    ": '{}' and '{}'".format(
+                        _frequency.shape[0],
+                        luminosity.shape[0])
+                    )
+        self._frequency = _frequency.to('Hz', u.spectral())
+        self.luminosity = luminosity.to('erg / s')
 
-        self.delta_frequency = frequency[1] - frequency[0]
+    @property
+    def frequency(self):
+        return self._frequency[:-1]
 
-        self._flux_nu = np.zeros_like(frequency.value) * u.Unit('erg / (s Hz cm^2)')
-        self._flux_lambda = np.zeros_like(frequency.value) * u.Unit('erg / (s Angstrom cm^2)')
+    @property
+    def delta_frequency(self):
+        return self.frequency[1] - self.frequency[0]
 
-        self.luminosity_density_nu = np.zeros_like(self.frequency) * u.Unit('erg / (s Hz)')
-        self.luminosity_density_lambda = np.zeros_like(self.frequency) * u.Unit('erg / (s Angstrom)')
+    @property
+    def wavelength(self):
+        return self.frequency.to('angstrom', u.spectral())
 
     @property
-    def frequency(self):
-        return self._frequency[:-1]
+    def luminosity_density_nu(self):
+        return (self.luminosity / self.delta_frequency).to('erg / (s Hz)')
 
+    @property
+    def luminosity_density_lambda(self):
+        return self.f_nu_to_f_lambda(
+                self.luminosity_density_nu,
+                )
 
     @property
+    @require_field('distance')
     def flux_nu(self):
-        if self.distance is None:
-            raise AttributeError('supernova distance not supplied - flux calculation impossible')
-        else:
-            return self._flux_nu
+        return self.luminosity_to_flux(
+                self.luminosity_density_nu,
+                self.distance)
 
     @property
+    @require_field('distance')
     def flux_lambda(self):
-        if self.distance is None:
-            raise AttributeError('supernova distance not supplied - flux calculation impossible')
-        return self._flux_lambda
-
-    def update_luminosity(self, spectrum_luminosity):
-        self.luminosity_density_nu = (spectrum_luminosity / self.delta_frequency).to('erg / (s Hz)')
-        self.luminosity_density_lambda = self.f_nu_to_f_lambda(self.luminosity_density_nu.value) \
-                                         * u.Unit('erg / (s Angstrom)')
+        return self.luminosity_to_flux(
+                self.luminosity_density_lambda,
+                self.distance
+                )
 
-        if self.distance is not None:
-            self._flux_nu = (self.luminosity_density_nu / (4 * np.pi * self.distance.to('cm')**2))
-
-            self._flux_lambda = self.f_nu_to_f_lambda(self.flux_nu.value) * u.Unit('erg / (s Angstrom cm^2)')
+    @staticmethod
+    def luminosity_to_flux(luminosity, distance):
+        return luminosity / (4 * np.pi * distance.to('cm')**2)
 
     def f_nu_to_f_lambda(self, f_nu):
-        return f_nu * self.frequency.value**2 / constants.c.cgs.value / 1e8
-
+        return f_nu * self.frequency / self.wavelength
 
     def plot(self, ax, mode='wavelength'):
         if mode == 'wavelength':
             ax.plot(self.wavelength.value, self.flux_lambda.value)
-            ax.set_xlabel('Wavelength [%s]' % self.wavelength.unit._repr_latex_())
+            ax.set_xlabel('Wavelength [{}]'.format(
+                self.wavelength.unit._repr_latex_())
+                )
             ax.set_ylabel('Flux [%s]' % self.flux_lambda.unit._repr_latex_())
 
     def to_ascii(self, fname, mode='luminosity_density'):
         if mode == 'luminosity_density':
-            np.savetxt(fname, zip(self.wavelength.value, self.luminosity_density_lambda.value))
+            np.savetxt(
+                    fname, zip(
+                        self.wavelength.value,
+                        self.luminosity_density_lambda.value))
         elif mode == 'flux':
-            np.savetxt(fname, zip(self.wavelength.value, self.flux_lambda.value))
+            np.savetxt(
+                    fname,
+                    zip(self.wavelength.value, self.flux_lambda.value))
         else:
-            raise NotImplementedError('only mode "luminosity_density" and "flux" are implemented')
-
-    def to_hdf(self, path_or_buf, path='', name=''):
-        """
-        Store the spectrum to an HDF structure.
-
-        Parameters
-        ----------
-        path_or_buf
-            Path or buffer to the HDF store
-        path : str
-            Path inside the HDF store to store the spectrum
-        name : str, optional
-            A different name than 'spectrum', if needed.
-
-        Returns
-        -------
-        None
-
-        """
-        if not name:
-            name = 'spectrum'
-        spectrum_path = os.path.join(path, name)
-        properties = ['luminosity_density_nu', 'delta_frequency', 'wavelength',
-                      'luminosity_density_lambda']
-        to_hdf(path_or_buf, spectrum_path, {name: getattr(self, name) for name
-                                            in properties})
+            raise NotImplementedError(
+                    'only mode "luminosity_density"'
+                    'and "flux" are implemented')
+
+    def to_hdf(self, path_or_buf, path):
+        pass
+
+    @classmethod
+    def from_hdf(cls, path_or_buf, path):
+        pass