append dummy index to references and fix docs (tardis-sn#3315)

* append dummy index to references

* remove fill for source since they should should never be empty

* remove whitespace change

* change solve in montecarlo tutorial

* add hydrogen only test

* add tolerances

* rtol 1e-12

Author: jvshields

tardis/conftest.py

@@ -4,6 +4,7 @@
 
 import pytest
 from astropy.version import version as astropy_version
+from astropy import units as u
 
 from tardis import run_tardis
 from tardis.io.configuration.config_reader import Configuration
@@ -168,11 +169,17 @@ def generate_reference(request):
 
 @pytest.fixture(scope="session")
 def tardis_regression_path(request):
-    tardis_regression_path = request.config.getoption("--tardis-regression-data")
+    tardis_regression_path = request.config.getoption(
+        "--tardis-regression-data"
+    )
     if tardis_regression_path is None:
         pytest.skip("--tardis-regression-data was not specified")
     else:
-        return Path(os.path.expandvars(tardis_regression_path)).expanduser().resolve()
+        return (
+            Path(os.path.expandvars(tardis_regression_path))
+            .expanduser()
+            .resolve()
+        )
 
 
 @pytest.fixture(scope="function")
@@ -182,6 +189,7 @@ def tardis_config_verysimple():
         YAMLLoader,
     )
 
+
 @pytest.fixture(scope="session")
 def config_verysimple_for_simulation_one_loop(
     config_verysimple, atomic_data_fname
@@ -194,6 +202,17 @@ def config_verysimple_for_simulation_one_loop(
     return config
 
 
+@pytest.fixture(scope="function")
+def config_verysimple_hydrogen_only(config_verysimple):
+    config = deepcopy(config_verysimple)
+    config.model.abundances = {
+        "type": "uniform",
+        "H": 1.0,
+        "model_isotope_time_0": 0.0 * u.s,
+    }
+    return config
+
+
 @pytest.fixture(scope="function")
 def tardis_config_verysimple_nlte():
     return yaml_load_file(

tardis/opacities/macro_atom/macroatom_solver.py

@@ -441,10 +441,8 @@ def _solve_first_macroatom_iteration(
         )
 
         macro_atom_transition_metadata["source_level_idx"] = (
-            (macro_atom_transition_metadata.source.map(source_to_index))
-            .fillna(0)
-            .astype(np.int64)
-        )
+            macro_atom_transition_metadata.source.map(source_to_index)
+        ).astype(np.int64)
 
         macro_block_references = create_macro_block_references(
             macro_atom_transition_metadata
@@ -596,8 +594,21 @@ def create_macro_block_references(macro_atom_transition_metadata):
         .groupby("source")
         .apply(lambda x: x.index[0])
     )
+
+    # Append a dummy index so that the interactions can access a "block end" if a packet activates the macroatom highest level of the heaviest element in the montecarlo.
+    # Without this the kernel will crash trying to access an index that doesn't exist.
+    macro_data = np.append(
+        macro_data.values, len(macro_atom_transition_metadata)
+    )
+    unique_source_multi_index = unique_source_multi_index.append(
+        pd.MultiIndex.from_tuples(
+            [(-99, -99, -99)],
+            names=["atomic_number", "ion_number", "level_number"],
+        )
+    )
+
     macro_block_references = pd.Series(
-        data=macro_data.values,
+        data=macro_data,
         index=unique_source_multi_index,
         name="macro_block_references",
     )

tardis/transport/montecarlo/tests/test_montecarlo_main_loop.py

@@ -137,3 +137,23 @@ def test_montecarlo_main_loop_vpacket_log(
         rtol=1e-12,
         atol=1e-15,
     )
+
+
+def test_montecarlo_main_loop_hydrogen_only(
+    config_verysimple_hydrogen_only, atomic_dataset, regression_data
+):
+    atomic_dataset = deepcopy(atomic_dataset)
+    montecarlo_main_loop_simulation = Simulation.from_config(
+        config_verysimple_hydrogen_only,
+        atom_data=atomic_dataset,
+        virtual_packet_logging=False,
+    )
+    montecarlo_main_loop_simulation.run_convergence()
+    montecarlo_main_loop_simulation.run_final()
+
+    transport = montecarlo_main_loop_simulation.transport.transport_state
+    assert transport.j_estimator is not None
+    expected_j_estimator = regression_data.sync_ndarray(transport.j_estimator)
+    npt.assert_allclose(
+        transport.j_estimator, expected_j_estimator, atol=0, rtol=1e-12
+    )