Add parametric doppler broadening to Sed objects (#1058)

Author: WillJRoper

docs/source/emissions/emission_objects/sed_example.ipynb

@@ -981,6 +981,82 @@
     "# Get attenuation at 5500 angstrom\n",
     "att_at_5500 = get_attenuation_at_5500(int_sed, att_sed)"
    ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "9d38967c",
+   "metadata": {},
+   "source": [
+    "### Doppler broadening\n",
+    "\n",
+    "`Sed` objects also provide tools for producing Doppler-broadened versions of an SED. The method `doppler_broaden` returns a new SED broadened using a specified velocity dispersion. This can either operate in-place or return a new `Sed` object.\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "0942c4b6",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from unyt import km, s\n",
+    "\n",
+    "log10age = 6.0  # log10(age/yr)\n",
+    "metallicity = 0.01\n",
+    "spectra_type = \"nebular\"\n",
+    "grid_point = grid.get_grid_point(log10ages=log10age, metallicity=metallicity)\n",
+    "\n",
+    "# Extract SED\n",
+    "sed = grid.get_sed_at_grid_point(grid_point, spectra_type=spectra_type)\n",
+    "\n",
+    "# Apply Doppler broadening specifying a velocity\n",
+    "sed_broadened1 = sed.doppler_broaden(1000.0 * km / s)\n",
+    "\n",
+    "# Plot to compare\n",
+    "fig, ax = sed.plot_spectra(label=\"original\")\n",
+    "fig, ax = sed_broadened1.plot_spectra(\n",
+    "    label=\"broadened\", fig=fig, ax=ax, xlimits=(1000, 1400)\n",
+    ")\n",
+    "ax.set_xscale(\"linear\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ea1d0b35",
+   "metadata": {},
+   "source": [
+    "### Thermal Broadening\n",
+    "\n",
+    "There is also a variant of this method for thermal broadening, called `thermally_broaden`, which computes the broadening from a given temperature and mean molecular mass (which defaults to `1*amu`). This can either operate in-place or return a new `Sed` object."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "236895aa",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from unyt import K\n",
+    "\n",
+    "sed_thermal = sed.thermally_broaden(1e4 * K)\n",
+    "fig, ax = sed.plot_spectra(label=\"original\")\n",
+    "fig, ax = sed_thermal.plot_spectra(\n",
+    "    label=\"thermally broadened\", fig=fig, ax=ax, xlimits=(1000, 1400)\n",
+    ")\n",
+    "ax.set_xscale(\"linear\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "82167512",
+   "metadata": {},
+   "source": [
+    "These methods are particularly useful when modelling environments where thermal motions contribute to line broadening, such as H II regions.\n",
+    "\n",
+    "Note that, if you have a particle based component, self-consistent Doppler broadening based on particle velocities can be applied during emission generation by enabling `vel_shift` on an `EmissionModel`."
+   ]
   }
  ],
  "metadata": {},

src/synthesizer/emissions/sed.py

@@ -19,15 +19,20 @@
 import numpy as np
 from matplotlib.colors import LogNorm
 from scipy.interpolate import interp1d
+from scipy.signal import fftconvolve
 from scipy.stats import linregress
 from spectres import spectres
 from unyt import (
     Hz,
+    K,
+    amu,
     angstrom,
     c,
     erg,
     eV,
     h,
+    kb,
+    km,
     pc,
     s,
     unyt_array,
@@ -1598,6 +1603,88 @@ def calculate_ionising_photon_production_rate(
 
         return ion_photon_prod_rate
 
+    @accepts(sigma_v=km / s)
+    def doppler_broaden(self, sigma_v, inplace=False):
+        """Doppler broaden the spectra.
+
+        Doppler broadens the spectra either in place or returning a new Sed.
+
+        Args:
+            sigma_v (unyt_array):
+                The velocity dispersion to broaden the spectra by.
+            inplace (bool):
+                Flag for whether to modify self or return a new Sed.
+
+        Returns:
+            Sed:
+                A new Sed containing the rest frame spectra of self broadened
+                by the provided velocity dispersion. Only returned if
+                in_place is False.
+        """
+        # Convert wavelength grid to log-lambda space
+        x = np.log(self.lam)
+
+        # Regrid to uniform log-lambda spacing, using a very fine grid
+        x_uniform = np.linspace(x.min(), x.max(), 100000)
+        lnu_uniform = spectres(x_uniform, x, self.lnu, fill=0.0, verbose=False)
+
+        # Convert velocity sigma to log-lambda sigma
+        # Δx = ln(λ) gives Δv = c*Δx
+        sigma_x = sigma_v / c
+
+        # Gaussian kernel in log-lambda
+        dx = x_uniform[1] - x_uniform[0]
+        N = len(x_uniform)
+        half = N // 2
+
+        # Define and normalise the broadening kernel
+        kernel_x = (np.arange(N) - half) * dx
+        kernel = np.exp(-(kernel_x**2) / (2 * sigma_x**2))
+        kernel /= kernel.sum()
+
+        # Convolution in velocity/log-lambda space
+        lnu_broad = fftconvolve(lnu_uniform, kernel, mode="same")
+
+        # Re-interpolate back onto original wavelength grid and update the sed
+        new_lnu = (
+            spectres(x, x_uniform, lnu_broad, fill=0.0, verbose=False)
+            * self.lnu.units
+        )
+
+        # Return new Sed or modify in place
+        if inplace:
+            self.lnu = new_lnu
+            return self
+        else:
+            return Sed(self.lam, new_lnu)
+
+    @accepts(temperature=K, mu=amu)
+    def thermally_broaden(self, temperature, mu=1.0 * amu, inplace=False):
+        """Create a spectra including the thermal broadening.
+
+        This simply calculates the velocity dispersion from the temperature
+        and mean molecular weight.
+
+        Args:
+            temperature (unyt_array):
+                The temperature to broaden the spectra by.
+            mu (unyt_array):
+                The mean molecular weight of the gas. By default assumed to be
+                1 amu.
+            inplace (bool):
+                Flag for whether to modify self or return a new Sed.
+
+        Returns:
+            Sed:
+                A new Sed containing the rest frame spectra of self
+                broadened by the provided velocity dispersion. Only returned if
+                inplace is False.
+        """
+        # Calculate the velocity dispersion
+        sigma_v = np.sqrt(kb * temperature / mu)
+
+        return self.doppler_broaden(sigma_v, inplace=inplace)
+
     def plot_spectra(self, **kwargs):
         """Plot the spectra.