Merge pull request #750 from tardis-sn/fi_spectrum_integrated

[TEP005][QOL] Add MontecarloRunner.spectrum_integrated

scripts/tardis

@@ -67,28 +67,13 @@ tardis_config = config_reader.Configuration.from_yaml(args.config_fname)
 simulation = Simulation.from_config(tardis_config)
 
 
-def get_integrated_spectrum():
-    if (
-            tardis_config.plasma.disable_electron_scattering and
-            tardis_config.plasma.line_interaction_type == 'downbranch'
-            ):
-        return simulation.runner.integrator.calculate_spectrum(
-                simulation.runner.spectrum_frequency[:-1],
-                1000)
-    else:
-        raise config_reader.ConfigurationError(
-                'The formal integral module currently works only for '
-                'plasma settings "downbranch" and disabled electron scattering'
-                '.\nWe got: {}'.format(str(tardis_config.plasma)))
-
-
 def get_virtual_spectrum():
     # Catch warning when acessing invalid spectrum_virtual
     with warnings.catch_warnings(record=True) as w:
         spectrum = simulation.runner.spectrum_virtual
 
         if len(w) > 0 and w[-1]._category_name == 'UserWarning':
-            print(
+            warnings.warn(
                     'Virtual spectrum is not available, using the '
                     'real packet spectrum instead.')
             spectrum = simulation.runner.spectrum
@@ -99,20 +84,22 @@ print('Saving the {} spectrum.'.format(tardis_config.spectrum.method))
 # Believe it or not, that is a fancy switch-case statement
 # We need functions so that only one spectrum is accessed so we can catch
 # the warning properly
-spectrum = {
+get_spectrum = {
         'real': lambda: simulation.runner.spectrum,
         'virtual': get_virtual_spectrum,
-        'integrated': get_integrated_spectrum,
+        'integrated': lambda: simulation.runner.spectrum_integrated,
         }[tardis_config.spectrum.method]
 
 
 if args.gdb:
-    import os; print(os.getpid()); raw_input() # Workaround to attach gdb
+    import os
+    print(os.getpid())
+    raw_input()  # Workaround to attach gdb
 if args.profile:
     import cProfile
     cProfile.runctx('simulation.run()', locals(),
                     globals(), filename=args.profiler_log_file)
 else:
     simulation.run()
 
-spectrum().to_ascii(args.spectrum_fname)
+get_spectrum().to_ascii(args.spectrum_fname)

tardis/montecarlo/base.py

@@ -124,6 +124,11 @@ def spectrum_virtual(self):
                 self.spectrum_frequency,
                 self.montecarlo_virtual_luminosity)
 
+    @property
+    def spectrum_integrated(self):
+        return self.integrator.calculate_spectrum(
+                self.spectrum_frequency[:-1])
+
     @property
     def integrator(self):
         if self._integrator is None:

tardis/montecarlo/formal_integral.py

@@ -1,3 +1,4 @@
+import warnings
 import numpy as np
 import pandas as pd
 from astropy import units as u
@@ -6,25 +7,72 @@
 from tardis.montecarlo.spectrum import TARDISSpectrum
 
 
+class IntegrationError(Exception):
+    pass
+
+
 class FormalIntegrator(object):
 
-    def __init__(self, model, plasma, runner):
+    def __init__(self, model, plasma, runner, points=1000):
         self.model = model
         self.plasma = plasma
         self.runner = runner
+        self.points = points
+
+    def check(self, raises=True):
+        '''
+        A method that determines if the formal integral can be performed with
+        the current configuration settings
+
+        The function returns False if the configuration conflicts with the
+        required settings. If raises evaluates to True, then a
+        IntegrationError is raised instead
+        '''
+        def raise_or_return(message):
+            if raises:
+                raise IntegrationError(message)
+            else:
+                warnings.warn(message)
+                return False
+
+        for obj in (self.model, self.plasma, self.runner):
+            if obj is None:
+                return raise_or_return(
+                        'The integrator is missing either model, plasma or '
+                        'runner. Please make sure these are provided to the '
+                        'FormalIntegrator.'
+                        )
+
+        if not self.runner.line_interaction_type == 'downbranch':
+            return raise_or_return(
+                    'The FormalIntegrator currently only works for '
+                    'line_interaction_type == "downbranch"'
+                    )
 
-    def calculate_spectrum(self, frequency, N=1000):
+        if self.runner.sigma_thomson.value > 1e-100:
+            return raise_or_return(
+                    'The FormalIntegrator currently only works with '
+                    'disable_electron_scattering turned on.'
+                    )
+
+        return True
+
+    def calculate_spectrum(self, frequency, points=None, raises=True):
         # Very crude implementation
         # The c extension needs bin centers (or something similar)
         # while TARDISSpectrum needs bin edges
+        self.check(raises)
+        N = points or self.points
         frequency = frequency.to('Hz', u.spectral())
+
         luminosity = u.Quantity(
                 formal_integral(
                     self,
                     frequency,
                     N),
                 'erg'
                 ) * (frequency[1] - frequency[0])
+
         # Ugly hack to convert to 'bin edges'
         frequency = u.Quantity(
                 np.concatenate([
@@ -33,6 +81,7 @@ def calculate_spectrum(self, frequency, N=1000):
                         frequency.value[-1] + np.diff(frequency.value)[-1]
                         ]]),
                     frequency.unit)
+
         return TARDISSpectrum(
                 frequency,
                 luminosity