Merge pull request #515 from flatironinstitute/development

Development

Author: BalzaniEdoardo

src/nemos/basis/_basis.py

@@ -94,15 +94,18 @@ def min_max_rescale_samples(
     else:
         nanmin, nanmax, asarray, where = np.nanmin, np.nanmax, np.asarray, np.where
 
-    # if not normalize all array
-    vmin = nanmin(sample_pts, axis=0) if bounds is None else bounds[0]
-    vmax = nanmax(sample_pts, axis=0) if bounds is None else bounds[1]
+    if bounds is None:
+        vmin = nanmin(sample_pts, axis=0)
+        vmax = nanmax(sample_pts, axis=0)
+    elif isinstance(bounds[0], (tuple, list)):
+        vmin = asarray([b[0] for b in bounds])
+        vmax = asarray([b[1] for b in bounds])
+    else:
+        vmin, vmax = bounds[0], bounds[1]
+
     scaling = asarray(vmax - vmin)
-    # do not normalize if samples contain a single value (in which case vmax=vmin)
     scaling = where(scaling == 0, 1.0, scaling)
-    sample_pts -= vmin
-    sample_pts /= scaling
-
+    sample_pts = (sample_pts - vmin) / scaling
     return sample_pts, scaling
 
 
@@ -115,7 +118,7 @@ def _is_single_bound(bounds) -> bool:
     if len(bounds) != 2:
         return False
     # Single bound has numeric elements; multiple bounds have tuple/None elements
-    return all(isinstance(b, (int, float, np.number)) for b in bounds)
+    return all(isinstance(b, (int, float, np.number)) or b is None for b in bounds)
 
 
 def _fill_bounds(b):
@@ -184,7 +187,7 @@ class Basis(Base, abc.ABC, BasisTransformerMixin):
     def __init__(
         self,
     ) -> None:
-        self._n_input_dimensionality = getattr(self, "_n_input_dimensionality", 0)
+        self._n_inputs = getattr(self, "_n_inputs", 0)
 
         # specified only after inputs/input shapes are provided
         self._input_shape_product = getattr(self, "_input_shape_product", None)
@@ -635,8 +638,8 @@ def evaluate(self, *xi: ArrayLike | Tsd | TsdFrame | TsdTensor) -> FeatureMatrix
         """
         X = np.hstack(
             (
-                self.basis1.evaluate(*xi[: self.basis1._n_input_dimensionality]),
-                self.basis2.evaluate(*xi[self.basis1._n_input_dimensionality :]),
+                self.basis1.evaluate(*xi[: self.basis1._n_inputs]),
+                self.basis2.evaluate(*xi[self.basis1._n_inputs :]),
             )
         )
         return X
@@ -691,8 +694,8 @@ def _compute_features(
         )
         X = hstack_pynapple(
             (
-                comp_feature_1(*xi[: self.basis1._n_input_dimensionality]),
-                comp_feature_2(*xi[self.basis1._n_input_dimensionality :]),
+                comp_feature_1(*xi[: self.basis1._n_inputs]),
+                comp_feature_2(*xi[self.basis1._n_inputs :]),
             ),
         )
         return X
@@ -1048,8 +1051,8 @@ def evaluate(self, *xi: ArrayLike | Tsd | TsdFrame | TsdTensor) -> FeatureMatrix
         """
         # evaluate preserves the shape of the input arrays
         shape = xi[0].shape
-        x1 = self.basis1.evaluate(*xi[: self.basis1._n_input_dimensionality])
-        x2 = self.basis2.evaluate(*xi[self.basis1._n_input_dimensionality :])
+        x1 = self.basis1.evaluate(*xi[: self.basis1._n_inputs])
+        x2 = self.basis2.evaluate(*xi[self.basis1._n_inputs :])
         # Required in case xi.shape[-1] == 0
         # For example, in a multiplication with Zero basis
         x1_shape = math.prod(x1.shape[:-1])
@@ -1098,8 +1101,8 @@ def _compute_features(
         comp_feature_2 = getattr(
             self.basis2, "_compute_features", self.basis2.compute_features
         )
-        x1 = comp_feature_1(*xi[: self.basis1._n_input_dimensionality])
-        x2 = comp_feature_2(*xi[self.basis1._n_input_dimensionality :])
+        x1 = comp_feature_1(*xi[: self.basis1._n_inputs])
+        x2 = comp_feature_2(*xi[self.basis1._n_inputs :])
         # multiplicative basis inputs are of the same shape, checked and
         # set just before the call to this method
         n_samples = x1.shape[0]

src/nemos/basis/_basis_mixin.py

@@ -10,7 +10,15 @@
 from contextlib import contextmanager
 from functools import wraps
 from itertools import chain
-from typing import TYPE_CHECKING, Any, Generator, Literal, Optional, Tuple, Union
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Generator,
+    List,
+    Literal,
+    Optional,
+    Tuple,
+)
 
 import jax
 import jax.numpy as jnp
@@ -548,21 +556,19 @@ def _compute_features(self, *xi: ArrayLike | Tsd | TsdFrame | TsdTensor):
     @support_pynapple(conv_type="jax")
     def _apply_fill_value(self, *xi: ArrayLike, out: NDArray) -> jax.Array:
         """Apply fill value to out-of-bounds samples."""
-        # Use jnp.where for JAX compatibility
         to_fill = jnp.any(
             jnp.stack(
                 [
                     jnp.any(
-                        (jnp.reshape(x, (x.shape[0], -1)) < self.bounds[0])
-                        | (jnp.reshape(x, (x.shape[0], -1)) > self.bounds[1]),
+                        (jnp.reshape(x, (x.shape[0], -1)) < lo)
+                        | (jnp.reshape(x, (x.shape[0], -1)) > hi),
                         axis=1,
                     )
-                    for x in xi
+                    for x, (lo, hi) in zip(xi, self._get_bounds_per_dim(), strict=True)
                 ]
             ),
             axis=0,
         )
-        # Reshape to_fill to broadcast correctly: (n_samples,) -> (n_samples, 1, 1, ...)
         to_fill_broadcast = to_fill.reshape(to_fill.shape[0], *([1] * (out.ndim - 1)))
         return jnp.where(to_fill_broadcast, self.fill_value, out)
 
@@ -605,10 +611,46 @@ def _set_input_independent_states(self) -> EvalBasisMixin:
         return self
 
     @property
-    def bounds(self):
-        """Range of values covered by the basis."""
+    def bounds(self) -> List[Tuple[float, float]] | Tuple[float, float] | None:
+        """Returns bounds, as provided."""
         return self._bounds
 
+    def _get_bounds_per_dim(self) -> List[Tuple[float, float]]:
+        """Return bounds,  broadcast to one pair per input dimension."""
+        if self._bounds is None or isinstance(self._bounds[0], (int, float)):
+            return [self._bounds] * self._n_inputs
+        return list(self._bounds)
+
+    @bounds.setter
+    def bounds(self, values):
+        if values is None:
+            self._bounds = None
+            return
+
+        if isinstance(values, np.ndarray):
+            values = values.tolist()
+
+        if not isinstance(values, (list, tuple)):
+            raise TypeError(
+                f"Invalid bounds ``{values}`` provided, "
+                "bounds should be one or multiple tuples of 2 floats, "
+                "matching the inputs of the basis.\n"
+            )
+
+        # Validate: single (lo, hi) pair or one per dimension
+        if len(values) == 2 and all(not isinstance(v, (list, tuple)) for v in values):
+            # Single pair
+            self._bounds = self._format_bounds(values)
+        else:
+            # One pair per dimension
+            if len(values) != self._n_inputs:
+                raise ValueError(
+                    f"Invalid bounds ``{values}`` provided, "
+                    "bounds should be one or multiple tuples of 2 floats, "
+                    "matching the inputs of the basis.\n"
+                )
+            self._bounds = tuple(self._format_bounds(v) for v in values)
+
     @staticmethod
     def _format_bounds(values: Any) -> Tuple[Any, Exception | None]:
         """Check bounds and cast to tuple."""
@@ -640,19 +682,6 @@ def _format_bounds(values: Any) -> Tuple[Any, Exception | None]:
 
         return values
 
-    @bounds.setter
-    def bounds(self, values: Union[None, Tuple[float, float]]):
-        """Setter for bounds."""
-        if values is None:
-            self._bounds = None
-            return
-        values = self._format_bounds(values)
-        if values is not None and len(values) != 2:
-            raise ValueError(
-                f"The provided `bounds` must be of length two. Length {len(values)} provided instead!"
-            )
-        self._bounds = values
-
 
 class ConvBasisMixin:
     """Mixin class for convolutional basis."""
@@ -883,7 +912,7 @@ def __init__(
         self, basis1: BasisMixin, basis2: BasisMixin, label: Optional[str] = None
     ):
         # number of input arrays that the basis receives
-        self._n_input_dimensionality = infer_input_dimensionality(
+        self._n_inputs = infer_input_dimensionality(
             basis1
         ) + infer_input_dimensionality(basis2)
 

src/nemos/basis/_composition_utils.py

@@ -165,8 +165,8 @@ def _composite_basis_setter_logic(new: BasisMixin, current: BasisMixin):
     # Carry-on input shape info if dimensions match
     for attr in ("_input_shape_product", "_input_shape_"):
         if getattr(new, attr, None) is None and getattr(
-            new, "_n_input_dimensionality", None
-        ) == getattr(current, "_n_input_dimensionality", None):
+            new, "_n_inputs", None
+        ) == getattr(current, "_n_inputs", None):
             setattr(new, attr, getattr(current, attr, None))
     return new
 
@@ -226,10 +226,10 @@ def _atomic_basis_label_setter_logic(
 def infer_input_dimensionality(bas: BasisMixin) -> int:
     """Infer input dimensionality from compute_features signature.
 
-    If `_n_input_dimensionality` return the attribute, otherwise return
+    If `_n_inputs` return the attribute, otherwise return
     the number of fixed arguments in `compute_features`.
     """
-    n_input_dim = getattr(bas, "_n_input_dimensionality", None)
+    n_input_dim = getattr(bas, "_n_inputs", None)
     if n_input_dim is None:
         # infer from compute_features (facilitate custom basis compatibility).
         # assume compute_features is always implemented.
@@ -398,13 +398,13 @@ def set_input_shape(bas, *xi):
         else 1
     )
     # get the attribute if available
-    n_input_dim = getattr(bas, "_n_input_dimensionality", n_args)
+    n_input_dim = getattr(bas, "_n_inputs", n_args)
 
     if len(xi) == 1 and xi[0] is None:
         xi = (None,) * n_input_dim
 
     elif len(xi) != n_input_dim:
-        expected_inputs = getattr(bas, "_n_input_dimensionality", 1)
+        expected_inputs = getattr(bas, "_n_inputs", 1)
         raise ValueError(
             f"set_input_shape expects {expected_inputs} input"
             f"{'s' if expected_inputs > 1 else ''}, but {len(xi)} were provided."
@@ -436,7 +436,7 @@ def set_input_shape(bas, *xi):
         if hasattr(bas.basis1, "compute_features")
         else 1
     )
-    n_input_dim_1 = getattr(bas.basis1, "_n_input_dimensionality", n_args_1)
+    n_input_dim_1 = getattr(bas.basis1, "_n_inputs", n_args_1)
 
     out1 = set_input_shape(bas.basis1, *xi[:n_input_dim_1])
     out2 = set_input_shape(bas.basis2, *xi[n_input_dim_1:])

src/nemos/basis/_custom_basis.py

@@ -219,7 +219,7 @@ def __init__(
 
         self._input_shape_product = None
 
-        self._n_input_dimensionality = infer_input_dimensionality(self)
+        self._n_inputs = infer_input_dimensionality(self)
         self._n_basis_funcs = len(self.funcs)
 
         self.basis_kwargs = basis_kwargs
@@ -253,9 +253,9 @@ def funcs(self, val: Iterable[Callable[[NDArray, ...], NDArray]]):
         if not all(isinstance(f, Callable) for f in val):
             raise ValueError("User must provide an iterable of callable.")
 
-        if hasattr(self, "_n_input_dimensionality"):
+        if hasattr(self, "_n_inputs"):
             inp_dim = sum(count_positional_and_var_args(f)[0] for f in val)
-            if inp_dim != self._n_input_dimensionality:
+            if inp_dim != self._n_inputs:
                 raise ValueError(
                     "The number of input time series required by the CustomBasis must be consistent. "
                     "Redefine a CustomBasis for a different number of inputs."
@@ -718,10 +718,10 @@ def input_shape(
         """
         input_shape = self._input_shape_
         if input_shape is None:
-            if self._n_input_dimensionality == 1:
+            if self._n_inputs == 1:
                 return None
             else:
-                return [None] * self._n_input_dimensionality
-        if self._n_input_dimensionality == 1:
+                return [None] * self._n_inputs
+        if self._n_inputs == 1:
             return input_shape[0]
-        return input_shape * self._n_input_dimensionality
+        return input_shape * self._n_inputs

src/nemos/basis/_decaying_exponential.py

@@ -39,13 +39,11 @@ def __init__(
         decay_rates: NDArray[np.floating],
         label: Optional[str] = "OrthExponentialBasis",
     ):
+        self._n_inputs = 1
         AtomicBasisMixin.__init__(self, n_basis_funcs=n_basis_funcs, label=label)
-        Basis.__init__(
-            self,
-        )
+        Basis.__init__(self)
         self.decay_rates = decay_rates
         self._check_rates()
-        self._n_input_dimensionality = 1
 
     @property
     def decay_rates(self):

src/nemos/basis/_fourier_basis.py

@@ -353,7 +353,7 @@ def __init__(
         ) = "no-intercept",
         label: Optional[str] = None,
     ) -> None:
-        self._n_input_dimensionality = self._check_ndim(ndim)
+        self._n_inputs = self._check_ndim(ndim)
         self.frequencies = frequencies
         self.frequency_mask = frequency_mask
         Basis.__init__(
@@ -462,8 +462,8 @@ def frequency_mask(
 
             values = values.astype(bool)
 
-            if not values.ndim == self._n_input_dimensionality:
-                ndim = self._n_input_dimensionality
+            if not values.ndim == self._n_inputs:
+                ndim = self._n_inputs
                 raise ValueError(
                     f"The frequency mask for a {ndim}-dimensional Fourier basis "
                     f"must be an {ndim}-dimensional array of 0s and 1s. "
@@ -513,21 +513,19 @@ def frequencies(
         self,
         frequencies: int | tuple[int, int] | list[int] | list[tuple[int, int]],
     ) -> None:
-        ndim = self._n_input_dimensionality
+        ndim = self._n_inputs
 
         if isinstance(frequencies, Number) and (frequencies == int(frequencies)):
             frequencies = [arange_constructor(frequencies)] * ndim
 
         elif isinstance(frequencies, tuple):
-            frequencies = _process_tuple_frequencies(
-                frequencies, self._n_input_dimensionality
-            )
+            frequencies = _process_tuple_frequencies(frequencies, self._n_inputs)
 
         elif is_at_least_1d_numpy_array_like(frequencies):
             frequencies = _check_and_sort_frequencies(*([frequencies] * ndim))
 
         elif isinstance(frequencies, list):
-            if len(frequencies) != self._n_input_dimensionality:
+            if len(frequencies) != self._n_inputs:
                 raise ValueError(
                     "Length of frequencies list must match input dimensionality."
                 )
@@ -589,7 +587,7 @@ def masked_frequencies(self) -> jnp.ndarray:
     @property
     def ndim(self):
         """The dimensionality of the basis."""
-        return self._n_input_dimensionality
+        return self._n_inputs
 
     @staticmethod
     def _check_ndim(ndim: int) -> int:
@@ -630,25 +628,24 @@ def evaluate(  # call these _evaluate
 
         """
         shape = sample_pts[0].shape
-
-        bounds = getattr(self, "bounds", None)
-        if bounds is None:
-            bounds = (None,) * self._n_input_dimensionality
+        bounds = self._get_bounds_per_dim()
 
         # min/max rescale to [0,1]:
         # The function does so over the time axis (each extra dim is
         # normalized independently)
         def _flat_samples_to_angles(xs):
             scaled_samples = jax.tree_util.tree_map(
-                lambda x, b: 2
-                * jnp.pi
-                * self._shift_angles(min_max_rescale_samples(x, b)[0].reshape(-1)),
+                lambda x, b: (
+                    2
+                    * jnp.pi
+                    * self._shift_angles(min_max_rescale_samples(x, b)[0].reshape(-1))
+                ),
                 xs,
                 bounds,
             )
             return jnp.stack(scaled_samples, axis=-1)
 
-        sample_pts = _flat_samples_to_angles(sample_pts)
+        sample_pts = _flat_samples_to_angles(list(sample_pts))
         angles = sample_pts @ self._freq_combinations
         out = jnp.concatenate(
             [jnp.cos(angles), jnp.sin(angles[..., self._has_zero_phase :])], axis=1