Add model and tidy up code

changelog/161.feature.rst

@@ -0,0 +1 @@
+Add function `~sunkit_spex.models.physical.albedo.get_albedo_matrix` to calculate Albedo correction for given input spectrum and add model `~sunkit_spex.models.physical.albedo.Albedo` to correct modeled photon spectrum for albedo.

docs/conf.py

@@ -48,6 +48,7 @@
     "sphinx_automodapi.smart_resolver",
     "sphinx_changelog",
     "sphinx_gallery.gen_gallery",
+    "matplotlib.sphinxext.plot_directive"
 ]
 
 # Add any paths that contain templates here, relative to this directory.

docs/index.rst

@@ -16,7 +16,7 @@ will help you know where to look for certain things:
 * :doc:`Topic guides </discussion/index>` discuss key topics and concepts at a
   fairly high level and provide useful background information and explanation.
 
-* :doc:`Reference guides </reference/index>` contain technical reference for APIs and
+* :doc:`Reference </reference/index>` contain technical reference for APIs and
   other aspects of sunkit-spex machinery. They describe how it works and how to
   use it but assume that you have a basic understanding of key concepts.
 

docs/reference/index.rst

@@ -4,10 +4,12 @@ Reference
 
 Software and API.
 
+.. automodapi:: sunkit_spex
 
 .. toctree::
    :maxdepth: 2
 
-   fitting/index
-   models/index
-   legacy/index
+
+   fitting
+   models
+   legacy

docs/reference/legacy.rst

@@ -1,13 +1,15 @@
 Legacy (`sunkit_spex.legacy`)
 *****************************
 
+.. warning::
+   The legacy module contains legacy code which will no longer be maintained and will be removed in the near future.
 
 .. automodapi:: sunkit_spex.legacy
 .. automodapi:: sunkit_spex.legacy.thermal
 .. automodapi:: sunkit_spex.legacy.integrate
 .. automodapi:: sunkit_spex.legacy.io
 .. automodapi:: sunkit_spex.legacy.emission
 .. automodapi:: sunkit_spex.legacy.constants
-.. automodapi:: sunkit_spex.legacy.fitting_legacy
-.. automodapi:: sunkit_spex.legacy.fitting_legacy.io
+.. automodapi:: sunkit_spex.legacy.fitting
+.. automodapi:: sunkit_spex.legacy.fitting.io
 .. automodapi:: sunkit_spex.legacy.photon_power_law

docs/reference/models.rst

@@ -5,4 +5,8 @@ Models (`sunkit_spex.models`)
 
 
 .. automodapi:: sunkit_spex.models
+.. automodapi:: sunkit_spex.models.models
+.. automodapi:: sunkit_spex.models.instrument_response
 .. automodapi:: sunkit_spex.models.physical
+.. automodapi:: sunkit_spex.models.physical.thermal
+.. automodapi:: sunkit_spex.models.physical.albedo

sunkit_spex/models/physical/albedo.py

@@ -18,7 +18,7 @@ class Albedo(FittableModel):
     r"""
     Aldedo model which adds albdeo correction to input spectrum.
 
-    Following [Kontar2006]_ using precomputed green matrices distributed as part of SSW_.
+    Following [Kontar2006]_ using precomputed green matrices distributed as part of [SSW]_.
 
     .. [Kontar2006] https://doi.org/10.1051/0004-6361:20053672
     .. [SSW] https://www.lmsal.com/solarsoft/
@@ -34,41 +34,61 @@ class Albedo(FittableModel):
 
     Examples
     ========
-     >>> import astropy.units as u
-     >>> import numpy as np
-     >>> from astropy.modeling.powerlaws import PowerLaw1D
-     >>> from sunkit_spex.models.physical.albedo import Albedo
-     >>> e_edges = np.linspace(10, 300, 10) * u.keV
-     >>> e_centers = e_edges[0:-1] + (0.5 * np.diff(e_edges))
-     >>> source = PowerLaw1D(amplitude=100*u.ph, x_0=10*u.keV, alpha=4)
-     >>> observed = source | Albedo(energy_edges=e_edges)
-     >>> observed(e_centers)
-     <Quantity [6.09547957e+00, 9.75099578e-02, 1.91466233e-02, 5.27952213e-03,
-           1.88744944e-03, 8.34602265e-04, 4.31433126e-04, 2.49285636e-04,
-           1.53959474e-04] ph>
+    .. plot::
+        :include-source:
+
+        import astropy.units as u
+        import numpy as np
+        import matplotlib.pyplot as plt
+
+        from astropy.modeling.powerlaws import PowerLaw1D
+        from astropy.visualization import quantity_support
+
+        from sunkit_spex.models.physical.albedo import Albedo
+
+        e_edges = np.linspace(5, 550, 600) * u.keV
+        e_centers = e_edges[0:-1] + (0.5 * np.diff(e_edges))
+        source = PowerLaw1D(amplitude=1*u.ph/(u.cm*u.s), x_0=5*u.keV, alpha=3)
+        albedo = Albedo(energy_edges=e_edges)
+        observed = source | albedo
+
+        with quantity_support():
+            plt.figure()
+            plt.plot(e_centers,  source(e_centers), 'k', label='Source')
+            for i, t in enumerate([0, 45, 90]*u.deg):
+                albedo.theta = t
+                plt.plot(e_centers,  observed(e_centers), '--', label=f'Observed, theta={t}', color=f'C{i+1}')
+                plt.plot(e_centers,  observed(e_centers) - source(e_centers), ':',
+                         label=f'Reflected, theta={t}', color=f'C{i+1}')
+
+            plt.ylim(1e-6,  1)
+            plt.xlim(5, 550)
+            plt.loglog()
+            plt.legend()
+            plt.show()
+
     """
 
+    n_inputs = 1
+    n_outputs = 1
     theta = Parameter(
         name="theta",
         default=0,
         unit=u.deg,
         min=-90,
         max=90,
         description="Angle between the observer and the source",
-        fixed=True,
+        fixed=False,
     )
     anisotropy = Parameter(default=1, description="The anisotropy used for albedo correction", fixed=True)
 
-    n_inputs = 1
-    n_outputs = 1
-
     def __init__(self, *args, **kwargs):
         self.energy_edges = kwargs.pop("energy_edges")
-        return super().__init__(*args, **kwargs)
+        super().__init__(*args, **kwargs)
 
     def evaluate(self, spectrum, theta, anisotropy):
-        albedo_matrix_T = get_albedo_matrix(self.energy_edges, theta, anisotropy)
-        return spectrum + spectrum @ albedo_matrix_T
+        albedo_matrix = get_albedo_matrix(self.energy_edges, theta, anisotropy)
+        return spectrum + spectrum @ albedo_matrix
 
 
 @lru_cache
@@ -124,9 +144,7 @@ def _get_green_matrix(theta: float) -> RegularGridInterpolator:
     energy_grid_edges = green["p"].edges[0]
     energy_grid_centers = energy_grid_edges[:, 0] + (np.diff(energy_grid_edges, axis=1) / 2).reshape(-1)
 
-    green_matrix_interpolator = RegularGridInterpolator((energy_grid_centers, energy_grid_centers), albedo)
-
-    return green_matrix_interpolator
+    return RegularGridInterpolator((energy_grid_centers, energy_grid_centers), albedo)
 
 
 @lru_cache
@@ -157,21 +175,16 @@ def _calculate_albedo_matrix(energy_edges: tuple[float], theta: float, anisotrop
     albedo_interp = (albedo_interp * de) / anisotropy
 
     # Take a transpose
-    albedo_interp_T = albedo_interp.T
-
-    return albedo_interp_T
+    return albedo_interp.T
 
 
 @u.quantity_input
 def get_albedo_matrix(energy_edges: Quantity[u.keV], theta: Quantity[u.deg], anisotropy=1):
     r"""
     Get albedo correction matrix.
 
-    Matrix used to correct a photon spectrum for the component reflected by the solar atmosphere following
-    [Kontar2006]_ using precomputed green matrices distributed as part of SSW_.
-
-    .. [Kontar2006] https://doi.org/10.1051/0004-6361:20053672
-    .. [SSW] https://www.lmsal.com/solarsoft/
+    Matrix used to correct a photon spectrum for the component reflected by the solar atmosphere following interpolated
+    to given angle and energy indices.
 
     Parameters
     ----------
@@ -201,7 +214,4 @@ def get_albedo_matrix(energy_edges: Quantity[u.keV], theta: Quantity[u.deg], ani
     if np.abs(theta) > 90 * u.deg:
         raise ValueError(f"Theta must be between -90 and 90 degrees: {theta}.")
     anisotropy = np.array(anisotropy).squeeze()
-    albedo_matrix = _calculate_albedo_matrix(
-        tuple(energy_edges.to_value(u.keV)), theta.to_value(u.deg), anisotropy.item()
-    )
-    return albedo_matrix
+    return _calculate_albedo_matrix(tuple(energy_edges.to_value(u.keV)), theta.to_value(u.deg), anisotropy.item())