From 9d5544c1650b43703bb8f3cdea2b57b25bef00d2 Mon Sep 17 00:00:00 2001
From: Joshua Shields <54691495+jvshields@users.noreply.github.com>
Date: Mon, 24 Jun 2024 16:42:05 -0400
Subject: [PATCH] add rotation broadening (#198)

* add rotation broadening

* fix test and units to output

* small regression test
---
 .../opacities/opacities_solvers/broadening.py | 60 ++++++++++++++++++-
 .../tests/test_broadening.py                  | 34 ++++++++++-
 2 files changed, 92 insertions(+), 2 deletions(-)

diff --git a/stardis/radiation_field/opacities/opacities_solvers/broadening.py b/stardis/radiation_field/opacities/opacities_solvers/broadening.py
index a07b3780..6dd2dc77 100644
--- a/stardis/radiation_field/opacities/opacities_solvers/broadening.py
+++ b/stardis/radiation_field/opacities/opacities_solvers/broadening.py
@@ -1,8 +1,9 @@
 import numpy as np
-from astropy import constants as const
+from astropy import constants as const, units as u
 import math
 import numba
 from numba import cuda
+from scipy.ndimage import convolve1d
 
 GPUs_available = cuda.is_available()
 
@@ -19,6 +20,7 @@
 BOHR_RADIUS = float(const.a0.cgs.value)
 VACUUM_ELECTRIC_PERMITTIVITY = 1.0 / (4.0 * PI)
 H_MASS = float(const.m_p.cgs.value)
+C_KMS = float(const.c.to(u.km / u.s).value)
 
 
 @numba.njit
@@ -700,3 +702,59 @@ def calculate_broadening(
     )
 
     return gammas, doppler_widths
+
+
+def rotation_broadening(
+    velocity_per_pix, wavelength, flux, v_rot=0 * u.km / u.s, limb_darkening=0.6
+):
+    """Convolve a spectrum with a rotational broadening profile. Only accurate if the velocity_per_pix is constant.
+    Taken from starkit https://github.com/starkit/starkit/blob/57b919a79c1fd10e61af6036ec9ac56f38a6f883/starkit/base/operations/stellar.py#L14
+    Originally adapted from Observations of Stellar Photospheres by David Gray.
+
+    Parameters:
+        velocity_per_pix : astropy.units.Quantity
+            Velocity resolution element of the spectrum
+        wavelength : astropy.units.Quantity
+            Wavelengths of the spectrum
+        flux : astropy.units.Quantity
+            Fluxes of the spectrum
+        v_rot : astropy.units.Quantity
+            Rotational velocity, v sin(i) of the star
+        limb_darkening : float (0.6)
+            Limb darkening coefficient of the star. Default is 0.6.
+
+    Returns:
+        wavelengths: astropy.units.Quantity
+            Wavelengths of the convolved spectrum
+        fluxes: astropy.units.Quantity
+            Fluxes of the convolved spectrum
+    """
+
+    velocity_per_pix = velocity_per_pix.to(u.km / u.s).value
+    v_rot = v_rot.to(u.km / u.s).value
+
+    v_rot_by_c = np.maximum(1e-5, np.abs(v_rot)) / C_KMS
+
+    half_width_pix = np.round((v_rot / velocity_per_pix)).astype(int)
+    profile_velocity = (
+        np.linspace(-half_width_pix, half_width_pix, 2 * half_width_pix + 1)
+        * velocity_per_pix
+    )
+    profile = np.maximum(0.0, 1.0 - (profile_velocity / v_rot) ** 2)
+
+    rotational_profile = (
+        2 * (1 - limb_darkening) * profile**0.5 + 0.5 * PI * limb_darkening * profile
+    ) / (PI * v_rot_by_c * (1 - limb_darkening / 3))
+
+    if np.abs(v_rot) < 1e-5:
+        return (wavelength, flux)
+
+    broadened_fluxes = (
+        convolve1d(flux, rotational_profile / rotational_profile.sum())
+        * u.erg
+        / u.s
+        / u.cm**2
+        / u.Angstrom
+    )
+
+    return wavelength, broadened_fluxes
diff --git a/stardis/radiation_field/opacities/opacities_solvers/tests/test_broadening.py b/stardis/radiation_field/opacities/opacities_solvers/tests/test_broadening.py
index 921be7e1..1708e5ab 100644
--- a/stardis/radiation_field/opacities/opacities_solvers/tests/test_broadening.py
+++ b/stardis/radiation_field/opacities/opacities_solvers/tests/test_broadening.py
@@ -1,6 +1,6 @@
 import pytest
 import numpy as np
-from astropy import constants as const
+from astropy import constants as const, units as u
 
 from stardis.radiation_field.opacities.opacities_solvers.broadening import (
     calc_doppler_width,
@@ -18,6 +18,7 @@
     calc_gamma_van_der_waals,
     _calc_gamma_van_der_waals_cuda,
     calc_gamma_van_der_waals_cuda,
+    rotation_broadening,
 )
 
 GPUs_available = False  # cuda.is_available()
@@ -613,3 +614,34 @@ def test_calc_gamma_van_der_waals_cuda_wrapped_sample_cuda_values(
         calc_gamma_van_der_waals_cuda(*map(cp.asarray, arg_list)),
         calc_gamma_van_der_waals_sample_values_expected_result,
     )
+
+
+def test_rotational_broadening(example_stardis_output):
+    actual_wavelengths, actual_fluxes_no_broadening = rotation_broadening(
+        20 * u.km / u.s,
+        example_stardis_output.lambdas,
+        example_stardis_output.spectrum_lambda,
+        v_rot=0 * u.km / u.s,
+    )
+
+    expected_broadening_fluxes = [
+        34069895.2932162,
+        34069822.27391726,
+        34069676.41567875,
+        34069458.07948008,
+        34069167.80744836,
+        34068806.32366434,
+        34068374.53527089,
+        34067873.53390926,
+    ]
+    actual_wavelengths, actual_fluxes = rotation_broadening(
+        20 * u.km / u.s,
+        example_stardis_output.lambdas,
+        example_stardis_output.spectrum_lambda,
+        v_rot=500 * u.km / u.s,
+    )
+    assert np.allclose(actual_wavelengths, example_stardis_output.lambdas)
+    assert np.allclose(
+        actual_fluxes_no_broadening, example_stardis_output.spectrum_lambda
+    )
+    assert np.allclose(actual_fluxes[:8].value, expected_broadening_fluxes)