fix test and units to output

jvshields · Jun 21, 2024 · c544eb1 · c544eb1
1 parent 57e164d
commit c544eb1
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 5 deletions.
diff --git a/stardis/radiation_field/opacities/opacities_solvers/broadening.py b/stardis/radiation_field/opacities/opacities_solvers/broadening.py
@@ -749,4 +749,12 @@ def rotation_broadening(
     if np.abs(v_rot) < 1e-5:
         return (wavelength, flux)
 
-    return wavelength, convolve1d(flux, rotational_profile / rotational_profile.sum())
+    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
@@ -616,10 +616,11 @@ def test_calc_gamma_van_der_waals_cuda_wrapped_sample_cuda_values(
     )
 
 
-def test_rotational_broadening(example_stellar_output):
+def test_rotational_broadening(example_stardis_output):
     actual_wavelengths, actual_fluxes = rotation_broadening(
         20 * u.km / u.s,
-        example_stellar_output.lambdas,
-        example_stellar_output.spectrum_lambda,
+        example_stardis_output.lambdas,
+        example_stardis_output.spectrum_lambda,
+        v_rot=0 * u.km / u.s,
     )
-    assert np.allclose(actual_wavelengths, example_stellar_output.lambdas)
+    assert np.allclose(actual_fluxes, example_stardis_output.spectrum_lambda)