Opacity state slice

tardis/opacities/opacity_state.py

@@ -106,6 +106,41 @@
         self.photo_ion_activation_idx = photo_ion_activation_idx
         self.k_packet_idx = k_packet_idx
 
+    def __getitem__(self, i):
+        """Get a shell or slice of shells of the attributes of the opacity state
+
+        Args:
+            i (int, slice): shell index or slice
+
+        Returns:
+            OpacityState : a shallow copy of the current instance
+        """
+        # NOTE: This currently will not work with continuum processes since it does not slice those arrays
+        return OpacityState(
+            self.electron_density[i],
+            self.t_electrons[i],
+            self.line_list_nu,
+            self.tau_sobolev[:, i],
+            self.transition_probabilities[:, i],
+            self.line2macro_level_upper,
+            self.macro_block_references,
+            self.transition_type,
+            self.destination_level_id,
+            self.transition_line_id,
+            self.bf_threshold_list_nu,
+            self.p_fb_deactivation,
+            self.photo_ion_nu_threshold_mins,
+            self.photo_ion_nu_threshold_maxs,
+            self.photo_ion_block_references,
+            self.chi_bf,
+            self.x_sect,
+            self.phot_nus,
+            self.ff_opacity_factor,
+            self.emissivities,
+            self.photo_ion_activation_idx,
+            self.k_packet_idx,
+        )
+
 
 def opacity_state_initialize(
     plasma,

tardis/transport/montecarlo/base.py

@@ -20,7 +20,7 @@
 from tardis.transport.montecarlo.montecarlo_transport_state import (
     MonteCarloTransportState,
 )
-from tardis.transport.montecarlo.numba_interface import (
+from tardis.opacities.opacity_state import ( 
     opacity_state_initialize,
 )
 from tardis.transport.montecarlo.packet_trackers import (
@@ -103,16 +103,19 @@ def initialize_transport_state(
         packet_collection = self.packet_source.create_packets(
             no_of_packets, seed_offset=iteration
         )
-        estimators = initialize_estimator_statistics(
-            plasma.tau_sobolevs.shape, gamma_shape
-        )
 
         geometry_state = simulation_state.geometry.to_numba()
         opacity_state = opacity_state_initialize(
             plasma,
             self.line_interaction_type,
             self.montecarlo_configuration.DISABLE_LINE_SCATTERING,
         )
+        opacity_state = opacity_state[simulation_state.geometry.v_inner_boundary_index: simulation_state.geometry.v_outer_boundary_index]
+
+        estimators = initialize_estimator_statistics(
+            opacity_state.tau_sobolev.shape, gamma_shape
+        )
+
         transport_state = MonteCarloTransportState(
             packet_collection,
             estimators,

tardis/transport/montecarlo/numba_interface.py

@@ -112,6 +112,41 @@
         self.photo_ion_activation_idx = photo_ion_activation_idx
         self.k_packet_idx = k_packet_idx
 
+    def __getitem__(self, i):
+        """Get a shell or slice of shells of the attributes of the opacity state
+
+        Args:
+            i (int, slice): shell index or slice
+
+        Returns:
+            OpacityState : a shallow copy of the current instance
+        """
+        # NOTE: This currently will not work with continuum processes since it does not slice those arrays
+        return OpacityState(
+            self.electron_density[i],
+            self.t_electrons[i],
+            self.line_list_nu,
+            self.tau_sobolev[:, i],
+            self.transition_probabilities[:, i],
+            self.line2macro_level_upper,
+            self.macro_block_references,
+            self.transition_type,
+            self.destination_level_id,
+            self.transition_line_id,
+            self.bf_threshold_list_nu,
+            self.p_fb_deactivation,
+            self.photo_ion_nu_threshold_mins,
+            self.photo_ion_nu_threshold_maxs,
+            self.photo_ion_block_references,
+            self.chi_bf,
+            self.x_sect,
+            self.phot_nus,
+            self.ff_opacity_factor,
+            self.emissivities,
+            self.photo_ion_activation_idx,
+            self.k_packet_idx,
+        )
+
 
 def opacity_state_initialize(
     plasma,

tardis/transport/montecarlo/tests/test_numba_interface.py

@@ -4,8 +4,19 @@
 import numpy as np
 
 
-@pytest.mark.parametrize("input_params", ["scatter", "macroatom", "downbranch"])
-def test_opacity_state_initialize(nb_simulation_verysimple, input_params):
+@pytest.mark.parametrize(
+    "input_params,sliced",
+    [
+        ("scatter", False),
+        ("macroatom", False),
+        ("macroatom", True),
+        ("downbranch", False),
+        ("downbranch", True),
+    ],
+)
+def test_opacity_state_initialize(
+    nb_simulation_verysimple, input_params, sliced
+):
     line_interaction_type = input_params
     plasma = nb_simulation_verysimple.plasma
     actual = numba_interface.opacity_state_initialize(
@@ -14,11 +25,20 @@ def test_opacity_state_initialize(nb_simulation_verysimple, input_params):
         disable_line_scattering=False,
     )
 
+    if sliced:
+        index = slice(2, 5)
+        actual = actual[index]
+    else:
+        index = ...
+
     npt.assert_allclose(
-        actual.electron_density, plasma.electron_densities.values
+        actual.electron_density, plasma.electron_densities.values[index]
     )
     npt.assert_allclose(actual.line_list_nu, plasma.atomic_data.lines.nu.values)
-    npt.assert_allclose(actual.tau_sobolev, plasma.tau_sobolevs.values)
+    print(actual.tau_sobolev.shape, plasma.tau_sobolevs.values[:, index].shape)
+    npt.assert_allclose(
+        actual.tau_sobolev, plasma.tau_sobolevs.values[:, index]
+    )
     if line_interaction_type == "scatter":
         empty = np.zeros(1, dtype=np.int64)
         npt.assert_allclose(
@@ -32,7 +52,7 @@ def test_opacity_state_initialize(nb_simulation_verysimple, input_params):
     else:
         npt.assert_allclose(
             actual.transition_probabilities,
-            plasma.transition_probabilities.values,
+            plasma.transition_probabilities.values[:, index],
         )
         npt.assert_allclose(
             actual.line2macro_level_upper,