From b9217dbbdc13ac31bc3d478823d8b16031643e26 Mon Sep 17 00:00:00 2001 From: natsat0919 Date: Thu, 19 Dec 2024 19:38:38 +0000 Subject: [PATCH] Speed up albedo correction, add AlbedoModel --- sunkit_spex/models/physical/albedo.py | 153 +++++++++++++----- .../models/physical/tests/test_albedo.py | 7 +- 2 files changed, 121 insertions(+), 39 deletions(-) diff --git a/sunkit_spex/models/physical/albedo.py b/sunkit_spex/models/physical/albedo.py index a5c532f1..63db9b50 100644 --- a/sunkit_spex/models/physical/albedo.py +++ b/sunkit_spex/models/physical/albedo.py @@ -1,47 +1,62 @@ +import warnings +import functools + import numpy as np from scipy.interpolate import RegularGridInterpolator from scipy.io import readsav -import astropy.units as u -from astropy.units import Quantity +# from astropy.units import Quantity +from astropy.modeling import Fittable1DModel, Parameter from sunpy.data import cache +__all__ = ["AlbedoModel"] -@u.quantity_input -def albedo(spec, energy: Quantity[u.keV], theta: Quantity[u.deg, 0, 90], anisotropy=1): - r""" - Add albedo correction to input spectrum - Correct input model spectrum for the component reflected by the solar atmosphere following [Kontar20006]_ using - precomputed green matrices from SSW. +class AlbedoModel(Fittable1DModel): + def __init__(self, energy, anisotropy, theta): + self.energy = energy + self.anisotropy = Parameter( + default=anisotropy, description="The anisotropy used for albedo correction", fixed=True + ) + self.theta = Parameter( + default=theta, description="Angle between the flare and the telescope in radians", fixed=True + ) + super().__init__() - .. [Kontar20006] https://doi.org/10.1051/0004-6361:20053672 + def evaluate(self, count_model): + """Corrects the composite count model for albedo. To be used as: ct_model = (ph_model|srm)|albedo""" + albedo_matrix_T = albedo(self.energy, self.theta.value, self.anisotropy.value) - Parameters - ---------- - spec : - Input spectrum - energy : - Energy edges associated with the spectrum - theta : - Angle between Sun-observer line and X-ray source - anisotropy : - Ratio of the flux in observer direction to the flux downwards, anisotropy=1 (default) the source is isotropic + return count_model + count_model @ albedo_matrix_T - Example - ------- - >>> import astropy.units as u - >>> import numpy as np - >>> from sunkit_spex.models.physical import albedo - >>> e = np.linspace(5, 500, 1000) - >>> e_c = e[:-1] + np.diff(e) - >>> s = 125*e_c**-3 - >>> corrected = albedo(s, e, theta=45*u.deg) - """ - base_url = "https://soho.nascom.nasa.gov/solarsoft/packages/xray/dbase/albedo/" + +# Wrapper for chaching function with numpy array args +def np_cache(function): + @functools.cache + def cached_wrapper(*args, **kwargs): + args = [np.array(a) if isinstance(a, tuple) else a for a in args] + kwargs = {k: np.array(v) if isinstance(v, tuple) else v for k, v in kwargs.items()} + + return function(*args, **kwargs) + + @functools.wraps(function) + def wrapper(*args, **kwargs): + args = [tuple(a) if isinstance(a, np.ndarray) else a for a in args] + kwargs = {k: tuple(v) if isinstance(v, np.ndarray) else v for k, v in kwargs.items()} + return cached_wrapper(*args, **kwargs) + + wrapper.cache_info = cached_wrapper.cache_info + wrapper.cache_clear = cached_wrapper.cache_clear + + return wrapper + + +@functools.cache +def load_green_matrices(theta): mu = np.cos(theta) + base_url = "https://soho.nascom.nasa.gov/solarsoft/packages/xray/dbase/albedo/" # what about 0 and 1 assume so close to 05 and 95 that it doesn't matter # load precomputed green matrices if 0.5 <= mu <= 0.95: @@ -70,17 +85,83 @@ def albedo(spec, energy: Quantity[u.keV], theta: Quantity[u.deg, 0, 90], anisotr albedo = green["p"].albedo[0] albedo = albedo.T - energy_grid_edges = green["p"].edges[0] * u.keV + + # energy_grid_edges = green["p"].edges[0] * u.keV + energy_grid_edges = green["p"].edges[0] + energy_grid_centers = energy_grid_edges[:, 0] + (np.diff(energy_grid_edges, axis=1) / 2).reshape(-1) - interp = RegularGridInterpolator((energy_grid_centers.to_value(u.keV), energy_grid_centers.to_value(u.keV)), albedo) + # interp_green_matrix = RegularGridInterpolator((energy_grid_centers.to_value(u.keV), energy_grid_centers.to_value(u.keV)), albedo) + interp_green_matrix = RegularGridInterpolator((energy_grid_centers, energy_grid_centers), albedo) + + return interp_green_matrix + +@np_cache +def calc_albedo_matrix(energy, interp_green_matrix, anisotropy): + # Re-introduce units + # energy = energy * u.keV de = np.diff(energy) energy_centers = energy[:-1] + de / 2 - X, Y = np.meshgrid(energy_centers.to_value(u.keV), energy_centers.to_value(u.keV)) - albedo_interp = interp((X, Y)) + # X, Y = np.meshgrid(energy_centers.to_value(u.keV), energy_centers.to_value(u.keV)) + X, Y = np.meshgrid(energy_centers, energy_centers) + + albedo_interp = interp_green_matrix((X, Y)) + + # Scale by anisotropy + # albedo_interp = albedo_interp * de.value / anisotropy + albedo_interp = albedo_interp * de / anisotropy + + # Take a transpose + albedo_interp_T = albedo_interp.T + + return albedo_interp_T + + +# @u.quantity_input +# def albedo(energy: Quantity[u.keV], theta: Quantity[u.deg, 0, 90], anisotropy=1): +def albedo(energy, theta, anisotropy=1): + r""" + Add albedo correction to input spectrum + + Correct input model spectrum for the component reflected by the solar atmosphere following [Kontar20006]_ using + precomputed green matrices from SSW. + + .. [Kontar20006] https://doi.org/10.1051/0004-6361:20053672 + + Parameters + ---------- + spec : + Input count spectrum + energy : + Energy edges associated with the spectrum + theta : + Angle between Sun-observer line and X-ray source + anisotropy : + Ratio of the flux in observer direction to the flux downwards, anisotropy=1 (default) the source is isotropic + + Example + ------- + >>> import astropy.units as u + >>> import numpy as np + >>> from sunkit_spex.models.physical.albedo import albedo + >>> e = np.linspace(5, 500, 1000) + >>> e_c = e[:-1] + np.diff(e) + >>> s = 125*e_c**-3 + >>> albedo_matrix = albedo(e, theta=0.2) + >>> corrected = s + s@albedo_matrix + """ + + # Add check for energy range restricted by the Green matrices + if energy[0] < 3 or energy[-1] > 600: + warnings.warn("\nCount energies are required to be >= 3keV and <=600 keV ") + + # Green matrix for a given theta + interp_green = load_green_matrices(theta) - albedo_interp = albedo_interp * de.value / anisotropy + # Transpose of a matrix used for albedo correction, need to strip energy of units otherwise problem with caching + # albedo_matrix_T = calc_albedo_matrix(energy.value, interp_green, anisotropy) + albedo_matrix_T = calc_albedo_matrix(energy, interp_green, anisotropy) - return spec + spec @ albedo_interp.T + return albedo_matrix_T diff --git a/sunkit_spex/models/physical/tests/test_albedo.py b/sunkit_spex/models/physical/tests/test_albedo.py index 58911aca..5d33964c 100644 --- a/sunkit_spex/models/physical/tests/test_albedo.py +++ b/sunkit_spex/models/physical/tests/test_albedo.py @@ -1,7 +1,5 @@ import numpy as np -import astropy.units as u - from sunkit_spex.models.physical.albedo import albedo @@ -9,5 +7,8 @@ def test_albedo(): e = np.linspace(4, 600, 597) e_c = e[:-1] + np.diff(e) s = 125 * e_c**-3 - out = albedo(s, e * u.keV, 45 * u.deg) + # e = e *u.keV + theta = 0.2 # u.rad + albedo_matrix = albedo(e, theta=theta) + out = s + s @ albedo_matrix print(out)