Folded the velocity shift flag into EmissionModels (#823)

examples/particle/plot_velocity_broadened_spectra.py

@@ -29,8 +29,8 @@
 grid_dir = "../../tests/test_grid/"
 grid = Grid(grid_name, grid_dir=grid_dir)
 
-# Define the model
-model = NebularEmission(grid)
+# Define the model with velocity shift
+model = NebularEmission(grid, vel_shift=True)
 
 # Create galaxy object
 galaxy = load_CAMELS_IllustrisTNG(
@@ -45,10 +45,9 @@
     np.random.normal(100, 500, galaxy.stars.coordinates.shape) * km / s
 )
 
-# Get the spectra (this will automatically use the tau_vs we just calculated
-# since the emission model has tau_v="tau_v")
+# Get the spectra
 start_with_shift = time.time()
-galaxy.stars.get_spectra(model, vel_shift=True)
+galaxy.stars.get_spectra(model)
 print(
     "Time to get spectra with velocity shift: "
     f"{time.time() - start_with_shift}"
@@ -60,7 +59,10 @@
 # Clear the spectra
 galaxy.clear_all_emissions()
 
-# Get the spectra without the velocity broadening
+# Get the spectra without the velocity broadening (we could use the
+# set_vel_shift method on an EmissionModel to turn off the velocity shift
+# but here we demonstrate how to do it with the get_spectra method's
+# overides)
 start_without_shift = time.time()
 galaxy.stars.get_spectra(model, vel_shift=False)
 print(

src/synthesizer/base_galaxy.py

@@ -791,7 +791,7 @@ def get_spectra(
         fesc=None,
         covering_fraction=None,
         mask=None,
-        vel_shift=False,
+        vel_shift=None,
         verbose=True,
         **kwargs,
     ):

src/synthesizer/components/component.py

@@ -244,7 +244,7 @@ def get_spectra(
         tau_v=None,
         fesc=None,
         mask=None,
-        vel_shift=False,
+        vel_shift=None,
         verbose=True,
         **kwargs,
     ):

src/synthesizer/emission_models/base_model.py

@@ -151,6 +151,10 @@ class EmissionModel(Extraction, Generation, DustAttenuation, Combination):
             "per particle". If True, the spectra and lines will be stored per
             particle. Integrated spectra are made automatically by summing the
             per particle spectra. Default is False.
+        vel_shift (bool):
+            A flag for whether the emission produced by this model should take
+            into account the velocity shift due to peculiar velocities.
+            Only applicable to particle based emitters. Default is False.
     """
 
     # Define quantities
@@ -180,6 +184,7 @@ def __init__(
         post_processing=(),
         save=True,
         per_particle=False,
+        vel_shift=False,
         **kwargs,
     ):
         """
@@ -269,6 +274,10 @@ def __init__(
                 be "per particle". If True, the spectra and lines will be
                 stored per particle. Integrated spectra are made automatically
                 by summing the per particle spectra. Default is False.
+            vel_shift (bool):
+                A flag for whether the emission produced by this model should
+                take into account the velocity shift due to peculiar
+                velocities. Default is False.
             **kwargs:
                 Any additional keyword arguments to store. These can be used
                 to store additional information needed by the model.
@@ -332,6 +341,7 @@ def __init__(
             lum_attenuated_model=lum_attenuated_model,
             fesc=fesc,
             lam_mask=lam_mask,
+            vel_shift=vel_shift,
         )
 
         # Containers for children and parents
@@ -391,6 +401,7 @@ def _init_operations(
         lum_attenuated_model,
         fesc,
         lam_mask,
+        vel_shift,
     ):
         """
         Initialise the correct parent operation.
@@ -423,10 +434,14 @@ def _init_operations(
                 The escape fraction.
             lam_mask (ndarray):
                 The mask to apply to the wavelength array.
+            vel_shift (bool):
+                A flag for whether the emission produced by this model should
+                take into account the velocity shift due to peculiar
+                velocities. Particle emitters only.
         """
         # Which operation are we doing?
         if self._is_extracting:
-            Extraction.__init__(self, grid, extract, fesc, lam_mask)
+            Extraction.__init__(self, grid, extract, fesc, lam_mask, vel_shift)
         elif self._is_combining:
             Combination.__init__(self, combine)
         elif self._is_dust_attenuating:
@@ -1029,6 +1044,32 @@ def set_per_particle(self, per_particle):
         # Unpack the model now we're done
         self.unpack_model()
 
+    @property
+    def vel_shift(self):
+        """Get the velocity shift flag."""
+        if hasattr(self, "_use_vel_shift"):
+            return self._use_vel_shift
+        else:
+            return False
+
+    def set_vel_shift(self, vel_shift, set_all=False):
+        """
+        Set whether we should apply velocity shifts to the spectra.
+
+        Only applicable to particle emitters.
+
+        Args:
+            vel_shift (bool):
+                Whether to set the velocity shift flag.
+            set_all (bool):
+                Whether to set the emitter on all models.
+        """
+        if not set_all:
+            self._use_vel_shift = vel_shift
+        else:
+            for model in self._models.values():
+                model.set_vel_shift(vel_shift)
+
     @property
     def lum_intrinsic_model(self):
         """Get the intrinsic model for computing dust luminosity."""
@@ -2029,7 +2070,14 @@ def plot_emission_tree(
         return fig, ax
 
     def _apply_overrides(
-        self, emission_model, dust_curves, tau_v, fesc, covering_fraction, mask
+        self,
+        emission_model,
+        dust_curves,
+        tau_v,
+        fesc,
+        covering_fraction,
+        mask,
+        vel_shift,
     ):
         """
         Apply overrides to an emission model copy.
@@ -2096,6 +2144,13 @@ def _apply_overrides(
                     - A dictionary of the form:
                       {<label>: {"attr": <attr>, "thresh": <thresh>, "op":<op>}
                       to add a specific mask to a particular model.
+            vel_shift (dict/bool):
+                Overide the models flag for using peculiar velocities to apply
+                doppler shift to the generated spectra. Only applicable for
+                particle spectra. Can be a boolean to apply to all models or a
+                dictionary of the form:
+                    {<label>: bool}
+                to apply to specific models.
         """
         # If we have dust curves to apply, apply them
         if dust_curves is not None:
@@ -2142,6 +2197,16 @@ def _apply_overrides(
             for label, mask in mask.items():
                 emission_model[label].add_mask(**mask)
 
+        # If we have velocity shifts to apply, apply them
+        if vel_shift is not None:
+            if isinstance(vel_shift, dict):
+                for label, value in vel_shift.items():
+                    emission_model._models[label].set_vel_shift(value)
+            else:
+                for model in emission_model._models.values():
+                    if model._is_extracting:
+                        model.set_vel_shift(vel_shift)
+
     def _get_spectra(
         self,
         emitters,
@@ -2150,7 +2215,7 @@ def _get_spectra(
         fesc=None,
         covering_fraction=None,
         mask=None,
-        vel_shift=False,
+        vel_shift=None,
         verbose=True,
         spectra=None,
         particle_spectra=None,
@@ -2229,7 +2294,9 @@ def _get_spectra(
                 A dictionary of particle spectra to add to. This is used for
                 recursive calls to this function.
             vel_shift (bool)
-                Flags whether to apply doppler shift to the spectrum.
+                Overide the models flag for using peculiar velocities to apply
+                doppler shift to the generated spectra. Only applicable for
+                particle spectra.
             _is_related (bool)
                 Are we generating related model spectra? If so we don't want
                 to apply any post processing functions or delete any spectra,
@@ -2262,7 +2329,13 @@ def _get_spectra(
 
         # Apply any overides we have
         self._apply_overrides(
-            emission_model, dust_curves, tau_v, fesc, covering_fraction, mask
+            emission_model,
+            dust_curves,
+            tau_v,
+            fesc,
+            covering_fraction,
+            mask,
+            vel_shift,
         )
 
         # Make a spectra dictionary if we haven't got one yet
@@ -2286,7 +2359,6 @@ def _get_spectra(
             emitters,
             spectra,
             particle_spectra,
-            vel_shift=vel_shift,
             verbose=verbose,
             **kwargs,
         )
@@ -2570,7 +2642,13 @@ def _get_lines(
 
         # Apply any overides we have
         self._apply_overrides(
-            emission_model, dust_curves, tau_v, fesc, covering_fraction, mask
+            emission_model,
+            dust_curves,
+            tau_v,
+            fesc,
+            covering_fraction,
+            mask,
+            None,
         )
 
         # If we haven't got a lines dictionary yet we'll make one

src/synthesizer/emission_models/operations.py

@@ -32,7 +32,7 @@ class Extraction:
             The escape fraction.
     """
 
-    def __init__(self, grid, extract, fesc, lam_mask):
+    def __init__(self, grid, extract, fesc, lam_mask, vel_shift):
         """
         Initialise the extraction model.
 
@@ -45,6 +45,9 @@ def __init__(self, grid, extract, fesc, lam_mask):
                 The escape fraction.
             lam_mask (ndarray):
                 The wavelength mask to apply to the spectra.
+            vel_shift (bool):
+                Should the emission take into account the velocity shift due
+                to peculiar velocities? (Particle Only!)
         """
         # Attach the grid
         self._grid = grid
@@ -58,13 +61,16 @@ def __init__(self, grid, extract, fesc, lam_mask):
         # Attach the wavelength mask
         self._lam_mask = lam_mask
 
+        # Should the emission take into account the velocity shift due to
+        # peculiar velocities? (Particle Only!)
+        self._use_vel_shift = vel_shift
+
     def _extract_spectra(
         self,
         emission_model,
         emitters,
         spectra,
         particle_spectra,
-        vel_shift,
         verbose,
         **kwargs,
     ):
@@ -80,8 +86,6 @@ def _extract_spectra(
                 The dictionary to store the extracted spectra in.
             particle_spectra (dict):
                 The dictionary to store the extracted particle spectra in.
-            vel_shift (bool):
-                Flags whether to apply doppler shift to the spectra
             verbose (bool):
                 Are we talking?
             kwargs (dict):
@@ -126,6 +130,7 @@ def _extract_spectra(
                 generator_func = emitter.generate_lnu
 
             # Get this base spectra
+            print(f"Extracting {label} from {spectra_key}:", self.vel_shift)
             sed = Sed(
                 emission_model.lam,
                 generator_func(
@@ -135,7 +140,7 @@ def _extract_spectra(
                     if isinstance(this_model.fesc, str)
                     else this_model.fesc,
                     mask=this_mask,
-                    vel_shift=vel_shift,
+                    vel_shift=self.vel_shift,
                     lam_mask=this_model._lam_mask,
                     verbose=verbose,
                     **kwargs,
@@ -349,6 +354,7 @@ def _extract_summary(self):
         summary.append(f"  Grid: {self._grid.grid_name}")
         summary.append(f"  Extract key: {self._extract}")
         summary.append(f"  Escape fraction: {self._fesc}")
+        summary.append(f"  Use velocity shift: {self._use_vel_shift}")
 
         return summary
 

src/synthesizer/particle/blackholes.py

@@ -929,7 +929,7 @@ def get_particle_spectra(
         tau_v=None,
         covering_fraction=None,
         mask=None,
-        vel_shift=False,
+        vel_shift=None,
         verbose=True,
         **kwargs,
     ):