Initial implementation of scaled lookup table for TOSCA

This is an interpolated lookup approach to the TOSCA resolution
function, using McStas data from Adrien Perrichon

Author: ajjackson

pyproject.toml

@@ -68,4 +68,4 @@ include-package-data = true
 where =  ["src"]
 
 [tool.setuptools.package-data]
-resolution_functions = ["instrument_data/*.yaml"]
+resolution_functions = ["instrument_data/*.yaml", "instrument_data/*.npz"]

src/resolution_functions/instrument.py

@@ -12,14 +12,16 @@
 import yaml
 from typing import Iterator, Iterable, Optional, Union, TYPE_CHECKING
 
+
+INSTRUMENT_DATA_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'instrument_data')
+
+
 from .models import MODELS
 
 if TYPE_CHECKING:
     from .models.model_base import ModelData, InstrumentModel
     from inspect import Signature
 
-INSTRUMENT_DATA_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'instrument_data')
-
 INSTRUMENT_MAP: dict[str, tuple[str, None | str]] = {
     'ARCS': ('arcs.yaml', None),
     'CNCS': ('cncs.yaml', None),

src/resolution_functions/instrument_data/tosca.yaml

@@ -146,3 +146,10 @@ version:
                             <<: *tosca_configurations_forward
                             # TODO: Recompute below once average_secondary_flight_path for Forward has been recomputed
                             distance_sample_analyzer: 0.2312830382624369  # m (as above)
+            perrichon2022: "perrichon2022_v1"
+            perrichon2022_v1:
+                function: "scaled_tabulated"
+                citation: ["https://doi.org/10.1016/j.nima.2022.167401"]
+                parameters:
+                    npz: "tosca_perrichon2022_v1.npz"
+                configurations: {}

src/resolution_functions/instrument_data/tosca_perrichon2022_v1.npz

@@ -13,6 +13,7 @@
 
 from typing import TYPE_CHECKING
 
+from .lookuptables import ScaledTabulatedModel
 from .polynomial import PolynomialModel1D, DiscontinuousPolynomialModel1D
 from .panther_abins import PantherAbINSModel
 from .pychop import PyChopModelFermi, PyChopModelNonFermi, PyChopModelCNCS, PyChopModelLET
@@ -32,5 +33,6 @@
     'pychop_fit_fermi': PyChopModelFermi,
     'pychop_fit_cncs': PyChopModelCNCS,
     'pychop_fit_let': PyChopModelLET,
+    'scaled_tabulated': ScaledTabulatedModel,
 }
 """A dictionary mapping the unique name of a model to the corresponding class."""

src/resolution_functions/models/__init__.py

@@ -0,0 +1,143 @@
+"""
+A model based on intepolated lookup tables
+
+All classes within are exposed for reference only and should not be instantiated directly. For
+obtaining the :term:`resolution function` of an :term:`instrument`, please use the
+`resolution_functions.instrument.Instrument.get_resolution_function` method.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import ClassVar, TYPE_CHECKING
+
+import numpy as np
+from numpy.polynomial.polynomial import Polynomial
+
+from .model_base import InstrumentModel, ModelData
+from .mixins import SimpleBroaden1DMixin
+from ..instrument import INSTRUMENT_DATA_PATH
+
+if TYPE_CHECKING:
+    from jaxtyping import Float
+
+
+@dataclass(init=True, repr=True, frozen=True, slots=True, kw_only=True)
+class ScaledTabulatedModelData(ModelData):
+    """
+    Data for the `ScaledTabulatedModel` :term:`model`.
+
+    Attributes
+    ----------
+    function
+        The name of the function, i.e. the alias for `ScaledTabulatedModel`.
+    citation
+        The citation for the model. Please use this to look up more details and cite the model.
+    npz
+        Relative path from Instrument yaml files to lookup table file
+    restrictions
+    defaults
+    """
+    npz: str
+
+
+class ScaledTabulatedModel(SimpleBroaden1DMixin, InstrumentModel):
+    """
+    Model using a lookup table to model a 1D :term:`instrument`.
+
+    This allows non-Gaussian shapes to be produced. For smooth interpolation
+    and data efficiency, the x-axis is scaled in proportion to an approximated
+    standard deviation. This standard deviation is fitted to a polynomial
+    function of energy.
+
+    Parameters
+    ----------
+    model_data
+        The data associated with the model for a given version of a given instrument.
+
+    Attributes
+    ----------
+    input
+        The names of the columns in the ``omega_q`` array expected by all computation methods, i.e.
+        the names of the independent variables ([Q, w]) that the model models.
+    data_class
+        Reference to the `PolynomialModelData` type.
+    npz
+        The .npz file containing the model data
+    citation
+    """
+    input = ('energy_transfer',)
+
+    data_class: ClassVar[type[ScaledTabultedModelData]] = ScaledTabulatedModelData
+
+    def __init__(self, model_data: ScaledTabulatedModelData, **_):
+        from pathlib import Path
+
+        from numpy.polynomial import Polynomial
+        from scipy.interpolate import RegularGridInterpolator
+
+        super().__init__(model_data)
+        self.data = np.load(Path(INSTRUMENT_DATA_PATH) / model_data.npz)
+
+        self.polynomial = Polynomial(coef=self.data["coef"],
+                                     domain=self.data["domain"],
+                                                   window=self.data["window"])
+        self._interp = RegularGridInterpolator((self.data["energy_transfer"], self.data["kernel_energies"]),
+                                               self.data["table"],
+                                               method="linear",
+                                               bounds_error=False,
+                                               fill_value=0.)
+
+    def get_characteristics(self, omega_q: Float[np.ndarray, 'energy_transfer dimension=1']
+                            ) -> dict[str, Float[np.ndarray, 'sigma']]:
+        """
+        Computes the broadening width at each value of energy transfer (`omega_q`).
+
+        The model approximates the broadening using the Gaussian distribution, so the returned
+        widths are in the form of the standard deviation (sigma).
+
+        Parameters
+        ----------
+        omega_q
+            The energy transfer in meV at which to compute the width in sigma of the kernel.
+            This *must* be a ``sample`` x 1 2D array where ``sample`` is the number of energy
+            transfers.
+
+        Returns
+        -------
+        characteristics
+            The characteristics of the broadening function, i.e. the Gaussian width as sigma.
+        """
+        return {'sigma': self.polynomial(omega_q[:, 0])}
+
+    def get_kernel(self,
+                   omega_q: Float[np.ndarray, 'sample dimension=1'],
+                   mesh: Float[np.ndarray, 'mesh'],
+                   ) -> Float[np.ndarray, 'sample mesh']:
+
+        assert len(omega_q.shape) == 2 and omega_q.shape[1] == 1
+        energy = omega_q
+
+        scale_factors = self.polynomial(energy)
+        scaled_x_values = mesh / scale_factors
+
+        # Clip lookup energies to known maximum; width scaling should give a
+        # reasonable extrapolation from there
+        energy = np.minimum(energy, max(self.data["energy_transfer"]))
+
+        energy_expanded = np.meshgrid(energy[:, None], mesh, indexing="ij")[0]
+        lookup_mesh = np.stack([energy_expanded, scaled_x_values], axis=-1)
+        interp_kernels = self._interp(lookup_mesh) / scale_factors
+        return interp_kernels
+
+    def get_peak(self,
+                 omega_q: Float[np.ndarray, 'sample dimension=1'],
+                 mesh: Float[np.ndarray, 'mesh'],
+                 ) -> Float[np.ndarray, 'sample mesh']:
+        shifted_meshes = [mesh - energy for energy in omega_q[:, 0]]
+
+        shifted_kernels = [
+            self.get_kernel(np.array([omega_q]), shifted_mesh)
+            for omega_q, shifted_mesh in zip(omega_q, shifted_meshes)
+        ]
+
+        return np.array(np.vstack(shifted_kernels))