Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Defer thermal model setup until actually needed #184

Draft
wants to merge 2 commits into
base: main
Choose a base branch
from
Draft

Defer thermal model setup until actually needed #184

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
fb1d3ba
First attempt using globals - a class will probably be better
samaloney Jan 30, 2025
d1eccdb
Better apprach using class and cached properites
samaloney Jan 30, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
128 changes: 90 additions & 38 deletions sunkit_spex/models/physical/thermal.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -190,11 +190,36 @@ def setup_default_abundances(filename=None):
return load_xray_abundances()


# Read line, continuum and abundance data into global variables.
CONTINUUM_GRID = setup_continuum_parameters()
LINE_GRID = setup_line_parameters()
DEFAULT_ABUNDANCES = setup_default_abundances()
DEFAULT_ABUNDANCE_TYPE = "sun_coronal_ext"
class ThermalParameters:
_continuum_grid = None
_line_grid = None
_default_abundances = None
_default_abundance_type = "sun_coronal_ext"

@property
def continuum_grid(self):
if not self._continuum_grid:
self._continuum_grid = setup_continuum_parameters()
return self._continuum_grid

@property
def line_grid(self):
if not self._line_grid:
self._line_grid = setup_line_parameters()
return self._line_grid

@property
def default_abundances(self):
if not self._default_abundances:
self._default_abundances = setup_default_abundances()
return self._default_abundances

@property
def default_abundance_type(self):
return self._default_abundance_type


THERMAL_PARAMETERS = ThermalParameters()


@u.quantity_input(
Expand All @@ -204,7 +229,7 @@ def thermal_emission(
energy_edges,
temperature,
emission_measure,
abundance_type=DEFAULT_ABUNDANCE_TYPE,
abundance_type=THERMAL_PARAMETERS.default_abundance_type,
relative_abundances=None,
observer_distance=(1 * u.AU).to(u.cm),
):
Expand All @@ -219,13 +244,25 @@ def thermal_emission(
# Convert inputs to known units and confirm they are within range.
energy_edges_keV, temperature_K = _sanitize_inputs(energy_edges, temperature)
energy_range = (
min(CONTINUUM_GRID["energy range keV"][0], LINE_GRID["energy range keV"][0]),
max(CONTINUUM_GRID["energy range keV"][1], LINE_GRID["energy range keV"][1]),
min(
THERMAL_PARAMETERS.continuum_grid["energy range keV"][0],
THERMAL_PARAMETERS.line_grid["energy range keV"][0],
),
max(
THERMAL_PARAMETERS.continuum_grid["energy range keV"][1],
THERMAL_PARAMETERS.line_grid["energy range keV"][1],
),
)
_error_if_input_outside_valid_range(energy_edges_keV, energy_range, "energy", "keV")
temp_range = (
min(CONTINUUM_GRID["temperature range K"][0], LINE_GRID["temperature range K"][0]),
max(CONTINUUM_GRID["temperature range K"][1], LINE_GRID["temperature range K"][1]),
min(
THERMAL_PARAMETERS.continuum_grid["temperature range K"][0],
THERMAL_PARAMETERS.line_grid["temperature range K"][0],
),
max(
THERMAL_PARAMETERS.continuum_grid["temperature range K"][1],
THERMAL_PARAMETERS.line_grid["temperature range K"][1],
),
)
_error_if_input_outside_valid_range(temperature_K, temp_range, "temperature", "K")
# Calculate abundances
Expand All @@ -246,7 +283,7 @@ def continuum_emission(
energy_edges,
temperature,
emission_measure,
abundance_type=DEFAULT_ABUNDANCE_TYPE,
abundance_type=THERMAL_PARAMETERS.default_abundance_type,
relative_abundances=None,
observer_distance=(1 * u.AU).to(u.cm),
):
Expand All @@ -260,8 +297,12 @@ def continuum_emission(
{doc_string_params}"""
# Convert inputs to known units and confirm they are within range.
energy_edges_keV, temperature_K = _sanitize_inputs(energy_edges, temperature)
_error_if_input_outside_valid_range(energy_edges_keV, CONTINUUM_GRID["energy range keV"], "energy", "keV")
_error_if_input_outside_valid_range(temperature_K, CONTINUUM_GRID["temperature range K"], "temperature", "K")
_error_if_input_outside_valid_range(
energy_edges_keV, THERMAL_PARAMETERS.continuum_grid["energy range keV"], "energy", "keV"
)
_error_if_input_outside_valid_range(
temperature_K, THERMAL_PARAMETERS.continuum_grid["temperature range K"], "temperature", "K"
)
# Calculate abundances
abundances = _calculate_abundances(abundance_type, relative_abundances)
# Calculate flux.
Expand All @@ -279,7 +320,7 @@ def line_emission(
energy_edges,
temperature,
emission_measure,
abundance_type=DEFAULT_ABUNDANCE_TYPE,
abundance_type=THERMAL_PARAMETERS.default_abundance_type,
relative_abundances=None,
observer_distance=(1 * u.AU).to(u.cm),
):
Expand All @@ -289,8 +330,12 @@ def line_emission(
{docstring_params}"""
# Convert inputs to known units and confirm they are within range.
energy_edges_keV, temperature_K = _sanitize_inputs(energy_edges, temperature)
_warn_if_input_outside_valid_range(energy_edges_keV, LINE_GRID["energy range keV"], "energy", "keV")
_error_if_input_outside_valid_range(temperature_K, LINE_GRID["temperature range K"], "temperature", "K")
_warn_if_input_outside_valid_range(
energy_edges_keV, THERMAL_PARAMETERS.line_grid["energy range keV"], "energy", "keV"
)
_error_if_input_outside_valid_range(
temperature_K, THERMAL_PARAMETERS.line_grid["temperature range K"], "temperature", "K"
)
# Calculate abundances
abundances = _calculate_abundances(abundance_type, relative_abundances)

Expand Down Expand Up @@ -330,7 +375,7 @@ def _continuum_emission(energy_edges_keV, temperature_K, abundances):

# Mask Unwanted Abundances
abundance_mask = np.zeros(len(abundances))
abundance_mask[CONTINUUM_GRID["abundance index"]] = 1.0
abundance_mask[THERMAL_PARAMETERS.continuum_grid["abundance index"]] = 1.0
abundances *= abundance_mask

# Calculate Continuum Intensity Summed Over All Elements
Expand All @@ -341,23 +386,24 @@ def _continuum_emission(energy_edges_keV, temperature_K, abundances):
# temperatures here. If only a few temperatures are input, do this step only
# when looping over input temperatures. This minimizes computation.
n_tband = 3
n_t_grid = len(CONTINUUM_GRID["log10T"])
n_t_grid = len(THERMAL_PARAMETERS.continuum_grid["log10T"])
n_temperature_K = len(temperature_K)
n_thresh = n_temperature_K * n_tband
if n_thresh >= n_t_grid:
intensity_per_em_at_source_allT = np.zeros(CONTINUUM_GRID["intensity"].shape[1:])
intensity_per_em_at_source_allT = np.zeros(THERMAL_PARAMETERS.continuum_grid["intensity"].shape[1:])
for i in range(0, n_t_grid):
intensity_per_em_at_source_allT[i] = np.matmul(
abundances[CONTINUUM_GRID["sorted abundance index"]], CONTINUUM_GRID["intensity"][:, i]
abundances[THERMAL_PARAMETERS.continuum_grid["sorted abundance index"]],
THERMAL_PARAMETERS.continuum_grid["intensity"][:, i],
)
# 2. Add dummy axes to energy and temperature grid arrays for later vectorized operations.
repeat_E_grid = CONTINUUM_GRID["E_keV"][np.newaxis, :]
repeat_T_grid = CONTINUUM_GRID["T_keV"][:, np.newaxis]
dE_grid_keV = CONTINUUM_GRID["energy bin widths keV"][np.newaxis, :]
repeat_E_grid = THERMAL_PARAMETERS.continuum_grid["E_keV"][np.newaxis, :]
repeat_T_grid = THERMAL_PARAMETERS.continuum_grid["T_keV"][:, np.newaxis]
dE_grid_keV = THERMAL_PARAMETERS.continuum_grid["energy bin widths keV"][np.newaxis, :]
# 3. Identify the indices of the temperature bins containing each input temperature and
# the bins above and below them. For each input temperature, these three bins will
# act as a temperature band over which we'll interpolate the continuum emission.
selt = np.digitize(log10T_in, CONTINUUM_GRID["log10T"]) - 1
selt = np.digitize(log10T_in, THERMAL_PARAMETERS.continuum_grid["log10T"]) - 1
tband_idx = selt[:, np.newaxis] + np.arange(n_tband)[np.newaxis, :]

# Finally, loop over input temperatures and calculate continuum emission for each.
Expand All @@ -366,11 +412,12 @@ def _continuum_emission(energy_edges_keV, temperature_K, abundances):
# If not already done above, calculate continuum intensity summed over
# all elements as a function of energy/wavelength over the temperature band.
if n_thresh < n_t_grid:
element_intensities_per_em_at_source = CONTINUUM_GRID["intensity"][:, tband_idx[j]]
element_intensities_per_em_at_source = THERMAL_PARAMETERS.continuum_grid["intensity"][:, tband_idx[j]]
intensity_per_em_at_source = np.zeros(element_intensities_per_em_at_source.shape[1:])
for i in range(0, n_tband):
intensity_per_em_at_source[i] = np.matmul(
abundances[CONTINUUM_GRID["sorted abundance index"]], element_intensities_per_em_at_source[:, i]
abundances[THERMAL_PARAMETERS.continuum_grid["sorted abundance index"]],
element_intensities_per_em_at_source[:, i],
)
else:
intensity_per_em_at_source = intensity_per_em_at_source_allT[tband_idx[j]]
Expand All @@ -385,13 +432,17 @@ def _continuum_emission(energy_edges_keV, temperature_K, abundances):
# Interpolate the normalized temperature component of the intensity grid the the
# input temperature.
flux[j] = _interpolate_continuum_intensities(
gaunt, CONTINUUM_GRID["log10T"][tband_idx[j]], CONTINUUM_GRID["E_keV"], energy_gmean_keV, logt
gaunt,
THERMAL_PARAMETERS.continuum_grid["log10T"][tband_idx[j]],
THERMAL_PARAMETERS.continuum_grid["E_keV"],
energy_gmean_keV,
logt,
)
# Rescale the interpolated intensity.
flux = flux * np.exp(-(energy_gmean_keV[np.newaxis, :] / T_in_keV[:, np.newaxis]))

# Put intensity into correct units.
return flux * CONTINUUM_GRID["intensity unit"]
return flux * THERMAL_PARAMETERS.continuum_grid["intensity unit"]


def _line_emission(energy_edges_keV, temperature_K, abundances):
Expand All @@ -417,27 +468,28 @@ def _line_emission(energy_edges_keV, temperature_K, abundances):

# Find indices of lines within user input energy range.
energy_roi_indices = np.logical_and(
LINE_GRID["line peaks keV"] >= energy_edges_keV.min(), LINE_GRID["line peaks keV"] <= energy_edges_keV.max()
THERMAL_PARAMETERS.line_grid["line peaks keV"] >= energy_edges_keV.min(),
THERMAL_PARAMETERS.line_grid["line peaks keV"] <= energy_edges_keV.max(),
)
n_energy_roi_indices = energy_roi_indices.sum()
# If there are emission lines within the energy range of interest, compile spectrum.
if n_energy_roi_indices > 0:
# Mask Unwanted Abundances
abundance_mask = np.zeros(len(abundances))
abundance_mask[LINE_GRID["abundance index"]] = 1.0
abundance_mask[THERMAL_PARAMETERS.line_grid["abundance index"]] = 1.0
abundances *= abundance_mask
# Extract only the lines within the energy range of interest.
line_abundances = abundances[LINE_GRID["line atomic numbers"][energy_roi_indices] - 2]
line_abundances = abundances[THERMAL_PARAMETERS.line_grid["line atomic numbers"][energy_roi_indices] - 2]
# Above magic number of of -2 is comprised of:
# a -1 to account for the fact that index is atomic number -1, and
# another -1 because abundance index is offset from abundance index by 1.

##### Calculate Line Intensities within the Input Energy Range #####
# Calculate abundance-normalized intensity of each line in energy range of
# interest as a function of energy and temperature.
line_intensity_grid = LINE_GRID["intensity"][energy_roi_indices]
line_intensity_grid = THERMAL_PARAMETERS.line_grid["intensity"][energy_roi_indices]
line_intensities = _calculate_abundance_normalized_line_intensities(
np.log10(temperature_K), line_intensity_grid, LINE_GRID["log10T"]
np.log10(temperature_K), line_intensity_grid, THERMAL_PARAMETERS.line_grid["log10T"]
)
# Scale line intensities by abundances to get true line intensities.
line_intensities *= line_abundances
Expand All @@ -448,7 +500,7 @@ def _line_emission(energy_edges_keV, temperature_K, abundances):
# when averaged over neighboring bins.
# This has the effect of appearing to double the number of lines as regards
# the dimensionality of the line_intensities array.
line_peaks_keV = LINE_GRID["line peaks keV"][energy_roi_indices]
line_peaks_keV = THERMAL_PARAMETERS.line_grid["line peaks keV"][energy_roi_indices]
split_line_intensities, line_spectrum_bins = _weight_emission_bins_to_line_centroid(
line_peaks_keV, energy_edges_keV, line_intensities
)
Expand All @@ -466,7 +518,7 @@ def _line_emission(energy_edges_keV, temperature_K, abundances):
# Scale flux by observer distance, emission measure and spectral bin width
# and put into correct units.
energy_bin_widths = (energy_edges_keV[1:] - energy_edges_keV[:-1]) * u.keV
return flux * LINE_GRID["intensity unit"] / energy_bin_widths
return flux * THERMAL_PARAMETERS.line_grid["intensity unit"] / energy_bin_widths


def _interpolate_continuum_intensities(data_grid, log10T_grid, energy_grid_keV, energy_keV, log10T):
Expand Down Expand Up @@ -753,15 +805,15 @@ def _warn_if_input_outside_valid_range(input_values, grid_range, param_name, par


def _calculate_abundances(abundance_type, relative_abundances):
abundances = DEFAULT_ABUNDANCES[abundance_type].data
abundances = THERMAL_PARAMETERS.default_abundances[abundance_type].data
if relative_abundances:
# Convert input relative abundances to array where
# first axis is atomic number, i.e == index + 1
# Second axis is relative abundance value.
rel_abund_array = np.array(relative_abundances).T
# Confirm relative abundances are for valid elements and positive.
min_abundance_z = DEFAULT_ABUNDANCES["atomic number"].min()
max_abundance_z = DEFAULT_ABUNDANCES["atomic number"].max()
min_abundance_z = THERMAL_PARAMETERS.default_abundances["atomic number"].min()
max_abundance_z = THERMAL_PARAMETERS.default_abundances["atomic number"].max()
if rel_abund_array[0].min() < min_abundance_z or rel_abund_array[0].max() > max_abundance_z:
raise ValueError(
"Relative abundances can only be set for elements with "
Expand Down