Modernise: importlib, pathlib, omega_q becomes "points"

Refactored instrument tests a bit to patch importlib.resources.files()
as a fixture. This is a bit more risky than changing our own
INSTRUMENT_DATA_PATH variable; if we run into conflict with other uses
for this function then we might need to reintroduce the variable.

Alternatively we could consider using an environment variable so
additional paths can be set at runtime.

Author: ajjackson

src/resins/instrument.py

@@ -6,7 +6,7 @@
 
 from collections import ChainMap
 import dataclasses
-import os
+import importlib
 
 import numpy as np
 import yaml
@@ -18,8 +18,6 @@
     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),
@@ -360,7 +358,7 @@ def _get_file(instrument_name: str) -> tuple[str, Union[str, None]]:
                 f'"{instrument_name}" is not a valid instrument name. Only the following instruments are '
                 f'supported: {list(INSTRUMENT_MAP)}')
 
-        return os.path.join(INSTRUMENT_DATA_PATH, file_name), implied_version
+        return str(importlib.resources.files("resins.instrument_data") / file_name), implied_version
 
     def get_model_data(self, model_name: Optional[str] = None, **kwargs) -> ModelData:
         """

src/resins/models/lookuptables.py

@@ -8,14 +8,13 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
+import importlib
 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
@@ -67,16 +66,14 @@ class ScaledTabulatedModel(SimpleBroaden1DMixin, InstrumentModel):
     """
     input = ('energy_transfer',)
 
-    data_class: ClassVar[type[ScaledTabultedModelData]] = ScaledTabulatedModelData
+    data_class: ClassVar[type[ScaledTabulatedModelData]] = 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.data = np.load(importlib.resources.files("resins.instrument_data") / model_data.npz)
 
         self.polynomial = Polynomial(coef=self.data["coef"],
                                      domain=self.data["domain"],
@@ -110,7 +107,7 @@ def get_characteristics(self, omega_q: Float[np.ndarray, 'energy_transfer dimens
         return {'sigma': self.polynomial(omega_q[:, 0])}
 
     def get_kernel(self,
-                   omega_q: Float[np.ndarray, 'sample dimension=1'],
+                   points: Float[np.ndarray, 'sample dimension=1'],
                    mesh: Float[np.ndarray, 'mesh'],
                    ) -> Float[np.ndarray, 'sample mesh']:
 
@@ -130,14 +127,14 @@ def get_kernel(self,
         return interp_kernels
 
     def get_peak(self,
-                 omega_q: Float[np.ndarray, 'sample dimension=1'],
+                 points: 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)
+            self.get_kernel(np.array([point]), shifted_mesh)
+            for point, shifted_mesh in zip(points, shifted_meshes)
         ]
 
         return np.array(np.vstack(shifted_kernels))

tests/unit_tests/test_instrument.py

@@ -1,6 +1,6 @@
 from dataclasses import dataclass
 import inspect
-import os
+from pathlib import Path
 import typing
 
 import pytest
@@ -11,8 +11,8 @@
 from resins.models.model_base import ModelData, InstrumentModel
 
 
-TEST_DIR = os.path.dirname(os.path.abspath(__file__))
-FAKE_YAML = os.path.join(TEST_DIR, 'fake_instrument.yaml')
+TEST_DIR = Path(__file__).parent
+FAKE_YAML = TEST_DIR / 'fake_instrument.yaml'
 
 
 @dataclass(init=True, repr=True, frozen=True, slots=True)
@@ -68,13 +68,6 @@ def broaden(self, omega_q, data, mesh):
 def mock_models():
     return {'mock': MockModel}
 
-@pytest.fixture
-def mock_instrument_map():
-    return {
-        'TEST': ('fake_instrument.yaml', None),
-        'ALIAS': ('fake_instrument.yaml', 'VERSION1'),
-    }
-
 
 @pytest.fixture(scope='module')
 def data():
@@ -127,35 +120,45 @@ def test_instrument(data):
     )
 
 
-def test_available_instruments(mock_instrument_map, mocker):
-    mocker.patch('resins.instrument.INSTRUMENT_MAP', mock_instrument_map)
-    assert i.Instrument.available_instruments() == ['TEST', 'ALIAS']
+@pytest.fixture
+def patch_importlib_files(mocker):
+    """Make all calls to importlib.resources.files() return TEST_DIR
+
+    This is used to replace resins.instrument_data with test data
+    """
+    mocker.patch('importlib.resources.files', lambda _: Path(TEST_DIR))    
 
 
-def test_private_available_versions(mock_instrument_map, mocker):
-    mocker.patch('resins.instrument.INSTRUMENT_MAP', mock_instrument_map)
-    mocker.patch('resins.instrument.INSTRUMENT_DATA_PATH', TEST_DIR)
+@pytest.fixture
+def patch_instrument_map(mocker):
+    """Replace main INSTRUMENT_MAP with test examples"""
 
+    instrument_map = {
+        'TEST': ('fake_instrument.yaml', None),
+        'ALIAS': ('fake_instrument.yaml', 'VERSION1'),
+    }
+    mocker.patch('resins.instrument.INSTRUMENT_MAP', instrument_map)
+
+
+def test_available_instruments(patch_instrument_map):
+    assert i.Instrument.available_instruments() == ['TEST', 'ALIAS']
+
+
+def test_private_available_versions(patch_instrument_map, patch_importlib_files):
     actual_versions, actual_default = i.Instrument._available_versions(FAKE_YAML)
 
     assert actual_default == 'TEST'
     assert actual_versions == ['VERSION1', 'TEST']
 
 
-def test_available_versions(mock_instrument_map, mocker):
-    mocker.patch('resins.instrument.INSTRUMENT_MAP', mock_instrument_map)
-    mocker.patch('resins.instrument.INSTRUMENT_DATA_PATH', TEST_DIR)
-
+def test_available_versions(patch_instrument_map, patch_importlib_files):
     actual_versions, actual_default = i.Instrument.available_versions('TEST')
 
     assert actual_default == 'TEST'
     assert actual_versions == ['VERSION1', 'TEST']
 
 
-def test_available_versions_alias(mock_instrument_map, mocker):
-    mocker.patch('resins.instrument.INSTRUMENT_MAP', mock_instrument_map)
-    mocker.patch('resins.instrument.INSTRUMENT_DATA_PATH', TEST_DIR)
-
+def test_available_versions_alias(patch_instrument_map, patch_importlib_files):
     actual_versions, actual_default = i.Instrument.available_versions('ALIAS')
 
     assert actual_default == 'VERSION1'
@@ -185,10 +188,7 @@ def test_from_file_invalid_version():
         i.Instrument.from_file(FAKE_YAML, 'INVALID_VERSION')
 
 
-def test_from_default(data, mock_instrument_map, mocker):
-    mocker.patch('resins.instrument.INSTRUMENT_MAP', mock_instrument_map)
-    mocker.patch('resins.instrument.INSTRUMENT_DATA_PATH', TEST_DIR)
-
+def test_from_default(data, patch_instrument_map, patch_importlib_files):
     instrument = i.Instrument.from_default('TEST', 'VERSION1')
 
     assert isinstance(instrument, i.Instrument)
@@ -197,10 +197,7 @@ def test_from_default(data, mock_instrument_map, mocker):
     assert instrument._models == data['version']['VERSION1']['models']
 
 
-def test_from_default_default(data, mock_instrument_map, mocker):
-    mocker.patch('resins.instrument.INSTRUMENT_MAP', mock_instrument_map)
-    mocker.patch('resins.instrument.INSTRUMENT_DATA_PATH', TEST_DIR)
-
+def test_from_default_default(data, patch_instrument_map, patch_importlib_files):
     instrument = i.Instrument.from_default('TEST')
 
     assert isinstance(instrument, i.Instrument)
@@ -211,20 +208,14 @@ def test_from_default_default(data, mock_instrument_map, mocker):
 
 @pytest.mark.parametrize("name,expected_path,implied_ver",
                          [('TEST', FAKE_YAML, None), ('ALIAS', FAKE_YAML, 'VERSION1')])
-def test_get_file(name, expected_path, implied_ver, mock_instrument_map, mocker):
-    mocker.patch('resins.instrument.INSTRUMENT_MAP', mock_instrument_map)
-    mocker.patch('resins.instrument.INSTRUMENT_DATA_PATH', TEST_DIR)
-
+def test_get_file(name, expected_path, implied_ver, patch_instrument_map, patch_importlib_files):
     actual_path, actual_version = i.Instrument._get_file(name)
 
-    assert actual_path == expected_path
+    assert Path(actual_path) == expected_path
     assert actual_version == implied_ver
 
 
-def test_get_file_invalid(mock_instrument_map, mocker):
-    mocker.patch('resins.instrument.INSTRUMENT_MAP', mock_instrument_map)
-    mocker.patch('resins.instrument.INSTRUMENT_DATA_PATH', TEST_DIR)
-
+def test_get_file_invalid(patch_instrument_map, patch_importlib_files):
     with pytest.raises(i.InvalidInstrumentError):
         i.Instrument._get_file('INVALID_INSTRUMENT')