Implement kernels and convolutions (#39)

* Implement get_kernel

* Rewrite get_characteristics interface

* Position kernels at energy transfer

* Implement convolve method

* Move the mesh argument to last

* Add missing docstrings

* Change __call__ to call convolve

* Update the input class variable docs

* Add corresponding documentation page

* Update InstrumentModel docstring

* Update README

* Fix convolve

* Use two different get kernel methods

* Fix integration tests: compare with "sigma" characteristic

* Add mandatory methods to InstrumentModel mock class

* Tweak pychop test for more convenient formatting

Was getting some false negatives from error message formatting
depending if the float was displayed as a bare float or as np.float64.

That seems unimportant to the test, so cast to float()

* Fix convolve to use get_peak instead of get_kernel

SimpleConvolve expects the mesh to span the full energy transfer range and so does the get_peak method (get_kernel does not)

* Add unit tests for mixins

* Update quickstart.rst

* Update README.md with get_peak changes

* Update how-to guides

* Fix test_mixins convolve test data generation

* Add convolve result plot to docs

* Fix test_convolve and add test data

* Switch convolve to use einsum

* Rename convolve into broaden

* Clarify mixin output array shapes in docs

* Update InstrumentModel docstring according to review

Co-authored-by: Adam J. Jackson <a.j.jackson@physics.org>

* Update InstrumentModel method docstrings to clarify array shapes

* Refactor meshes to be a list of arrays

* Use reshape for obtaining a 1D array of energies

* Add a nice exception text

* Revert to slicing

* Apply suggestions from code review

Co-authored-by: Adam J. Jackson <a.j.jackson@physics.org>

---------

Co-authored-by: Rastislav Turanyi <rastislav.turanyi@stfc.ac.uk>
Co-authored-by: Adam J. Jackson <a.j.jackson@physics.org>

Author: 3 people

.gitignore

@@ -1,5 +1,6 @@
 # Documentation build files
 docs/include/auto-instruments
+tests/unit_tests/data/*.png
 
 # Byte-compiled / optimized / DLL files
 __pycache__/

README.md

@@ -9,7 +9,7 @@ literature. The main purposes are:
 
 1. Provide one, central place implementing various models for INS instruments (resolution functions)
 2. Provide a simple way to obtain the broadening at a given frequency, for a given instrument and settings
-3. (in progress) Provide a way to apply broadening to a spectrum
+3. Provide a way to apply broadening to a spectrum
 
 ## Quick Start
 
@@ -18,31 +18,53 @@ class and get the instrument object of your choice:
 
 ```
 >>> from resolution_functions import Instrument
->>> tosca = Instrument.from_default('TOSCA')
->>> tosca
-Instrument(name='TOSCA', version='TOSCA', models=['AbINS', 'book', 'vision'])
+>>> maps = Instrument.from_default('MAPS', 'MAPS')
+>>> print(maps)
+Instrument(name=MAPS, version=MAPS)
 ```
 
-To get the resolution function, call the `get_resolution_function` method, which returns a callable 
-that can be called to get the broadening at specified frequencies. However, you will need to know 
-which model you want to use, as well as any model-specific parameters.
+To get the resolution function, call the `get_resolution_function` method (providing all your 
+choices for the required settings and configurations), which returns a callable that can be called 
+to broaden the data.
 
 ```
 >>> # The available models for a given instrument can be queried:
->>> tosca.available_models
-['AbINS', 'book', 'vision']
+>>> maps.available_models
+['PyChop_fit']
 >>> # There are multiple ways of querying the model-specific parameters, but the most comprehensive is
->>> tosca.get_model_signature('book')
-<Signature (model_name: Optional[str] = 'book', *, detector_bank: Literal['Backward', 'Forward'] = 'Backward', _) -> resolution_functions.models.tosca_book.ToscaBookModel>
+>>> maps.get_model_signature('PyChop_fit')
+<Signature (model_name: Optional[str] = 'PyChop_fit_v1', *, chopper_package: Literal['A', 'B', 'S'] = 'A', e_init: Annotated[ForwardRef('Optional[float]'), 'restriction=[0, 2000]'] = 500, chopper_frequency: Annotated[ForwardRef('Optional[int]'), 'restriction=[50, 601, 50]'] = 400, fitting_order: 'int' = 4, _) -> resolution_functions.models.pychop.PyChopModelFermi>
 >>> # Now we can get the resolution function
->>> book = tosca.get_resolution_function('book', detector_bank='Forward')
->>> print(book)
-ToscaBookModel(citation=[''])
->>> book(100)
-0.81802604002035
+>>> pychop = maps.get_resolution_function('PyChop_fit', chopper_package='B', e_init=500, chopper_frequency=300)
+>>> print(pychop)
+PyChopModelFermi(citation=[''])
+```
+
+Calling the model (like a function) broadens the data at the provided combinations of energy 
+transfer and momentum ([w, Q]), using a mesh and the corresponding data:
+
+```
 >>> import numpy as np
->>> book(np.array([100, 200, 300]))
-array([0.81802604, 1.34222267, 1.88255039])
+>>> energy_transfer = np.array([100, 200, 300])[:, np.newaxis]
+>>> data = np.array([0.6, 1.5, 0.9])
+>>> mesh = np.linspace(0, 500, 1000)
+>>> pychop(energy_transfer, data, mesh)
+array([3.43947518e-028, ... 5.99877942e-002, ... 7.31766110e-249])
+```
+
+However, the model also provides methods that go lower; 
+
+- `get_kernel` computes the broadening kernel at each [w, Q] (centered on 0)
+- `get_peak` computes the broadening peak at each [w, Q] (centered on the [w, Q])
+- `get_characteristics` returns only the characteristic parameters of the kernel 
+  at each [w, Q] (such as the standard deviation of the normal distribution)
+
+```
+>>> peaks = pychop.get_peak(energy_transfer, mesh)
+>>> mesh_centered_on_0 = np.linspace(-100, 100, 1000)
+>>> kernels = pychop.get_kernel(energy_transfer, mesh_centered_on_0)
+>>> pychop.get_characteristics(energy_transfer)
+{'sigma': array([9.15987016, 7.38868127, 5.93104319])}
 ```
 
 ## Installation

docs/source/api/models/mixins.rst

@@ -0,0 +1,6 @@
+resolution\_functions.models.mixins
+-------------------------------------------
+
+.. automodule:: resolution_functions.models.mixins
+   :members:
+   :show-inheritance:

docs/source/figures/example_convolve.png

@@ -39,7 +39,7 @@ which can be used as a starting point to generate some resolution data:
 >>> import numpy as np
 >>> model = maps.get_resolution_function('PyChop_fit', chopper_package='B')
 >>> energies = np.arange(0, data.defaults['e_init'], 0.5)
->>> resolution = model(energies)
+>>> resolution = model.get_characteristics(energies)['sigma']
 
 Plotting with matplotlib:
 
@@ -91,7 +91,7 @@ All the data can then be generated using by loopiing over these variables:
 >>> for i, chopper_frequency in enumerate(test_choppers):
 ...     for j, e_init in enumerate(test_e_init):
 ...         model = maps.get_resolution_function('PyChop_fit', chopper_package='B', e_init=e_init, chopper_frequency=chopper_frequency)
-...         results[i, j, :] = model(energy_transfer)
+...         results[i, j, :] = model.get_characteristics(energy_transfer)['sigma']
 
 This can then be plotted in various ways. For example, we can check the effect
 of ``e_init`` by plotting the resolution for its different values at a constant

docs/source/figures/example_get_kernel.png

@@ -5,9 +5,9 @@ With the package :doc:`installed<installation>`, the first step is to create an
 instance of a particular :term:`version` of a particular :term:`instrument`:
 
 >>> from resolution_functions import Instrument
->>> tosca = Instrument('TOSCA')
->>> print(tosca)
-Instrument(name='TOSCA', version='TOSCA', models=['AbINS', 'book', 'vision'])
+>>> maps = Instrument.from_default('MAPS', 'MAPS')
+>>> print(maps)
+Instrument(name=MAPS, version=MAPS)
 
 To get the :term:`resolution function`, a couple choices might be necessary:
 
@@ -20,29 +20,55 @@ If we don't know what the possibilities are for the chosen instrument, the
 information can be found either in the :doc:`documentation<instruments>` or
 programmatically:
 
->>> tosca.available_models
-['AbINS', 'book', 'vision']
->>> tosca.get_model_signature('book')
-<Signature (model_name: Optional[str] = 'book', *, detector_bank: Literal['Backward', 'Forward'] = 'Backward', _) -> resolution_functions.models.tosca_book.ToscaBookModel>
+>>> maps.available_models
+['PyChop_fit']
+>>> maps.get_model_signature('PyChop_fit')
+<Signature (model_name: Optional[str] = 'PyChop_fit_v1', *, chopper_package: Literal['A', 'B', 'S'] = 'A', e_init: Annotated[ForwardRef('Optional[float]'), 'restriction=[0, 2000]'] = 500, chopper_frequency: Annotated[ForwardRef('Optional[int]'), 'restriction=[50, 601, 50]'] = 400, fitting_order: 'int' = 4, _) -> resolution_functions.models.pychop.PyChopModelFermi>
 
 With this, it is possible to make the choices and obtain the resolution function
 via the
 :py:meth:`~resolution_functions.instrument.Instrument.get_resolution_function`
 method:
 
->>> book = tosca.get_resolution_function('book', detector_bank='Forward')
+>>> pychop = maps.get_resolution_function('PyChop_fit', chopper_package='B', e_init=500, chopper_frequency=300)
 >>> print(book)
-ToscaBookModel(citation=[''])
+PyChopModelFermi(citation=[''])
 
 .. note::
 
-    The settings *must* be passed in as keyword arguments.
+    The settings and configurations *must* be passed in as keyword arguments.
 
-The resolution function is a callable object that computes the resolution at the
-given energy transfer (in meV) and/or momentum:
+The obtained model can be called (like a function) to broaden data at the
+provided combinations of energy transfer and momentum ([w, Q]), using a mesh and
+the corresponding data:
 
->>> book(100)
-0.81802604002035
 >>> import numpy as np
->>> book(np.array([100, 200, 300]))
-array([0.81802604, 1.34222267, 1.88255039])
+>>> energy_transfer = np.array([100, 200, 300])[:, np.newaxis]
+>>> data = np.array([0.6, 1.5, 0.9])
+>>> mesh = np.linspace(0, 500, 1000)
+>>> result = pychop(energy_transfer, data, mesh)
+
+which can be plotted as:
+
+.. image:: /figures/example_convolve.png
+
+However, the model also provides methods for more fundamental operations: `get_kernel` computes
+the broadening kernel at each [w, Q] (centered on 0), `get_peak` computes the
+broadening peak at each [w, Q] (centered on the [w, Q]), and
+`get_characteristics` returns only the characteristic parameters of the kernel
+at each [w, Q] (such as the standard deviation of the normal distribution):
+
+>>> pychop.get_characteristics(energy_transfer)
+{'sigma': array([9.15987016, 7.38868127, 5.93104319])}
+>>> peaks = pychop.get_peak(energy_transfer, mesh)
+
+Can be plotted as:
+
+.. image:: /figures/example_get_peak.png
+
+>>> mesh_centered_on_0 = np.linspace(-100, 100, 1000)
+>>> kernels = pychop.get_kernel(energy_transfer, mesh_centered_on_0)
+
+Can be plotted as:
+
+.. image:: /figures/example_get_kernel.png

docs/source/figures/example_get_peak.png

@@ -0,0 +1,140 @@
+"""
+Mixins providing generic implementations for
+`~resolution_functions.models.model_base.InstrumentModel` methods.
+
+The classes defined here are mixins to be used by specific models via multiple inheritance, allowing
+common code to be shared between models. Please note, however, that when doing this, the mixin
+**must** be the first base class (i.e. ``class Foo(Mixin, InstrumentModel)``) so that its
+implementation of a method overrides the abstract declaration in ``InstrumentModel``.
+"""
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import numpy as np
+from scipy.stats import norm
+
+if TYPE_CHECKING:
+    from jaxtyping import Float
+    from .model_base import InstrumentModel
+
+
+class GaussianKernel1DMixin:
+    """
+    A mixin providing the implementation for the Gaussian kernel ``get_kernel`` method.
+
+    Implements `resolution_functions.models.model_base.InstrumentModel.get_kernel` method of models
+    whose broadening can be represented by a 1D Gaussian distribution. Any model that satisfies this
+    condition should inherit the ``get_kernel`` method from this mixin instead of writing its own
+    implementation.
+
+    Technically, any model that implements the
+    `resolution_functions.models.model_base.InstrumentModel.get_characteristics` method and which
+    returns the ``sigma`` parameter in its dictionary can use this mixin to inherit the Gaussian
+    ``get_kernel`` method. However, it is recommended that only models that actually model a
+    Gaussian kernel should use this mixin.
+    """
+    def get_kernel(self,
+                   points: Float[np.ndarray, 'sample dimension=1'],
+                   mesh: Float[np.ndarray, 'mesh'],
+                   ) -> Float[np.ndarray, 'sample mesh']:
+        """
+        Computes the Gaussian kernel centered on zero on the provided `mesh` at each (energy
+        transfer or momentum scalar) point in `points`.
+
+        Parameters
+        ----------
+        points
+            The energy transfer or momentum scalar for which to compute the kernel. This *must* be
+            a Nx1 2D array where N is the number of w/Q values.
+        mesh
+            The mesh on which to evaluate the kernel. A 1D array which must encompass 0.
+
+        Returns
+        -------
+        kernels
+            The Gaussian kernel at w/Q point as given by this model, computed on the
+            `mesh` and centered on zero. This is a 2D N x M array where N is the number of w/Q
+            values and M is the length of the `mesh` array.
+        """
+        return self._get_kernel(points, mesh, 0.)
+
+    def get_peak(self,
+                 points: Float[np.ndarray, 'sample dimension=1'],
+                 mesh: Float[np.ndarray, 'mesh'],
+                 ) -> Float[np.ndarray, 'sample mesh']:
+        """
+        Computes the Gaussian broadening peak on the provided `mesh` at each (energy transfer or
+        momentum scalar) point in `points`.
+
+        Parameters
+        ----------
+        points
+            The energy transfer or momentum scalar for which to compute the peak. This *must* be a
+            Nx1 2D array where N is the number of w/Q values.
+        mesh
+            The mesh on which to evaluate the kernel. This is a 1D array which *must* span the
+            `points` w\Q space of interest.
+
+        Returns
+        -------
+        peaks
+            The Gaussian peak at each w/Q point in `points` as given by this model, computed on the
+            `mesh` and centered on the corresponding w/Q. This is a 2D N x M array where
+            N is the number of w/Q values and M is the length of the `mesh` array.
+        """
+        return self._get_kernel(points, mesh, points)
+
+    def _get_kernel(self: InstrumentModel,
+                    points: Float[np.ndarray, 'sample dimension=1'],
+                    mesh: Float[np.ndarray, 'mesh'],
+                    displacement: float | Float[np.ndarray, 'sample'] = 0.
+                    ) -> Float[np.ndarray, 'sample mesh']:
+        """Computes the kernel using the specified `displacement`."""
+        new_mesh = np.zeros((len(points), len(mesh)))
+        new_mesh[:, :] = mesh
+
+        sigma = self.get_characteristics(points)['sigma']
+        return norm.pdf(new_mesh, loc=displacement, scale=sigma[:, np.newaxis])
+
+
+class SimpleBroaden1DMixin:
+    """
+    A mixin providing the most simple implementation for the ``convolve`` method.
+
+    Implements `resolution_functions.models.model_base.InstrumentModel.convolve` method in the
+    most simple and basic way - the dot product between the matrix of kernels (obtained from the
+    ``get_kernel`` method) and the intensities.
+
+    This implementation should be mostly used as a reference method given that it is correct but
+    inefficient. It should be able to work with any model, so it may be used when other
+    implementations are unavailable.
+    """
+    def broaden(self: InstrumentModel,
+                points: Float[np.ndarray, 'sample dimension=1'],
+                data: Float[np.ndarray, 'data'],
+                mesh: Float[np.ndarray, 'mesh'],
+                ) -> Float[np.ndarray, 'spectrum']:
+        """
+        Broadens the `data` on the full `mesh` using the straightforward scheme.
+
+        Parameters
+        ----------
+        points
+            The independent variable (energy transfer or momentum scalar) whose `data` to broaden.
+            This *must* be a ``sample`` x 1 2D array where ``sample`` is the number of w/Q values
+            for which there is `data`. Therefore, the ``sample`` dimension *must* match the length
+            of the `data` array.
+        data
+            The intensities at the w/Q `points`.
+        mesh
+            The mesh to use for the broadening. This is a 1D array which *must* span the entire
+            `points` space of interest.
+
+        Returns
+        -------
+        spectrum
+            The broadened spectrum. This is a 1D array of the same length as `mesh`.
+        """
+        kernels = self.get_peak(points, mesh)
+        return np.einsum('i,ij...->j...', data, kernels)