Merge pull request #262 from FlorianPfaff/feature/hypertoroidal_grid

Added HypertoroidalGridDistribution

Author: FlorianPfaff

pyrecest/_backend/__init__.py

@@ -234,6 +234,7 @@ def get_backend_name():
         "sqrtm",
         "svd",
         "matrix_rank",
+        "block_diag",  # For PyRecEst
     ],
     "random": [
         "choice",

pyrecest/_backend/jax/linalg.py

@@ -18,6 +18,7 @@
     expm,
     sqrtm,
     polar,
+    block_diag,  # For PyRecEst
 )
 
 unsupported_functions = [

pyrecest/_backend/numpy/linalg.py

@@ -14,7 +14,10 @@
     norm,
     svd,
 )
-from scipy.linalg import expm
+from scipy.linalg import (
+    expm,
+    block_diag,  # For PyRecEst
+)
 
 from .._shared_numpy.linalg import (
     fractional_matrix_power,

pyrecest/_backend/pytorch/__init__.py

@@ -63,8 +63,6 @@
     isinf,
     deg2rad,
     argsort,
-    max,
-    min,
     roll,
     dstack,
     vmap,
@@ -357,11 +355,12 @@ def shape(val):
     return val.shape
 
 
-def amax(a, axis=None):
+def max(a, axis=None):
     if axis is None:
         return _torch.max(array(a))
     return _torch.max(array(a), dim=axis).values
 
+amax=max
 
 def maximum(a, b):
     return _torch.max(array(a), array(b))
@@ -811,10 +810,11 @@ def sort(a, axis=-1):
     return sorted_a
 
 
-def amin(a, axis=-1):
+def min(a, axis=-1):
     (values, _) = _torch.min(a, dim=axis)
     return values
 
+amin = min
 
 def take(a, indices, axis=0):
     if not _torch.is_tensor(indices):

pyrecest/_backend/pytorch/linalg.py

@@ -15,6 +15,7 @@
     solve,
 )
 from torch.linalg import matrix_exp as expm
+from torch import block_diag  # For PyRecEst
 
 from .._backend_config import np_atol as atol
 from ..numpy import linalg as _gsnplinalg

pyrecest/distributions/abstract_grid_distribution.py

@@ -6,6 +6,7 @@
 
 # pylint: disable=redefined-builtin,no-name-in-module,no-member
 from pyrecest.backend import abs, any, mean
+from math import prod
 
 from .abstract_distribution_type import AbstractDistributionType
 
@@ -24,7 +25,10 @@ def __init__(
         assert (
             not grid_type == "custom" or grid is not None
         )  # if grid_type is custom, grid needs to be given
-        assert grid is None or grid.shape == () or grid.shape[0] == grid_values.shape[0]
+        assert (
+            # Use builtin prod because .shape is a tuple of ints
+            grid is None or grid.shape == () or grid.shape[0] == prod(grid_values.shape)
+        )
         assert (
             grid is None or grid.shape == () or grid.ndim == 1 or grid.shape[1] == dim
         )