fix bug and test additional case

tardis-sn · Jun 13, 2024 · 6c2ea1a · 6c2ea1a
1 parent 7cfa55f
commit 6c2ea1a
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 3 deletions.
diff --git a/stardis/io/model/tests/test_model_io.py b/stardis/io/model/tests/test_model_io.py
@@ -72,9 +72,14 @@ def test_marcs_model(marcs_model):
 
 def test_rescale_nuclide_mass_fraction(example_stellar_model):
     rescaled = rescale_nuclide_mass_fractions(
-        example_stellar_model.composition.nuclide_mass_fraction, [5], [0.5]
+        example_stellar_model.composition.nuclide_mass_fraction, [4, 5], [1.1, 0.8]
     )
     assert np.allclose(
         rescaled.loc[5].values,
-        example_stellar_model.composition.nuclide_mass_fraction.loc[5].values * 0.5,
+        example_stellar_model.composition.nuclide_mass_fraction.loc[5].values * 0.8,
+    )
+
+    assert np.allclose(
+        rescaled.loc[5].values,
+        example_stellar_model.composition.nuclide_mass_fraction.loc[4].values * 1.1,
     )
diff --git a/stardis/io/model/util.py b/stardis/io/model/util.py
@@ -83,7 +83,7 @@ def rescale_nuclide_mass_fractions(nuclide_mass_fraction, nuclides, scale_factor
         if nuclide not in new_mass_fractions.columns:
             raise ValueError(f"{nuclide} not available in the simulation")
 
-        new_mass_fractions[nuclides] = new_mass_fractions[nuclides] * scale_factor
+        new_mass_fractions[nuclide] = new_mass_fractions[nuclide] * scale_factor
 
     return new_mass_fractions.T.div(
         new_mass_fractions.T.sum(axis=0)