Restructure/separate transport state (#2492)

* restructure of geometry

* add radial1d boundary logic

* black format

* several fixes

* fix epsilon

* add testing of boundaries

* change the r_inner_active

* first integration with `from_config` working

* hunting down density indexing bug

* all model tests (without csvy) pass

* more fixes

* fix of model to simulation_state

* fix inner boundary packet error

* fix some leftovers

* final fix for csvy

* blackify

* restructure to readers and remove some leftover code

* further cleanup

* first start of the restructure

* add comment about removing quantitiness

* add velocity check

* add new abundance functions

* remove default units

* add new matter module

* several restructures. move decay from io to tardis/model/matter

* mid restructure

* slow progress on abundance refactor

* include effective_element_masses

* further updates

* fix of csvy readers

* last fixes to custom abundance widget

* restructuring the radiation field

* some cleanup

* cleanup

* further cleanup

* clearnup

* changes

* remove matter base.py

* removing matter

* fixed last tests

* first commit for regression_data

* add several parsers

* add composition

* remove isotopemassfractions

* some small fixes

* fixed w, t_rad to dilution_factor, t_radiative

* add setters for w, t_rad

* several fixes for instantiating the right t_rads.
also ensuring that they are positive and len(t_rad) == dilution factor

* add the from_config atom_data

* fixing model.ipynb

* fix some more tests

* start of montecarlo restructure

* large scale changes

* further disentangling state and solver

* fix tests

* fix some of the issues wdigets

* assert mass fractions positive.
Change IsotopeAbundances to Massfractions

* update estimators

* add packet collection

* move spectrum to mct_state

* add radiation_field_state

* Refactor packet initialization in
MontecarloTransportSolver

* Update model_nparray_attrs in test_base.py and fix
output_energy in base.py

* Fix temperature assertion in
test_initial_temperature() and update variable
names in test_montecarlo_main_loop()

* Fix Integrator

* Fix SDEC plot for MCTState

* Fix variable names in transport module

* Refactor the to_hdf code.

* Refactor Monte Carlo transport solver
initialization and transport state

* Fix isotope mass fraction typo and t_radiatve duplicate definition

* Document remove functionality for effective element masses

* Fix variable name in
convert_to_nuclide_mass_fraction function

* Refactor test names and variable names in
test_base.py

* restructuring montecarlo state

* Refactor gamma ray simulation setup and test
functions

* fix gamma ray tests

* Add model_isotope_time_0 property to abundances
schema

* Add test for montecarlo main loop with vpacket
tracking

* Refactor config_reader and config_validator
modules

* Fix model_isotope_time_0 property in
model_definitions.yml

* Refactor simulation module to use pathlib

* remove bad comment and refactor with ruff

* Fix variable name in test_simulation_state_mass()
function

* Fix unit conversion in calculate_cell_masses
function

* Fix logger debug message in CSV reader and update
method names in Composition class

* black formatting

* Refactor code to use composition variable in
test_gamma_ray_transport.py

* Refactor imports in atom_web_download.py and
update return statement in config_internal.py

* fix for documentation not builiding

* add fix for model_isotope_time_0

* fix grid test

* restructure the grid

* final change

* add stuff

* add vpacket_log

* transfer over to VPacketCollection constructor

* Refactor model_reader.py and shell_info.py
changing w to dilution_factor

* Update damping constants and variable names
in doc notebook

* Refactor vpacket tracking in montecarlo_radial1d
function

* fix w in several places in tardis

* slowly fixing radiation_field changes

* fixing montecarlo tests

* slowly fixing vpacket_tracker

* last fix

* change hdf from model to simulation_state

* fixing model->simulation-state

* Fix model_reader.py imports and formatting issues

* Refactor variable names for clarity and
consistency and being correct

* Update variable names in model.ipynb

* Update variable name for average temperature
calculation

* Fix file paths in test_grid.py, test_base.py,
test_density.py, and test_base.py

* Fix transport dictionary comparison

* Implement regression data configurations

* Refactor snapshot fixtures and add regression data
handling

* Update to new testing framework

* Refactor imports and fixtures in conftest.py and
test_nlte_solver.py

* Make regression-data optional

* Fix regression data synchronization in plasma
tests

* Refactor variable names in model reader and
simulation base

* Refactor store_model_to_hdf to
store_simulation_state_to_hdf

* Refactor HDF file handling and model reader
imports

* add docstr

* Remove syrupy_extensions.py

* Refactor code and remove unused imports and
deleted file

* Fix test_base and test_simulation

* setting up regression with github actions

* Update LFS cache IDs for refdata and regression
data

* Fix test_montecarlo_radial1d function

* Update sync_hdf_store method in RegressionData class

* Refactor test_base.py for improved accuracy and precision

* update docstrings

* Refactor transport state access in test_base.py

* Fix warnings and add HDF output in sdec_plot.ipynb and cleanup regression_data.py

* Update demo.hdf file

Delete demo.hdf file

Update file paths in code

* Fix negative values in last_interaction_types

* Update URL for atomic data repository

* Refactor code for improved readability and maintainability

* Refactor code to remove unnecessary array conversion

* Refactor code to remove unnecessary conversion to numpy array in test_hdf_scalars

* Refactor vpacket tracking in montecarlo_main_loop

* Update packet_collection and vpacket_tracker

* fixup of code

* Add packet functionality to VPacketCollection

* ruffify test_sdec_plot

* Remove unused import in parse_input.py

* Update VPACKET_TAU_RUSSIAN constant

* Refactor montecarlo_main_loop and FormalIntegrator classes

* Fixes most tests

* Removes old montecarlo_state file added in merge

* Fixes docs build

* Fix virtual packet tracking for main loop

* Restructure vpacket log storage

Now located in the MCTState. Tests are still failing!

* Remove some data from testing

Spectra and scalars need a different testing method, out of scope for this PR.

* Applied black formatting

---------

Co-authored-by: Andrew Fullard <[email protected]>

docs/io/output/rpacket_tracking.ipynb

@@ -202,7 +202,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "type(sim.transport.rpacket_tracker)"
+    "type(sim.transport.transport_state.rpacket_tracker)"
    ]
   },
   {
@@ -243,7 +243,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "len(sim.transport.rpacket_tracker)"
+    "len(sim.transport.transport_state.rpacket_tracker)"
    ]
   },
   {
@@ -281,7 +281,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "sim.transport.rpacket_tracker[10].index"
+    "sim.transport.transport_state.rpacket_tracker[10].index"
    ]
   },
   {
@@ -300,7 +300,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "sim.transport.rpacket_tracker[10].seed"
+    "sim.transport.transport_state.rpacket_tracker[10].seed"
    ]
   },
   {
@@ -319,7 +319,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "sim.transport.rpacket_tracker[10].status"
+    "sim.transport.transport_state.rpacket_tracker[10].status"
    ]
   },
   {
@@ -347,7 +347,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "len(sim.transport.rpacket_tracker[10].shell_id)"
+    "len(sim.transport.transport_state.rpacket_tracker[10].shell_id)"
    ]
   },
   {
@@ -403,7 +403,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "type(sim.transport.rpacket_tracker_df)"
+    "type(sim.transport.transport_state.rpacket_tracker_df)"
    ]
   },
   {
@@ -461,7 +461,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "sim.transport.rpacket_tracker_df"
+    "sim.transport.transport_state.rpacket_tracker_df"
    ]
   },
   {
@@ -490,7 +490,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "sim.transport.rpacket_tracker_df.loc[10]"
+    "sim.transport.transport_state.rpacket_tracker_df.loc[10]"
    ]
   },
   {
@@ -509,7 +509,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "sim.transport.rpacket_tracker_df.loc[10][\"energy\"]"
+    "sim.transport.transport_state.rpacket_tracker_df.loc[10][\"energy\"]"
    ]
   },
   {
@@ -528,7 +528,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "sim.transport.rpacket_tracker_df.loc[10,5][\"energy\"]"
+    "sim.transport.transport_state.rpacket_tracker_df.loc[10,5][\"energy\"]"
    ]
   },
   {
@@ -538,7 +538,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "sim.transport.rpacket_tracker_df.loc[10][\"energy\"][5]"
+    "sim.transport.transport_state.rpacket_tracker_df.loc[10][\"energy\"][5]"
    ]
   },
   {
@@ -567,7 +567,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.13"
+   "version": "3.11.5"
   },
   "vscode": {
    "interpreter": {

docs/io/visualization/montecarlo_packet_visualization.ipynb

@@ -89,7 +89,7 @@
    "outputs": [],
    "source": [
     "# accessing the rpacket_tracker dataframe\n",
-    "sim.transport.rpacket_tracker_df"
+    "sim.transport.transport_state.rpacket_tracker_df"
    ]
   },
   {
@@ -191,11 +191,11 @@
    "source": [
     "# function for getting coordinates of all packets\n",
     "def get_coordinates_multiple_packets_ints(r_packet_tracker,time):\n",
-    "    thetas = np.linspace(0, 2*math.pi,sim.transport.rpacket_tracker_df.index[-1][0]+2)\n",
+    "    thetas = np.linspace(0, 2*math.pi,sim.transport.transport_state.rpacket_tracker_df.index[-1][0]+2)\n",
     "    x = []\n",
     "    y = []\n",
     "    inters = []\n",
-    "    for i in range(sim.transport.rpacket_tracker_df.index[-1][0]+1):\n",
+    "    for i in range(sim.transport.transport_state.rpacket_tracker_df.index[-1][0]+1):\n",
     "        xs,ys,ints = get_x_y_ints_with_theta_init(r_packet_tracker.loc[i][\"r\"],r_packet_tracker.loc[i][\"mu\"],time,r_packet_tracker.loc[i][\"interaction_type\"],'n',thetas[i])\n",
     "        x.append(xs)\n",
     "        y.append(ys)\n",
@@ -236,7 +236,7 @@
     "# animated plot\n",
     "\n",
     "v_shells = sim.simulation_state.radius.value * 1e-5 / sim.simulation_state.time_explosion.value\n",
-    "xs,ys,ints = get_coordinates_multiple_packets_ints(sim.transport.rpacket_tracker_df,sim.simulation_state.time_explosion.value)\n",
+    "xs,ys,ints = get_coordinates_multiple_packets_ints(sim.transport.transport_state.rpacket_tracker_df,sim.simulation_state.time_explosion.value)\n",
     "xs,ys,ints,max_size = get_equal_array_size(xs,ys,ints)\n",
     "\n",
     "fig = go.Figure()\n",
@@ -273,7 +273,7 @@
     "\n",
     "#Add the packet trajectory\n",
     "\n",
-    "df = sim.transport.rpacket_tracker_df\n",
+    "df = sim.transport.transport_state.rpacket_tracker_df\n",
     "interaction_from_num = {0: \"No Interaction\", 1: \"EScattering\", 2: \"Line\"}\n",
     "interaction_color_from_num = {0: \"darkslategrey\", 1: \"#3366FF\", 2: \"#FF3300\"}\n",
     "interaction_opacity_from_num = {0: 0, 1: 1, 2: 1}\n",

tardis/analysis.py

@@ -3,25 +3,26 @@
 """
 
 import re
-import os
 
-from astropy import units as u
-from tardis import constants
 import numpy as np
 import pandas as pd
+from astropy import units as u
+
+from tardis import constants
 
 INVALID_ION_ERROR_MSG = "Atomic number, ion_number pair not present in model"
 
 
-class LastLineInteraction(object):
+class LastLineInteraction:
     @classmethod
     def from_simulation(cls, simulation, packet_filter_mode="packet_out_nu"):
+        transport_state = simulation.transport.transport_state
         return cls(
-            simulation.transport.last_line_interaction_in_id,
-            simulation.transport.last_line_interaction_out_id,
-            simulation.transport.last_line_interaction_shell_id,
-            simulation.transport.transport_state.packet_collection.output_nus,
-            simulation.transport.last_interaction_in_nu,
+            transport_state.last_line_interaction_in_id,
+            transport_state.last_line_interaction_out_id,
+            transport_state.last_line_interaction_shell_id,
+            transport_state.packet_collection.output_nus,
+            transport_state.last_interaction_in_nu,
             simulation.plasma.atomic_data.lines,
             packet_filter_mode,
         )
@@ -260,7 +261,7 @@ def onpress(event):
         fig.canvas.mpl_connect("on_press", onpress)
 
 
-class TARDISHistory(object):
+class TARDISHistory:
     """
     Records the history of the model
     """

tardis/conftest.py

@@ -232,3 +232,14 @@ def simulation_verysimple(config_verysimple, atomic_dataset):
     sim = Simulation.from_config(config_verysimple, atom_data=atomic_data)
     sim.iterate(4000)
     return sim
+
+
+@pytest.fixture(scope="session")
+def simulation_verysimple_vpacket_tracking(config_verysimple, atomic_dataset):
+    atomic_data = deepcopy(atomic_dataset)
+    sim = Simulation.from_config(
+        config_verysimple, atom_data=atomic_data, virtual_packet_logging=True
+    )
+    sim.last_no_of_packets = 4000
+    sim.run_final()
+    return sim

tardis/grid/tests/test_grid.py

@@ -4,8 +4,11 @@
 import pandas as pd
 
 import tardis
+
+from pathlib import Path
 import tardis.grid as grid
 
+
 DATA_PATH = Path(tardis.__path__[0]) / "grid" / "tests" / "data"
 
 

tardis/model/parse_input.py

@@ -598,6 +598,14 @@ def initialize_packet_source(config, geometry, packet_source):
     ValueError
         If both t_inner and luminosity_requested are None.
     """
+    if config.montecarlo.enable_full_relativity:
+        packet_source = BlackBodySimpleSourceRelativistic(
+            base_seed=config.montecarlo.seed,
+            time_explosion=config.supernova.time_explosion,
+        )
+    else:
+        packet_source = BlackBodySimpleSource(base_seed=config.montecarlo.seed)
+
     luminosity_requested = config.supernova.luminosity_requested
     if config.plasma.initial_t_inner > 0.0 * u.K:
         packet_source.radius = geometry.r_inner[0]