From 140bb96a6c0d162900e6106d62348bcaef987e2d Mon Sep 17 00:00:00 2001 From: Andrew Date: Wed, 26 Jun 2024 13:15:26 -0400 Subject: [PATCH] I/O Refactoring (#2652) * Moves parse_input to IO and removes unused file * Removes unused file and imports * Separates parse_input into chunks Also cleans up imports * Deprecates old CMFGen reader and moves new one * Fixes logging import * Fixes import sorting * Consistent file naming * Consistent names and ordering for parsing functions _from_config and _from_csvy to match typical class methods Config first, CSVY second * Abundances renamed to mass fractions in IO With small extensions outside that area * Pluralise * Fixes mass fraction and density dimension mismatch for CSVY models * Adds docstring * black * Adds dummy type variable to schema for convergence Should fix read/write issue. * Cleanup of structure config reading Not really happy with it, but helps show some of the current problems with the config arrangement compared to the code * Fixes error with density time point and black * Fixes packet source re-init for custom sources Also adjusts notebook to deal with quantity-based packet frequencies * Resolve some escape sequence errors * Plural mass fractions * black * Clarify docstrings --- docs/io/optional/how_to_custom_source.ipynb | 3 +- tardis/io/__init__.py | 17 +- .../schemas/montecarlo_definitions.yml | 12 + tardis/io/model/__init__.py | 2 +- tardis/io/model/cmfgen.py | 75 -- tardis/io/model/density.py | 187 ---- .../model/parse_composition_configuration.py | 100 +++ .../io/model/parse_density_configuration.py | 110 ++- .../io/model/parse_geometry_configuration.py | 173 ++++ .../parse_mass_fraction_configuration.py | 227 +++++ .../parse_packet_source_configuration.py | 77 ++ .../parse_radiation_field_configuration.py | 188 ++++ tardis/io/model/readers/base.py | 42 +- tardis/io/model/readers/cmfgen.py | 128 ++- tardis/io/model/readers/cmfgen_deprecated.py | 91 ++ tardis/io/model/readers/csvy.py | 22 +- tardis/io/model/readers/generic_readers.py | 58 +- tardis/io/model/readers/stella.py | 12 +- .../model/readers/tests/test_ascii_readers.py | 4 +- .../model/readers/tests/test_cmfgen_reader.py | 2 +- tardis/io/model/readers/util.py | 19 +- tardis/io/model_reader.py | 846 ------------------ tardis/io/tests/test_model_reader.py | 14 +- tardis/model/base.py | 63 +- tardis/model/parse_input.py | 722 --------------- tardis/model/tests/test_csvy_model.py | 2 +- tardis/simulation/base.py | 12 +- .../visualization/widgets/custom_abundance.py | 8 +- 28 files changed, 1141 insertions(+), 2075 deletions(-) delete mode 100644 tardis/io/model/cmfgen.py delete mode 100644 tardis/io/model/density.py create mode 100644 tardis/io/model/parse_composition_configuration.py create mode 100644 tardis/io/model/parse_geometry_configuration.py create mode 100644 tardis/io/model/parse_mass_fraction_configuration.py create mode 100644 tardis/io/model/parse_packet_source_configuration.py create mode 100644 tardis/io/model/parse_radiation_field_configuration.py create mode 100644 tardis/io/model/readers/cmfgen_deprecated.py delete mode 100644 tardis/io/model_reader.py delete mode 100644 tardis/model/parse_input.py diff --git a/docs/io/optional/how_to_custom_source.ipynb b/docs/io/optional/how_to_custom_source.ipynb index 99bd5da7915..046faf00780 100644 --- a/docs/io/optional/how_to_custom_source.ipynb +++ b/docs/io/optional/how_to_custom_source.ipynb @@ -131,7 +131,6 @@ " truncation_frequency = (\n", " u.Quantity(self.truncation_wavelength, u.Angstrom)\n", " .to(u.Hz, equivalencies=u.spectral())\n", - " .value\n", " )\n", "\n", " # Draw nus from blackbody distribution and reject based on truncation_frequency.\n", @@ -220,7 +219,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.11.5" + "version": "3.12.3" } }, "nbformat": 4, diff --git a/tardis/io/__init__.py b/tardis/io/__init__.py index e4e678f32e4..59f55b94a85 100644 --- a/tardis/io/__init__.py +++ b/tardis/io/__init__.py @@ -1,18 +1,3 @@ """ -A collection of subpackages and submodules to handle input and output data. +A collection of subpackages and submodules to handle input and output data. """ - -from tardis.io.configuration.config_internal import ( - get_internal_configuration, - get_data_dir, -) - -# readin model_data - -from tardis.io.model.readers.generic_readers import ( - read_simple_ascii_abundances, -) -from tardis.io.model.readers.generic_readers import ( - read_simple_ascii_density, -) -from tardis.io.model.readers.base import read_density_file diff --git a/tardis/io/configuration/schemas/montecarlo_definitions.yml b/tardis/io/configuration/schemas/montecarlo_definitions.yml index ea1233ac3b1..1adf22599ae 100644 --- a/tardis/io/configuration/schemas/montecarlo_definitions.yml +++ b/tardis/io/configuration/schemas/montecarlo_definitions.yml @@ -55,6 +55,10 @@ definitions: description: specifies the threshold that is taken as convergence (i.e. 0.05 means that the value does not change more than 5%) minimum: 0 + type: + type: string + default: 'damped' + description: THIS IS A DUMMY VARIABLE DO NOT USE t_rad: type: object additionalProperties: false @@ -69,6 +73,10 @@ definitions: description: specifies the threshold that is taken as convergence (i.e. 0.05 means that the value does not change more than 5%) minimum: 0 + type: + type: string + default: 'damped' + description: THIS IS A DUMMY VARIABLE DO NOT USE required: - threshold w: @@ -85,6 +93,10 @@ definitions: description: specifies the threshold that is taken as convergence (i.e. 0.05 means that the value does not change more than 5%) minimum: 0 + type: + type: string + default: 'damped' + description: THIS IS A DUMMY VARIABLE DO NOT USE required: - threshold lock_t_inner_cycles: diff --git a/tardis/io/model/__init__.py b/tardis/io/model/__init__.py index 25d3a7257a6..9a877f17b98 100644 --- a/tardis/io/model/__init__.py +++ b/tardis/io/model/__init__.py @@ -1,4 +1,4 @@ from tardis.io.model.readers.stella import read_stella_model # from tardis.io.model.stella import read_stella_model -from tardis.io.model.cmfgen import read_cmfgen_model +from tardis.io.model.readers.cmfgen import read_cmfgen_model diff --git a/tardis/io/model/cmfgen.py b/tardis/io/model/cmfgen.py deleted file mode 100644 index 0c4364541d3..00000000000 --- a/tardis/io/model/cmfgen.py +++ /dev/null @@ -1,75 +0,0 @@ -import re -import pandas as pd -from astropy import units as u -from pathlib import Path -import dataclasses - - -@dataclasses.dataclass -class CMFGENModel: - metadata: dict - data: pd.DataFrame - - -HEADER_RE_STR = [ - ("t0:\s+(\d+\.\d+)+\s+day", "t0"), -] - -COLUMN_ROW = 1 -UNIT_ROW = 2 -DATA_START_ROW = 3 - - -def read_cmfgen_model(fname): - """ - Read in a CMFGEN model file and return the data and model - - Parameters - ---------- - - fname : str - - Returns - ------- - model : CMFGENModel - - """ - header_re = [re.compile(re_str[0]) for re_str in HEADER_RE_STR] - metadata = {} - with open(fname) as fh: - for i, line in enumerate(fh): - if i < len(HEADER_RE_STR): - header_re_match = header_re[i].match(line) - metadata[HEADER_RE_STR[i][1]] = header_re_match.group(1) - elif i == COLUMN_ROW: - if "Index" in line: - column_names = re.split(r"\s", line.strip()) - column_names = [ - col.lower().replace(" ", "_") for col in column_names - ] - column_names = column_names[ - 1: - ] # Remove Index from column names - else: - raise ValueError( - '"Index" is required in the Cmfgen input file to infer columns' - ) - elif i == UNIT_ROW: - units = re.split(r"\s", line.strip()) - units = units[1:] # Remove index column - for col, unit in zip(column_names, units): - if u.Unit(unit) == "": # dimensionless - continue - metadata[f"{col}_unit"] = u.Unit(unit) - break - - metadata["t0"] = float(metadata["t0"]) * u.day - data = pd.read_csv( - fname, - delim_whitespace=True, - skiprows=DATA_START_ROW, - header=None, - index_col=0, - ) - data.columns = column_names - return CMFGENModel(metadata, data) diff --git a/tardis/io/model/density.py b/tardis/io/model/density.py deleted file mode 100644 index c9048e0112c..00000000000 --- a/tardis/io/model/density.py +++ /dev/null @@ -1,187 +0,0 @@ -import numpy as np -from astropy import units as u - -from tardis.io.configuration.config_reader import Configuration -from tardis.util.base import quantity_linspace - - -def parse_density_section( - density_configuration: Configuration, - v_middle: u.Quantity, - time_explosion: u.Quantity, -): - if density_configuration.type == "branch85_w7": - density_0 = calculate_power_law_density( - v_middle, - density_configuration.w7_v_0, - density_configuration.w7_rho_0, - -7, - ) - time_0 = density_configuration.w7_time_0 - elif density_configuration.type == "uniform": - density_0 = density_configuration.value.to("g cm^-3") * np.ones_like( - v_middle.value - ) - time_0 = density_configuration.get("time_0", time_explosion) - elif density_configuration.type == "power_law": - density_0 = calculate_power_law_density( - v_middle, - density_configuration.v_0, - density_configuration.rho_0, - density_configuration.exponent, - ) - time_0 = density_configuration.get("time_0", time_explosion) - elif density_configuration.type == "exponential": - density_0 = calculate_exponential_density( - v_middle, density_configuration.v_0, density_configuration.rho_0 - ) - time_0 = density_configuration.get("time_0", time_explosion) - else: - raise ValueError( - f"Unrecognized density type " f"'{density_configuration.type}'" - ) - return density_0, time_0 - - -def parse_config_v1_density(config: Configuration) -> u.Quantity: - """ - Create a new HomologousDensity instance from a Configuration object. - - Parameters - ---------- - config : tardis.io.config_reader.Configuration - - Returns - ------- - HomologousDensity - - """ - - velocity = quantity_linspace( - config.model.structure.velocity.start, - config.model.structure.velocity.stop, - config.model.structure.velocity.num + 1, - ).cgs - - adjusted_velocity = velocity.insert(0, 0) - v_middle = adjusted_velocity[1:] * 0.5 + adjusted_velocity[:-1] * 0.5 - time_explosion = config.supernova.time_explosion.cgs - d_conf = config.model.structure.density - density_0, time_0 = parse_density_section(d_conf, v_middle, time_explosion) - - return calculate_density_after_time(density_0, time_0, time_explosion) - - -def parse_csvy_density(csvy_model_config, time_explosion: u.Quantity): - """ - Create a new HomologousDensity instance from a base - Configuration object and a csvy model Configuration object. - - Parameters - ---------- - config : tardis.io.config_reader.Configuration - csvy_model_config : tardis.io.config_reader.Configuration - - Returns - ------- - HomologousDensity - - """ - if hasattr(csvy_model_config, "velocity"): - velocity = quantity_linspace( - csvy_model_config.velocity.start, - csvy_model_config.velocity.stop, - csvy_model_config.velocity.num + 1, - ).cgs - else: - velocity_field_index = [ - field.name for field in csvy_model_config.datatype.fields - ].index("velocity") - velocity_unit = u.Unit( - csvy_model_config.datatype.fields[velocity_field_index].unit - ) - velocity = csvy_model_config.velocity.values * velocity_unit - - adjusted_velocity = velocity.insert(0, 0) - v_middle = adjusted_velocity[1:] * 0.5 + adjusted_velocity[:-1] * 0.5 - no_of_shells = len(adjusted_velocity) - 1 - - if hasattr(csvy_model_config, "density"): - density_0, time_0 = parse_density_section( - csvy_model_config.density, v_middle, time_explosion - ) - return calculate_density_after_time(density_0, time_0, time_explosion) - - -def calculate_power_law_density(velocities, velocity_0, rho_0, exponent): - """ - - This function computes a descret exponential density profile. - :math:`\\rho = \\rho_0 \\times \\left( \\frac{v}{v_0} \\right)^n` - - Parameters - ---------- - velocities : astropy.Quantity - Array like velocity profile - velocity_0 : astropy.Quantity - reference velocity - rho_0 : astropy.Quantity - reference density - exponent : float - exponent used in the powerlaw - - Returns - ------- - densities : astropy.Quantity - - """ - densities = rho_0 * np.power((velocities / velocity_0), exponent) - return densities - - -def calculate_exponential_density(velocities, velocity_0, rho_0): - """ - This function computes the exponential density profile. - :math:`\\rho = \\rho_0 \\times \\exp \\left( -\\frac{v}{v_0} \\right)` - - Parameters - ---------- - velocities : astropy.Quantity - Array like velocity profile - velocity_0 : astropy.Quantity - reference velocity - rho_0 : astropy.Quantity - reference density - - Returns - ------- - densities : astropy.Quantity - - """ - densities = rho_0 * np.exp(-(velocities / velocity_0)) - return densities - - -def calculate_density_after_time( - densities: u.Quantity, time_0: u.Quantity, time_explosion: u.Quantity -) -> u.Quantity: - """ - scale the density from an initial time of the model to the - time of the explosion by ^-3 - - Parameters - ---------- - densities : astropy.units.Quantity - densities - time_0 : astropy.units.Quantity - time of the model - time_explosion : astropy.units.Quantity - time to be scaled to - - Returns - ------- - scaled_density : astropy.units.Quantity - in g / cm^3 - """ - - return (densities * (time_explosion / time_0) ** -3).to(u.g / (u.cm**3)) diff --git a/tardis/io/model/parse_composition_configuration.py b/tardis/io/model/parse_composition_configuration.py new file mode 100644 index 00000000000..898f136e05f --- /dev/null +++ b/tardis/io/model/parse_composition_configuration.py @@ -0,0 +1,100 @@ +from tardis.io.model.parse_density_configuration import ( + parse_density_from_config, + parse_density_from_csvy, +) +from tardis.io.model.parse_mass_fraction_configuration import ( + parse_mass_fractions_from_config, + parse_mass_fractions_from_csvy, +) +from tardis.model.matter.composition import Composition + + +def parse_composition_from_config(atom_data, config, time_explosion, geometry): + """Parse the composition data from a CSVY model. + + Parameters + ---------- + atom_data : object + The atom data used for parsing. + config : object + The configuration data. + time_explosion : float + The time of the explosion. + geometry : object + The geometry of the model. + + Returns + ------- + Composition + The parsed composition + array + Electron densities. + """ + density, electron_densities = parse_density_from_config(config) + + ( + nuclide_mass_fractions, + raw_isotope_mass_fractions, + ) = parse_mass_fractions_from_config(config, geometry, time_explosion) + + return ( + Composition( + density, + nuclide_mass_fractions, + raw_isotope_mass_fractions, + atom_data.atom_data.mass.copy(), + ), + electron_densities, + ) + + +def parse_composition_from_csvy( + atom_data, csvy_model_config, csvy_model_data, time_explosion, geometry +): + """ + Parse the composition data from a CSVY model. + + Parameters + ---------- + atom_data : object + The atom data used for parsing. + csvy_model_config : object + The configuration data of the CSVY model. + csvy_model_data : object + The data of the CSVY model. + time_explosion : float + The time of the explosion. + geometry : object + The geometry of the model. + + Returns + ------- + Composition : object + The parsed composition + + Raises + ------ + None. + + Notes + ----- + This function parses the composition data from a CSVY model. It calls the 'parse_density_csvy' + function to parse the density data, and the 'parse_mass_fraction_csvy' function to parse the mass fraction + and isotope mass fraction data. The parsed data is returned as a Composition object. + """ + density = parse_density_from_csvy( + csvy_model_config, csvy_model_data, time_explosion + ) + + ( + nuclide_mass_fractions, + raw_isotope_mass_fractions, + ) = parse_mass_fractions_from_csvy( + csvy_model_config, csvy_model_data, geometry, time_explosion + ) + return Composition( + density, + nuclide_mass_fractions, + raw_isotope_mass_fractions, + atom_data.atom_data.mass.copy(), + ) diff --git a/tardis/io/model/parse_density_configuration.py b/tardis/io/model/parse_density_configuration.py index d70046b5e79..29e054c4e72 100644 --- a/tardis/io/model/parse_density_configuration.py +++ b/tardis/io/model/parse_density_configuration.py @@ -1,16 +1,22 @@ +import logging +from typing import Tuple + import numpy as np from astropy import units as u from tardis.io.configuration.config_reader import ( - ConfigurationNameSpace, Configuration, + ConfigurationNameSpace, +) +from tardis.io.model.parse_geometry_configuration import ( + parse_structure_from_config, ) from tardis.util.base import quantity_linspace -from typing import Tuple +logger = logging.getLogger(__name__) -def parse_density_section( +def parse_density_section_config( density_configuration: ConfigurationNameSpace, v_middle: u.Quantity, time_explosion: u.Quantity, @@ -21,7 +27,6 @@ def parse_density_section( Parameters ---------- - density_configuration : tardis.io.config_reader.Configuration v_middle : astropy.Quantity middle of the velocity bins @@ -63,12 +68,12 @@ def parse_density_section( time_0 = density_configuration.get("time_0", time_explosion) else: raise ValueError( - f"Unrecognized density type " f"'{density_configuration.type}'" + f"Unrecognized density type '{density_configuration.type}'" ) return density_0, time_0 -def parse_config_v1_density(config: Configuration) -> u.Quantity: +def parse_density_from_config(config: Configuration) -> u.Quantity: """ Parse the configuration file and produce a density at time_explosion. @@ -82,23 +87,39 @@ def parse_config_v1_density(config: Configuration) -> u.Quantity: density: u.Quantity """ + time_explosion = config.supernova.time_explosion.cgs + ( + density_time, + velocity, + density, + electron_densities, + temperature, + ) = parse_structure_from_config(config) + + if density is None: + adjusted_velocity = velocity.insert(0, 0) + v_middle = adjusted_velocity[1:] * 0.5 + adjusted_velocity[:-1] * 0.5 + d_conf = config.model.structure.density + density, density_time = parse_density_section_config( + d_conf, v_middle, time_explosion + ) - velocity = quantity_linspace( - config.model.structure.velocity.start, - config.model.structure.velocity.stop, - config.model.structure.velocity.num + 1, - ).cgs + density = calculate_density_after_time( + density, density_time, time_explosion + ) - adjusted_velocity = velocity.insert(0, 0) - v_middle = adjusted_velocity[1:] * 0.5 + adjusted_velocity[:-1] * 0.5 - time_explosion = config.supernova.time_explosion.cgs - d_conf = config.model.structure.density - density_0, time_0 = parse_density_section(d_conf, v_middle, time_explosion) + # Note: This is the number of shells *without* taking in mind the + # v boundaries. + if len(density) == len(velocity): + logger.warning( + "Number of density points larger than number of shells. Assuming inner point irrelevant" + ) + density = density[1:] - return calculate_density_after_time(density_0, time_0, time_explosion) + return density, electron_densities -def parse_csvy_density( +def parse_density_section_csvy( csvy_model_config: Configuration, time_explosion: u.Quantity ) -> u.Quantity: """ @@ -135,12 +156,62 @@ def parse_csvy_density( no_of_shells = len(adjusted_velocity) - 1 if hasattr(csvy_model_config, "density"): - density_0, time_0 = parse_density_section( + density_0, time_0 = parse_density_section_config( csvy_model_config.density, v_middle, time_explosion ) return calculate_density_after_time(density_0, time_0, time_explosion) +def parse_density_from_csvy(csvy_model_config, csvy_model_data, time_explosion): + """ + Parse the density data from a CSVY model. + + Parameters + ---------- + csvy_model_config : object + The configuration data of the CSVY model. + csvy_model_data : object + The data of the CSVY model. + time_explosion : float + The time of the explosion. + + Returns + ------- + density : object + The parsed density data. + + Raises + ------ + None. + + Notes + ----- + This function parses the density data from a CSVY model. If the CSVY model configuration + contains a 'density' attribute, it uses the 'parse_csvy_density' function to parse the + density data. Otherwise, it calculates the density data using the 'calculate_density_after_time' + function. The parsed density data is returned. + """ + if hasattr(csvy_model_config, "density"): + density = parse_density_section_csvy(csvy_model_config, time_explosion) + else: + time_0 = csvy_model_config.model_density_time_0 + density_field_index = [ + field["name"] for field in csvy_model_config.datatype.fields + ].index("density") + density_unit = u.Unit( + csvy_model_config.datatype.fields[density_field_index]["unit"] + ) + density_0 = csvy_model_data["density"].values * density_unit + density_0 = density_0.to("g/cm^3")[1:] + # Removing as the new architecture removes the 0th shell already + # density_0 = density_0.insert(0, 0) + density = calculate_density_after_time( + density_0, time_0, time_explosion + ) + + return density + + def calculate_power_law_density( velocities: u.Quantity, velocity_0: u.Quantity, @@ -218,5 +289,4 @@ def calculate_density_after_time( scaled_density : astropy.units.Quantity in g / cm^3 """ - return (densities * (time_explosion / time_0) ** -3).to(u.g / (u.cm**3)) diff --git a/tardis/io/model/parse_geometry_configuration.py b/tardis/io/model/parse_geometry_configuration.py new file mode 100644 index 00000000000..c517a4a7ea4 --- /dev/null +++ b/tardis/io/model/parse_geometry_configuration.py @@ -0,0 +1,173 @@ +import os + +from astropy import units as u + +from tardis.io.model.readers.base import read_density_file +from tardis.model.geometry.radial1d import HomologousRadial1DGeometry +from tardis.util.base import quantity_linspace + + +def parse_structure_from_config(config): + """Parses the structure section from a config object + + Parameters + ---------- + config : object + The configuration to parse + + Returns + ------- + Quantity + Time at which densities are valid + Quantity + Velocities + Quantity + Densities + Quantity + Electron densities + Quantity + Temperatures + + Raises + ------ + NotImplementedError + For structure types that are not "specific" or "file" + """ + density_time = None + velocity = None + density = None + electron_densities = None + temperature = None + structure_config = config.model.structure + if structure_config.type == "specific": + velocity = quantity_linspace( + structure_config.velocity.start, + structure_config.velocity.stop, + structure_config.velocity.num + 1, + ).cgs + + elif structure_config.type == "file": + if os.path.isabs(structure_config.filename): + structure_config_fname = structure_config.filename + else: + structure_config_fname = os.path.join( + config.config_dirname, structure_config.filename + ) + + ( + density_time, + velocity, + density, + electron_densities, + temperature, + ) = read_density_file(structure_config_fname, structure_config.filetype) + density = density.insert(0, 0) + else: + raise NotImplementedError + + return density_time, velocity, density, electron_densities, temperature + + +def parse_geometry_from_config(config, time_explosion): + """ + Parse the geometry data from a TARDIS config. + + Parameters + ---------- + config : object + Configuration object. + time_explosion : float + The time of the explosion + + Returns + ------- + HomologousRadial1DGeometry + The parsed geometry + """ + ( + density_time, + velocity, + density, + electron_densities, + temperature, + ) = parse_structure_from_config(config) + + return HomologousRadial1DGeometry( + velocity[:-1], # v_inner + velocity[1:], # v_outer + v_inner_boundary=config.model.structure.get("v_inner_boundary", None), + v_outer_boundary=config.model.structure.get("v_outer_boundary", None), + time_explosion=time_explosion, + ) + + +def parse_geometry_from_csvy( + config, csvy_model_config, csvy_model_data, time_explosion +): + """ + Parse the geometry data from a CSVY model. + + Parameters + ---------- + config : object + The configuration data. + csvy_model_config : object + The configuration data of the CSVY model. + csvy_model_data : object + The data of the CSVY model. + time_explosion : float + The time of the explosion. + + Returns + ------- + geometry : object + The parsed geometry. + + Raises + ------ + None. + + Notes + ----- + This function parses the geometry data from a CSVY model. It extracts the velocity + information from the CSVY model configuration or data. The parsed velocity data is + used to create a homologous radial 1D geometry object, which is returned. + """ + if hasattr(config, "model"): + if hasattr(config.model, "v_inner_boundary"): + v_boundary_inner = config.model.v_inner_boundary + else: + v_boundary_inner = None + + if hasattr(config.model, "v_outer_boundary"): + v_boundary_outer = config.model.v_outer_boundary + else: + v_boundary_outer = None + else: + v_boundary_inner = None + v_boundary_outer = None + + if hasattr(csvy_model_config, "velocity"): + velocity = quantity_linspace( + csvy_model_config.velocity.start, + csvy_model_config.velocity.stop, + csvy_model_config.velocity.num + 1, + ).cgs + else: + velocity_field_index = [ + field["name"] for field in csvy_model_config.datatype.fields + ].index("velocity") + velocity_unit = u.Unit( + csvy_model_config.datatype.fields[velocity_field_index]["unit"] + ) + velocity = csvy_model_data["velocity"].values * velocity_unit + velocity = velocity.to("cm/s") + + geometry = HomologousRadial1DGeometry( + velocity[:-1], # v_inner + velocity[1:], # v_outer + v_inner_boundary=v_boundary_inner, + v_outer_boundary=v_boundary_outer, + time_explosion=time_explosion, + ) + return geometry diff --git a/tardis/io/model/parse_mass_fraction_configuration.py b/tardis/io/model/parse_mass_fraction_configuration.py new file mode 100644 index 00000000000..e55de12a3e6 --- /dev/null +++ b/tardis/io/model/parse_mass_fraction_configuration.py @@ -0,0 +1,227 @@ +import logging +import os + +import astropy.units as u +import numpy as np +import pandas as pd + +from tardis.io.model.readers.base import read_mass_fractions_file +from tardis.io.model.readers.csvy import parse_csv_mass_fractions +from tardis.io.model.readers.generic_readers import read_uniform_mass_fractions +from tardis.model.matter.composition import Composition +from tardis.model.matter.decay import IsotopicMassFraction + +logger = logging.getLogger(__name__) + + +def parse_mass_fractions_from_config(config, geometry, time_explosion): + """ + Parse the mass fraction configuration data. + + Parameters + ---------- + config : object + The configuration data. + geometry : object + The geometry of the model. + time_explosion : float + The time of the explosion. + + Returns + ------- + nuclide_mass_fractions : object + The parsed nuclide mass fraction. + + raw_isotope_mass_fractions : object + The parsed raw isotope mass fraction. This is the isotope mass fraction data before decay. + + Raises + ------ + None. + + Notes + ----- + This function parses the abundance configuration data and returns the parsed nuclide + mass fraction. The mass fraction configuration can be of type 'uniform' or 'file'. If it + is of type 'uniform', the mass fraction and isotope mass fraction are read using the + 'read_uniform_mass fractions' function. If it is of type 'file', the mass fraction and + isotope mass fraction are read from a file using the 'read_abundances_file' function. + The parsed data is then processed to replace NaN values with 0.0, remove rows with + zero sum, and normalize the data if necessary. The resulting nuclide mass fraction + is returned. + """ + mass_fractions_section = config.model.abundances + isotope_mass_fractions = pd.DataFrame() + + if mass_fractions_section.type == "uniform": + mass_fractions, isotope_mass_fractions = read_uniform_mass_fractions( + mass_fractions_section, geometry.no_of_shells + ) + + elif mass_fractions_section.type == "file": + if os.path.isabs(mass_fractions_section.filename): + mass_fractions_fname = mass_fractions_section.filename + else: + mass_fractions_fname = os.path.join( + config.config_dirname, mass_fractions_section.filename + ) + + ( + index, + mass_fractions, + isotope_mass_fractions, + ) = read_mass_fractions_file( + mass_fractions_fname, mass_fractions_section.filetype + ) + + mass_fractions = mass_fractions.replace(np.nan, 0.0) + mass_fractions = mass_fractions[mass_fractions.sum(axis=1) > 0] + + norm_factor = mass_fractions.sum(axis=0) + isotope_mass_fractions.sum( + axis=0 + ) + + if np.any(np.abs(norm_factor - 1) > 1e-12): + logger.warning( + "Mass fractions have not been normalized to 1. - normalizing" + ) + mass_fractions /= norm_factor + isotope_mass_fractions /= norm_factor + # The next line is if the mass_fractions are given via dict + # and not gone through the schema validator + raw_isotope_mass_fractions = isotope_mass_fractions + model_isotope_time_0 = config.model.abundances.get( + "model_isotope_time_0", 0.0 * u.day + ) + isotope_mass_fractions = IsotopicMassFraction( + isotope_mass_fractions, time_0=model_isotope_time_0 + ).decay(time_explosion) + + nuclide_mass_fractions = convert_to_nuclide_mass_fractions( + isotope_mass_fractions, mass_fractions + ) + return nuclide_mass_fractions, raw_isotope_mass_fractions + + +def parse_mass_fractions_from_csvy( + csvy_model_config, csvy_model_data, geometry, time_explosion +): + """ + Parse the mass fraction data from a CSVY model. + + Parameters + ---------- + csvy_model_config : object + The configuration data of the CSVY model. + csvy_model_data : object + The data of the CSVY model. + geometry : object + The geometry of the model. + + Returns + ------- + mass_fractions : pd.DataFrame + The parsed mass fraction data. + isotope_mass_fractions : pandas.DataFrame + The parsed isotope mass fraction data. + + Raises + ------ + None. + + Notes + ----- + This function parses the mass fraction data from a CSVY model. If the CSVY model + configuration contains an 'abundance' attribute, it uses the 'read_uniform_mass_fractions' + function to parse the mass fraction and isotope mass fraction data. Otherwise, it uses the + 'parse_csv_mass_fractions' function to parse the data. The parsed data is then processed + to replace NaN values with 0.0, remove rows with zero sum, and normalize the data + if necessary. The resulting mass fraction and isotope mass fraction arrays are returned. + """ + if hasattr(csvy_model_config, "abundance"): + mass_fractions_section = csvy_model_config.abundance + mass_fractions, isotope_mass_fractions = read_uniform_mass_fractions( + mass_fractions_section, geometry.no_of_shells + ) + else: + _, mass_fractions, isotope_mass_fractions = parse_csv_mass_fractions( + csvy_model_data + ) + mass_fractions = mass_fractions.loc[:, 1:] + mass_fractions.columns = np.arange(mass_fractions.shape[1]) + isotope_mass_fractions = isotope_mass_fractions.loc[:, 1:] + isotope_mass_fractions.columns = np.arange( + isotope_mass_fractions.shape[1] + ) + + mass_fractions = mass_fractions.replace(np.nan, 0.0) + mass_fractions = mass_fractions[mass_fractions.sum(axis=1) > 0] + isotope_mass_fractions = isotope_mass_fractions.replace(np.nan, 0.0) + isotope_mass_fractions = isotope_mass_fractions[ + isotope_mass_fractions.sum(axis=1) > 0 + ] + norm_factor = mass_fractions.sum(axis=0) + isotope_mass_fractions.sum( + axis=0 + ) + + if np.any(np.abs(norm_factor - 1) > 1e-12): + logger.warning( + "Mass fractions have not been normalized to 1. - normalizing" + ) + mass_fractions /= norm_factor + isotope_mass_fractions /= norm_factor + + raw_isotope_mass_fraction = isotope_mass_fractions + isotope_mass_fractions = IsotopicMassFraction( + isotope_mass_fractions, time_0=csvy_model_config.model_isotope_time_0 + ).decay(time_explosion) + return ( + convert_to_nuclide_mass_fractions( + isotope_mass_fractions, mass_fractions + ), + raw_isotope_mass_fraction, + ) + + +def convert_to_nuclide_mass_fractions(isotopic_mass_fractions, mass_fractions): + """ + Convert the mass fraction and isotope mass fraction data to nuclide mass fraction. + + Parameters + ---------- + isotope_mass_fraction : pandas.DataFrame + The isotope mass fraction data. + mass_fractions : pandas.DataFrame + The mass fraction data. + + Returns + ------- + nuclide_mass_fraction : pandas.DataFrame + The converted nuclide mass fraction. + + Raises + ------ + None. + + Notes + ----- + This function converts the mass fraction and isotope mass fraction data to nuclide mass fraction. + If the mass fraction data is not None, it is converted to nuclide mass fraction by mapping + the mass fraction index to nuclide indices using the 'convert_element2nuclide_index' function. + The resulting mass fraction data is then concatenated with the isotope mass fraction data to + obtain the final nuclide mass fraction. + """ + nuclide_mass_fractions = pd.DataFrame() + if mass_fractions is not None: + mass_fractions.index = Composition.convert_element2nuclide_index( + mass_fractions.index + ) + nuclide_mass_fractions = mass_fractions + else: + nuclide_mass_fractions = pd.DataFrame() + + if isotopic_mass_fractions is not None: + nuclide_mass_fractions = pd.concat( + [nuclide_mass_fractions, isotopic_mass_fractions] + ) + return nuclide_mass_fractions diff --git a/tardis/io/model/parse_packet_source_configuration.py b/tardis/io/model/parse_packet_source_configuration.py new file mode 100644 index 00000000000..90aab02e556 --- /dev/null +++ b/tardis/io/model/parse_packet_source_configuration.py @@ -0,0 +1,77 @@ +from astropy import units as u + +from tardis.transport.montecarlo.packet_source import ( + BlackBodySimpleSource, + BlackBodySimpleSourceRelativistic, +) + + +def initialize_packet_source(packet_source, config, geometry): + """ + Initialize the packet source based on config and geometry + + Parameters + ---------- + config : Config + The configuration object containing the supernova and plasma settings. + geometry : Geometry + The geometry object containing the inner radius information. + packet_source : BasePacketSource + The packet source object based on the configuration and geometry. + + Returns + ------- + packet_source : BasePacketSource + The packet source object based on the configuration and geometry. + + Raises + ------ + ValueError + If both t_inner and luminosity_requested are None. + """ + luminosity_requested = config.supernova.luminosity_requested + if config.plasma.initial_t_inner > 0.0 * u.K: + packet_source.radius = geometry.r_inner_active[0] + packet_source.temperature = config.plasma.initial_t_inner + + elif (config.plasma.initial_t_inner < 0.0 * u.K) and ( + luminosity_requested is not None + ): + packet_source.radius = geometry.r_inner_active[0] + packet_source.set_temperature_from_luminosity(luminosity_requested) + else: + raise ValueError( + "Both t_inner and luminosity_requested cannot be None." + ) + return packet_source + + +def parse_packet_source_from_config(config, geometry, legacy_mode_enabled): + """ + Parse the packet source based on the given configuration and geometry. + + Parameters + ---------- + config : Config + The configuration object containing the supernova and plasma settings. + geometry : Geometry + The geometry object containing the inner radius information. + + Returns + ------- + packet_source : BlackBodySimpleSource + The packet source object based on the configuration and geometry. + """ + if config.montecarlo.enable_full_relativity: + packet_source = BlackBodySimpleSourceRelativistic( + base_seed=config.montecarlo.seed, + time_explosion=config.supernova.time_explosion, + legacy_mode_enabled=legacy_mode_enabled, + ) + else: + packet_source = BlackBodySimpleSource( + base_seed=config.montecarlo.seed, + legacy_mode_enabled=legacy_mode_enabled, + ) + + return initialize_packet_source(packet_source, config, geometry) diff --git a/tardis/io/model/parse_radiation_field_configuration.py b/tardis/io/model/parse_radiation_field_configuration.py new file mode 100644 index 00000000000..e12c319feb2 --- /dev/null +++ b/tardis/io/model/parse_radiation_field_configuration.py @@ -0,0 +1,188 @@ +import logging +import os + +import numpy as np +from astropy import units as u + +from tardis import constants as const +from tardis.io.model.parse_geometry_configuration import ( + parse_structure_from_config, +) +from tardis.model.radiation_field_state import ( + DiluteBlackBodyRadiationFieldState, +) + +logger = logging.getLogger(__name__) + + +def parse_radiation_field_state_from_config( + config, geometry, dilution_factor=None, packet_source=None +): + """ + Parses the radiation field state based on the provided configuration, radiative temperature, geometry, dilution factor, and packet source. + + Parameters + ---------- + config : Config + The configuration object. + geometry : Geometry + The geometry object. + dilution_factor : {None, ndarray}, optional + The dilution factor. If None, it is calculated based on the geometry. + packet_source : {None, PacketSource}, optional + The packet source object. + + Returns + ------- + DiluteThermalRadiationFieldState + The parsed radiation field state. + + Raises + ------ + AssertionError + If the length of t_radiative or dilution_factor is not compatible with the geometry. + """ + ( + density_time, + velocity, + density, + electron_densities, + temperature, + ) = parse_structure_from_config(config) + + if temperature is None: + if config.plasma.initial_t_rad > 0 * u.K: + temperature = ( + np.ones(geometry.no_of_shells) * config.plasma.initial_t_rad + ) + else: + temperature = calculate_t_radiative_from_t_inner( + geometry, packet_source + ) + + assert len(temperature) == geometry.no_of_shells + + if dilution_factor is None: + dilution_factor = calculate_geometric_dilution_factor(geometry) + + assert len(dilution_factor) == geometry.no_of_shells + + return DiluteBlackBodyRadiationFieldState( + temperature, dilution_factor, geometry + ) + + +def parse_radiation_field_state_from_csvy( + config, csvy_model_config, csvy_model_data, geometry, packet_source +): + """Parses the radiation field state for CSVY model inputs. + + Parameters + ---------- + config : Config + The configuration object. + csvy_model_config : Config + CSVY model configuration. + csvy_model_data : Config + CSVY model data + geometry : Geometry + The geometry object. + packet_source : {None, PacketSource}, optional + The packet source object. + + Returns + ------- + DiluteThermalRadiationFieldState + The parsed radiation field state. + """ + t_radiative = None + dilution_factor = None + + if hasattr(csvy_model_data, "columns") and ( + "t_rad" in csvy_model_data.columns + ): + t_rad_field_index = [ + field["name"] for field in csvy_model_config.datatype.fields + ].index("t_rad") + t_rad_unit = u.Unit( + csvy_model_config.datatype.fields[t_rad_field_index]["unit"] + ) + t_radiative = csvy_model_data["t_rad"].iloc[1:].values * t_rad_unit + + elif config.plasma.initial_t_rad > 0 * u.K: + t_radiative = ( + np.ones(geometry.no_of_shells) * config.plasma.initial_t_rad + ) + else: + t_radiative = calculate_t_radiative_from_t_inner( + geometry, packet_source + ) + + if np.any(t_radiative < 1000 * u.K): + logging.critical( + "Radiative temperature is too low in some of the shells, temperatures below 1000K " + f"(e.g., T_rad = {t_radiative[np.argmin(t_radiative)]} in shell {np.argmin(t_radiative)} in your model) " + "are not accurately handled by TARDIS.", + ) + + if hasattr(csvy_model_data, "columns") and ( + "dilution_factor" in csvy_model_data.columns + ): + dilution_factor = csvy_model_data["dilution_factor"].iloc[1:].values + else: + dilution_factor = calculate_geometric_dilution_factor(geometry) + + return DiluteBlackBodyRadiationFieldState( + t_radiative, dilution_factor, geometry + ) + + +def calculate_t_radiative_from_t_inner(geometry, packet_source): + """ + Calculates the radiative temperature based on the inner temperature and the geometry of the system. + + Parameters + ---------- + geometry : Geometry + The geometry object. + packet_source : PacketSource + The packet source object. + + Returns + ------- + Quantity + The calculated radiative temperature. + """ + lambda_wien_inner = const.b_wien / packet_source.temperature + t_radiative = const.b_wien / ( + lambda_wien_inner + * (1 + (geometry.v_middle - geometry.v_inner_boundary) / const.c) + ) + return t_radiative + + +def calculate_geometric_dilution_factor(geometry): + """Calculates the geometric dilution factor w for models without a defined w. + + Parameters + ---------- + geometry : Geometry + The geometry object + + Returns + ------- + np.array + The dilution factors for the inout geometry + """ + return 0.5 * ( + 1 + - np.sqrt( + 1 + - ( + geometry.r_inner[geometry.v_inner_boundary_index] ** 2 + / geometry.r_middle**2 + ) + .to(1) + .value + ) + ) diff --git a/tardis/io/model/readers/base.py b/tardis/io/model/readers/base.py index c66da4b2b03..ec0ac37dd54 100644 --- a/tardis/io/model/readers/base.py +++ b/tardis/io/model/readers/base.py @@ -1,11 +1,11 @@ -from tardis.io.model.readers.cmfgen import ( +from tardis.io.model.readers.cmfgen_deprecated import ( read_cmfgen_composition, read_cmfgen_density, ) from tardis.io.model.readers.generic_readers import ( ConfigurationError, read_csv_composition, - read_simple_ascii_abundances, + read_simple_ascii_mass_fractions, read_simple_ascii_density, ) @@ -16,9 +16,9 @@ from tardis.io.model.readers.artis import read_artis_density -def read_abundances_file( - abundance_filename, - abundance_filetype, +def read_mass_fractions_file( + mass_fractions_filename, + mass_fractions_filetype, inner_boundary_index=None, outer_boundary_index=None, ): @@ -27,9 +27,9 @@ def read_abundances_file( Parameters ---------- - abundance_filename : str + mass_fractions_filename : str filename or path of the density file - abundance_filetype : str + mass_fractions_filetype : str type of the density file inner_boundary_index : int index of the inner shell, default None @@ -38,30 +38,32 @@ def read_abundances_file( """ file_parsers = { - "simple_ascii": read_simple_ascii_abundances, - "artis": read_simple_ascii_abundances, + "simple_ascii": read_simple_ascii_mass_fractions, + "artis": read_simple_ascii_mass_fractions, "cmfgen_model": read_cmfgen_composition, "custom_composition": read_csv_composition, } - isotope_abundance = pd.DataFrame() - if abundance_filetype in ["cmfgen_model", "custom_composition"]: - index, abundances, isotope_abundance = file_parsers[abundance_filetype]( - abundance_filename - ) + isotope_mass_fractions = pd.DataFrame() + if mass_fractions_filetype in ["cmfgen_model", "custom_composition"]: + index, mass_fractions, isotope_mass_fractions = file_parsers[ + mass_fractions_filetype + ](mass_fractions_filename) else: - index, abundances = file_parsers[abundance_filetype](abundance_filename) + index, mass_fractions = file_parsers[mass_fractions_filetype]( + mass_fractions_filename + ) if outer_boundary_index is not None: outer_boundary_index_m1 = outer_boundary_index - 1 else: outer_boundary_index_m1 = None index = index[inner_boundary_index:outer_boundary_index] - abundances = abundances.loc[ + mass_fractions = mass_fractions.loc[ :, slice(inner_boundary_index, outer_boundary_index_m1) ] - abundances.columns = np.arange(len(abundances.columns)) - return index, abundances, isotope_abundance + mass_fractions.columns = np.arange(len(mass_fractions.columns)) + return index, mass_fractions, isotope_mass_fractions def read_density_file(filename, filetype): @@ -83,6 +85,10 @@ def read_density_file(filename, filetype): the array containing the velocities unscaled_mean_densities : np.ndarray the array containing the densities + electron_densities : np.ndarray + The array containing electron densities + temperature : np.ndarray + The array containing temperatures """ file_parsers = { "artis": read_artis_density, diff --git a/tardis/io/model/readers/cmfgen.py b/tardis/io/model/readers/cmfgen.py index d11abbb0c9b..77865a3a5d4 100644 --- a/tardis/io/model/readers/cmfgen.py +++ b/tardis/io/model/readers/cmfgen.py @@ -1,91 +1,75 @@ -from tardis.io.model.readers.util import read_csv_isotope_abundances -from tardis.util.base import parse_quantity - - +import re import pandas as pd from astropy import units as u +from pathlib import Path +import dataclasses -import warnings +@dataclasses.dataclass +class CMFGENModel: + metadata: dict + data: pd.DataFrame -def read_cmfgen_density(fname: str): - """ - Reading a density file of the following structure (example; lines starting with a hash will be ignored): - The first density describes the mean density in the center of the model and is not used. - The file consists of a header row and next row contains unit of the respective attributes - Note that the first column has to contain a running index +HEADER_RE_STR = [ + (r"t0:\s+(\d+\.\d+)+\s+day", "t0"), +] - Example: +COLUMN_ROW = 1 +UNIT_ROW = 2 +DATA_START_ROW = 3 - index velocity densities electron_densities temperature - - km/s g/cm^3 /cm^3 K - 0 871.66905 4.2537191e-09 2.5953807e+14 7.6395577 - 1 877.44269 4.2537191e-09 2.5953807e+14 7.6395577 - Rest columns contain abundances of elements and isotopes +def read_cmfgen_model(fname): + """ + Read in a CMFGEN model file and return the data and model Parameters ---------- + fname : str - filename or path with filename Returns ------- - time_of_model : astropy.units.Quantity - time at which the model is valid - velocity : np.ndarray - mean_density : np.ndarray - electron_densities : np.ndarray - temperature : np.ndarray - """ - warnings.warn( - "The current CMFGEN model parser is deprecated", DeprecationWarning - ) - - df = pd.read_csv(fname, comment="#", delimiter=r"\s+", skiprows=[0, 2]) - - with open(fname) as fh: - for row_index, line in enumerate(fh): - if row_index == 0: - time_of_model_string = line.strip().replace("t0:", "") - time_of_model = parse_quantity(time_of_model_string) - elif row_index == 2: - quantities = line.split() - - velocity = u.Quantity(df["velocity"].values, quantities[1]).to("cm/s") - temperature = u.Quantity(df["temperature"].values, quantities[2])[1:] - mean_density = u.Quantity(df["densities"].values, quantities[3])[1:] - electron_densities = u.Quantity( - df["electron_densities"].values, quantities[4] - )[1:] + model : CMFGENModel - return ( - time_of_model, - velocity, - mean_density, - electron_densities, - temperature, - ) - - -def read_cmfgen_composition(fname, delimiter=r"\s+"): - """Read composition from a CMFGEN model file - - The CMFGEN file format contains information about the ejecta state in the - first four columns and the following ones contain elemental and isotopic - abundances. - - WARNING : deprecated - - fname : str - filename of the csv file """ - - warnings.warn( - "The current CMFGEN model parser is deprecated", DeprecationWarning - ) - - return read_csv_isotope_abundances( - fname, delimiter=delimiter, skip_columns=4, skip_rows=[0, 2, 3] + header_re = [re.compile(re_str[0]) for re_str in HEADER_RE_STR] + metadata = {} + with open(fname) as fh: + for i, line in enumerate(fh): + if i < len(HEADER_RE_STR): + header_re_match = header_re[i].match(line) + metadata[HEADER_RE_STR[i][1]] = header_re_match.group(1) + elif i == COLUMN_ROW: + if "Index" in line: + column_names = re.split(r"\s", line.strip()) + column_names = [ + col.lower().replace(" ", "_") for col in column_names + ] + column_names = column_names[ + 1: + ] # Remove Index from column names + else: + raise ValueError( + '"Index" is required in the Cmfgen input file to infer columns' + ) + elif i == UNIT_ROW: + units = re.split(r"\s", line.strip()) + units = units[1:] # Remove index column + for col, unit in zip(column_names, units): + if u.Unit(unit) == "": # dimensionless + continue + metadata[f"{col}_unit"] = u.Unit(unit) + break + + metadata["t0"] = float(metadata["t0"]) * u.day + data = pd.read_csv( + fname, + delim_whitespace=True, + skiprows=DATA_START_ROW, + header=None, + index_col=0, ) + data.columns = column_names + return CMFGENModel(metadata, data) diff --git a/tardis/io/model/readers/cmfgen_deprecated.py b/tardis/io/model/readers/cmfgen_deprecated.py new file mode 100644 index 00000000000..898f7a5dc2c --- /dev/null +++ b/tardis/io/model/readers/cmfgen_deprecated.py @@ -0,0 +1,91 @@ +from tardis.io.model.readers.util import read_csv_isotope_mass_fractions +from tardis.util.base import parse_quantity + + +import pandas as pd +from astropy import units as u + + +import warnings + + +def read_cmfgen_density(fname: str): + """ + Reading a density file of the following structure (example; lines starting with a hash will be ignored): + The first density describes the mean density in the center of the model and is not used. + The file consists of a header row and next row contains unit of the respective attributes + Note that the first column has to contain a running index + + Example: + + index velocity densities electron_densities temperature + - km/s g/cm^3 /cm^3 K + 0 871.66905 4.2537191e-09 2.5953807e+14 7.6395577 + 1 877.44269 4.2537191e-09 2.5953807e+14 7.6395577 + + Rest columns contain abundances of elements and isotopes + + Parameters + ---------- + fname : str + filename or path with filename + + Returns + ------- + time_of_model : astropy.units.Quantity + time at which the model is valid + velocity : np.ndarray + mean_density : np.ndarray + electron_densities : np.ndarray + temperature : np.ndarray + """ + warnings.warn( + "The current CMFGEN model parser is deprecated", DeprecationWarning + ) + + df = pd.read_csv(fname, comment="#", delimiter=r"\s+", skiprows=[0, 2]) + + with open(fname) as fh: + for row_index, line in enumerate(fh): + if row_index == 0: + time_of_model_string = line.strip().replace("t0:", "") + time_of_model = parse_quantity(time_of_model_string) + elif row_index == 2: + quantities = line.split() + + velocity = u.Quantity(df["velocity"].values, quantities[1]).to("cm/s") + temperature = u.Quantity(df["temperature"].values, quantities[2])[1:] + mean_density = u.Quantity(df["densities"].values, quantities[3])[1:] + electron_densities = u.Quantity( + df["electron_densities"].values, quantities[4] + )[1:] + + return ( + time_of_model, + velocity, + mean_density, + electron_densities, + temperature, + ) + + +def read_cmfgen_composition(fname, delimiter=r"\s+"): + """Read composition from a CMFGEN model file + + The CMFGEN file format contains information about the ejecta state in the + first four columns and the following ones contain elemental and isotopic + abundances. + + WARNING : deprecated + + fname : str + filename of the csv file + """ + + warnings.warn( + "The current CMFGEN model parser is deprecated", DeprecationWarning + ) + + return read_csv_isotope_mass_fractions( + fname, delimiter=delimiter, skip_columns=4, skip_rows=[0, 2, 3] + ) diff --git a/tardis/io/model/readers/csvy.py b/tardis/io/model/readers/csvy.py index fa6f350497e..2c03272b5ec 100644 --- a/tardis/io/model/readers/csvy.py +++ b/tardis/io/model/readers/csvy.py @@ -96,9 +96,9 @@ def load_csv_from_csvy(fname): return data -def parse_csv_abundances(csvy_data): +def parse_csv_mass_fractions(csvy_data): """ - A parser for the csv data part of a csvy model file. This function filters out columns that are not abundances. + A parser for the csv data part of a csvy model file. This function filters out columns that are not mass fractions. Parameters ---------- @@ -107,18 +107,18 @@ def parse_csv_abundances(csvy_data): Returns ------- index : np.ndarray - abundances : pandas.DataFrame - isotope_abundance : pandas.MultiIndex + mass_fractions : pandas.DataFrame + isotope_mass_fraction : pandas.MultiIndex """ - abundance_col_names = [ + mass_fraction_col_names = [ name for name in csvy_data.columns if is_valid_nuclide_or_elem(name) ] - df = csvy_data.loc[:, abundance_col_names] + df = csvy_data.loc[:, mass_fraction_col_names] df = df.transpose() - abundance = pd.DataFrame( + mass_fractions = pd.DataFrame( columns=np.arange(df.shape[1]), index=pd.Index([], name="atomic_number"), dtype=np.float64, @@ -127,20 +127,20 @@ def parse_csv_abundances(csvy_data): isotope_index = pd.MultiIndex( [[]] * 2, [[]] * 2, names=["atomic_number", "mass_number"] ) - isotope_abundance = pd.DataFrame( + isotope_mass_fractions = pd.DataFrame( columns=np.arange(df.shape[1]), index=isotope_index, dtype=np.float64 ) for element_symbol_string in df.index[0:]: if element_symbol_string in Z_DICT.values(): z = elem_to_Z(element_symbol_string) - abundance.loc[z, :] = df.loc[element_symbol_string].tolist() + mass_fractions.loc[z, :] = df.loc[element_symbol_string].tolist() else: nuc = Nuclide(element_symbol_string) z = nuc.Z mass_no = nuc.A - isotope_abundance.loc[(z, mass_no), :] = df.loc[ + isotope_mass_fractions.loc[(z, mass_no), :] = df.loc[ element_symbol_string ].tolist() - return abundance.index, abundance, isotope_abundance + return mass_fractions.index, mass_fractions, isotope_mass_fractions diff --git a/tardis/io/model/readers/generic_readers.py b/tardis/io/model/readers/generic_readers.py index 6f037522127..38acb03fec3 100644 --- a/tardis/io/model/readers/generic_readers.py +++ b/tardis/io/model/readers/generic_readers.py @@ -1,14 +1,14 @@ -from astropy import units as u +from pathlib import Path from typing import Any, Tuple -from numpy import recfromtxt + +import numpy as np import pandas as pd +from astropy import units as u +from numpy import recfromtxt from radioactivedecay import Nuclide from radioactivedecay.utils import Z_DICT, elem_to_Z -from pathlib import Path - -import numpy as np -from tardis.io.model.readers.util import read_csv_isotope_abundances +from tardis.io.model.readers.util import read_csv_isotope_mass_fractions from tardis.util.base import parse_quantity @@ -41,7 +41,6 @@ def read_simple_ascii_density( mean_density: u.Quantity mean density """ - with open(fname) as fh: time_of_model_string = fh.readline().strip() time_of_model = parse_quantity(time_of_model_string) @@ -61,10 +60,10 @@ def read_simple_ascii_density( def read_csv_composition(fname, delimiter=r"\s+"): """Read composition from a simple CSV file - The CSV file can contain specific isotopes or elemental abundances in the + The CSV file can contain specific isotopes or elemental mass fractions in the different columns. The first row must contain the header in which the contents of each column is specified by the elemental symbol (for elemental - abundances) or by the symbol plus mass number (for isotopic abundances). + mass fractions) or by the symbol plus mass number (for isotopic mass fractions). Example: C O Fe Ni56 Co @@ -73,16 +72,15 @@ def read_csv_composition(fname, delimiter=r"\s+"): fname : str filename of the csv file """ - - return read_csv_isotope_abundances( + return read_csv_isotope_mass_fractions( fname, delimiter=delimiter, skip_columns=0, skip_rows=[1] ) -def read_simple_ascii_abundances(fname): +def read_simple_ascii_mass_fractions(fname): """ - Reading an abundance file of the following structure (example; lines starting with hash will be ignored): - The first line of abundances describe the abundances in the center of the model and are not used. + Reading a mass fraction file of the following structure (example; lines starting with hash will be ignored): + The first line of mass fractions describe the mass fractions in the center of the model and are not used. #index element1, element2, ..., element30 0 0.4 0.3, .. 0.2 @@ -95,32 +93,32 @@ def read_simple_ascii_abundances(fname): ------- index : np.ndarray containing the indices - abundances : pandas.DataFrame + mass_fractions : pandas.DataFrame data frame containing index, element1 - element30 and columns according to the shells """ data = np.loadtxt(fname) index = data[1:, 0].astype(int) - abundances = pd.DataFrame( + mass_fractions = pd.DataFrame( data[1:, 1:].transpose(), index=np.arange(1, data.shape[1]) ) - return index, abundances + return index, mass_fractions -def read_uniform_abundances(abundances_section, no_of_shells): +def read_uniform_mass_fractions(mass_fractions_section, no_of_shells): """ Parameters ---------- - abundances_section : config.model.abundances + mass_fractions_section : config.model.abundances no_of_shells : int Returns ------- - abundance : pandas.DataFrame - isotope_abundance : pandas.DataFrame + mass_fractions : pandas.DataFrame + isotope_mass_fractions : pandas.DataFrame """ - abundance = pd.DataFrame( + mass_fractions = pd.DataFrame( columns=np.arange(no_of_shells), index=pd.Index(np.arange(1, 120), name="atomic_number"), dtype=np.float64, @@ -129,30 +127,30 @@ def read_uniform_abundances(abundances_section, no_of_shells): isotope_index = pd.MultiIndex( [[]] * 2, [[]] * 2, names=["atomic_number", "mass_number"] ) - isotope_abundance = pd.DataFrame( + isotope_mass_fractions = pd.DataFrame( columns=np.arange(no_of_shells), index=isotope_index, dtype=np.float64 ) - for element_symbol_string in abundances_section: + for element_symbol_string in mass_fractions_section: if element_symbol_string in ["type", "model_isotope_time_0"]: continue try: if element_symbol_string in Z_DICT.values(): z = elem_to_Z(element_symbol_string) - abundance.loc[z] = float( - abundances_section[element_symbol_string] + mass_fractions.loc[z] = float( + mass_fractions_section[element_symbol_string] ) else: nuc = Nuclide(element_symbol_string) mass_no = nuc.A z = nuc.Z - isotope_abundance.loc[(z, mass_no), :] = float( - abundances_section[element_symbol_string] + isotope_mass_fractions.loc[(z, mass_no), :] = float( + mass_fractions_section[element_symbol_string] ) except RuntimeError as err: raise RuntimeError( - f"Abundances are not defined properly in config file : {err.args}" + f"mass_fractions are not defined properly in config file : {err.args}" ) - return abundance, isotope_abundance + return mass_fractions, isotope_mass_fractions diff --git a/tardis/io/model/readers/stella.py b/tardis/io/model/readers/stella.py index e314e716e3b..41074c8c725 100644 --- a/tardis/io/model/readers/stella.py +++ b/tardis/io/model/readers/stella.py @@ -12,17 +12,17 @@ class STELLAModel: HEADER_RE_STR = [ - ("\s+days post max Lbol\s+(\d+\.\d*)", "t_max"), - ("\s+zones\s+(\d+)", "zones"), + (r"\s+days post max Lbol\s+(\d+\.\d*)", "t_max"), + (r"\s+zones\s+(\d+)", "zones"), ( - "\s+inner boundary mass\s+(\d+\.\d+E[+-]\d+)\s+\d+\.\d+E[+-]\d+", - "inner_boundary_mass", + r"\s+inner boundary mass\s+(\d+\.\d+E[+-]\d+)\s+\d+\.\d+E[+-]\d+", + r"inner_boundary_mass", ), - ("\s+total mass\s+(\d+\.\d+E[+-]\d+)\s+\d+\.\d+E[+-]\d+", "total_mass"), + (r"\s+total mass\s+(\d+\.\d+E[+-]\d+)\s+\d+\.\d+E[+-]\d+", "total_mass"), ] DATA_START_ROW = 5 -COLUMN_WITH_UNIT_RE = re.compile("(.+)\s+\((.+)\)") +COLUMN_WITH_UNIT_RE = re.compile(r"(.+)\s+\((.+)\)") def read_stella_model(fname): diff --git a/tardis/io/model/readers/tests/test_ascii_readers.py b/tardis/io/model/readers/tests/test_ascii_readers.py index 751a84fb824..10fb6bf5bca 100644 --- a/tardis/io/model/readers/tests/test_ascii_readers.py +++ b/tardis/io/model/readers/tests/test_ascii_readers.py @@ -13,7 +13,7 @@ from tardis.io.model.readers.base import read_density_file from tardis.io.model.readers.generic_readers import ( - read_simple_ascii_abundances, + read_simple_ascii_mass_fractions, read_simple_ascii_density, ) @@ -32,7 +32,7 @@ def test_simple_ascii_density_reader_time(example_model_file_dir): def test_simple_ascii_abundance_reader(example_model_file_dir): - (index, abundances,) = read_simple_ascii_abundances( + (index, abundances,) = read_simple_ascii_mass_fractions( example_model_file_dir / "artis_abundances.dat" ) diff --git a/tardis/io/model/readers/tests/test_cmfgen_reader.py b/tardis/io/model/readers/tests/test_cmfgen_reader.py index 6a31952ca42..25ba9fefc87 100644 --- a/tardis/io/model/readers/tests/test_cmfgen_reader.py +++ b/tardis/io/model/readers/tests/test_cmfgen_reader.py @@ -4,7 +4,7 @@ from pytest import fixture from astropy import units as u -from tardis.io.model.cmfgen import read_cmfgen_model +from tardis.io.model.readers.cmfgen import read_cmfgen_model MODEL_DATA_PATH = Path(__file__).parent / "data" diff --git a/tardis/io/model/readers/util.py b/tardis/io/model/readers/util.py index 8597af839bf..bc46b224607 100644 --- a/tardis/io/model/readers/util.py +++ b/tardis/io/model/readers/util.py @@ -4,13 +4,13 @@ from radioactivedecay.utils import Z_DICT, elem_to_Z -def read_csv_isotope_abundances( +def read_csv_isotope_mass_fractions( fname, delimiter=r"\s+", skip_columns=0, skip_rows=[1] ): """ A generic parser for a TARDIS composition stored as a CSV file - The parser can read in both elemental and isotopic abundances. The first + The parser can read in both elemental and isotopic mass fractions. The first column is always expected to contain a running index, labelling the grid cells. The parser also allows for additional information to be stored in the first skip_columns columns. These will be ignored if skip_columns > 0. @@ -39,16 +39,15 @@ def read_csv_isotope_abundances( Returns ------- index : np.ndarray - abundances : pandas.DataFrame - isotope_abundance : pandas.MultiIndex + mass_fractions : pandas.DataFrame + isotope_mass_fraction : pandas.MultiIndex """ - df = pd.read_csv( fname, comment="#", sep=delimiter, skiprows=skip_rows, index_col=0 ) df = df.transpose() - abundance = pd.DataFrame( + mass_fractions = pd.DataFrame( columns=np.arange(df.shape[1]), index=pd.Index([], name="atomic_number"), dtype=np.float64, @@ -57,20 +56,20 @@ def read_csv_isotope_abundances( isotope_index = pd.MultiIndex( [[]] * 2, [[]] * 2, names=["atomic_number", "mass_number"] ) - isotope_abundance = pd.DataFrame( + isotope_mass_fractions = pd.DataFrame( columns=np.arange(df.shape[1]), index=isotope_index, dtype=np.float64 ) for element_symbol_string in df.index[skip_columns:]: if element_symbol_string in Z_DICT.values(): z = elem_to_Z(element_symbol_string) - abundance.loc[z, :] = df.loc[element_symbol_string].tolist() + mass_fractions.loc[z, :] = df.loc[element_symbol_string].tolist() else: nuc = Nuclide(element_symbol_string) z = nuc.Z mass_no = nuc.A - isotope_abundance.loc[(z, mass_no), :] = df.loc[ + isotope_mass_fractions.loc[(z, mass_no), :] = df.loc[ element_symbol_string ].tolist() - return abundance.index, abundance, isotope_abundance + return mass_fractions.index, mass_fractions, isotope_mass_fractions diff --git a/tardis/io/model_reader.py b/tardis/io/model_reader.py deleted file mode 100644 index 3e1bc37c3d5..00000000000 --- a/tardis/io/model_reader.py +++ /dev/null @@ -1,846 +0,0 @@ -# reading different model files - -import logging -import warnings - -import h5py -import numpy as np -import pandas as pd -from astropy import units as u -from numpy import recfromtxt -from radioactivedecay import Nuclide -from radioactivedecay.utils import Z_DICT, elem_to_Z, NuclideStrError - -from tardis.io.configuration.config_reader import ConfigurationNameSpace -from tardis.transport.montecarlo.base import MonteCarloTransportSolver -from tardis.transport.montecarlo.packet_source import ( - BlackBodySimpleSource, - BlackBodySimpleSourceRelativistic, -) -from tardis.util.base import is_valid_nuclide_or_elem, parse_quantity - -# Adding logging support -logger = logging.getLogger(__name__) - - -class ConfigurationError(Exception): - pass - - -def read_density_file(filename, filetype): - """ - read different density file formats - - Parameters - ---------- - filename : str - filename or path of the density file - filetype : str - type of the density file - - Returns - ------- - time_of_model : astropy.units.Quantity - time at which the model is valid - velocity : np.ndarray - the array containing the velocities - unscaled_mean_densities : np.ndarray - the array containing the densities - """ - file_parsers = { - "artis": read_artis_density, - "simple_ascii": read_simple_ascii_density, - "cmfgen_model": read_cmfgen_density, - } - - electron_densities = None - temperature = None - if filetype == "cmfgen_model": - ( - time_of_model, - velocity, - unscaled_mean_densities, - electron_densities, - temperature, - ) = read_cmfgen_density(filename) - else: - (time_of_model, velocity, unscaled_mean_densities) = file_parsers[ - filetype - ](filename) - - v_inner = velocity[:-1] - v_outer = velocity[1:] - - invalid_volume_mask = (v_outer - v_inner) <= 0 - if invalid_volume_mask.sum() > 0: - message = "\n".join( - [ - f"cell {i:d}: v_inner {v_inner_i:s}, v_outer {v_outer_i:s}" - for i, v_inner_i, v_outer_i in zip( - np.arange(len(v_outer))[invalid_volume_mask], - v_inner[invalid_volume_mask], - v_outer[invalid_volume_mask], - ) - ] - ) - raise ConfigurationError( - "Invalid volume of following cell(s):\n" f"{message:s}" - ) - - return ( - time_of_model, - velocity, - unscaled_mean_densities, - electron_densities, - temperature, - ) - - -def read_abundances_file( - abundance_filename, - abundance_filetype, - inner_boundary_index=None, - outer_boundary_index=None, -): - """ - read different density file formats - - Parameters - ---------- - abundance_filename : str - filename or path of the density file - abundance_filetype : str - type of the density file - inner_boundary_index : int - index of the inner shell, default None - outer_boundary_index : int - index of the outer shell, default None - """ - file_parsers = { - "simple_ascii": read_simple_ascii_abundances, - "artis": read_simple_ascii_abundances, - "cmfgen_model": read_cmfgen_composition, - "custom_composition": read_csv_composition, - } - - isotope_abundance = pd.DataFrame() - if abundance_filetype in ["cmfgen_model", "custom_composition"]: - index, abundances, isotope_abundance = file_parsers[abundance_filetype]( - abundance_filename - ) - else: - index, abundances = file_parsers[abundance_filetype](abundance_filename) - - if outer_boundary_index is not None: - outer_boundary_index_m1 = outer_boundary_index - 1 - else: - outer_boundary_index_m1 = None - index = index[inner_boundary_index:outer_boundary_index] - abundances = abundances.loc[ - :, slice(inner_boundary_index, outer_boundary_index_m1) - ] - abundances.columns = np.arange(len(abundances.columns)) - return index, abundances, isotope_abundance - - -def read_uniform_abundances(abundances_section, no_of_shells): - """ - Parameters - ---------- - abundances_section : config.model.abundances - no_of_shells : int - - Returns - ------- - abundance : pandas.DataFrame - isotope_abundance : pandas.DataFrame - """ - abundance = pd.DataFrame( - columns=np.arange(no_of_shells), - index=pd.Index(np.arange(1, 120), name="atomic_number"), - dtype=np.float64, - ) - - isotope_index = pd.MultiIndex( - [[]] * 2, [[]] * 2, names=["atomic_number", "mass_number"] - ) - isotope_abundance = pd.DataFrame( - columns=np.arange(no_of_shells), index=isotope_index, dtype=np.float64 - ) - - for element_symbol_string in abundances_section: - if ( - element_symbol_string == "type" - or element_symbol_string == "model_isotope_time_0" - ): - continue - try: - if element_symbol_string in Z_DICT.values(): - z = elem_to_Z(element_symbol_string) - abundance.loc[z] = float( - abundances_section[element_symbol_string] - ) - else: - nuc = Nuclide(element_symbol_string) - mass_no = nuc.A - z = nuc.Z - isotope_abundance.loc[(z, mass_no), :] = float( - abundances_section[element_symbol_string] - ) - - except NuclideStrError as err: - raise RuntimeError( - f"Abundances are not defined properly in config file : {err.args}" - ) - - return abundance, isotope_abundance - - -def read_simple_ascii_density(fname): - """ - Reading a density file of the following structure (example; lines starting with a hash will be ignored): - The first density describes the mean density in the center of the model and is not used. - 5 s - #index velocity [km/s] density [g/cm^3] - 0 1.1e4 1.6e8 - 1 1.2e4 1.7e8 - - Parameters - ---------- - fname : str - filename or path with filename - - Returns - ------- - time_of_model : astropy.units.Quantity - time at which the model is valid - data : pandas.DataFrame - data frame containing index, velocity (in km/s) and density - """ - with open(fname) as fh: - time_of_model_string = fh.readline().strip() - time_of_model = parse_quantity(time_of_model_string) - - data = recfromtxt( - fname, - skip_header=1, - names=("index", "velocity", "density"), - dtype=(int, float, float), - ) - velocity = (data["velocity"] * u.km / u.s).to("cm/s") - mean_density = (data["density"] * u.Unit("g/cm^3"))[1:] - - return time_of_model, velocity, mean_density - - -def read_artis_density(fname): - """ - Reading a density file of the following structure (example; lines starting with a hash will be ignored): - The first density describes the mean density in the center of the model and is not used. - 5 - #index velocity [km/s] log10(density) [log10(g/cm^3)] - 0 1.1e4 1.6e8 - 1 1.2e4 1.7e8 - - Parameters - ---------- - fname : str - filename or path with filename - - Returns - ------- - time_of_model : astropy.units.Quantity - time at which the model is valid - data : pandas.DataFrame - data frame containing index, velocity (in km/s) and density - """ - with open(fname) as fh: - for i, line in enumerate(open(fname)): - if i == 0: - no_of_shells = np.int64(line.strip()) - elif i == 1: - time_of_model = u.Quantity(float(line.strip()), "day").to("s") - elif i == 2: - break - - artis_model_columns = [ - "index", - "velocities", - "mean_densities_0", - "ni56_fraction", - "co56_fraction", - "fe52_fraction", - "cr48_fraction", - ] - artis_model = recfromtxt( - fname, - skip_header=2, - usecols=(0, 1, 2, 4, 5, 6, 7), - unpack=True, - dtype=[(item, np.float64) for item in artis_model_columns], - ) - - velocity = u.Quantity(artis_model["velocities"], "km/s").to("cm/s") - mean_density = u.Quantity(10 ** artis_model["mean_densities_0"], "g/cm^3")[ - 1: - ] - - return time_of_model, velocity, mean_density - - -def read_cmfgen_density(fname): - """ - Reading a density file of the following structure (example; lines starting with a hash will be ignored): - The first density describes the mean density in the center of the model and is not used. - The file consists of a header row and next row contains unit of the respective attributes - Note that the first column has to contain a running index - - Example: - - index velocity densities electron_densities temperature - - km/s g/cm^3 /cm^3 K - 0 871.66905 4.2537191e-09 2.5953807e+14 7.6395577 - 1 877.44269 4.2537191e-09 2.5953807e+14 7.6395577 - - Rest columns contain abundances of elements and isotopes - - Parameters - ---------- - fname : str - filename or path with filename - - Returns - ------- - time_of_model : astropy.units.Quantity - time at which the model is valid - velocity : np.ndarray - mean_density : np.ndarray - electron_densities : np.ndarray - temperature : np.ndarray - """ - warnings.warn( - "The current CMFGEN model parser is deprecated", DeprecationWarning - ) - - df = pd.read_csv(fname, comment="#", delimiter=r"\s+", skiprows=[0, 2]) - - with open(fname) as fh: - for row_index, line in enumerate(fh): - if row_index == 0: - time_of_model_string = line.strip().replace("t0:", "") - time_of_model = parse_quantity(time_of_model_string) - elif row_index == 2: - quantities = line.split() - - velocity = u.Quantity(df["velocity"].values, quantities[1]).to("cm/s") - temperature = u.Quantity(df["temperature"].values, quantities[2])[1:] - mean_density = u.Quantity(df["densities"].values, quantities[3])[1:] - electron_densities = u.Quantity( - df["electron_densities"].values, quantities[4] - )[1:] - - return ( - time_of_model, - velocity, - mean_density, - electron_densities, - temperature, - ) - - -def read_simple_ascii_abundances(fname): - """ - Reading an abundance file of the following structure (example; lines starting with hash will be ignored): - The first line of abundances describe the abundances in the center of the model and are not used. - #index element1, element2, ..., element30 - 0 0.4 0.3, .. 0.2 - - Parameters - ---------- - fname : str - filename or path with filename - - Returns - ------- - index : np.ndarray - containing the indices - abundances : pandas.DataFrame - data frame containing index, element1 - element30 and columns according to the shells - """ - data = np.loadtxt(fname) - - index = data[1:, 0].astype(int) - abundances = pd.DataFrame( - data[1:, 1:].transpose(), index=np.arange(1, data.shape[1]) - ) - - return index, abundances - - -def read_cmfgen_composition(fname, delimiter=r"\s+"): - """Read composition from a CMFGEN model file - - The CMFGEN file format contains information about the ejecta state in the - first four columns and the following ones contain elemental and isotopic - abundances. - - WARNING : deprecated - - fname : str - filename of the csv file - """ - warnings.warn( - "The current CMFGEN model parser is deprecated", DeprecationWarning - ) - - return read_csv_isotope_abundances( - fname, delimiter=delimiter, skip_columns=4, skip_rows=[0, 2, 3] - ) - - -def read_csv_composition(fname, delimiter=r"\s+"): - """Read composition from a simple CSV file - - The CSV file can contain specific isotopes or elemental abundances in the - different columns. The first row must contain the header in which the - contents of each column is specified by the elemental symbol (for elemental - abundances) or by the symbol plus mass number (for isotopic abundances). - - Example: C O Fe Ni56 Co - - The i-th row specifies the composition in the i-th shell - - fname : str - filename of the csv file - """ - return read_csv_isotope_abundances( - fname, delimiter=delimiter, skip_columns=0, skip_rows=[1] - ) - - -def read_csv_isotope_abundances( - fname, delimiter=r"\s+", skip_columns=0, skip_rows=[1] -): - """ - A generic parser for a TARDIS composition stored as a CSV file - - The parser can read in both elemental and isotopic abundances. The first - column is always expected to contain a running index, labelling the grid - cells. The parser also allows for additional information to be stored in - the first skip_columns columns. These will be ignored if skip_columns > 0. - Note that the first column, containing the cell index is not taken into - account here. - - Specific header lines can be skipped by the skip_rows keyword argument - - It is expected that the first row of the date block (after skipping the - rows specified in skip_rows) specifies the different elements and isotopes. - Each row after contains the composition in the corresponding grid shell. - The first composition row describes the composition of the photosphere and - is essentially ignored (for the default value of skip_rows). - - Example: - - Index C O Ni56 - 0 1 1 1 - 1 0.4 0.3 0.2 - - Parameters - ---------- - fname : str - filename or path with filename - - Returns - ------- - index : np.ndarray - abundances : pandas.DataFrame - isotope_abundance : pandas.MultiIndex - """ - df = pd.read_csv( - fname, comment="#", sep=delimiter, skiprows=skip_rows, index_col=0 - ) - df = df.transpose() - - abundance = pd.DataFrame( - columns=np.arange(df.shape[1]), - index=pd.Index([], name="atomic_number"), - dtype=np.float64, - ) - - isotope_index = pd.MultiIndex( - [[]] * 2, [[]] * 2, names=["atomic_number", "mass_number"] - ) - isotope_abundance = pd.DataFrame( - columns=np.arange(df.shape[1]), index=isotope_index, dtype=np.float64 - ) - - for element_symbol_string in df.index[skip_columns:]: - if element_symbol_string in Z_DICT.values(): - z = elem_to_Z(element_symbol_string) - abundance.loc[z, :] = df.loc[element_symbol_string].tolist() - else: - nuc = Nuclide(element_symbol_string) - z = nuc.Z - mass_no = nuc.A - isotope_abundance.loc[(z, mass_no), :] = df.loc[ - element_symbol_string - ].tolist() - - return abundance.index, abundance, isotope_abundance - - -def parse_csv_abundances(csvy_data): - """ - A parser for the csv data part of a csvy model file. This function filters out columns that are not abundances. - - Parameters - ---------- - csvy_data : pandas.DataFrame - - Returns - ------- - index : np.ndarray - abundances : pandas.DataFrame - isotope_abundance : pandas.MultiIndex - """ - abundance_col_names = [ - name for name in csvy_data.columns if is_valid_nuclide_or_elem(name) - ] - df = csvy_data.loc[:, abundance_col_names] - - df = df.transpose() - - abundance = pd.DataFrame( - columns=np.arange(df.shape[1]), - index=pd.Index([], name="atomic_number"), - dtype=np.float64, - ) - - isotope_index = pd.MultiIndex( - [[]] * 2, [[]] * 2, names=["atomic_number", "mass_number"] - ) - isotope_abundance = pd.DataFrame( - columns=np.arange(df.shape[1]), index=isotope_index, dtype=np.float64 - ) - - for element_symbol_string in df.index[0:]: - if element_symbol_string in Z_DICT.values(): - z = elem_to_Z(element_symbol_string) - abundance.loc[z, :] = df.loc[element_symbol_string].tolist() - else: - nuc = Nuclide(element_symbol_string) - z = nuc.Z - mass_no = nuc.A - isotope_abundance.loc[(z, mass_no), :] = df.loc[ - element_symbol_string - ].tolist() - - return abundance.index, abundance, isotope_abundance - - -def transport_to_dict(transport): - """ - Retrieves all the data from a transport object and returns a dictionary. - - Parameters - ---------- - transport : tardis.transport.montecarlo.MontecarloTransport - - Returns - ------- - transport_dict : dict - integrator_settings : dict - virtual_spectrum_spawn_range : dict - """ - transport_dict = { - "Edotlu_estimator": transport.Edotlu_estimator, - "bf_heating_estimator": transport.bf_heating_estimator, - "disable_electron_scattering": transport.disable_electron_scattering, - "enable_full_relativity": transport.enable_full_relativity, - "input_energy": transport.input_energy, - "input_mu": transport.input_mu, - "input_nu": transport.input_nu, - "input_r_cgs": transport.input_r, - "j_blue_estimator": transport.j_blue_estimator, - "j_estimator": transport.j_estimator, - "last_interaction_in_nu": transport.last_interaction_in_nu, - "last_interaction_type": transport.last_interaction_type, - "last_line_interaction_in_id": transport.last_line_interaction_in_id, - "last_line_interaction_out_id": transport.last_line_interaction_out_id, - "last_line_interaction_shell_id": transport.last_line_interaction_shell_id, - "line_interaction_type": transport.line_interaction_type, - "nu_bar_estimator": transport.nu_bar_estimator, - "photo_ion_estimator": transport.photo_ion_estimator, - "photo_ion_estimator_statistics": transport.photo_ion_estimator_statistics, - "r_inner": transport.r_inner_cgs, - "r_outer": transport.r_outer_cgs, - "packet_source_base_seed": transport.packet_source.base_seed, - "spectrum_frequency_cgs": transport.spectrum_frequency, - "spectrum_method": transport.spectrum_method, - "stim_recomb_cooling_estimator": transport.stim_recomb_cooling_estimator, - "stim_recomb_estimator": transport.stim_recomb_estimator, - "time_of_simulation_cgs": transport.time_of_simulation, - "use_gpu": transport.use_gpu, - "v_inner": transport.v_inner_cgs, - "v_outer": transport.v_outer_cgs, - "nthreads": transport.nthreads, - "virt_logging": transport.virt_logging, - "virt_packet_energies": transport.virt_packet_energies, - "virt_packet_initial_mus": transport.virt_packet_initial_mus, - "virt_packet_initial_rs": transport.virt_packet_initial_rs, - "virt_packet_last_interaction_in_nu": transport.virt_packet_last_interaction_in_nu, - "virt_packet_last_interaction_type": transport.virt_packet_last_interaction_type, - "virt_packet_last_line_interaction_in_id": transport.virt_packet_last_line_interaction_in_id, - "virt_packet_last_line_interaction_out_id": transport.virt_packet_last_line_interaction_out_id, - "virt_packet_last_line_interaction_shell_id": transport.virt_packet_last_line_interaction_shell_id, - "virt_packet_nus": transport.virt_packet_nus, - "volume_cgs": transport.volume, - } - - for key, value in transport_dict.items(): - if key.endswith("_cgs"): - transport_dict[key] = [value.cgs.value, value.unit.to_string()] - - integrator_settings = transport.integrator_settings - virtual_spectrum_spawn_range = transport.virtual_spectrum_spawn_range - - return ( - transport_dict, - integrator_settings, - virtual_spectrum_spawn_range, - ) - - -def store_transport_to_hdf(transport, fname): - """ - Stores data from MontecarloTransport object into a hdf file. - - Parameters - ---------- - transport : tardis.transport.montecarlo.MontecarloTransport - filename : str - """ - with h5py.File(fname, "a") as f: - transport_group = f.require_group("transport") - transport_group.clear() - - ( - transport_data, - integrator_settings, - virtual_spectrum_spawn_range, - ) = transport_to_dict(transport) - - for key, value in transport_data.items(): - if key.endswith("_cgs"): - transport_group.create_dataset(key, data=value[0]) - transport_group.create_dataset(key + "_unit", data=value[1]) - else: - if value is not None: - transport_group.create_dataset(key, data=value) - - integrator_settings_group = transport_group.create_group( - "integrator_settings" - ) - for key, value in integrator_settings.items(): - integrator_settings_group.create_dataset(key, data=value) - - virtual_spectrum_spawn_range_group = transport_group.create_group( - "virtual_spectrum_spawn_range" - ) - for key, value in virtual_spectrum_spawn_range.items(): - virtual_spectrum_spawn_range_group.create_dataset(key, data=value) - - -def transport_from_hdf(fname): - """ - Creates a MontecarloTransport object using data stored in a hdf file. - - Parameters - ---------- - fname : str - - Returns - ------- - new_transport : tardis.transport.montecarlo.MontecarloTransport - """ - d = {} - - # Loading data from hdf file - with h5py.File(fname, "r") as f: - transport_group = f["transport"] - for key, value in transport_group.items(): - if not key.endswith("_unit"): - if isinstance(value, h5py.Dataset): - if isinstance(value[()], bytes): - d[key] = value[()].decode("utf-8") - else: - d[key] = value[()] - else: - data_inner = {} - for key_inner, value_inner in value.items(): - if isinstance(value_inner[()], bytes): - data_inner[key] = value_inner[()].decode("utf-8") - else: - data_inner[key] = value_inner[()] - data_inner[key_inner] = value_inner[()] - d[key] = data_inner - - for key, value in transport_group.items(): - if key.endswith("_unit"): - d[key[:-5]] = [d[key[:-5]], value[()]] - - # Converting cgs data to astropy quantities - for key, value in d.items(): - if key.endswith("_cgs"): - d[key] = u.Quantity(value[0], unit=u.Unit(value[1].decode("utf-8"))) - - # Using packet source seed to packet source - if not d["enable_full_relativity"]: - d["packet_source"] = BlackBodySimpleSource(d["packet_source_base_seed"]) - else: - d["packet_source"] = BlackBodySimpleSourceRelativistic( - d["packet_source_base_seed"] - ) - - # Converting virtual spectrum spawn range values to astropy quantities - vssr = d["virtual_spectrum_spawn_range"] - d["virtual_spectrum_spawn_range"] = { - "start": u.Quantity(vssr["start"], unit=u.Unit("Angstrom")), - "end": u.Quantity(vssr["end"], unit=u.Unit("Angstrom")), - } - - # Converting dictionaries to ConfigurationNameSpace - d["integrator_settings"] = ConfigurationNameSpace(d["integrator_settings"]) - d["virtual_spectrum_spawn_range"] = ConfigurationNameSpace( - d["virtual_spectrum_spawn_range"] - ) - - # Creating a transport object and storing data - new_transport = MonteCarloTransportSolver( - spectrum_frequency=d["spectrum_frequency_cgs"], - virtual_spectrum_spawn_range=d["virtual_spectrum_spawn_range"], - disable_electron_scattering=d["disable_electron_scattering"], - enable_full_relativity=d["enable_full_relativity"], - line_interaction_type=d["line_interaction_type"], - integrator_settings=d["integrator_settings"], - spectrum_method=d["spectrum_method"], - packet_source=d["packet_source"], - nthreads=d["nthreads"], - enable_virtual_packet_logging=d["virt_logging"], - use_gpu=d["use_gpu"], - montecarlo_configuration=d["montecarlo_configuration"], - ) - - new_transport.Edotlu_estimator = d["Edotlu_estimator"] - new_transport.bf_heating_estimator = d["bf_heating_estimator"] - new_transport.input_energy = d["input_energy"] - new_transport.input_mu = d["input_mu"] - new_transport.input_nu = d["input_nu"] - new_transport.input_r = d["input_r_cgs"] - new_transport.j_blue_estimator = d["j_blue_estimator"] - new_transport.j_estimator = d["j_estimator"] - new_transport.last_interaction_in_nu = d["last_interaction_in_nu"] - new_transport.last_interaction_type = d["last_interaction_type"] - new_transport.last_line_interaction_in_id = d["last_line_interaction_in_id"] - new_transport.last_line_interaction_out_id = d[ - "last_line_interaction_out_id" - ] - new_transport.last_line_interaction_shell_id = d[ - "last_line_interaction_shell_id" - ] - new_transport.nu_bar_estimator = d["nu_bar_estimator"] - new_transport.photo_ion_estimator = d["photo_ion_estimator"] - new_transport.photo_ion_estimator_statistics = d[ - "photo_ion_estimator_statistics" - ] - new_transport.r_inner_cgs = d["r_inner"] - new_transport.r_outer_cgs = d["r_outer"] - new_transport.stim_recomb_cooling_estimator = d[ - "stim_recomb_cooling_estimator" - ] - new_transport.stim_recomb_estimator = d["stim_recomb_estimator"] - new_transport.time_of_simulation = d["time_of_simulation_cgs"] - new_transport.v_inner_cgs = d["v_inner"] - new_transport.v_outer_cgs = d["v_outer"] - new_transport.virt_packet_energies = d["virt_packet_energies"] - new_transport.virt_packet_initial_mus = d["virt_packet_initial_mus"] - new_transport.virt_packet_initial_rs = d["virt_packet_initial_rs"] - new_transport.virt_packet_last_interaction_in_nu = d[ - "virt_packet_last_interaction_in_nu" - ] - new_transport.virt_packet_last_interaction_type = d[ - "virt_packet_last_interaction_type" - ] - new_transport.virt_packet_last_line_interaction_in_id = d[ - "virt_packet_last_line_interaction_in_id" - ] - new_transport.virt_packet_last_line_interaction_out_id = d[ - "virt_packet_last_line_interaction_out_id" - ] - new_transport.virt_packet_last_line_interaction_shell_id = d[ - "virt_packet_last_line_interaction_shell_id" - ] - new_transport.virt_packet_nus = d["virt_packet_nus"] - new_transport.volume = d["volume_cgs"] - - return new_transport - - -def model_to_dict(model): - """ - Retrieves all the data from a SimulationState object and returns a dictionary. - - Parameters - ---------- - transport : tardis.model.SimulationState - - Returns - ------- - model_dict : dict - isotope_abundance : dict - """ - model_dict = { - "velocity_cgs": model.velocity, - "abundance": model.abundance, - "time_explosion_cgs": model.time_explosion, - "t_inner_cgs": model.t_inner, - "t_radiative_cgs": model.t_radiative, - "dilution_factor": model.dilution_factor, - "v_boundary_inner_cgs": model.v_boundary_inner, - "v_boundary_outer_cgs": model.v_boundary_outer, - "w": model.w, - "t_rad_cgs": model.t_rad, - "r_inner_cgs": model.r_inner, - "density_cgs": model.density, - } - - for key, value in model_dict.items(): - if hasattr(value, "unit"): - model_dict[key] = [value.cgs.value, value.unit.to_string()] - - isotope_abundance = model.raw_isotope_abundance.__dict__ - - return model_dict, isotope_abundance - - -def store_model_to_hdf(model, fname): - """ - Stores data from SimulationState object into a hdf file. - - Parameters - ---------- - model : tardis.model.SimulationState - filename : str - """ - with h5py.File(fname, "a") as f: - model_group = f.require_group("model") - model_group.clear() - - model_dict, isotope_abundance = model_to_dict(model) - - for key, value in model_dict.items(): - if key.endswith("_cgs"): - model_group.create_dataset(key, data=value[0]) - model_group.create_dataset(key + "_unit", data=value[1]) - else: - model_group.create_dataset(key, data=value) diff --git a/tardis/io/tests/test_model_reader.py b/tardis/io/tests/test_model_reader.py index 97a6ac7fe63..d1ed9808a75 100644 --- a/tardis/io/tests/test_model_reader.py +++ b/tardis/io/tests/test_model_reader.py @@ -5,14 +5,14 @@ from tardis.io.configuration.config_reader import Configuration from tardis.io.model.readers.artis import read_artis_density -from tardis.io.model.readers.cmfgen import ( +from tardis.io.model.readers.cmfgen_deprecated import ( read_cmfgen_composition, read_cmfgen_density, ) from tardis.io.model.readers.generic_readers import ( read_csv_composition, - read_simple_ascii_abundances, - read_uniform_abundances, + read_simple_ascii_mass_fractions, + read_uniform_mass_fractions, ) @@ -61,8 +61,8 @@ def test_simple_read_artis_density(artis_density_fname: str): # Artis files are currently read with read ascii files function -def test_read_simple_ascii_abundances(artis_abundances_fname): - index, abundances = read_simple_ascii_abundances(artis_abundances_fname) +def test_read_simple_ascii_mass_fractions(artis_abundances_fname): + index, abundances = read_simple_ascii_mass_fractions(artis_abundances_fname) assert len(abundances.columns) == 69 assert np.isclose(abundances[23].loc[2], 2.672351e-08, atol=1.0e-12) @@ -91,8 +91,8 @@ def test_read_cmfgen_isotope_abundances(cmfgen_fname): assert isotope_abundance.shape == (2, 9) -def test_read_uniform_abundances(isotope_uniform_abundance): - abundances, isotope_abundance = read_uniform_abundances( +def test_read_uniform_mass_fractions(isotope_uniform_abundance): + abundances, isotope_abundance = read_uniform_mass_fractions( isotope_uniform_abundance, 20 ) assert np.isclose(abundances.loc[8, 2], 0.19, atol=1.0e-12) diff --git a/tardis/model/base.py b/tardis/model/base.py index adb86b99bd6..776e2f920fd 100644 --- a/tardis/model/base.py +++ b/tardis/model/base.py @@ -5,27 +5,27 @@ import numpy as np from astropy import units as u -from tardis import constants from tardis.io.configuration.config_reader import Configuration from tardis.io.configuration.config_validator import validate_dict +from tardis.io.model.parse_composition_configuration import ( + parse_composition_from_config, + parse_composition_from_csvy, +) +from tardis.io.model.parse_geometry_configuration import ( + parse_geometry_from_config, + parse_geometry_from_csvy, +) +from tardis.io.model.parse_packet_source_configuration import ( + parse_packet_source_from_config, +) +from tardis.io.model.parse_radiation_field_configuration import ( + parse_radiation_field_state_from_config, + parse_radiation_field_state_from_csvy, +) from tardis.io.model.readers.csvy import ( load_csvy, ) from tardis.io.util import HDFWriterMixin -from tardis.model.matter.composition import Composition -from tardis.model.parse_input import ( - parse_abundance_config, - parse_csvy_composition, - parse_csvy_geometry, - parse_csvy_radiation_field_state, - parse_radiation_field_state, - parse_structure_config, - parse_packet_source, -) -from tardis.transport.montecarlo.packet_source import BlackBodySimpleSource -from tardis.model.radiation_field_state import ( - DiluteBlackBodyRadiationFieldState, -) from tardis.util.base import is_valid_nuclide_or_elem logger = logging.getLogger(__name__) @@ -289,31 +289,18 @@ def from_config(cls, config, atom_data, legacy_mode_enabled=False): """ time_explosion = config.supernova.time_explosion.cgs - ( - electron_densities, - t_radiative, - geometry, - density, - ) = parse_structure_config(config, time_explosion) + geometry = parse_geometry_from_config(config, time_explosion) - nuclide_mass_fraction, raw_isotope_abundance = parse_abundance_config( - config, geometry, time_explosion + composition, electron_densities = parse_composition_from_config( + atom_data, config, time_explosion, geometry ) - # using atom_data.mass.copy() to ensure that the original atom_data is not modified - composition = Composition( - density, - nuclide_mass_fraction, - raw_isotope_abundance, - atom_data.atom_data.mass.copy(), - ) - - packet_source = parse_packet_source( + packet_source = parse_packet_source_from_config( config, geometry, legacy_mode_enabled ) - radiation_field_state = parse_radiation_field_state( + + radiation_field_state = parse_radiation_field_state_from_config( config, - t_radiative, geometry, dilution_factor=None, packet_source=packet_source, @@ -395,11 +382,11 @@ def from_csvy(cls, config, atom_data=None, legacy_mode_enabled=False): electron_densities = None - geometry = parse_csvy_geometry( + geometry = parse_geometry_from_csvy( config, csvy_model_config, csvy_model_data, time_explosion ) - composition = parse_csvy_composition( + composition = parse_composition_from_csvy( atom_data, csvy_model_config, csvy_model_data, @@ -407,11 +394,11 @@ def from_csvy(cls, config, atom_data=None, legacy_mode_enabled=False): geometry, ) - packet_source = parse_packet_source( + packet_source = parse_packet_source_from_config( config, geometry, legacy_mode_enabled ) - radiation_field_state = parse_csvy_radiation_field_state( + radiation_field_state = parse_radiation_field_state_from_csvy( config, csvy_model_config, csvy_model_data, geometry, packet_source ) diff --git a/tardis/model/parse_input.py b/tardis/model/parse_input.py deleted file mode 100644 index d1b5ba3766f..00000000000 --- a/tardis/model/parse_input.py +++ /dev/null @@ -1,722 +0,0 @@ -import logging -import os - -import numpy as np -import pandas as pd -from astropy import units as u - -from tardis import constants as const -from tardis.io.model.parse_density_configuration import ( - calculate_density_after_time, - parse_config_v1_density, - parse_csvy_density, -) -from tardis.io.model.readers.base import read_abundances_file, read_density_file -from tardis.io.model.readers.csvy import parse_csv_abundances -from tardis.io.model.readers.generic_readers import read_uniform_abundances -from tardis.model.geometry.radial1d import HomologousRadial1DGeometry -from tardis.model.matter.composition import Composition -from tardis.model.matter.decay import IsotopicMassFraction -from tardis.model.radiation_field_state import ( - DiluteBlackBodyRadiationFieldState, -) -from tardis.transport.montecarlo.packet_source import ( - BlackBodySimpleSource, - BlackBodySimpleSourceRelativistic, -) -from tardis.util.base import quantity_linspace - -logger = logging.getLogger(__name__) - - -def parse_structure_config(config, time_explosion, enable_homology=True): - """ - Parse the structure configuration data. - - Parameters - ---------- - config : object - The configuration data. - time_explosion : float - The time of the explosion. - enable_homology : bool, optional - Whether to enable homology (default is True). - - Returns - ------- - electron_densities : object - The parsed electron densities. - temperature : object - The parsed temperature. - geometry : object - The parsed geometry. - density : object - The parsed density. - - Raises - ------ - NotImplementedError - If the structure configuration type is not supported. - - Notes - ----- - This function parses the structure configuration data and returns the parsed electron - densities, temperature, geometry, and density. The structure configuration can be of - type 'specific' or 'file'. If it is of type 'specific', the velocity and density are - parsed from the configuration. If it is of type 'file', the velocity and density are - read from a file. The parsed data is used to create a homologous radial 1D geometry object. - """ - electron_densities = None - temperature = None - structure_config = config.model.structure - if structure_config.type == "specific": - velocity = quantity_linspace( - structure_config.velocity.start, - structure_config.velocity.stop, - structure_config.velocity.num + 1, - ).cgs - density = parse_config_v1_density(config) - - elif structure_config.type == "file": - if os.path.isabs(structure_config.filename): - structure_config_fname = structure_config.filename - else: - structure_config_fname = os.path.join( - config.config_dirname, structure_config.filename - ) - - ( - time_0, - velocity, - density_0, - electron_densities, - temperature, - ) = read_density_file(structure_config_fname, structure_config.filetype) - density_0 = density_0.insert(0, 0) - - density = calculate_density_after_time( - density_0, time_0, time_explosion - ) - - else: - raise NotImplementedError - - # Note: This is the number of shells *without* taking in mind the - # v boundaries. - if len(density) == len(velocity): - logger.warning( - "Number of density points larger than number of shells. Assuming inner point irrelevant" - ) - density = density[1:] - geometry = HomologousRadial1DGeometry( - velocity[:-1], # v_inner - velocity[1:], # v_outer - v_inner_boundary=structure_config.get("v_inner_boundary", None), - v_outer_boundary=structure_config.get("v_outer_boundary", None), - time_explosion=time_explosion, - ) - return electron_densities, temperature, geometry, density - - -def parse_csvy_geometry( - config, csvy_model_config, csvy_model_data, time_explosion -): - """ - Parse the geometry data from a CSVY model. - - Parameters - ---------- - config : object - The configuration data. - csvy_model_config : object - The configuration data of the CSVY model. - csvy_model_data : object - The data of the CSVY model. - time_explosion : float - The time of the explosion. - - Returns - ------- - geometry : object - The parsed geometry. - - Raises - ------ - None. - - Notes - ----- - This function parses the geometry data from a CSVY model. It extracts the velocity - information from the CSVY model configuration or data. The parsed velocity data is - used to create a homologous radial 1D geometry object, which is returned. - """ - if hasattr(config, "model"): - if hasattr(config.model, "v_inner_boundary"): - v_boundary_inner = config.model.v_inner_boundary - else: - v_boundary_inner = None - - if hasattr(config.model, "v_outer_boundary"): - v_boundary_outer = config.model.v_outer_boundary - else: - v_boundary_outer = None - else: - v_boundary_inner = None - v_boundary_outer = None - - if hasattr(csvy_model_config, "velocity"): - velocity = quantity_linspace( - csvy_model_config.velocity.start, - csvy_model_config.velocity.stop, - csvy_model_config.velocity.num + 1, - ).cgs - else: - velocity_field_index = [ - field["name"] for field in csvy_model_config.datatype.fields - ].index("velocity") - velocity_unit = u.Unit( - csvy_model_config.datatype.fields[velocity_field_index]["unit"] - ) - velocity = csvy_model_data["velocity"].values * velocity_unit - velocity = velocity.to("cm/s") - - geometry = HomologousRadial1DGeometry( - velocity[:-1], # v_inner - velocity[1:], # v_outer - v_inner_boundary=v_boundary_inner, - v_outer_boundary=v_boundary_outer, - time_explosion=time_explosion, - ) - return geometry - - -def parse_abundance_config(config, geometry, time_explosion): - """ - Parse the abundance configuration data. - - Parameters - ---------- - config : object - The configuration data. - geometry : object - The geometry of the model. - time_explosion : float - The time of the explosion. - - Returns - ------- - nuclide_mass_fraction : object - The parsed nuclide mass fraction. - - raw_isotope_abundance : object - The parsed raw isotope abundance. This is the isotope abundance data before decay. - - Raises - ------ - None. - - Notes - ----- - This function parses the abundance configuration data and returns the parsed nuclide - mass fraction. The abundance configuration can be of type 'uniform' or 'file'. If it - is of type 'uniform', the abundance and isotope abundance are read using the - 'read_uniform_abundances' function. If it is of type 'file', the abundance and - isotope abundance are read from a file using the 'read_abundances_file' function. - The parsed data is then processed to replace NaN values with 0.0, remove rows with - zero sum, and normalize the data if necessary. The resulting nuclide mass fraction - is returned. - """ - abundances_section = config.model.abundances - isotope_abundance = pd.DataFrame() - - if abundances_section.type == "uniform": - abundance, isotope_abundance = read_uniform_abundances( - abundances_section, geometry.no_of_shells - ) - - elif abundances_section.type == "file": - if os.path.isabs(abundances_section.filename): - abundances_fname = abundances_section.filename - else: - abundances_fname = os.path.join( - config.config_dirname, abundances_section.filename - ) - - index, abundance, isotope_abundance = read_abundances_file( - abundances_fname, abundances_section.filetype - ) - - abundance = abundance.replace(np.nan, 0.0) - abundance = abundance[abundance.sum(axis=1) > 0] - - norm_factor = abundance.sum(axis=0) + isotope_abundance.sum(axis=0) - - if np.any(np.abs(norm_factor - 1) > 1e-12): - logger.warning( - "Abundances have not been normalized to 1. - normalizing" - ) - abundance /= norm_factor - isotope_abundance /= norm_factor - # The next line is if the abundances are given via dict - # and not gone through the schema validator - raw_isotope_abundance = isotope_abundance - model_isotope_time_0 = config.model.abundances.get( - "model_isotope_time_0", 0.0 * u.day - ) - isotope_abundance = IsotopicMassFraction( - isotope_abundance, time_0=model_isotope_time_0 - ).decay(time_explosion) - - nuclide_mass_fraction = convert_to_nuclide_mass_fraction( - isotope_abundance, abundance - ) - return nuclide_mass_fraction, raw_isotope_abundance - - -def convert_to_nuclide_mass_fraction(isotopic_mass_fraction, mass_fraction): - """ - Convert the abundance and isotope abundance data to nuclide mass fraction. - - Parameters - ---------- - isotope_abundance : pandas.DataFrame - The isotope abundance data. - abundance : pandas.DataFrame - The abundance data. - - Returns - ------- - nuclide_mass_fraction : pandas.DataFrame - The converted nuclide mass fraction. - - Raises - ------ - None. - - Notes - ----- - This function converts the abundance and isotope abundance data to nuclide mass fraction. - If the abundance data is not None, it is converted to nuclide mass fraction by mapping - the abundance index to nuclide indices using the 'convert_element2nuclide_index' function. - The resulting abundance data is then concatenated with the isotope abundance data to - obtain the final nuclide mass fraction. - """ - nuclide_mass_fraction = pd.DataFrame() - if mass_fraction is not None: - mass_fraction.index = Composition.convert_element2nuclide_index( - mass_fraction.index - ) - nuclide_mass_fraction = mass_fraction - else: - nuclide_mass_fraction = pd.DataFrame() - - if isotopic_mass_fraction is not None: - nuclide_mass_fraction = pd.concat( - [nuclide_mass_fraction, isotopic_mass_fraction] - ) - return nuclide_mass_fraction - - -def parse_csvy_composition( - atom_data, csvy_model_config, csvy_model_data, time_explosion, geometry -): - """ - Parse the composition data from a CSVY model. - - Parameters - ---------- - atom_data : object - The atom data used for parsing. - csvy_model_config : object - The configuration data of the CSVY model. - csvy_model_data : object - The data of the CSVY model. - time_explosion : float - The time of the explosion. - geometry : object - The geometry of the model. - - Returns - ------- - density : object - The parsed density data. - abundance : object - The parsed abundance data. - isotope_abundance : object - The parsed isotope abundance data. - elemental_mass : object - The elemental mass data. - - Raises - ------ - None. - - Notes - ----- - This function parses the composition data from a CSVY model. It calls the 'parse_density_csvy' - function to parse the density data, and the 'parse_abundance_csvy' function to parse the abundance - and isotope abundance data. The parsed data is returned as density, abundance, isotope_abundance, - and elemental_mass. - """ - density = parse_density_csvy( - csvy_model_config, csvy_model_data, time_explosion - ) - - nuclide_mass_fraction, raw_isotope_mass_fraction = parse_abundance_csvy( - csvy_model_config, csvy_model_data, geometry, time_explosion - ) - return Composition( - density, - nuclide_mass_fraction, - raw_isotope_mass_fraction, - atom_data.atom_data.mass.copy(), - ) - - -def parse_abundance_csvy( - csvy_model_config, csvy_model_data, geometry, time_explosion -): - """ - Parse the abundance data from a CSVY model. - - Parameters - ---------- - csvy_model_config : object - The configuration data of the CSVY model. - csvy_model_data : object - The data of the CSVY model. - geometry : object - The geometry of the model. - - Returns - ------- - abundance : pd.DataFrame - The parsed abundance data. - isotope_abundance : pandas.DataFrame - The parsed isotope abundance data. - - Raises - ------ - None. - - Notes - ----- - This function parses the abundance data from a CSVY model. If the CSVY model - configuration contains an 'abundance' attribute, it uses the 'read_uniform_abundances' - function to parse the abundance and isotope abundance data. Otherwise, it uses the - 'parse_csv_abundances' function to parse the data. The parsed data is then processed - to replace NaN values with 0.0, remove rows with zero sum, and normalize the data - if necessary. The resulting abundance and isotope abundance arrays are returned. - """ - if hasattr(csvy_model_config, "abundance"): - abundances_section = csvy_model_config.abundance - mass_fraction, isotope_mass_fraction = read_uniform_abundances( - abundances_section, geometry.no_of_shells - ) - else: - _, mass_fraction, isotope_mass_fraction = parse_csv_abundances( - csvy_model_data - ) - mass_fraction = mass_fraction.loc[:, 1:] - mass_fraction.columns = np.arange(mass_fraction.shape[1]) - isotope_mass_fraction = isotope_mass_fraction.loc[:, 1:] - isotope_mass_fraction.columns = np.arange( - isotope_mass_fraction.shape[1] - ) - - mass_fraction = mass_fraction.replace(np.nan, 0.0) - mass_fraction = mass_fraction[mass_fraction.sum(axis=1) > 0] - isotope_mass_fraction = isotope_mass_fraction.replace(np.nan, 0.0) - isotope_mass_fraction = isotope_mass_fraction[ - isotope_mass_fraction.sum(axis=1) > 0 - ] - norm_factor = mass_fraction.sum(axis=0) + isotope_mass_fraction.sum(axis=0) - - if np.any(np.abs(norm_factor - 1) > 1e-12): - logger.warning( - "Abundances have not been normalized to 1. - normalizing" - ) - mass_fraction /= norm_factor - isotope_mass_fraction /= norm_factor - - raw_isotope_mass_fraction = isotope_mass_fraction - isotope_mass_fraction = IsotopicMassFraction( - isotope_mass_fraction, time_0=csvy_model_config.model_isotope_time_0 - ).decay(time_explosion) - return ( - convert_to_nuclide_mass_fraction(isotope_mass_fraction, mass_fraction), - raw_isotope_mass_fraction, - ) - - -def parse_density_csvy(csvy_model_config, csvy_model_data, time_explosion): - """ - Parse the density data from a CSVY model. - - Parameters - ---------- - csvy_model_config : object - The configuration data of the CSVY model. - csvy_model_data : object - The data of the CSVY model. - time_explosion : float - The time of the explosion. - - Returns - ------- - density : object - The parsed density data. - - Raises - ------ - None. - - Notes - ----- - This function parses the density data from a CSVY model. If the CSVY model configuration - contains a 'density' attribute, it uses the 'parse_csvy_density' function to parse the - density data. Otherwise, it calculates the density data using the 'calculate_density_after_time' - function. The parsed density data is returned. - """ - if hasattr(csvy_model_config, "density"): - density = parse_csvy_density(csvy_model_config, time_explosion) - else: - time_0 = csvy_model_config.model_density_time_0 - density_field_index = [ - field["name"] for field in csvy_model_config.datatype.fields - ].index("density") - density_unit = u.Unit( - csvy_model_config.datatype.fields[density_field_index]["unit"] - ) - density_0 = csvy_model_data["density"].values * density_unit - # Removing as thee new architecture removes the 0th shell already - # density_0 = density_0.to("g/cm^3")[1:] - # density_0 = density_0.insert(0, 0) - density = calculate_density_after_time( - density_0, time_0, time_explosion - ) - - return density - - -def parse_radiation_field_state( - config, t_radiative, geometry, dilution_factor=None, packet_source=None -): - """ - Parses the radiation field state based on the provided configuration, radiative temperature, geometry, dilution factor, and packet source. - - Parameters - ---------- - config : Config - The configuration object. - t_radiative : {None, Quantity}, optional - The radiative temperature. If None, it is calculated based on the initial_t_rad value in the plasma configuration. - geometry : Geometry - The geometry object. - dilution_factor : {None, ndarray}, optional - The dilution factor. If None, it is calculated based on the geometry. - packet_source : {None, PacketSource}, optional - The packet source object. - - Returns - ------- - DiluteThermalRadiationFieldState - The parsed radiation field state. - - Raises - ------ - AssertionError - If the length of t_radiative or dilution_factor is not compatible with the geometry. - """ - if t_radiative is None: - if config.plasma.initial_t_rad > 0 * u.K: - t_radiative = ( - np.ones(geometry.no_of_shells) * config.plasma.initial_t_rad - ) - else: - t_radiative = calculate_t_radiative_from_t_inner( - geometry, packet_source - ) - - assert len(t_radiative) == geometry.no_of_shells - - if dilution_factor is None: - dilution_factor = calculate_geometric_dilution_factor(geometry) - - assert len(dilution_factor) == geometry.no_of_shells - - return DiluteBlackBodyRadiationFieldState( - t_radiative, dilution_factor, geometry - ) - - -def initialize_packet_source( - config, geometry, packet_source, legacy_mode_enabled -): - """ - Initialize the packet source based on config and geometry - - Parameters - ---------- - config : Config - The configuration object containing the supernova and plasma settings. - geometry : Geometry - The geometry object containing the inner radius information. - packet_source : BasePacketSource - The packet source object based on the configuration and geometry. - - Returns - ------- - packet_source : BasePacketSource - The packet source object based on the configuration and geometry. - - Raises - ------ - 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, - legacy_mode_enabled=legacy_mode_enabled, - ) - else: - packet_source = BlackBodySimpleSource( - base_seed=config.montecarlo.seed, - legacy_mode_enabled=legacy_mode_enabled, - ) - - luminosity_requested = config.supernova.luminosity_requested - if config.plasma.initial_t_inner > 0.0 * u.K: - packet_source.radius = geometry.r_inner_active[0] - packet_source.temperature = config.plasma.initial_t_inner - - elif (config.plasma.initial_t_inner < 0.0 * u.K) and ( - luminosity_requested is not None - ): - packet_source.radius = geometry.r_inner_active[0] - packet_source.set_temperature_from_luminosity(luminosity_requested) - else: - raise ValueError( - "Both t_inner and luminosity_requested cannot be None." - ) - - return packet_source - - -def parse_packet_source(config, geometry, legacy_mode_enabled): - """ - Parse the packet source based on the given configuration and geometry. - - Parameters - ---------- - config : Config - The configuration object containing the supernova and plasma settings. - geometry : Geometry - The geometry object containing the inner radius information. - - Returns - ------- - packet_source : BlackBodySimpleSource - The packet source object based on the configuration and geometry. - """ - if config.montecarlo.enable_full_relativity: - packet_source = BlackBodySimpleSourceRelativistic( - base_seed=config.montecarlo.seed, - time_explosion=config.supernova.time_explosion, - legacy_mode_enabled=legacy_mode_enabled, - ) - else: - packet_source = BlackBodySimpleSource( - base_seed=config.montecarlo.seed, - legacy_mode_enabled=legacy_mode_enabled, - ) - - return initialize_packet_source( - config, geometry, packet_source, legacy_mode_enabled - ) - - -def parse_csvy_radiation_field_state( - config, csvy_model_config, csvy_model_data, geometry, packet_source -): - t_radiative = None - dilution_factor = None - - if hasattr(csvy_model_data, "columns") and ( - "t_rad" in csvy_model_data.columns - ): - t_rad_field_index = [ - field["name"] for field in csvy_model_config.datatype.fields - ].index("t_rad") - t_rad_unit = u.Unit( - csvy_model_config.datatype.fields[t_rad_field_index]["unit"] - ) - t_radiative = csvy_model_data["t_rad"].iloc[1:].values * t_rad_unit - - elif config.plasma.initial_t_rad > 0 * u.K: - t_radiative = ( - np.ones(geometry.no_of_shells) * config.plasma.initial_t_rad - ) - else: - t_radiative = calculate_t_radiative_from_t_inner( - geometry, packet_source - ) - - if np.any(t_radiative < 1000 * u.K): - logging.critical( - "Radiative temperature is too low in some of the shells, temperatures below 1000K " - f"(e.g., T_rad = {t_radiative[np.argmin(t_radiative)]} in shell {np.argmin(t_radiative)} in your model) " - "are not accurately handled by TARDIS.", - ) - - if hasattr(csvy_model_data, "columns") and ( - "dilution_factor" in csvy_model_data.columns - ): - dilution_factor = csvy_model_data["dilution_factor"].iloc[1:].values - else: - dilution_factor = calculate_geometric_dilution_factor(geometry) - - return DiluteBlackBodyRadiationFieldState( - t_radiative, dilution_factor, geometry - ) - - -def calculate_t_radiative_from_t_inner(geometry, packet_source): - """ - Calculates the radiative temperature based on the inner temperature and the geometry of the system. - - Parameters - ---------- - geometry : Geometry - The geometry object. - packet_source : PacketSource - The packet source object. - - Returns - ------- - Quantity - The calculated radiative temperature. - """ - lambda_wien_inner = const.b_wien / packet_source.temperature - t_radiative = const.b_wien / ( - lambda_wien_inner - * (1 + (geometry.v_middle - geometry.v_inner_boundary) / const.c) - ) - return t_radiative - - -def calculate_geometric_dilution_factor(geometry): - return 0.5 * ( - 1 - - np.sqrt( - 1 - - ( - geometry.r_inner[geometry.v_inner_boundary_index] ** 2 - / geometry.r_middle**2 - ) - .to(1) - .value - ) - ) diff --git a/tardis/model/tests/test_csvy_model.py b/tardis/model/tests/test_csvy_model.py index c82c23b8e26..7463a475628 100644 --- a/tardis/model/tests/test_csvy_model.py +++ b/tardis/model/tests/test_csvy_model.py @@ -178,7 +178,7 @@ def test_read_csvy_abundances( # to be decayed after being loaded into SimulationState. # Now the abundances are decayed before # assert model_isotopes.shape == reference_input_isotopes_shape - # Same applies to the comparison - we are summing up the mass_fraction to compare pre/post decay + # Same applies to the comparison - we are summing up the mass_fractions to compare pre/post decay npt.assert_array_almost_equal( reference_input_isotopes.to_numpy()[0], model_isotopes.sum(axis=0).to_numpy(), diff --git a/tardis/simulation/base.py b/tardis/simulation/base.py index 219e20f363f..7efa377a314 100644 --- a/tardis/simulation/base.py +++ b/tardis/simulation/base.py @@ -12,12 +12,14 @@ 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.util import HDFWriterMixin from tardis.model import SimulationState -from tardis.model.parse_input import initialize_packet_source -from tardis.transport.montecarlo.base import MonteCarloTransportSolver from tardis.plasma.standard_plasmas import assemble_plasma from tardis.simulation.convergence import ConvergenceSolver +from tardis.transport.montecarlo.base import MonteCarloTransportSolver from tardis.util.base import is_notebook from tardis.visualization import ConvergencePlots @@ -677,12 +679,10 @@ def from_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( - config, - simulation_state.geometry, - packet_source, - legacy_mode_enabled, + packet_source, config, simulation_state.geometry ) if "plasma" in kwargs: plasma = kwargs["plasma"] diff --git a/tardis/visualization/widgets/custom_abundance.py b/tardis/visualization/widgets/custom_abundance.py index 82a986cf41f..490aaa8c7ac 100644 --- a/tardis/visualization/widgets/custom_abundance.py +++ b/tardis/visualization/widgets/custom_abundance.py @@ -11,7 +11,7 @@ from pathlib import Path import tardis -from tardis.io.model.readers.generic_readers import read_uniform_abundances +from tardis.io.model.readers.generic_readers import read_uniform_mass_fractions from tardis.util.base import ( quantity_linspace, is_valid_nuclide_or_elem, @@ -27,7 +27,7 @@ from tardis.io.configuration.config_validator import validate_dict from tardis.io.model.readers.csvy import load_csvy from tardis.io.model.readers.csvy import ( - parse_csv_abundances, + parse_csv_mass_fractions, ) from tardis.util.base import atomic_number2element_symbol, quantity_linspace from tardis.visualization.tools.convergence_plot import transition_colors @@ -156,11 +156,11 @@ def from_csvy(cls, fpath): if hasattr(csvy_model_config, "abundance"): abundances_section = csvy_model_config.abundance - abundance, isotope_abundance = read_uniform_abundances( + abundance, isotope_abundance = read_uniform_mass_fractions( abundances_section, no_of_shells ) else: - _, abundance, isotope_abundance = parse_csv_abundances( + _, abundance, isotope_abundance = parse_csv_mass_fractions( csvy_model_data ) abundance = abundance.loc[:, 1:]