diff --git a/docs/physics/setup/model.ipynb b/docs/physics/setup/model.ipynb index 426d8010ebd..e9ca9fffda8 100644 --- a/docs/physics/setup/model.ipynb +++ b/docs/physics/setup/model.ipynb @@ -81,7 +81,7 @@ "id": "cee054e9", "metadata": {}, "source": [ - "In the cell below, we set up a model. We use the [specific structure](../../io/configuration/components/models/index.rst#specific-structure) where we supply $t_\\mathrm{explosion}$, the velocity of the inner and outer boundaries of the supernova (labeled `start` and `stop`), and the number of shells (labeled `num`). The shells are then evenly spaced between the inner and outer boundaries of the supernova. The time after the explosion, the inner and outer velocities, and the number of shells can be varied to get different shell structures. The `SimulationState` object stores information about the model in the following attributes: `velocity` shows the velocity of the shell boundaries, `v_inner` shows the velocities of the inner boundaries of each shell, `v_outer` shows the velocity of the outer boundaries of each shell, and `v_middle` shows the velocity of the middle of each shell. Similarly, `radius`, `r_inner`, `r_outer`, and `r_middle` show the radii of each shell boundary, the inner boundaries, the outer boundaries, and the middles of each shell, respectively. `v_boundary_inner` shows the velocity of the inner boundary of the computational domain, and `v_boundary_outer` shows the velocity of the outer boundary of the computational domain. Finally, `volume` shows the volume of each shell, calculated via the formula of the volume of a spherical shell: $V=\\frac{4}{3}\\pi (r_\\mathrm{outer}^3-r_\\mathrm{inner}^3)$." + "In the cell below, we set up a model. We use the [specific structure](../../io/configuration/components/models/index.rst#specific-structure) where we supply $t_\\mathrm{explosion}$, the velocity of the inner and outer boundaries of the supernova (labeled `start` and `stop`), and the number of shells (labeled `num`). The shells are then evenly spaced between the inner and outer boundaries of the supernova. The time after the explosion, the inner and outer velocities, and the number of shells can be varied to get different shell structures. The `SimulationState` object stores information about the model in the following attributes: `velocity` shows the velocity of the shell boundaries, `v_inner` shows the velocities of the inner boundaries of each shell, `v_outer` shows the velocity of the outer boundaries of each shell, and `v_middle` shows the velocity of the middle of each shell. Similarly, `radius`, `r_inner`, `r_outer`, and `r_middle` show the radii of each shell boundary, the inner boundaries, the outer boundaries, and the middles of each shell, respectively. `v_inner_boundary` shows the velocity of the inner boundary of the computational domain, and `v_outer_boundary` shows the velocity of the outer boundary of the computational domain. Finally, `volume` shows the volume of each shell, calculated via the formula of the volume of a spherical shell: $V=\\frac{4}{3}\\pi (r_\\mathrm{outer}^3-r_\\mathrm{inner}^3)$." ] }, { @@ -111,8 +111,8 @@ "print('v_inner:\\n', shell_simulation_state.v_inner)\n", "print('v_outer:\\n', shell_simulation_state.v_outer)\n", "print('v_middle:\\n', shell_simulation_state.v_middle)\n", - "print('v_boundary_inner:\\n', shell_simulation_state.v_boundary_inner)\n", - "print('v_boundary_outer:\\n', shell_simulation_state.v_boundary_outer)\n", + "print('v_inner_boundary:\\n', shell_simulation_state.v_inner_boundary)\n", + "print('v_outer_boundary:\\n', shell_simulation_state.v_outer_boundary)\n", "print('radius:\\n', shell_simulation_state.radius)\n", "print('r_inner:\\n', shell_simulation_state.r_inner)\n", "print('r_outer:\\n', shell_simulation_state.r_outer)\n", @@ -125,7 +125,7 @@ "id": "1ee56110", "metadata": {}, "source": [ - "Notice that `radius = velocity*time_explosion`, and similarly for `r_inner`, `r_outer`, and `r_middle`. You can get the radius of the photosphere via `v_boundary_inner*time_explosion` and outer edge of the supernova via `v_boundary_outer*time_explosion`.\n", + "Notice that `radius = velocity*time_explosion`, and similarly for `r_inner`, `r_outer`, and `r_middle`. You can get the radius of the photosphere via `v_inner_boundary*time_explosion` and outer edge of the supernova via `v_boundary_outer*time_explosion`.\n", "\n", "
\n", " \n", diff --git a/docs/workflows/v_inner_solver_workflow.ipynb b/docs/workflows/v_inner_solver_workflow.ipynb new file mode 100644 index 00000000000..f52d347223f --- /dev/null +++ b/docs/workflows/v_inner_solver_workflow.ipynb @@ -0,0 +1,120 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from tardis.workflows.v_inner_solver import InnerVelocitySolverWorkflow\n", + "from tardis.io.configuration.config_reader import Configuration" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "config = Configuration.from_yaml('../tardis_example.yml')" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "scrolled": true + }, + "outputs": [], + "source": [ + "from astropy import units as u\n", + "\n", + "config.montecarlo.convergence_strategy['v_inner_boundary'] = {\n", + " 'damping_constant' : 0.5,\n", + " 'threshold' : 0.01,\n", + " 'type' : 'damped'\n", + " }\n", + "\n", + "config.montecarlo.convergence_strategy.stop_if_converged = True\n", + "config.model.structure.velocity.start = 5000 * u.km/u.s # Larger window over which to search\n", + "config.model.structure.velocity.num = 50 # Increase number of shells\n", + "\n", + "workflow = InnerVelocitySolverWorkflow(\n", + " config, tau=2.0/3,\n", + " mean_optical_depth=\"rosseland\"\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "workflow.run()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import matplotlib.pyplot as plt" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "spectrum = workflow.spectrum_solver.spectrum_real_packets\n", + "spectrum_virtual = workflow.spectrum_solver.spectrum_virtual_packets\n", + "spectrum_integrated = workflow.spectrum_solver.spectrum_integrated" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "%matplotlib inline\n", + "plt.figure(figsize=(10, 6.5))\n", + "\n", + "spectrum.plot(label=\"Normal packets\")\n", + "spectrum_virtual.plot(label=\"Virtual packets\")\n", + "spectrum_integrated.plot(label='Formal integral')\n", + "\n", + "plt.xlim(500, 9000)\n", + "plt.title(\"TARDIS example model spectrum\")\n", + "plt.xlabel(r\"Wavelength [$\\AA$]\")\n", + "plt.ylabel(r\"Luminosity density [erg/s/$\\AA$]\")\n", + "plt.legend()\n", + "plt.show()" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "tardis", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.4" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/tardis/io/configuration/config_reader.py b/tardis/io/configuration/config_reader.py index c4efb949031..bfed003e583 100644 --- a/tardis/io/configuration/config_reader.py +++ b/tardis/io/configuration/config_reader.py @@ -431,7 +431,7 @@ def parse_convergence_section(convergence_section_dict): """ convergence_parameters = ["damping_constant", "threshold", "type"] - for convergence_variable in ["t_inner", "t_rad", "w"]: + for convergence_variable in ["t_inner", "t_rad", "w", "v_inner_boundary"]: if convergence_variable not in convergence_section_dict: convergence_section_dict[convergence_variable] = {} convergence_variable_section = convergence_section_dict[ diff --git a/tardis/io/configuration/schemas/montecarlo_definitions.yml b/tardis/io/configuration/schemas/montecarlo_definitions.yml index 1adf22599ae..b79a4b7da00 100644 --- a/tardis/io/configuration/schemas/montecarlo_definitions.yml +++ b/tardis/io/configuration/schemas/montecarlo_definitions.yml @@ -5,7 +5,6 @@ definitions: title: 'Damped Convergence Strategy' type: object additionalProperties: false - properties: properties: type: enum: @@ -59,6 +58,24 @@ definitions: type: string default: 'damped' description: THIS IS A DUMMY VARIABLE DO NOT USE + v_inner_boundary: + type: object + additionalProperties: false + properties: + damping_constant: + type: number + default: 0.0 + description: damping constant + minimum: 0 + threshold: + type: number + 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 diff --git a/tardis/io/model/parse_geometry_configuration.py b/tardis/io/model/parse_geometry_configuration.py index c517a4a7ea4..e2f0c8fd20f 100644 --- a/tardis/io/model/parse_geometry_configuration.py +++ b/tardis/io/model/parse_geometry_configuration.py @@ -135,17 +135,17 @@ def parse_geometry_from_csvy( """ if hasattr(config, "model"): if hasattr(config.model, "v_inner_boundary"): - v_boundary_inner = config.model.v_inner_boundary + v_inner_boundary = config.model.v_inner_boundary else: - v_boundary_inner = None + v_inner_boundary = None if hasattr(config.model, "v_outer_boundary"): - v_boundary_outer = config.model.v_outer_boundary + v_outer_boundary = config.model.v_outer_boundary else: - v_boundary_outer = None + v_outer_boundary = None else: - v_boundary_inner = None - v_boundary_outer = None + v_inner_boundary = None + v_outer_boundary = None if hasattr(csvy_model_config, "velocity"): velocity = quantity_linspace( @@ -166,8 +166,8 @@ def parse_geometry_from_csvy( geometry = HomologousRadial1DGeometry( velocity[:-1], # v_inner velocity[1:], # v_outer - v_inner_boundary=v_boundary_inner, - v_outer_boundary=v_boundary_outer, + v_inner_boundary=v_inner_boundary, + v_outer_boundary=v_outer_boundary, time_explosion=time_explosion, ) return geometry diff --git a/tardis/model/base.py b/tardis/model/base.py index b3f3ae98851..960c47a9455 100644 --- a/tardis/model/base.py +++ b/tardis/model/base.py @@ -58,11 +58,11 @@ class SimulationState(HDFWriterMixin): dilution_factor : np.ndarray If None, the dilution_factor will be initialized with the geometric dilution factor. - v_boundary_inner : astropy.units.Quantity - v_boundary_outer : astropy.units.Quantity + v_inner_boundary : astropy.units.Quantity + v_outer_boundary : astropy.units.Quantity raw_velocity : np.ndarray The complete array of the velocities, without being cut by - `v_boundary_inner` and `v_boundary_outer` + `v_inner_boundary` and `v_outer_boundary` electron_densities : astropy.units.quantity.Quantity Attributes @@ -81,8 +81,8 @@ class SimulationState(HDFWriterMixin): density : astropy.units.quantity.Quantity volume : astropy.units.quantity.Quantity no_of_shells : int - The number of shells as formed by `v_boundary_inner` and - `v_boundary_outer` + The number of shells as formed by `v_inner_boundary` and + `v_outer_boundary` no_of_raw_shells : int """ @@ -206,11 +206,11 @@ def radius(self): return self.time_explosion * self.velocity @property - def v_boundary_inner(self): + def v_inner_boundary(self): return self.geometry.v_inner_boundary @property - def v_boundary_outer(self): + def v_outer_boundary(self): return self.geometry.v_outer_boundary @property diff --git a/tardis/spectrum/formal_integral.py b/tardis/spectrum/formal_integral.py index fdbef19ebf8..02955947c71 100644 --- a/tardis/spectrum/formal_integral.py +++ b/tardis/spectrum/formal_integral.py @@ -416,9 +416,27 @@ def make_source_function(self): ------- Numpy array containing ( 1 - exp(-tau_ul) ) S_ul ordered by wavelength of the transition u -> l """ + simulation_state = self.simulation_state + # slice for the active shells + local_slice = slice( + simulation_state.geometry.v_inner_boundary_index, + simulation_state.geometry.v_outer_boundary_index, + ) + transport = self.transport montecarlo_transport_state = transport.transport_state + transition_probabilities = ( + self.original_plasma.transition_probabilities.iloc[:, local_slice] + ) + tau_sobolevs = self.original_plasma.tau_sobolevs.iloc[:, local_slice] + electron_densities = self.original_plasma.electron_densities.values[ + local_slice + ] + columns = range(simulation_state.no_of_shells) + + tau_sobolevs.columns = columns + transition_probabilities.columns = columns # macro_ref = self.atomic_data.macro_atom_references macro_ref = self.atomic_data.macro_atom_references @@ -431,9 +449,7 @@ def make_source_function(self): if transport.line_interaction_type == "macroatom": internal_jump_mask = (macro_data.transition_type >= 0).values ma_int_data = macro_data[internal_jump_mask] - internal = self.original_plasma.transition_probabilities[ - internal_jump_mask - ] + internal = transition_probabilities[internal_jump_mask] source_level_idx = ma_int_data.source_level_idx.values destination_level_idx = ma_int_data.destination_level_idx.values @@ -442,7 +458,7 @@ def make_source_function(self): montecarlo_transport_state.packet_collection.time_of_simulation * simulation_state.volume ) - exptau = 1 - np.exp(-self.original_plasma.tau_sobolevs) + exptau = 1 - np.exp(-tau_sobolevs) Edotlu = ( Edotlu_norm_factor * exptau @@ -475,14 +491,14 @@ def make_source_function(self): upper_level_index = self.atomic_data.lines.index.droplevel( "level_number_lower" ) - e_dot_lu = pd.DataFrame(Edotlu.values, index=upper_level_index) + e_dot_lu = pd.DataFrame( + Edotlu.values, index=upper_level_index, columns=columns + ) e_dot_u = e_dot_lu.groupby(level=[0, 1, 2]).sum() e_dot_u_src_idx = macro_ref.loc[e_dot_u.index].references_idx.values if transport.line_interaction_type == "macroatom": - C_frame = pd.DataFrame( - columns=np.arange(no_shells), index=macro_ref.index - ) + C_frame = pd.DataFrame(columns=columns, index=macro_ref.index) q_indices = (source_level_idx, destination_level_idx) for shell in range(no_shells): Q = sp.coo_matrix( @@ -505,7 +521,7 @@ def make_source_function(self): transitions_index = transitions.set_index( ["atomic_number", "ion_number", "source_level_number"] ).index.copy() - tmp = self.original_plasma.transition_probabilities[ + tmp = transition_probabilities[ (self.atomic_data.macro_atom_data.transition_type == -1).values ] q_ul = tmp.set_index(transitions_index) @@ -520,13 +536,15 @@ def make_source_function(self): att_S_ul = wave * (q_ul * e_dot_u) * t / (4 * np.pi) result = pd.DataFrame( - att_S_ul.values, index=transitions.transition_line_id.values + att_S_ul.values, + index=transitions.transition_line_id.values, + columns=columns, ) att_S_ul = result.loc[lines.index.values].values # Jredlu should already by in the correct order, i.e. by wavelength of # the transition l->u (similar to Jbluelu) - Jredlu = Jbluelu * np.exp(-self.original_plasma.tau_sobolevs) + att_S_ul + Jredlu = Jbluelu * np.exp(-tau_sobolevs) + att_S_ul if self.interpolate_shells > 0: ( att_S_ul, @@ -543,12 +561,8 @@ def make_source_function(self): transport.r_outer_i = ( montecarlo_transport_state.geometry_state.r_outer ) - transport.tau_sobolevs_integ = ( - self.original_plasma.tau_sobolevs.values - ) - transport.electron_densities_integ = ( - self.original_plasma.electron_densities.values - ) + transport.tau_sobolevs_integ = tau_sobolevs.values + transport.electron_densities_integ = electron_densities return att_S_ul, Jredlu, Jbluelu, e_dot_u @@ -575,7 +589,9 @@ def interpolate_integrator_quantities( transport.electron_densities_integ = interp1d( r_middle, - plasma.electron_densities, + plasma.electron_densities.iloc[ + self.simulation_state.geometry.v_inner_boundary_index : self.simulation_state.geometry.v_outer_boundary_index + ], fill_value="extrapolate", kind="nearest", )(r_middle_integ) @@ -583,7 +599,10 @@ def interpolate_integrator_quantities( # (as in the MC simulation) transport.tau_sobolevs_integ = interp1d( r_middle, - plasma.tau_sobolevs, + plasma.tau_sobolevs.iloc[ + :, + self.simulation_state.geometry.v_inner_boundary_index : self.simulation_state.geometry.v_outer_boundary_index, + ], fill_value="extrapolate", kind="nearest", )(r_middle_integ) diff --git a/tardis/workflows/simple_tardis_workflow.py b/tardis/workflows/simple_tardis_workflow.py index 1f5b3cbbf39..84d7225a5cc 100644 --- a/tardis/workflows/simple_tardis_workflow.py +++ b/tardis/workflows/simple_tardis_workflow.py @@ -256,6 +256,8 @@ def solve_simulation_state( "t_inner" ] + return next_values + def solve_plasma(self, estimated_radfield_properties): """Update the plasma solution with the new radiation field estimates diff --git a/tardis/workflows/util.py b/tardis/workflows/util.py new file mode 100644 index 00000000000..39b3c489c85 --- /dev/null +++ b/tardis/workflows/util.py @@ -0,0 +1,87 @@ +import numpy as np +from astropy import constants as const + + +def get_tau_integ(plasma, simulation_state, bin_size=10): + """Estimate the integrated mean optical depth at each velocity bin + + Parameters + ---------- + plasma : tardis.plasma.BasePlasma + The tardis legacy plasma + simulation_state : tardis.model.base.SimulationState + the current simulation state + bin_size : int, optional. Default : 10 + bin size for the aggregation of line opacities + + Returns + ------- + dict + rosseland : np.ndarray + Roassland Mean Optical Depth + planck : np.ndarray + Planck Mean Optical Depth + """ + index = plasma.atomic_data.lines.nu.index + taus = plasma.tau_sobolevs.loc[index] + freqs = plasma.atomic_data.lines.nu.values + order = np.argsort(freqs) + freqs = freqs[order] + taus = plasma.tau_sobolevs.values[order] + + extra = bin_size - len(freqs) % bin_size + extra_freqs = np.arange(extra + 1) + 1 + extra_taus = np.zeros((extra + 1, taus.shape[1])) + freqs = np.hstack((extra_freqs, freqs)) + taus = np.vstack((extra_taus, taus)) + + bins_low = freqs[:-bin_size:bin_size] + bins_high = freqs[bin_size::bin_size] + delta_nu = bins_high - bins_low + n_bins = len(delta_nu) + + taus = taus[1 : n_bins * bin_size + 1] + freqs = freqs[1 : n_bins * bin_size + 1] + + ct = simulation_state.time_explosion.cgs.value * const.c.cgs.value + t_rad = simulation_state.radiation_field_state.temperature.cgs.value + + h = const.h.cgs.value + c = const.c.cgs.value + kb = const.k_B.cgs.value + + def B(nu, T): + return 2 * h * nu**3 / c**2 / (np.exp(h * nu / (kb * T)) - 1) + + def U(nu, T): + return B(nu, T) ** 2 * (c / nu) ** 2 * (2 * kb * T**2) ** -1 + + kappa_exp = ( + (bins_low / delta_nu).reshape(-1, 1) + / ct + * (1 - np.exp(-taus.reshape(n_bins, bin_size, -1))).sum(axis=1) + ) + kappa_thom = plasma.electron_densities.values * const.sigma_T.cgs.value + Bdnu = B(bins_low.reshape(-1, 1), t_rad.reshape(1, -1)) * delta_nu.reshape( + -1, 1 + ) + kappa_planck = kappa_thom + (Bdnu * kappa_exp).sum(axis=0) / ( + Bdnu.sum(axis=0) + ) + + udnu = U(bins_low.reshape(-1, 1), t_rad.reshape(1, -1)) * delta_nu.reshape( + -1, 1 + ) + kappa_tot = kappa_thom + kappa_exp + kappa_rosseland = ( + (udnu * kappa_tot**-1).sum(axis=0) / (udnu.sum(axis=0)) + ) ** -1 + + dr = ( + simulation_state.geometry.r_outer - simulation_state.geometry.r_inner + ).cgs.value + dtau = kappa_planck * dr + planck_integ_tau = np.cumsum(dtau[::-1])[::-1] + rosseland_integ_tau = np.cumsum((kappa_rosseland * dr)[::-1])[::-1] + + return {"rosseland": rosseland_integ_tau, "planck": planck_integ_tau} diff --git a/tardis/workflows/v_inner_solver.py b/tardis/workflows/v_inner_solver.py new file mode 100644 index 00000000000..039574b4bd9 --- /dev/null +++ b/tardis/workflows/v_inner_solver.py @@ -0,0 +1,353 @@ +import logging + +import numpy as np +import pandas as pd +from astropy import units as u +from scipy.interpolate import interp1d + +from tardis.plasma.radiation_field import DilutePlanckianRadiationField +from tardis.simulation.convergence import ConvergenceSolver +from tardis.workflows.simple_tardis_workflow import SimpleTARDISWorkflow +from tardis.workflows.util import get_tau_integ + +# logging support +logger = logging.getLogger(__name__) + +# TODO: +# Option to estimate initial v_inner from electron opacity +# Add option for desired optical depth +# Think about csvy vs other formats for specifying grid +# Handle non-explicit formats when going out of the simulation + + +class InnerVelocitySolverWorkflow(SimpleTARDISWorkflow): + TAU_TARGET = np.log(2.0 / 3.0) + + def __init__(self, configuration, mean_optical_depth="rosseland", tau=None): + super().__init__(configuration) + self.mean_optical_depth = mean_optical_depth.lower() + + self.convergence_solvers["v_inner_boundary"] = ConvergenceSolver( + self.convergence_strategy.v_inner_boundary + ) + + if tau is not None: + self.TAU_TARGET = np.log(tau) + + initial_v_inner = self.estimate_v_inner() + + self.simulation_state.geometry.v_inner_boundary = initial_v_inner + self.simulation_state.blackbody_packet_source.radius = ( + self.simulation_state.r_inner[0] + ) + + def estimate_v_inner(self): + """ + Compute the Rosseland Mean Optical Depth, + Estimate location where v_inner makes t=2/3 (or target) + Extrapolate with exponential fits + + Need some way to return and inspect the optical depths for later logging + """ + tau_integ = np.log( + get_tau_integ( + self.plasma_solver, + self.simulation_state, + )[self.mean_optical_depth] + ) + + interpolator = interp1d( + tau_integ, + self.simulation_state.geometry.v_inner, # Only use the active values as we only need a numerical estimate, not an index + fill_value="extrapolate", + ) + + estimated_v_inner = interpolator(self.TAU_TARGET) * u.cm / u.s + + if estimated_v_inner < self.simulation_state.geometry.v_inner[0]: + estimated_v_inner = self.simulation_state.geometry.v_inner[0] + logger.warning( + "WARNING: v_inner_boundary outside of simulation, setting to first shell" + ) + elif estimated_v_inner > self.simulation_state.geometry.v_inner[-1]: + estimated_v_inner = self.simulation_state.geometry.v_inner[-1] + logger.warning( + "WARNING: v_inner_boundary outside of simulation, setting to last shell" + ) + + return estimated_v_inner + + @property + def property_mask(self): + mask = np.zeros( + (len(self.simulation_state.geometry.r_inner)), dtype=bool + ) + mask[ + self.simulation_state.geometry.v_inner_boundary_index : self.simulation_state.geometry.v_outer_boundary_index + ] = True + return mask + + def get_convergence_estimates(self, transport_state): + """Compute convergence estimates from the transport state + + Parameters + ---------- + transport_state : MonteCarloTransportState + Transport state object to compute estimates + + Returns + ------- + dict + Convergence estimates + EstimatedRadiationFieldProperties + Dilute radiation file and j_blues dataclass + """ + estimates = super().get_convergence_estimates(transport_state) + + estimated_v_inner = self.estimate_v_inner() + + estimates[0].update( + {"v_inner_boundary": estimated_v_inner, "mask": self.property_mask} + ) + + return estimates + + def reproject(self, a1, m1, a2, m2): + """Reprojects two sub_arrays defined by a set of masks onto an array where the masks of both objects are true + + let A1, A2 be arrays of size gemetry.no_of_shells and + a1 = A1[m1], + a2 = A2[m2] + find a1*, a2* s.t. + a1* = A1[m1 & m2], + a2* = A2[m1 & m2] + this is equivalent to + a1* = A1[m1][m2[m1]] = a1[m2[m1]], + a2* = A2[m2][m1[m2]] = a2[m1[m2]] + + Parameters + ---------- + a1 : np.ndarray + A sub array of an array with the shape of a geometry property + m1 : np.ndarray(bool) + Mask such that the parrent array accessed at this mask gives a1 + a2 : np.ndarray + A sub array of an array with the shape of a geometry property + m2 : np.ndarray(bool) + Mask such that the parrent array accessed at this mask gives a2 + + Returns + ------- + a1* + reprojection of a1 onto m1 & m2 + a2* + reprojection of a2 onto m1 & m2 + """ + return a1[m2[m1]], a2[m1[m2]] + + def print_mask(self, mask): + return "".join([{True: "-", False: "X"}[m] for m in mask]).join("[]") + + def check_convergence( + self, + estimated_values, + ): + """Check convergence status for a dict of estimated values + + Parameters + ---------- + estimated_values : dict + Estimates to check convergence + + Returns + ------- + bool + If convergence has occurred + """ + convergence_statuses = [] + + mask = estimated_values["mask"] + + if not np.all(mask == self.property_mask): + convergence_statuses.append(False) + # Need to set status to False if change in mask size + logger.info( + f"Resized Geometry, Convergence Suppressed\n" + f"\t Old Geometry: {self.print_mask(mask)}\n" + f"\t New Geometry: {self.print_mask(self.property_mask)}" + ) + + for key, solver in self.convergence_solvers.items(): + current_value = getattr(self.simulation_state, key) + estimated_value = estimated_values[key] + + if key in ["t_radiative", "dilution_factor"]: + current_value, estimated_value = self.reproject( + current_value, self.property_mask, estimated_value, mask + ) + no_of_shells = current_value.shape[0] + else: + no_of_shells = 1 + + convergence_statuses.append( + solver.get_convergence_status( + current_value, estimated_value, no_of_shells + ) + ) + + if np.all(convergence_statuses): + hold_iterations = self.convergence_strategy.hold_iterations + self.consecutive_converges_count += 1 + logger.info( + f"Iteration converged {self.consecutive_converges_count:d}/{(hold_iterations + 1):d} consecutive " + f"times." + ) + + return self.consecutive_converges_count == hold_iterations + 1 + + self.consecutive_converges_count = 0 + return False + + def solve_simulation_state(self, estimated_values): + """Update the simulation state with new inputs computed from previous + iteration estimates. + + Parameters + ---------- + estimated_values : dict + Estimated from the previous iterations + + Returns + ------- + next_values : dict + The next values assigned to the simulation state + """ + next_values = super().solve_simulation_state(estimated_values) + self.simulation_state.geometry.v_inner_boundary = next_values[ + "v_inner_boundary" + ] + self.simulation_state.blackbody_packet_source.radius = ( + self.simulation_state.r_inner[0] + ) + + return next_values + + def solve_plasma( + self, + estimated_radfield_properties, + mask, + ): + """Update the plasma solution with the new radiation field estimates + + Parameters + ---------- + estimated_radfield_properties : EstimatedRadiationFieldProperties + The radiation field properties to use for updating the plasma + radiation_field: tardis.plasma.radiation_field.RadiationField + Current radiation field object from the last iteration + + Raises + ------ + ValueError + If the plasma solver radiative rates type is unknown + """ + radiation_field = DilutePlanckianRadiationField( + temperature=self.simulation_state.radiation_field_state.temperature, + dilution_factor=self.simulation_state.radiation_field_state.dilution_factor, + ) + update_properties = dict( + dilute_planckian_radiation_field=radiation_field + ) + # A check to see if the plasma is set with JBluesDetailed, in which + # case it needs some extra kwargs. + if ( + self.plasma_solver.plasma_solver_settings.RADIATIVE_RATES_TYPE + == "blackbody" + ): + planckian_radiation_field = ( + radiation_field.to_planckian_radiation_field() + ) + j_blues = planckian_radiation_field.calculate_mean_intensity( + self.plasma_solver.atomic_data.lines.nu.values + ) + update_properties["j_blues"] = pd.DataFrame( + j_blues, index=self.plasma_solver.atomic_data.lines.index + ) + elif ( + self.plasma_solver.plasma_solver_settings.RADIATIVE_RATES_TYPE + == "dilute-blackbody" + ): + j_blues = radiation_field.calculate_mean_intensity( + self.plasma_solver.atomic_data.lines.nu.values + ) + update_properties["j_blues"] = pd.DataFrame( + j_blues, index=self.plasma_solver.atomic_data.lines.index + ) + elif ( + self.plasma_solver.plasma_solver_settings.RADIATIVE_RATES_TYPE + == "detailed" + ): + j_blues = radiation_field.calculate_mean_intensity( + self.plasma_solver.atomic_data.lines.nu.values + ) + j_blues[:, mask] = estimated_radfield_properties.j_blues + update_properties["j_blues"] = pd.DataFrame( + j_blues, + index=self.plasma_solver.atomic_data.lines.index, + ) + else: + raise ValueError( + f"radiative_rates_type type unknown - {self.plasma.plasma_solver_settings.RADIATIVE_RATES_TYPE}" + ) + + self.plasma_solver.update(**update_properties) + + def run(self): + """Run the TARDIS simulation until convergence is reached""" + converged = False + while self.completed_iterations < self.total_iterations - 1: + logger.info( + f"\n\tStarting iteration {(self.completed_iterations + 1):d} of {self.total_iterations:d}" + ) + + opacity_states = self.solve_opacity() + + transport_state, virtual_packet_energies = self.solve_montecarlo( + opacity_states, self.real_packet_count + ) + + ( + estimated_values, + estimated_radfield_properties, + ) = self.get_convergence_estimates(transport_state) + + self.solve_simulation_state(estimated_values) + + self.solve_plasma( + estimated_radfield_properties, + estimated_values["mask"], + ) + + converged = self.check_convergence(estimated_values) + + self.completed_iterations += 1 + + if converged and self.convergence_strategy.stop_if_converged: + break + + if converged: + logger.info("\n\tStarting final iteration") + else: + logger.error( + "\n\tITERATIONS HAVE NOT CONVERGED, starting final iteration" + ) + opacity_states = self.solve_opacity() + transport_state, virtual_packet_energies = self.solve_montecarlo( + opacity_states, + self.final_iteration_packet_count, + self.virtual_packet_count, + ) + self.initialize_spectrum_solver( + transport_state, + virtual_packet_energies, + )