[Restructure/plasma] Refactor the simulation state initialization (#2664

)

* change numberdensity to input

* fixed number density

* some fixes

* removing density

* remove atomic and isotope mass

* add isotopic_number_density

* add opacities package

* Update imports in property_collections.py, base.py, test_numba_interface.py, transport_montecarlo_numba_interface.py, conftest.py, formal_integral.py, base.py, and macro_atom.py

* Add calculate_transition_probabilities function to util.py in macro_atom package

* Add calculate_transition_probabilities function to util.py in macro_atom package

* Remove unused imports and update plasma properties

* add __init__ to macroatom

* blackify tardis

* blackified

* chore: Update imports and remove unused code

* chore: Add PlanckRadiationField and DilutePlanckRadiationField classes

* chore: Update imports and remove unused code

* removed density

* ruff output

* cleanup and adding object mode

* starting to make radiation_field a thing

* switched over to old tau_sobolev calculation

* renamed function to indicate numba use

* address comments

* added dilute planckian radiation field

* refactor: Convert species lists to proper format in assemble_plasma function

* moved radiation field into plasma. Resulting in some renames

* some fixes

* black montecarlo

* chore: Initialize atom data and simulation state in `initialization.py`

* updating the documentation

* remove parse_input.py

* chore: Refactor radiation field configuration parsing and state creation

Refactor the `parse_radiation_field_configuration.py` module to improve code organization and readability. Update the import statements to reflect the changes in the module structure. Replace the deprecated `DiluteBlackBodyRadiationFieldState` class with the new `DilutePlanckianRadiationField` class from the `tardis.plasma.radiation_field` module. This change ensures consistency and compatibility with the latest codebase.

Also, update the `tardis/simulation/base.py` module to import the `DilutePlanckianRadiationField` class from the `tardis.plasma.radiation_field` module. This change ensures that the correct class is used for creating the radiation field in the simulation.

* revert astropy_helpers

* remove astropy_helpers

* removed test.txt

* blackified code

* cleanup simulation from merges

* fix

* remove unused Input

* added description

* blackiefied codebase

* blackify code

* chore: Refactor atom data parsing and simulation state initialization

* restructure logger

* merge changes

* blackify

* ruffify

tardis/io/model/parse_atom_data.py

@@ -0,0 +1,53 @@
+import logging
+from pathlib import Path
+
+from tardis.io.atom_data.base import AtomData
+
+logger = logging.getLogger(__name__)
+
+
+def parse_atom_data(config, atom_data=None):
+    """
+    Parse atom data for the simulation.
+
+    Parameters
+    ----------
+    config : object
+        The configuration object containing information about the atom data.
+    atom_data : object, optional
+        Existing atom data to be used, if provided.
+
+    Returns
+    -------
+    object
+        The initialized atom data.
+
+    Raises
+    ------
+    ValueError
+        If no atom_data option is found in the configuration.
+    """
+    if atom_data is None:
+        if "atom_data" in config:
+            if Path(config.atom_data).is_absolute():
+                atom_data_fname = Path(config.atom_data)
+            else:
+                atom_data_fname = Path(config.config_dirname) / config.atom_data
+
+        else:
+            raise ValueError("No atom_data option found in the configuration.")
+
+        logger.info(f"\n\tReading Atomic Data from {atom_data_fname}")
+
+        try:
+            atom_data = AtomData.from_hdf(atom_data_fname)
+        except TypeError as e:
+            print(
+                e,
+                "Error might be from the use of an old-format of the atomic database, \n"
+                "please see https://github.com/tardis-sn/tardis-refdata/tree/master/atom_data"
+                " for the most recent version.",
+            )
+            raise
+
+    return atom_data

tardis/io/model/parse_packet_source_configuration.py

@@ -46,7 +46,7 @@ def initialize_packet_source(packet_source, config, geometry):
     return packet_source
 
 
-def parse_packet_source_from_config(config, geometry, legacy_mode_enabled):
+def parse_packet_source_from_config(config, geometry, enable_legacy_mode):
     """
     Parse the packet source based on the given configuration and geometry.
 
@@ -66,12 +66,12 @@ def parse_packet_source_from_config(config, geometry, legacy_mode_enabled):
         packet_source = BlackBodySimpleSourceRelativistic(
             base_seed=config.montecarlo.seed,
             time_explosion=config.supernova.time_explosion,
-            legacy_mode_enabled=legacy_mode_enabled,
+            legacy_mode_enabled=enable_legacy_mode,
         )
     else:
         packet_source = BlackBodySimpleSource(
             base_seed=config.montecarlo.seed,
-            legacy_mode_enabled=legacy_mode_enabled,
+            legacy_mode_enabled=enable_legacy_mode,
         )
 
     return initialize_packet_source(packet_source, config, geometry)

tardis/io/model/parse_simulation_state.py

@@ -0,0 +1,54 @@
+from tardis.io.model.parse_packet_source_configuration import (
+    initialize_packet_source,
+)
+from tardis.model import SimulationState
+
+
+def parse_simulation_state(
+    config, packet_source, enable_legacy_mode, kwargs, atom_data
+):
+    """
+    Initialize the simulation state.
+
+    Parameters
+    ----------
+    config : object
+        The configuration object for the simulation.
+    packet_source : object
+        The packet source for the simulation.
+    legacy_mode_enabled : bool
+        Flag indicating if legacy mode is enabled.
+    kwargs : dict
+        Additional keyword arguments.
+    atom_data : object
+        The atom data for the simulation.
+
+    Returns
+    -------
+    object
+        The initialized simulation state.
+    """
+    if "model" in kwargs:
+        simulation_state = kwargs["model"]
+    else:
+        if hasattr(config, "csvy_model"):
+            simulation_state = SimulationState.from_csvy(
+                config,
+                atom_data=atom_data,
+                legacy_mode_enabled=enable_legacy_mode,
+            )
+        else:
+            simulation_state = SimulationState.from_config(
+                config,
+                atom_data=atom_data,
+                legacy_mode_enabled=enable_legacy_mode,
+            )
+        if packet_source is not None:
+            simulation_state.packet_source = initialize_packet_source(
+                config,
+                simulation_state.geometry,
+                packet_source,
+                enable_legacy_mode,
+            )
+
+    return simulation_state

tardis/io/model/readers/csvy.py

@@ -23,7 +23,7 @@ def load_csvy(fname):
     -------
     yaml_dict : dictionary
                 YAML part of the csvy file
-    data : pandas.dataframe
+    data : pandas.DataFrame
             csv data from csvy file
     """
     with open(fname) as fh:

tardis/plasma/properties/property_collections.py

@@ -2,8 +2,8 @@
     NonMarkovChainTransitionProbabilities,
     TransitionProbabilities,
 )
-from tardis.plasma.properties import *
 from tardis.opacities.tau_sobolev import TauSobolev
+from tardis.plasma.properties import *
 
 
 class PlasmaPropertyCollection(list):

tardis/simulation/base.py

@@ -1,7 +1,6 @@
 import logging
 import time
 from collections import OrderedDict
-from pathlib import Path
 
 import numpy as np
 import pandas as pd
@@ -10,13 +9,12 @@
 
 import tardis
 from tardis import constants as const
-from tardis.io.atom_data.base import AtomData
 from tardis.io.configuration.config_reader import ConfigurationError
-from tardis.io.model.parse_packet_source_configuration import (
-    initialize_packet_source,
+from tardis.io.model.parse_atom_data import parse_atom_data
+from tardis.io.model.parse_simulation_state import (
+    parse_simulation_state,
 )
 from tardis.io.util import HDFWriterMixin
-from tardis.model import SimulationState
 from tardis.plasma.radiation_field import DilutePlanckianRadiationField
 from tardis.plasma.standard_plasmas import assemble_plasma
 from tardis.simulation.convergence import ConvergenceSolver
@@ -27,7 +25,6 @@
 from tardis.util.base import is_notebook
 from tardis.visualization import ConvergencePlots
 
-# Adding logging support
 logger = logging.getLogger(__name__)
 
 
@@ -515,12 +512,12 @@ def log_plasma_state(
         ----------
         t_rad : astropy.units.Quanity
             current t_rad
-        w : astropy.units.Quanity
-            current w
+        dilution_factor : np.ndarray
+            current dilution_factor
         next_t_rad : astropy.units.Quanity
             next t_rad
-        next_w : astropy.units.Quanity
-            next_w
+        next_dilution_factor : np.ndarray
+            next dilution_factor
         log_sampling : int
             the n-th shells to be plotted
 
@@ -636,89 +633,61 @@ def from_config(
         show_convergence_plots=False,
         show_progress_bars=True,
         legacy_mode_enabled=False,
+        atom_data=None,
+        plasma=None,
+        transport=None,
         **kwargs,
     ):
         """
-        Create a new Simulation instance from a Configuration object.
+        Create a simulation instance from the provided configuration.
 
         Parameters
         ----------
-        config : tardis.io.config_reader.Configuration
-
+        config : object
+            The configuration object for the simulation.
+        packet_source : object, optional
+            The packet source for the simulation.
+        virtual_packet_logging : bool, optional
+            Flag indicating virtual packet logging.
+        show_convergence_plots : bool, optional
+            Flag indicating whether to show convergence plots.
+        show_progress_bars : bool, optional
+            Flag indicating whether to show progress bars.
+        legacy_mode_enabled : bool, optional
+            Flag indicating if legacy mode is enabled.
+        atom_data : object, optional
+            The atom data for the simulation.
+        plasma : object, optional
+            The plasma object for the simulation.
+        transport : object, optional
+            The transport solver for the simulation.
         **kwargs
-            Allow overriding some structures, such as model, plasma, atomic data
-            and the transport, instead of creating them from the configuration
-            object.
+            Additional keyword arguments.
 
         Returns
         -------
-        Simulation
+        object
+            The created simulation instance.
         """
         # Allow overriding some config structures. This is useful in some
         # unit tests, and could be extended in all the from_config classmethods.
 
-        atom_data = kwargs.get("atom_data", None)
-        if atom_data is None:
-            if "atom_data" in config:
-                if Path(config.atom_data).is_absolute():
-                    atom_data_fname = Path(config.atom_data)
-                else:
-                    atom_data_fname = (
-                        Path(config.config_dirname) / config.atom_data
-                    )
-
-            else:
-                raise ValueError(
-                    "No atom_data option found in the configuration."
-                )
-
-            logger.info(f"\n\tReading Atomic Data from {atom_data_fname}")
-
-            try:
-                atom_data = AtomData.from_hdf(atom_data_fname)
-            except TypeError as e:
-                print(
-                    e,
-                    "Error might be from the use of an old-format of the atomic database, \n"
-                    "please see https://github.com/tardis-sn/tardis-refdata/tree/master/atom_data"
-                    " for the most recent version.",
-                )
-                raise
-        if "model" in kwargs:
-            simulation_state = kwargs["model"]
-        else:
-            if hasattr(config, "csvy_model"):
-                simulation_state = SimulationState.from_csvy(
-                    config,
-                    atom_data=atom_data,
-                    legacy_mode_enabled=legacy_mode_enabled,
-                )
-            else:
-                simulation_state = SimulationState.from_config(
-                    config,
-                    atom_data=atom_data,
-                    legacy_mode_enabled=legacy_mode_enabled,
-                )
-            # Override with custom packet source from function argument if present
-            if packet_source is not None:
-                simulation_state.packet_source = initialize_packet_source(
-                    packet_source, config, simulation_state.geometry
-                )
-        if "plasma" in kwargs:
-            plasma = kwargs["plasma"]
-        else:
+        atom_data = parse_atom_data(config, atom_data=atom_data)
+        simulation_state = parse_simulation_state(
+            config, packet_source, legacy_mode_enabled, kwargs, atom_data
+        )
+        if plasma is None:
             plasma = assemble_plasma(
                 config,
                 simulation_state,
                 atom_data=atom_data,
             )
-        if "transport" in kwargs:
-            if packet_source is not None:
-                raise ConfigurationError(
-                    "Cannot specify packet_source and transport at the same time."
-                )
-            transport = kwargs["transport"]
-        else:
+
+        if (transport is not None) and (packet_source is not None):
+            raise ConfigurationError(
+                "Cannot specify packet_source and transport at the same time."
+            )
+        if transport is None:
             transport = MonteCarloTransportSolver.from_config(
                 config,
                 packet_source=simulation_state.packet_source,