Convenience functions for Matrix, Vector and NormalizationParameters

ompy/matrix.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 import logging
+import warnings
 import copy
 import numpy as np
 import matplotlib.pyplot as plt
@@ -218,6 +219,10 @@ def save(self, path: Union[str, Path], filetype: Optional[str] = None,
             filetype = filetype_from_suffix(path)
         filetype = filetype.lower()
 
+        if self.std is not None:
+            warnings.warn(UserWarning(
+                "The `std` attribute of Matrix class cannot be saved to file."))
+
         if filetype == 'numpy':
             save_numpy_2D(self.values, self.Eg, self.Ex, path)
         elif filetype == 'txt':
@@ -307,7 +312,7 @@ def plot(self, *, ax: Any = None,
             ax.tick_params(axis='x', rotation=40)
             ax.yaxis.set_major_locator(MeshLocator(self.Ex))
         # ax.xaxis.set_major_locator(ticker.FixedLocator(self.Eg, nbins=10))
-        #fix_pcolormesh_ticks(ax, xvalues=self.Eg, yvalues=self.Ex)
+        # fix_pcolormesh_ticks(ax, xvalues=self.Eg, yvalues=self.Ex)
 
         ax.set_title(title if title is not None else self.state)
         ax.set_xlabel(r"$\gamma$-ray energy $E_{\gamma}$")
@@ -612,7 +617,8 @@ def line_mask(self, E1: Iterable[float],
     def trapezoid(self, Ex_min: float, Ex_max: float,
                   Eg_min: float, Eg_max: Optional[float] = None,
                   inplace: bool = True) -> Optional[Matrix]:
-        """Create a trapezoidal cut or mask delimited by the diagonal of the matrix
+        """Create a trapezoidal cut or mask delimited by the diagonal of the
+            matrix
 
         Args:
             Ex_min: The bottom edge of the trapezoid
@@ -729,8 +735,8 @@ def diagonal_elements(self) -> Iterator[Tuple[int, int]]:
             entries with `Eg > Ex + dE`.
         Args:
             mat: The matrix to iterate over
-            Iterator[Tuple[int, int]]: Indicies (i, j) over the last non-zero (=diagonal)
-            elements.
+            Iterator[Tuple[int, int]]: Indicies (i, j) over the last
+                non-zero(=diagonal) elements.
         """
         return diagonal_elements(self.values)
 
@@ -863,6 +869,7 @@ def __matmul__(self, other: Matrix) -> Matrix:
         result.values = [email protected]
         return result
 
+
 class MeshLocator(ticker.Locator):
     def __init__(self, locs, nbins=10):
         'place ticks on the i-th data points where (i-offset)%base==0'

ompy/models.py

@@ -391,6 +391,12 @@ class NormalizationParameters(Model):
     """Storage for normalization parameters + some convenience functions
     """
 
+    #: Element number of the nucleus
+    Z: Optional[int] = field(default=None,
+            metadata="Element number of the nucleus")  # noqa
+    #: Mass number of the nucleus
+    A: Optional[int] = field(default=None,
+            metadata="Mass number of the nucleus")  # noqa
     #: Average s-wave resonance spacing D0 [eV]
     D0: Optional[Tuple[float, float]] = field(default=None,
             metadata='Average s-wave resonance spacing D0 [eV]')  # noqa

ompy/vector.py

@@ -80,6 +80,9 @@ def __init__(self, values: Optional[Iterable[float]] = None,
             self.load(path)
         self.verify_integrity()
 
+    def __len__(self):
+        return len(self.values)
+
     def verify_integrity(self, check_equidistant: bool = False):
         """ Verify the internal consistency of the vector
 

tests/test_matrix.py

@@ -23,10 +23,12 @@ def ones(shape: Tuple[int, int]) -> om.Matrix:
     mat = np.tril(mat)
     return om.Matrix(values=mat)
 
+
 @pytest.fixture()
 def Si28():
     return om.example_raw('Si28')
 
+
 @pytest.mark.parametrize(
         "axis,Emin,Emax,shape",
         [('Ex', None, None, (10, 10)),
@@ -78,7 +80,7 @@ def test_index(E, index):
     assert mat.index_Ex(E) == index
 
 
-@pytest.mark.filterwarnings('ignore:divide by zero encountered in true_divide:RuntimeWarning')
+@pytest.mark.filterwarnings('ignore:divide by zero encountered in true_divide:RuntimeWarning')  # noqa
 def test_numericals():
     E = np.array([0, 1, 2])
     values1 = np.array([[0, 1, 2.], [-2, 1, 2.],  [2, 3, -10.]])
@@ -98,6 +100,7 @@ def test_numericals():
     assert_equal((matrix2@matrix1).values, values2@values1)
     assert_equal((matrix1@matrix2).values, values1@values2)
 
+
 @pytest.mark.parametrize(
         "Ex,Eg",
         [(np.linspace(0, 10., num=10), np.linspace(10, 20., num=15)),
@@ -120,6 +123,19 @@ def test_bin_shift(Ex, Eg):
     assert_almost_equal(Ex, mat.Ex)
     assert_almost_equal(Eg, mat.Eg)
 
+
+@pytest.mark.parametrize(
+        "Ex,Eg",
+        [(np.linspace(0, 10., num=10), np.linspace(10, 20., num=15)),
+         ([0, 1, 2, 3, 7, 10.], [0, 1, 2, 3, 80, 90.])
+         ])
+def test_save_warning(Ex, Eg):
+    values = np.ones((len(Ex), len(Eg)), dtype="float")
+    mat = om.Matrix(values=values, Ex=Ex, Eg=Eg, std=0.5*values)
+    with pytest.warns(UserWarning):
+        mat.save("/tmp/mat.npy")
+
+
 # This does not work as of now...
 # def test_mutable():
 #     E = np.array([0, 1, 2])

tests/test_vector.py

@@ -24,6 +24,14 @@ def test_init():
         om.Vector(vals, [1, 2, 3, 4, 5])
 
 
+def test_len():
+    N = 100
+    E = np.linspace(0, 1, N)
+    vals = np.linspace(2, 3.4, N)
+    vec = om.Vector(values=vals, E=E)
+    assert_equal(len(vec), N)
+
+
 def test_save_load_no_std():
     E = np.linspace(0, 1, 100)
     vals = np.linspace(2, 3.4, 100)
@@ -184,7 +192,8 @@ def test_cut():
     assert_equal(vector.E, Ecut)
     assert_equal(vector.values, valcut)
 
-@pytest.mark.filterwarnings('ignore:divide by zero encountered in true_divide:RuntimeWarning')
+
+@pytest.mark.filterwarnings('ignore:divide by zero encountered in true_divide:RuntimeWarning')  # noqa
 def test_numericals():
     E = np.array([0, 1, 2])
     values1 = np.array([0, 1, -2.])