Legacy albedo correction (#206)

* STIX spectrogram loader

* updated changes from pre-commit

* added changelog

* Moving STIX loader to extern

* Added the actual StixLoader

* code fix from review

* pre-commit fixes

* change STIXLoader srm placeholder name

* attenuator srm loading

* function description fix

* code for albedo correction

* fix automatic attenuation selector - att. state at end of file

* albedo speed-up

* fix mcmc residuals

* pre-commit changes

* fix .plot() when no models are defined (get_models to ouput 2 vals instead 1)

* add changelog

* fix albedo docstring

* add energy and angle checks

* remove duplicate functions

* remove duplicated albedo code

* fix mcmc plotting with albedo component

* return albedo for joint fit of different non-combinable spectra

---------

Co-authored-by: William Setterberg <[email protected]>

Author: natsat0919

changelog/206.feature.rst

@@ -0,0 +1 @@
+Add the possibility to perform albedo correction within the legacy code. The alebdo matrix is generated with :func:`sunkit_spex.legacy.fitting.albedo.get_albedo_matrix` and the albedo correction is performed in  :class:`sunkit_spex.legacy.fitting.fitter.Fitter`.

sunkit_spex/legacy/fitting/albedo.py

@@ -0,0 +1,35 @@
+"""
+The following code is used to generate matrices for albedo correction
+in the legacy portion of the package.
+"""
+
+import numpy as np
+
+import astropy.units as u
+
+from sunkit_spex.models.physical.albedo import get_albedo_matrix as _get_albedo_matrix
+
+
+@u.quantity_input
+def get_albedo_matrix(energy_edges: u.keV, theta: u.deg, anisotropy=1):
+    r"""
+    Get an albedo correction matrix.
+
+    Wraps `~sunkit_spex.models.physical.albedo.get_albedo_matrix`,
+    adding support for 2D energy edges.
+
+    Parameters
+    ----------
+    energy_edges :
+        Energy edges associated with the spectrum (2D array)
+    theta :
+        Angle between Sun-observer line and X-ray source
+    anisotropy :
+        Ratio of the flux in observer direction to the flux downwards, 1 for an isotropic source
+    """
+    # The physical model Albedo matrix function expects 1D
+    # energy edges, so flatten the 2D edges from the legacy side,
+    # and just call the other function.
+    energy_edges = energy_edges.to_value(u.keV)
+    flat_edges = np.concatenate((energy_edges[:, 0], [energy_edges[-1, -1]])) << u.keV
+    return _get_albedo_matrix(flat_edges, theta, anisotropy)

sunkit_spex/legacy/fitting/fitter.py

@@ -29,12 +29,15 @@
 import numpy as np
 from mpl_toolkits.axes_grid1 import make_axes_locatable
 from scipy.interpolate import interp1d
+from scipy.io import readsav
 from scipy.linalg import LinAlgError
 # to fit the model
 from scipy.optimize import minimize
 
+import astropy.units as u
 from astropy.table import Table
 
+from sunkit_spex.legacy.fitting.albedo import get_albedo_matrix
 from sunkit_spex.legacy.fitting.data_loader import LoadSpec
 from sunkit_spex.legacy.fitting.instruments import rebin_any_array
 from sunkit_spex.legacy.fitting.likelihoods import LogLikelihoods
@@ -439,6 +442,13 @@ def __init__(
         # attribute to determine how the class is pickled
         self._pickle_reason = "normal"
 
+        # for alebdo correction
+        self.albedo_corr = False
+
+        self.albedo_angle = 0
+
+        self.albedo_anisotropy = 1
+
     @property
     def model(self):
         """***Property*** Allows a model to be set and a parameter table to be generated straight away.
@@ -1453,6 +1463,7 @@ def _counts_model(self, **kwargs):
         # return counts s^-1 (keV^-1 has been multiplied out) photon_channel_widths=ph_e_binning
 
         cts_models = []
+        albedo_excess_models = []
         # loop through the parameter groups (params for spectrum1, then for spectrum2, etc)
         for s, pgs in enumerate(self._param_groups):
             # take the spectrum parameters (e.g., [p2_spectrum1,p1_spectrum1]) and order to be the same as the model parameters (e.g., [p1,p2])
@@ -1469,8 +1480,8 @@ def _counts_model(self, **kwargs):
             )  # np.diff(kwargs["photon_channels"][s]).flatten() # remove energy bin dependence
 
             # fold the photon model through the SRM to create the count rate model, [photon s^-1 cm^-2] * [count photon^-1 cm^2] = [count s^-1]
-            cts_model = make_model(
-                energies=kwargs["photon_channels"][s], photon_model=m, parameters=None, srm=kwargs["total_responses"][s]
+            cts_model, cts_albedo_exess = make_model(
+                energies=kwargs["photon_channels"][s], photon_model=m, parameters=None, srm=kwargs["total_responses"][s], albedo_corr=self.albedo_corr, albedo_angle =self.albedo_angle, albedo_anisotropy=self.albedo_anisotropy
             )
 
             if "scaled_background_spectrum" + str(s + 1) in self._scaled_backgrounds:
@@ -1496,7 +1507,20 @@ def _counts_model(self, **kwargs):
                 cts_model[None, :]
             )  # need [None, :] these lines get rid of a dimension meaning later concatenation fails
 
-        return cts_models
+            if cts_albedo_exess.size > 0:
+                cts_albedo_exess = (
+                    cts_albedo_exess
+                    if len(LL_CLASS.remove_non_numbers(cts_albedo_exess[cts_albedo_exess != 0])) != 0
+                    else np.zeros(cts_albedo_exess.shape)
+                )
+
+                albedo_excess_models.append(
+                    cts_albedo_exess[None, :]
+                )  # need [None, :] these lines get rid of a dimension meaning later concatenation fails
+            else:
+                albedo_excess_models.append(np.array([[]]))
+
+        return cts_models, albedo_excess_models
 
     def _pseudo_model(self, free_params_list, tied_or_frozen_params_list, param_name_list_order, **other_inputs):
         """Bridging method between the input args (free,other) and different ordered args for the model calculation.
@@ -1776,7 +1800,7 @@ def _fit_stat(
         """
 
         # make sure only the free parameters are getting varied so put them first
-        mu = self._pseudo_model(
+        mu, _ = self._pseudo_model(
             free_params_list,
             tied_or_frozen_params_list,
             param_name_list_order,
@@ -2492,7 +2516,7 @@ def _calculate_model(self, **kwargs):
         tied_or_frozen_params_list.extend(list(update_fixed_params.values()))
 
         # make sure only the free parameters are getting varied
-        mu = self._pseudo_model(
+        mu, alebdo_excess_count = self._pseudo_model(
             free_params_list,
             tied_or_frozen_params_list,
             param_name_list_order,
@@ -2502,18 +2526,22 @@ def _calculate_model(self, **kwargs):
             total_responses=srm,
             **kwargs,
         )
+
         # turn counts s^-1 into counts s^-1 keV^-1
         for m, e in enumerate(e_binning):
             mu[m][0] /= e
+            if np.array(alebdo_excess_count[m][0]).size > 0:
+                alebdo_excess_count[m][0] /= e
 
         self._energy_fitting_indices = _energy_fitting_indices_orig
         # numpy is better to store but all models need to be the same length, this might not be the case
         try:
-            return np.concatenate(tuple(mu))
+            return np.concatenate(tuple(mu)), np.concatenate(tuple(alebdo_excess_count))
+
         except ValueError:
-            return [mu[i][0] for i in range(len(mu))]
+            return [mu[i][0] for i in range(len(mu))], alebdo_excess_count
 
-    def _calc_counts_model(self, photon_model, parameters=None, spectrum="spectrum1", include_bg=False, **kwargs):
+    def _calc_counts_model(self, photon_model, parameters=None, spectrum="spectrum1", include_bg=False, for_plotting=False, **kwargs):
         """Easily calculate a spectrum's count model from the parameter table and user photon model.
 
         Given a model (a calculated model array or a photon model function) then calcutes the counts
@@ -2572,12 +2600,27 @@ def _calc_counts_model(self, photon_model, parameters=None, spectrum="spectrum1"
             )
             return
 
-        cts_model = make_model(
-            energies=photon_channel_bins,
-            photon_model=m * np.diff(photon_channel_bins).flatten(),
-            parameters=None,
-            srm=srm,
-        )
+        if for_plotting:
+            cts_model, _ = make_model(
+                energies=photon_channel_bins,
+                photon_model=m * np.diff(photon_channel_bins).flatten(),
+                parameters=None,
+                srm=srm,
+                albedo_corr=False,
+                albedo_angle=self.albedo_angle,
+                albedo_anisotropy=self.albedo_anisotropy
+            )
+
+        else:
+            cts_model, _ = make_model(
+                energies=photon_channel_bins,
+                photon_model=m * np.diff(photon_channel_bins).flatten(),
+                parameters=None,
+                srm=srm,
+                albedo_corr=self.albedo_corr,
+                albedo_angle=self.albedo_angle,
+                albedo_anisotropy=self.albedo_anisotropy
+            )
 
         if include_bg and ("scaled_background_" + spectrum in self._scaled_backgrounds):
             cts_model += self._scaled_backgrounds["scaled_background_" + spectrum]
@@ -2693,6 +2736,7 @@ def _calculate_submodels(self, spectrum=None):
                         photon_model=self._submod_functions[p],
                         parameters=self._submod_value_inputs[s - 1][p],
                         spectrum="spectrum" + str(s),
+                        for_plotting=True
                     )
                     spec_submods.append(cts_mod)
                 all_spec_submods.append(spec_submods)
@@ -3231,14 +3275,20 @@ def _plot_mcmc_mods(
                 name: (_params[mcmc_freepar_labels.index(val)] if type(val) == str else val)
                 for name, val in _spec_rpars.items()
             }
+            _rpars['for_plotting'] = True
             e_mids, ctr = self._calc_counts_model(
                 photon_model=self._model, parameters=_pars, spectrum="spectrum" + str(s + 1), include_bg=True, **_rpars
             )
             _randcts.append(ctr)
+
             if _rebin_info is not None:
                 ctr = self._bin_model(ctr, *_rebin_info)
                 e_mids = np.mean(_rebin_info[2], axis=1)
 
+            if self.albedo_corr:
+                for i in range(len(res_info[1])):
+                    ctr[i] = ctr[i] + res_info[-1][i]
+
             residuals = [
                 (res_info[0][i] - ctr[i]) / res_info[1][i] if res_info[1][i] > 0 else 0 for i in range(len(res_info[1]))
             ]
@@ -3559,7 +3609,7 @@ def _setup_rebin_plotting(self, rebin_and_spec, data_arrays, _return_cts_rate_mo
         count_rate_model), and a 1d array of energies for residual plotting (energy_channels_res), respectively.
         """
         rebin_val, rebin_spec = rebin_and_spec[0], rebin_and_spec[1]
-        energy_channels, energy_channel_error, count_rates, count_rate_errors, count_rate_model = data_arrays
+        energy_channels, energy_channel_error, count_rates, count_rate_errors, albedo_excess_count, count_rate_model = data_arrays
         if type(rebin_val) != type(None):
             if rebin_spec in list(self.data.loaded_spec_data.keys()):
                 (
@@ -3714,7 +3764,7 @@ def _plot_1spec(
         -------
         Spectrum axes and residuals axes.
         """
-        energy_channels, energy_channel_error, count_rates, count_rate_errors, count_rate_model = data_arrays
+        energy_channels, energy_channel_error, count_rates, count_rate_errors, albedo_excess_count, count_rate_model = data_arrays
 
         axs = axes if type(axes) != type(None) else plt.gca()
         fitting_range = fitting_range if type(fitting_range) != type(None) else self.energy_fitting_range
@@ -3795,14 +3845,17 @@ def _plot_1spec(
             axs.plot(energy_channels, count_rate_model, linewidth=2, color="k")
             res.plot(energy_channels_res, residuals, color="k", alpha=0.8)  # , drawstyle='steps-mid'
 
+            if self.albedo_corr and albedo_excess_count.size > 0:
+                axs.plot(energy_channels, albedo_excess_count, color="grey")
+
         if self._latest_fit_run == "mcmc":
             _rebin_info = (
                 [old_bin_width, old_bins, new_bins, new_bin_width] if type(rebin_and_spec[0]) != type(None) else None
             )
             self._plot_mcmc_mods(
                 axs,
                 res,
-                [count_rates, count_rate_errors, energy_channels_res],
+                [count_rates, count_rate_errors, energy_channels_res, albedo_excess_count],
                 spectrum=submod_spec,
                 num_of_samples=num_of_samples,
                 hex_grid=hex_grid,
@@ -4194,7 +4247,7 @@ def _get_models(self, number_of_models):
             return self._calculate_model()
         else:
             self._param_groups = [None] * int(number_of_models)
-            return [np.array([1])] * int(number_of_models)
+            return [np.array([1])] * int(number_of_models), [np.array([1])] * int(number_of_models) #empty models, empty array for albedo
 
     def plot(self, subplot_axes_grid=None, rebin=None, num_of_samples=100, hex_grid=False, plot_final_result=True):
         """Plots the latest fit or sampling result.
@@ -4268,7 +4321,7 @@ def plot(self, subplot_axes_grid=None, rebin=None, num_of_samples=100, hex_grid=
         self._scaled_backgrounds = self._scaled_background_rates_full
 
         # only need enough axes for the number of spectra to plot so doesn't matter if more axes are given
-        models = self._get_models(number_of_models=number_of_plots)
+        models, albedo_excess_count = self._get_models(number_of_models=number_of_plots)
 
         axes, res_axes = [], []
         _count_rates, _count_rate_errors = [], []
@@ -4282,6 +4335,7 @@ def plot(self, subplot_axes_grid=None, rebin=None, num_of_samples=100, hex_grid=
                         self.data.loaded_spec_data["spectrum" + str(s + 1)]["count_channel_binning"] / 2,
                         self.data.loaded_spec_data["spectrum" + str(s + 1)]["count_rate"],
                         self.data.loaded_spec_data["spectrum" + str(s + 1)]["count_rate_error"],
+                        albedo_excess_count[s],
                         models[s],
                     ),
                     axes=ax,
@@ -4308,6 +4362,7 @@ def plot(self, subplot_axes_grid=None, rebin=None, num_of_samples=100, hex_grid=
                         self.data.loaded_spec_data["spectrum" + str(s)]["count_channel_binning"] / 2,
                         np.mean(np.array(_count_rates), axis=0),
                         np.sqrt(np.sum(np.array(_count_rate_errors) ** 2, axis=0)) / len(_count_rate_errors),
+                        np.array([]),
                         np.mean(models, axis=0),
                     ),
                     axes=ax,
@@ -5600,7 +5655,7 @@ def imports():
     return _imps
 
 
-def make_model(energies=None, photon_model=None, parameters=None, srm=None):
+def make_model(energies=None, photon_model=None, parameters=None, srm=None, albedo_corr=False, albedo_angle=0, albedo_anisotropy=1):
     """Takes a photon model array ( or function if you provide the pinputs with parameters), the spectral response matrix and returns a model count spectrum.
 
     Parameters
@@ -5632,6 +5687,36 @@ def make_model(energies=None, photon_model=None, parameters=None, srm=None):
     else:
         photon_spec = photon_model(energies, *parameters)
 
+    if albedo_corr:
+        photon_spec, albedo_excess_phot  = albedo(photon_spec, energies, albedo_angle, anisotropy=albedo_anisotropy)
+        albedo_excess_count = np.matmul(albedo_excess_phot, srm)
+    else:
+        albedo_excess_count = np.array([])
+
     model_cts_spectrum = np.matmul(photon_spec, srm)
 
-    return model_cts_spectrum
+    return model_cts_spectrum, albedo_excess_count
+
+#####
+# Function to calculate the albedo correction, adapted from Shane's PR request
+#####
+
+def albedo(spec, energy, theta, anisotropy=1):
+    r"""
+        Gets the albedo matrix for given angle and anisotropy and returns the albedo corrected spectrum as well as the albedo component by itself.
+
+        Parameters
+        ----------
+        spec :
+            Spectrum object
+        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, 1 for an isotropic source
+        """
+
+    albedo_matrix = get_albedo_matrix(energy*u.keV, theta, anisotropy)
+
+    return spec + spec @ albedo_matrix, spec @ albedo_matrix