Migrate to tardis composition (#188)

* restructure high level code to allow modification of model after creation * remove comment * fix conftext import * apply black * add docstrings * apply black (again) * checkpoint * accept lists of elements * add config option for rescaling abundances * fix merge conflict base problem * apply black * make broadening test also rescale abundances * debug github tests * continue debugging github tests * first checkpoint to use tardis composition * migrate to tardis composition * apply black * fix broadening tests * continue debugging tests * more test debugging * more test debugging attempts * fix inconsistent spacing in parametrizations in broadening tests * turn off cupy tests * fix tests - change output of one test b/c hydrogen mass is no longer an put * cleanup unused stardis composition * fix benchmarks * unpin asv * fix incorrect atom data truncation * change typo for mass fraction index * fix typos * explicit composotion arg specifications * fix conftest
tardis-sn · Jun 12, 2024 · c440059 · c440059
1 parent 74ce810
commit c440059
Show file tree

Hide file tree

Showing 14 changed files with 66 additions and 91 deletions.
diff --git a/.github/workflows/benchmarks.yml b/.github/workflows/benchmarks.yml
@@ -67,7 +67,7 @@ jobs:
           key: atom-data
 
       - name: Install asv
-        run: pip install asv==0.5.*
+        run: pip install asv
 
       - name: Accept all asv questions
         run: asv machine --yes

diff --git a/benchmarks/run_stardis.py b/benchmarks/run_stardis.py
@@ -58,7 +58,7 @@ def setup(self):
                 np.min(
                     [
                         len(
-                            stellar_model.composition.atomic_mass_fraction.columns.tolist()
+                            stellar_model.composition.elemental_mass_fraction.columns.tolist()
                         ),
                         config.model.final_atomic_number,
                     ]
@@ -169,7 +169,7 @@ def setup(self):
                 np.min(
                     [
                         len(
-                            stellar_model.composition.atomic_mass_fraction.columns.tolist()
+                            stellar_model.composition.elemental_mass_fraction.columns.tolist()
                         ),
                         config.model.final_atomic_number,
                     ]

diff --git a/setup.cfg b/setup.cfg
@@ -28,7 +28,6 @@ test =
     pytest
     pytest-doctestplus
     pytest-cov
-    nbmake
 docs =
     sphinx
     sphinx-automodapi
@@ -37,7 +36,7 @@ docs =
 testpaths = "stardis" "docs"
 doctest_plus = enabled
 text_file_format = rst
-addopts = --doctest-rst --nbmake --nbmake-kernel=python3
+addopts = --doctest-rst 
 
 [coverage:run]
 omit =

diff --git a/stardis/base.py b/stardis/base.py
@@ -31,7 +31,6 @@ def run_stardis(config_fname, tracing_lambdas_or_nus):
 
     config, adata, stellar_model = parse_config_to_model(config_fname)
     set_num_threads(config.n_threads)
-
     stellar_plasma = create_stellar_plasma(stellar_model, adata, config)
     stellar_radiation_field = create_stellar_radiation_field(
         tracing_nus, stellar_model, stellar_plasma, config

diff --git a/stardis/config_schema.yml b/stardis/config_schema.yml
@@ -1,6 +1,5 @@
 $schema: http://json-schema.org/draft-04/schema#
 type: object
-additionalProperties: false
 properties:
     stardis_config_version:
         enum:
@@ -15,7 +14,6 @@ properties:
         description: Path to the atomic data file. The file should be in hdf5 format output by carsus.
     model:
         type: object
-        additionalProperties: false
         properties:
             type:
                 enum:
@@ -38,7 +36,10 @@ properties:
                 multipleOf: 1
                 default: -99
                 description: If -99, use all shells. If positive, use only specified outermost depth points. Only used for mesa models. 
-
+            elemental_rescaling_dict:
+              type: object
+              default: {}
+              description: Rescale the abundances of the elements in the model. The keys are the atomic numbers of the elements to rescale and the values are the rescaling factors.
         required:
         - type
         - fname

diff --git a/stardis/conftest.py b/stardis/conftest.py
@@ -77,7 +77,7 @@ def example_stellar_plasma(
             np.min(
                 [
                     len(
-                        example_stellar_model.composition.atomic_mass_fraction.columns.tolist()
+                        example_stellar_model.composition.elemental_mass_fraction.columns.tolist()
                     ),
                     example_config.model.final_atomic_number,
                 ]

diff --git a/stardis/io/base.py b/stardis/io/base.py
@@ -74,7 +74,7 @@ def parse_config_to_model(config_fname):
             np.min(
                 [
                     len(
-                        stellar_model.composition.atomic_mass_fraction.columns.tolist()
+                        stellar_model.composition.elemental_mass_fraction.columns.tolist()
                     ),
                     config.model.final_atomic_number,
                 ]
@@ -86,4 +86,14 @@ def parse_config_to_model(config_fname):
         continuum_interaction_species=[],
     )
 
+    # if (
+    #     not config.model.elemental_rescaling_dict
+    # ):  # Pass if no rescaling is requested, else rescale by dictionary values provided
+    #     pass
+    # else:
+    #     stellar_model.composition.rescale_elements(
+    #         list(config.model.elemental_rescaling_dict.keys()),
+    #         list(config.model.elemental_rescaling_dict.values()),
+    #     )
+
     return config, adata, stellar_model
diff --git a/stardis/io/model/marcs.py b/stardis/io/model/marcs.py
@@ -8,9 +8,8 @@
 
 
 from stardis.model.geometry.radial1d import Radial1DGeometry
-from stardis.model.composition.base import Composition
-
 from stardis.model.base import StellarModel
+from tardis.model.matter.composition import Composition
 
 
 @dataclass
@@ -57,7 +56,16 @@ def to_composition(self, atom_data, final_atomic_number):
         atomic_mass_fraction = self.convert_marcs_raw_abundances_to_mass_fractions(
             atom_data, final_atomic_number
         )
-        return Composition(density, atomic_mass_fraction)
+
+        atomic_mass_fraction["mass_number"] = -1
+        atomic_mass_fraction.set_index("mass_number", append=True, inplace=True)
+
+        return Composition(
+            density,
+            atomic_mass_fraction,
+            raw_isotope_abundance=None,
+            element_masses=atom_data.atom_data.mass.copy(),
+        )
 
     def convert_marcs_raw_abundances_to_mass_fractions(
         self, atom_data, final_atomic_number

diff --git a/stardis/io/model/mesa.py b/stardis/io/model/mesa.py
@@ -6,7 +6,7 @@
 import logging
 
 from stardis.model.geometry.radial1d import Radial1DGeometry
-from stardis.model.composition.base import Composition
+from tardis.model.matter.composition import Composition
 
 from stardis.model.base import StellarModel
 from stardis.io.model.util import (
@@ -85,8 +85,16 @@ def to_uniform_composition_from_solar(
             data=np.repeat(solar_profile.values, len(self.data), axis=1),
         )
 
+        atomic_mass_fraction["mass_number"] = -1
+        atomic_mass_fraction.set_index("mass_number", append=True, inplace=True)
+
         atomic_mass_fraction.index.name = "atomic_number"
-        return Composition(density, atomic_mass_fraction)
+        return Composition(
+            density,
+            atomic_mass_fraction,
+            raw_isotope_abundance=None,
+            element_masses=atom_data.atom_data.mass.copy(),
+        )
 
     def to_stellar_model(
         self,

diff --git a/stardis/model/composition/__init__.py b/stardis/model/composition/__init__.py
diff --git a/stardis/model/composition/base.py b/stardis/model/composition/base.py
diff --git a/stardis/plasma/base.py b/stardis/plasma/base.py
@@ -544,9 +544,9 @@ def create_stellar_plasma(
     return BasePlasma(
         plasma_properties=plasma_modules,
         t_rad=stellar_model.temperatures.value,
-        abundance=stellar_model.composition.atomic_mass_fraction,
+        abundance=stellar_model.composition.elemental_mass_fraction,
         atomic_data=atom_data,
-        density=stellar_model.composition.density.value,
+        number_density=stellar_model.composition.elemental_number_density,
         link_t_rad_t_electron=1.0,
         nlte_ionization_species=[],
         nlte_excitation_species=[],

diff --git a/stardis/radiation_field/opacities/opacities_solvers/broadening.py b/stardis/radiation_field/opacities/opacities_solvers/broadening.py
@@ -18,6 +18,7 @@
 ELEMENTARY_CHARGE = float(const.e.esu.value)
 BOHR_RADIUS = float(const.a0.cgs.value)
 VACUUM_ELECTRIC_PERMITTIVITY = 1.0 / (4.0 * PI)
+H_MASS = float(const.m_p.cgs.value)
 
 
 @numba.njit
@@ -410,7 +411,6 @@ def _calc_gamma_van_der_waals(
     n_eff_lower,
     temperature,
     h_density,
-    h_mass,
 ):
     """
     Calculates broadening parameter for van der Waals broadening.
@@ -428,21 +428,18 @@ def _calc_gamma_van_der_waals(
         Temperature of depth points being considered.
     h_density : float
         Number density of Hydrogen at depth point being considered.
-    h_mass : float
-        Atomic mass of Hydrogen in grams.
 
     Returns
     -------
     gamma_van_der_waals : float
         Broadening parameter for van der Waals broadening.
     """
-    ion_number, n_eff_upper, n_eff_lower, temperature, h_density, h_mass = (
+    ion_number, n_eff_upper, n_eff_lower, temperature, h_density = (
         int(ion_number),
         float(n_eff_upper),
         float(n_eff_lower),
         float(temperature),
         float(h_density),
-        float(h_mass),
     )
     c6 = (
         6.46e-34
@@ -455,7 +452,7 @@ def _calc_gamma_van_der_waals(
 
     gamma_van_der_waals = (
         17
-        * (8 * BOLTZMANN_CONSTANT * temperature / (PI * h_mass)) ** 0.3
+        * (8 * BOLTZMANN_CONSTANT * temperature / (PI * H_MASS)) ** 0.3
         * c6**0.4
         * h_density
     )
@@ -470,15 +467,13 @@ def calc_gamma_van_der_waals(
     n_eff_lower,
     temperature,
     h_density,
-    h_mass,
 ):
     return _calc_gamma_van_der_waals(
         ion_number,
         n_eff_upper,
         n_eff_lower,
         temperature,
         h_density,
-        h_mass,
     )
 
 
@@ -490,7 +485,6 @@ def _calc_gamma_van_der_waals_cuda(
     n_eff_lower,
     temperature,
     h_density,
-    h_mass,
 ):
     tid = cuda.grid(1)
     size = len(res)
@@ -502,7 +496,6 @@ def _calc_gamma_van_der_waals_cuda(
             n_eff_lower[tid],
             temperature[tid],
             h_density[tid],
-            h_mass[tid],
         )
 
 
@@ -512,7 +505,6 @@ def calc_gamma_van_der_waals_cuda(
     n_eff_lower,
     temperature,
     h_density,
-    h_mass,
     nthreads=256,
     ret_np_ndarray=False,
     dtype=float,
@@ -523,7 +515,6 @@ def calc_gamma_van_der_waals_cuda(
         n_eff_lower,
         temperature,
         h_density,
-        h_mass,
     )
 
     shortest_arg_idx = np.argmin(map(len, arg_list))
@@ -554,7 +545,6 @@ def calc_gamma(
     electron_density,
     temperature,
     h_density,
-    h_mass,
     linear_stark=True,
     quadratic_stark=True,
     van_der_waals=True,
@@ -584,8 +574,6 @@ def calc_gamma(
         Temperature of depth points being considered.
     h_density : float
         Number density of Hydrogen at depth point being considered.
-    h_mass : float
-        Atomic mass of Hydrogen in grams.
     linear_stark : bool, optional
         True if linear Stark broadening is to be considered, otherwise False.
         By default True.
@@ -636,7 +624,6 @@ def calc_gamma(
             n_eff_lower,
             temperature,
             h_density,
-            h_mass,
         )
     else:
         gamma_van_der_waals = np.zeros_like(gamma_linear_stark)
@@ -698,7 +685,6 @@ def calculate_broadening(
         electron_density=stellar_plasma.electron_densities.values,
         temperature=stellar_model.temperatures.value,
         h_density=stellar_plasma.ion_number_density.loc[1, 0].values,
-        h_mass=stellar_plasma.atomic_mass.loc[1],
         linear_stark=linear_stark,
         quadratic_stark=quadratic_stark,
         van_der_waals=van_der_waals,
@@ -708,7 +694,9 @@ def calculate_broadening(
     doppler_widths = calc_doppler_width(
         lines.nu.values[:, np.newaxis],
         stellar_model.temperatures.value,
-        stellar_plasma.atomic_mass.loc[lines.atomic_number].values[:, np.newaxis],
+        stellar_model.composition.nuclide_masses.loc[lines.atomic_number].values[
+            :, np.newaxis
+        ],
     )
 
     return gammas, doppler_widths