Skip to content

Commit

Permalink
Restructure/transport solver (#2476)
Browse files Browse the repository at this point in the history 
* restructure of geometry

* add radial1d boundary logic

* black format

* several fixes

* fix epsilon

* add testing of boundaries

* change the r_inner_active

* first integration with `from_config` working

* hunting down density indexing bug

* all model tests (without csvy) pass

* more fixes

* fix of model to simulation_state

* fix inner boundary packet error

* fix some leftovers

* final fix for csvy

* blackify

* restructure to readers and remove some leftover code

* further cleanup

* first start of the restructure

* add comment about removing quantitiness

* add velocity check

* add new abundance functions

* remove default units

* add new matter module

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

* mid restructure

* slow progress on abundance refactor

* include effective_element_masses

* further updates

* fix of csvy readers

* last fixes to custom abundance widget

* restructuring the radiation field

* some cleanup

* cleanup

* further cleanup

* clearnup

* changes

* remove matter base.py

* removing matter

* fixed last tests

* first commit for regression_data

* add several parsers

* add composition

* remove isotopemassfractions

* some small fixes

* fixed w, t_rad to dilution_factor, t_radiative

* add setters for w, t_rad

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

* add the from_config atom_data

* fixing model.ipynb

* fix some more tests

* start of montecarlo restructure

* large scale changes

* further disentangling state and solver

* fix tests

* fix some of the issues wdigets

* assert mass fractions positive.
Change IsotopeAbundances to Massfractions

* update estimators

* add packet collection

* move spectrum to mct_state

* add radiation_field_state

* Refactor packet initialization in
MontecarloTransportSolver

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

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

* Fix Integrator

* Fix SDEC plot for MCTState

* Fix variable names in transport module

* Refactor the to_hdf code.

* Fix isotope mass fraction typo and t_radiatve duplicate definition

* Document remove functionality for effective element masses

* Fix variable name in
convert_to_nuclide_mass_fraction function

* Refactor test names and variable names in
test_base.py

* Refactor gamma ray simulation setup and test
functions

* fix gamma ray tests

* Add model_isotope_time_0 property to abundances
schema

* Add test for montecarlo main loop with vpacket
tracking

* Refactor config_reader and config_validator
modules

* Fix model_isotope_time_0 property in
model_definitions.yml

* Refactor simulation module to use pathlib

* remove bad comment and refactor with ruff

* Fix variable name in test_simulation_state_mass()
function

* Fix unit conversion in calculate_cell_masses
function

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

* black formatting

* Refactor code to use composition variable in
test_gamma_ray_transport.py

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

* fix for documentation not builiding

* add fix for model_isotope_time_0

* fix grid test

* restructure the grid

* final change

* add stuff

* add vpacket_log

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

* Update damping constants and variable names
in doc notebook

* fix w in several places in tardis

* slowly fixing radiation_field changes

* last fix

* change hdf from model to simulation_state

* fixing model->simulation-state

* Fix model_reader.py imports and formatting issues

* Refactor variable names for clarity and
consistency and being correct

* Update variable names in model.ipynb

* Update variable name for average temperature
calculation

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

* Fix transport dictionary comparison

* Implement regression data configurations

* Refactor snapshot fixtures and add regression data
handling

* Update to new testing framework

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

* Make regression-data optional

* Fix regression data synchronization in plasma
tests

* Refactor variable names in model reader and
simulation base

* Refactor store_model_to_hdf to
store_simulation_state_to_hdf

* Refactor HDF file handling and model reader
imports

* add docstr

* Remove syrupy_extensions.py

* Refactor code and remove unused imports and
deleted file

* Fix test_base and test_simulation

* setting up regression with github actions

* Update LFS cache IDs for refdata and regression
data

* Fix test_montecarlo_radial1d function

* Update sync_hdf_store method in RegressionData class

* Refactor test_base.py for improved accuracy and precision

* update docstrings

* Refactor transport state access in test_base.py

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

* Update demo.hdf file

Delete demo.hdf file

Update file paths in code

* Fix negative values in last_interaction_types

* Update URL for atomic data repository

* Refactor code for improved readability and maintainability

* Refactor code to remove unnecessary array conversion

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

* fixed comments

* fix some of josh's requests

* Fix callback example doc and parse_packet_source missing import

* Fixes custom packet source through init function

* Fixes docs build

* black format on parse_input

* Fixes tests for packet source

Checks if the temperature is a quantity or not before trying .value

---------

Co-authored-by: Andrew Fullard <andrewgfullard@gmail.com>
  • Loading branch information
@wkerzendorf @andrewfullard
wkerzendorf and andrewfullard authored Jan 17, 2024
1 parent 5608438 commit 3dfcb52
Show file tree
Hide file tree
Showing 38 changed files with 1,116 additions and 1,634 deletions.
2 changes: 1 addition & 1 deletion docs/io/optional/callback_example.ipynb
View file
Original file line number Diff line number Diff line change
Expand Up @@ -78,7 +78,7 @@
"source": [
"def append_num_emitted_to_list(sim, lst):\n",
" if sim.iterations_executed < sim.iterations:\n",
" num_emitted_packets = len(sim.transport.emitted_packet_nu)\n",
" num_emitted_packets = len(sim.transport.transport_state.emitted_packet_nu)\n",
" lst.append(num_emitted_packets)"
]
},
Expand Down
36 changes: 24 additions & 12 deletions docs/io/optional/custom_source.ipynb
View file
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,9 @@
"from tardis import constants as const\n",
"from astropy import units as u\n",
"from tardis.montecarlo.packet_source import BlackBodySimpleSource\n",
"from tardis.montecarlo.montecarlo_numba.packet_collections import (\n",
" PacketCollection,\n",
")\n",
"from tardis import run_tardis\n",
"import matplotlib.pyplot as plt\n",
"from tardis.io.atom_data import download_atom_data"
Expand Down Expand Up @@ -87,7 +90,7 @@
" self.truncation_wavelength = truncation_wavelength\n",
" super().__init__(**kwargs)\n",
"\n",
" def create_packets(self, no_of_packets, drawing_sample_size=None):\n",
" def create_packets(self, no_of_packets, drawing_sample_size=None, seed_offset=0, *args, **kwargs):\n",
" \"\"\"\n",
" Packet source that generates a truncated Blackbody source.\n",
"\n",
Expand All @@ -108,12 +111,17 @@
" Packet energies\n",
" \"\"\"\n",
"\n",
" self._reseed(self.base_seed + seed_offset)\n",
" packet_seeds = self.rng.choice(\n",
" self.MAX_SEED_VAL, no_of_packets, replace=True\n",
" )\n",
"\n",
" # Makes uniform array of packet radii from blackbody source\n",
" radii = self.create_packet_radii(no_of_packets)\n",
" radii = self.create_packet_radii(no_of_packets, *args, **kwargs)\n",
"\n",
" # Use mus and energies from normal blackbody source.\n",
" mus = self.create_packet_mus(no_of_packets)\n",
" energies = self.create_packet_energies(no_of_packets)\n",
" mus = self.create_packet_mus(no_of_packets, *args, **kwargs)\n",
" energies = self.create_packet_energies(no_of_packets, *args, **kwargs)\n",
"\n",
" # If not specified, draw 2 times as many packets and reject any beyond no_of_packets.\n",
" if drawing_sample_size is None:\n",
Expand All @@ -128,18 +136,22 @@
"\n",
" # Draw nus from blackbody distribution and reject based on truncation_frequency.\n",
" # If more nus.shape[0] > no_of_packets use only the first no_of_packets.\n",
" nus = self.create_packet_nus(drawing_sample_size)\n",
" nus = self.create_packet_nus(drawing_sample_size, *args, **kwargs)\n",
" nus = nus[nus < truncation_frequency][:no_of_packets]\n",
"\n",
" # Only required if the truncation wavelength is too big compared to the maximum\n",
" # of the blackbody distribution. Keep sampling until nus.shape[0] > no_of_packets.\n",
" while nus.shape[0] < no_of_packets:\n",
" additional_nus = self.create_packet_nus(drawing_sample_size)\n",
" additional_nus = self.create_packet_nus(drawing_sample_size, *args, **kwargs)\n",
" mask = additional_nus < truncation_frequency\n",
" additional_nus = additional_nus[mask][:no_of_packets]\n",
" nus = np.hstack([nus, additional_nus])[:no_of_packets]\n",
"\n",
" return radii, nus, mus, energies"
" radiation_field_luminosity = (\n",
" self.calculate_radfield_luminosity().to(u.erg / u.s).value\n",
" )\n",
"\n",
" return PacketCollection(radii, nus, mus, energies, packet_seeds, radiation_field_luminosity)"
]
},
{
Expand Down Expand Up @@ -179,11 +191,11 @@
"outputs": [],
"source": [
"%matplotlib inline\n",
"plt.plot(mdl.transport.spectrum_virtual.wavelength,\n",
" mdl.transport.spectrum_virtual.luminosity_density_lambda,\n",
"plt.plot(mdl.transport.transport_state.spectrum_virtual.wavelength,\n",
" mdl.transport.transport_state.spectrum_virtual.luminosity_density_lambda,\n",
" color='red', label='truncated blackbody (custom packet source)')\n",
"plt.plot(mdl_norm.transport.spectrum_virtual.wavelength,\n",
" mdl_norm.transport.spectrum_virtual.luminosity_density_lambda,\n",
"plt.plot(mdl_norm.transport.transport_state.spectrum_virtual.wavelength,\n",
" mdl_norm.transport.transport_state.spectrum_virtual.luminosity_density_lambda,\n",
" color='blue', label='normal blackbody (default packet source)')\n",
"plt.xlabel('$\\lambda [\\AA]$')\n",
"plt.ylabel('$L_\\lambda$ [erg/s/$\\AA$]')\n",
Expand All @@ -208,7 +220,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.7.10"
"version": "3.11.5"
}
},
"nbformat": 4,
Expand Down
32 changes: 16 additions & 16 deletions docs/physics/montecarlo/initialization.ipynb
View file
Original file line number Diff line number Diff line change
Expand Up @@ -79,6 +79,9 @@
"source": [
"import numpy as np\n",
"from tardis.montecarlo.packet_source import BlackBodySimpleSource\n",
"from tardis.montecarlo.montecarlo_numba.packet_collections import (\n",
" PacketCollection,\n",
")\n",
"from astropy import units as u\n",
"from tardis import constants as const\n",
"import matplotlib.pyplot as plt"
Expand Down Expand Up @@ -164,24 +167,21 @@
"# We define our packet source\n",
"packet_source = BlackBodySimpleSource(base_seed=base_seed)\n",
"\n",
"# Create separate packet seeds to sample each packet's montecarlo trajectory\n",
"packet_seeds = packet_source.create_packet_seeds(n_packets, seed_offset)\n",
"\n",
"# Set radii and temperature from model\n",
"packet_source.radius = r_boundary_inner\n",
"packet_source.temperature = temperature_inner.value\n",
"packet_source.temperature = temperature_inner\n",
"\n",
"# Create packets\n",
"radii, nus, mus, energies = packet_source.create_packets(n_packets)\n",
"packet_collection = packet_source.create_packets(n_packets)\n",
"\n",
"# Sets the energies in units of ergs\n",
"energies *= u.erg\n",
"packet_collection.initial_energies *= u.erg\n",
"\n",
"# Sets the frequencies in units of Hz\n",
"nus *= u.Hz\n",
"packet_collection.initial_nus *= u.Hz\n",
"\n",
"print(\"Energies:\", energies)\n",
"print(\"Radii:\", radii)"
"print(\"Energies:\", packet_collection.initial_energies)\n",
"print(\"Radii:\", packet_collection.initial_radii)"
]
},
{
Expand All @@ -205,7 +205,7 @@
"print(\"Time of simulation:\", t_simulation)\n",
"\n",
"# Array of luminosity contribution by each packet\n",
"lumin_per_packet = energies / t_simulation\n",
"lumin_per_packet = packet_collection.initial_energies / t_simulation\n",
"print(\"Luminosity per packet:\", lumin_per_packet)"
]
},
Expand Down Expand Up @@ -263,15 +263,15 @@
"source": [
"# We set important quantites for making our histogram\n",
"bins = 200\n",
"nus_planck = np.linspace(min(nus), max(nus), bins)\n",
"nus_planck = np.linspace(min(packet_collection.initial_nus), max(packet_collection.initial_nus), bins).value\n",
"bin_width = nus_planck[1] - nus_planck[0]\n",
"\n",
"# In the histogram plot below, the weights argument is used\n",
"# to make sure our plotted spectrum has the correct y-axis scale\n",
"plt.hist(nus.value, bins=bins, weights=lumin_per_packet / bin_width)\n",
"plt.hist(packet_collection.initial_nus.value, bins=bins, weights=lumin_per_packet / bin_width)\n",
"\n",
"# We plot the planck function for comparison\n",
"plt.plot(nus_planck, planck_function(nus_planck))\n",
"plt.plot(nus_planck * u.Hz, planck_function(nus_planck * u.Hz))\n",
"\n",
"plt.xlabel(\"Frequency (Hz)\")\n",
"plt.ylabel(\"Luminosity density w.r.t. frequency (erg/s/Hz)\")\n",
Expand All @@ -296,7 +296,7 @@
"source": [
"x = np.linspace(0, 1, 1000)\n",
"\n",
"plt.hist(mus, bins=bins, density=True)\n",
"plt.hist(packet_collection.initial_mus, bins=bins, density=True)\n",
"plt.plot(x, 2 * x)\n",
"plt.xlabel(\"Propagation direction\")\n",
"plt.ylabel(\"Probability density\")\n",
Expand All @@ -312,7 +312,7 @@
"source": [
"thetas = np.linspace(0, np.pi / 2, 1000)\n",
"\n",
"plt.hist(np.arccos(mus), bins=bins, density=True)\n",
"plt.hist(np.arccos(packet_collection.initial_mus), bins=bins, density=True)\n",
"plt.plot(thetas, np.sin(2 * thetas))\n",
"plt.xlabel(\"Angle with normal (rad)\")\n",
"plt.ylabel(\"Probability density\")\n",
Expand Down Expand Up @@ -348,7 +348,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.8.13"
"version": "3.11.5"
}
},
"nbformat": 4,
Expand Down
12 changes: 6 additions & 6 deletions docs/physics/spectrum/basic.ipynb
View file
Original file line number Diff line number Diff line change
Expand Up @@ -47,7 +47,7 @@
"source": [
"from tardis.io.configuration.config_reader import Configuration\n",
"from tardis.simulation import Simulation\n",
"from tardis.montecarlo import TARDISSpectrum\n",
"from tardis.montecarlo.spectrum import TARDISSpectrum\n",
"from tardis.io.atom_data.util import download_atom_data\n",
"from astropy import units as u\n",
"import numpy as np\n",
Expand Down Expand Up @@ -118,7 +118,7 @@
"metadata": {},
"outputs": [],
"source": [
"nus = sim.transport.output_nu\n",
"nus = sim.transport.transport_state.output_nu\n",
"nus"
]
},
Expand All @@ -129,7 +129,7 @@
"metadata": {},
"outputs": [],
"source": [
"energies = sim.transport.output_energy\n",
"energies = sim.transport.transport_state.output_energy\n",
"energies"
]
},
Expand Down Expand Up @@ -183,7 +183,7 @@
"metadata": {},
"outputs": [],
"source": [
"luminosities = energies / sim.transport.time_of_simulation\n",
"luminosities = energies / sim.transport.transport_state.time_of_simulation\n",
"luminosities"
]
},
Expand All @@ -203,7 +203,7 @@
"metadata": {},
"outputs": [],
"source": [
"emitted_mask = sim.transport.emitted_packet_mask\n",
"emitted_mask = sim.transport.transport_state.emitted_packet_mask\n",
"emitted_mask"
]
},
Expand Down Expand Up @@ -460,7 +460,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.7.12"
"version": "3.11.5"
}
},
"nbformat": 4,
Expand Down
18 changes: 13 additions & 5 deletions docs/physics/update_and_conv/update_and_conv.ipynb
View file
Original file line number Diff line number Diff line change
Expand Up @@ -311,7 +311,7 @@
},
"outputs": [],
"source": [
"j_estimator = transport.j_estimator * (u.erg * u.cm) \n",
"j_estimator = transport.transport_state.estimators.j_estimator * (u.erg * u.cm) \n",
"j_estimator"
]
},
Expand All @@ -324,7 +324,7 @@
},
"outputs": [],
"source": [
"nu_bar_estimator = transport.nu_bar_estimator * (u.erg * u.cm * u.Hz)\n",
"nu_bar_estimator = transport.transport_state.estimators.nu_bar_estimator * (u.erg * u.cm * u.Hz)\n",
"nu_bar_estimator"
]
},
Expand All @@ -345,7 +345,7 @@
"outputs": [],
"source": [
"V = simulation_state.volume\n",
"Delta_t = transport.calculate_time_of_simulation(simulation_state)\n",
"Delta_t = transport.transport_state.time_of_simulation\n",
"prefactor = 1 / (4 * np.pi * V * Delta_t)\n",
"J = prefactor * j_estimator\n",
"J"
Expand Down Expand Up @@ -462,7 +462,7 @@
"#nu_lower = tardis_config.supernova.luminosity_wavelength_end.to(u.Hz, u.spectral)\n",
"#nu_upper = tardis_config.supernova.luminosity_wavelength_start.to(u.Hz, u.spectral)\n",
"\n",
"L_output = transport.calculate_emitted_luminosity(0,np.inf)\n",
"L_output = transport.transport_state.calculate_emitted_luminosity(0,np.inf)\n",
"L_output"
]
},
Expand Down Expand Up @@ -570,6 +570,14 @@
"source": [
"plasma.electron_densities"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "bd488259",
"metadata": {},
"outputs": [],
"source": []
}
],
"metadata": {
Expand All @@ -588,7 +596,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.8.13"
"version": "3.11.5"
},
"vscode": {
"interpreter": {
Expand Down
8 changes: 4 additions & 4 deletions docs/quickstart.ipynb
View file
Original file line number Diff line number Diff line change
Expand Up @@ -145,9 +145,9 @@
"metadata": {},
"outputs": [],
"source": [
"spectrum = sim.transport.spectrum\n",
"spectrum_virtual = sim.transport.spectrum_virtual\n",
"spectrum_integrated = sim.transport.spectrum_integrated"
"spectrum = sim.transport.transport_state.spectrum\n",
"spectrum_virtual = sim.transport.transport_state.spectrum_virtual\n",
"spectrum_integrated = sim.transport.transport_state.spectrum_integrated"
]
},
{
Expand Down Expand Up @@ -191,7 +191,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.8.17"
"version": "3.11.5"
}
},
"nbformat": 4,
Expand Down
14 changes: 7 additions & 7 deletions tardis/analysis.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,14 +15,14 @@

class LastLineInteraction(object):
@classmethod
def from_model(cls, model, packet_filter_mode="packet_out_nu"):
def from_simulation(cls, simulation, packet_filter_mode="packet_out_nu"):
return cls(
model.transport.last_line_interaction_in_id,
model.transport.last_line_interaction_out_id,
model.transport.last_line_interaction_shell_id,
model.transport.output_nu,
model.transport.last_interaction_in_nu,
model.plasma.atomic_data.lines,
simulation.transport.last_line_interaction_in_id,
simulation.transport.last_line_interaction_out_id,
simulation.transport.last_line_interaction_shell_id,
simulation.transport.transport_state.packet_collection.output_nus,
simulation.transport.last_interaction_in_nu,
simulation.plasma.atomic_data.lines,
packet_filter_mode,
)

Expand Down
9 changes: 4 additions & 5 deletions tardis/gui/widgets.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -406,7 +406,6 @@ def match_dicts(self, dict1, dict2): # dict1<=dict2

elif isinstance(dict2[key], list):
if isinstance(dict2[key][1], list):

# options = dict2[key][1] #This is passed by reference.
# So copy the list manually.
options = [
Expand Down Expand Up @@ -937,17 +936,17 @@ def __init__(self, parent, wavelength_start, wavelength_end, tablecreator):
f"Line Interaction: {wavelength_start:.2f} - {wavelength_end:.2f} (A) "
)
self.layout = QtWidgets.QVBoxLayout()
packet_nu_line_interaction = analysis.LastLineInteraction.from_model(
self.parent.model
packet_nu_line_interaction = (
analysis.LastLineInteraction.from_simulation(self.parent.model)
)
packet_nu_line_interaction.packet_filter_mode = "packet_nu"
packet_nu_line_interaction.wavelength_start = (
wavelength_start * u.angstrom
)
packet_nu_line_interaction.wavelength_end = wavelength_end * u.angstrom

line_in_nu_line_interaction = analysis.LastLineInteraction.from_model(
self.parent.model
line_in_nu_line_interaction = (
analysis.LastLineInteraction.from_simulation(self.parent.model)
)
line_in_nu_line_interaction.packet_filter_mode = "line_in_nu"
line_in_nu_line_interaction.wavelength_start = (
Expand Down
Loading

0 comments on commit 3dfcb52

Please sign in to comment.