Adds __repr__ for our custom classes

src/nemos/basis/_basis.py

@@ -15,7 +15,7 @@
 from ..base_class import Base
 from ..type_casting import support_pynapple
 from ..typing import FeatureMatrix
-from ..utils import row_wise_kron
+from ..utils import format_repr, row_wise_kron
 from ..validation import check_fraction_valid_samples
 from ._basis_mixin import BasisTransformerMixin, CompositeBasisMixin
 
@@ -521,6 +521,9 @@ def __pow__(self, exponent: int) -> MultiplicativeBasis:
             result = result * self
         return result
 
+    def __repr__(self):
+        return format_repr(self)
+
     def _get_feature_slicing(
         self,
         n_inputs: Optional[tuple] = None,
@@ -736,11 +739,23 @@ class AdditiveBasis(CompositeBasisMixin, Basis):
     >>> basis_1 = nmo.basis.BSplineEval(10)
     >>> basis_2 = nmo.basis.RaisedCosineLinearEval(15)
     >>> additive_basis = basis_1 + basis_2
-
+    >>> additive_basis
+    AdditiveBasis(
+        basis1=BSplineEval(n_basis_funcs=10, order=4),
+        basis2=RaisedCosineLinearEval(n_basis_funcs=15, width=2.0),
+    )
     >>> # can add another basis to the AdditiveBasis object
     >>> X = np.random.normal(size=(30, 3))
     >>> basis_3 = nmo.basis.RaisedCosineLogEval(100)
     >>> additive_basis_2 = additive_basis + basis_3
+    >>> additive_basis_2
+    AdditiveBasis(
+        basis1=AdditiveBasis(
+            basis1=BSplineEval(n_basis_funcs=10, order=4),
+            basis2=RaisedCosineLinearEval(n_basis_funcs=15, width=2.0),
+        ),
+        basis2=RaisedCosineLogEval(n_basis_funcs=100, width=2.0, time_scaling=50.0, enforce_decay_to_zero=True),
+    )
     """
 
     def __init__(self, basis1: Basis, basis2: Basis) -> None:
@@ -1158,11 +1173,23 @@ class MultiplicativeBasis(CompositeBasisMixin, Basis):
     >>> basis_1 = nmo.basis.BSplineEval(10)
     >>> basis_2 = nmo.basis.RaisedCosineLinearEval(15)
     >>> multiplicative_basis = basis_1 * basis_2
-
+    >>> multiplicative_basis
+    MultiplicativeBasis(
+        basis1=BSplineEval(n_basis_funcs=10, order=4),
+        basis2=RaisedCosineLinearEval(n_basis_funcs=15, width=2.0),
+    )
     >>> # Can multiply or add another basis to the AdditiveBasis object
     >>> # This will cause the number of output features of the result basis to grow accordingly
     >>> basis_3 = nmo.basis.RaisedCosineLogEval(100)
     >>> multiplicative_basis_2 = multiplicative_basis * basis_3
+    >>> multiplicative_basis_2
+    MultiplicativeBasis(
+        basis1=MultiplicativeBasis(
+            basis1=BSplineEval(n_basis_funcs=10, order=4),
+            basis2=RaisedCosineLinearEval(n_basis_funcs=15, width=2.0),
+        ),
+        basis2=RaisedCosineLogEval(n_basis_funcs=100, width=2.0, time_scaling=50.0, enforce_decay_to_zero=True),
+    )
     """
 
     def __init__(self, basis1: Basis, basis2: Basis) -> None:

src/nemos/basis/_basis_mixin.py

@@ -14,6 +14,7 @@
 from pynapple import Tsd, TsdFrame, TsdTensor
 
 from ..convolve import create_convolutional_predictor
+from ..utils import _get_terminal_size
 from ._transformer_basis import TransformerBasis
 
 if TYPE_CHECKING:
@@ -661,3 +662,25 @@ def set_input_shape(self, *xi: int | tuple[int, ...] | NDArray) -> Basis:
             )._input_shape_product,
         )
         return self
+
+    def __repr__(self, n=0):
+        _, rows = _get_terminal_size()
+        rows = rows // 4
+        # number of nested composite bases
+        n += 1
+        tab = "    "
+        try:
+            basis1 = self.basis1.__repr__(n=n)
+        except TypeError:
+            basis1 = self.basis1
+        try:
+            basis2 = self.basis2.__repr__(n=n)
+        except TypeError:
+            basis2 = self.basis2
+        if n < rows:
+            rep = f"{self.__class__.__name__}(\n{n*tab}basis1={basis1},\n{n*tab}basis2={basis2},\n{(n-1)*tab})"
+        elif n == rows:
+            rep = f"{self.__class__.__name__}(\n{n*tab}...\n{(n-1)*tab})"
+        else:
+            rep = None
+        return rep

src/nemos/basis/_transformer_basis.py

@@ -360,6 +360,9 @@ def __sklearn_clone__(self) -> TransformerBasis:
         cloned_obj = TransformerBasis(self.basis.__sklearn_clone__())
         return cloned_obj
 
+    def __repr__(self):
+        return f"Transformer({self.basis})"
+
     def set_params(self, **parameters) -> TransformerBasis:
         """
         Set TransformerBasis parameters.

src/nemos/basis/basis.py

@@ -75,6 +75,8 @@ class BSplineEval(EvalBasisMixin, BSplineBasis):
     >>> n_basis_funcs = 5
     >>> order = 3
     >>> bspline_basis = BSplineEval(n_basis_funcs, order=order)
+    >>> bspline_basis
+    BSplineEval(n_basis_funcs=5, order=3)
     >>> sample_points = linspace(0, 1, 100)
     >>> basis_functions = bspline_basis.compute_features(sample_points)
     """
@@ -226,6 +228,8 @@ class BSplineConv(ConvBasisMixin, BSplineBasis):
     >>> n_basis_funcs = 5
     >>> order = 3
     >>> bspline_basis = BSplineConv(n_basis_funcs, order=order, window_size=10)
+    >>> bspline_basis
+    BSplineConv(n_basis_funcs=5, window_size=10, order=3)
     >>> sample_points = linspace(0, 1, 100)
     >>> features = bspline_basis.compute_features(sample_points)
     """
@@ -365,6 +369,8 @@ class CyclicBSplineEval(EvalBasisMixin, CyclicBSplineBasis):
     >>> n_basis_funcs = 5
     >>> order = 3
     >>> cyclic_bspline_basis = CyclicBSplineEval(n_basis_funcs, order=order)
+    >>> cyclic_bspline_basis
+    CyclicBSplineEval(n_basis_funcs=5, order=3)
     >>> sample_points = linspace(0, 1, 100)
     >>> features = cyclic_bspline_basis.compute_features(sample_points)
     """
@@ -507,6 +513,8 @@ class CyclicBSplineConv(ConvBasisMixin, CyclicBSplineBasis):
     >>> n_basis_funcs = 5
     >>> order = 3
     >>> cyclic_bspline_basis = CyclicBSplineConv(n_basis_funcs, order=order, window_size=10)
+    >>> cyclic_bspline_basis
+    CyclicBSplineConv(n_basis_funcs=5, window_size=10, order=3)
     >>> sample_points = linspace(0, 1, 100)
     >>> features = cyclic_bspline_basis.compute_features(sample_points)
     """
@@ -670,6 +678,8 @@ class MSplineEval(EvalBasisMixin, MSplineBasis):
     >>> n_basis_funcs = 5
     >>> order = 3
     >>> mspline_basis = MSplineEval(n_basis_funcs, order=order)
+    >>> mspline_basis
+    MSplineEval(n_basis_funcs=5, order=3)
     >>> sample_points = linspace(0, 1, 100)
     >>> features = mspline_basis.compute_features(sample_points)
     """
@@ -836,6 +846,8 @@ class MSplineConv(ConvBasisMixin, MSplineBasis):
     >>> n_basis_funcs = 5
     >>> order = 3
     >>> mspline_basis = MSplineConv(n_basis_funcs, order=order, window_size=10)
+    >>> mspline_basis
+    MSplineConv(n_basis_funcs=5, window_size=10, order=3)
     >>> sample_points = linspace(0, 1, 100)
     >>> features = mspline_basis.compute_features(sample_points)
     """
@@ -982,6 +994,8 @@ class RaisedCosineLinearEval(EvalBasisMixin, RaisedCosineBasisLinear):
     >>> from nemos.basis import RaisedCosineLinearEval
     >>> n_basis_funcs = 5
     >>> raised_cosine_basis = RaisedCosineLinearEval(n_basis_funcs)
+    >>> raised_cosine_basis
+    RaisedCosineLinearEval(n_basis_funcs=5, width=2.0)
     >>> sample_points = np.random.randn(100)
     >>> # convolve the basis
     >>> features = raised_cosine_basis.compute_features(sample_points)
@@ -1125,6 +1139,8 @@ class RaisedCosineLinearConv(ConvBasisMixin, RaisedCosineBasisLinear):
     >>> from nemos.basis import RaisedCosineLinearConv
     >>> n_basis_funcs = 5
     >>> raised_cosine_basis = RaisedCosineLinearConv(n_basis_funcs, window_size=10)
+    >>> raised_cosine_basis
+    RaisedCosineLinearConv(n_basis_funcs=5, window_size=10, width=2.0)
     >>> sample_points = np.random.randn(100)
     >>> # convolve the basis
     >>> features = raised_cosine_basis.compute_features(sample_points)
@@ -1269,6 +1285,8 @@ class RaisedCosineLogEval(EvalBasisMixin, RaisedCosineBasisLog):
     >>> from nemos.basis import RaisedCosineLogEval
     >>> n_basis_funcs = 5
     >>> raised_cosine_basis = RaisedCosineLogEval(n_basis_funcs)
+    >>> raised_cosine_basis
+    RaisedCosineLogEval(n_basis_funcs=5, width=2.0, time_scaling=50.0, enforce_decay_to_zero=True)
     >>> sample_points = np.random.randn(100)
     >>> # convolve the basis
     >>> features = raised_cosine_basis.compute_features(sample_points)
@@ -1424,6 +1442,8 @@ class RaisedCosineLogConv(ConvBasisMixin, RaisedCosineBasisLog):
     >>> from nemos.basis import RaisedCosineLogConv
     >>> n_basis_funcs = 5
     >>> raised_cosine_basis = RaisedCosineLogConv(n_basis_funcs, window_size=10)
+    >>> raised_cosine_basis
+    RaisedCosineLogConv(n_basis_funcs=5, window_size=10, width=2.0, time_scaling=50.0, enforce_decay_to_zero=True)
     >>> sample_points = np.random.randn(100)
     >>> # convolve the basis
     >>> features = raised_cosine_basis.compute_features(sample_points)
@@ -1561,6 +1581,8 @@ class OrthExponentialEval(EvalBasisMixin, OrthExponentialBasis):
     >>> decay_rates = np.array([0.01, 0.02, 0.03, 0.04, 0.05])  # sample decay rates
     >>> window_size = 10
     >>> ortho_basis = OrthExponentialEval(n_basis_funcs, decay_rates)
+    >>> ortho_basis
+    OrthExponentialEval(n_basis_funcs=5)
     >>> sample_points = linspace(0, 1, 100)
     >>> # evaluate the basis
     >>> features = ortho_basis.compute_features(sample_points)
@@ -1701,6 +1723,8 @@ class OrthExponentialConv(ConvBasisMixin, OrthExponentialBasis):
     >>> decay_rates = np.array([0.01, 0.02, 0.03, 0.04, 0.05])  # sample decay rates
     >>> window_size = 10
     >>> ortho_basis = OrthExponentialConv(n_basis_funcs, window_size, decay_rates)
+    >>> ortho_basis
+    OrthExponentialConv(n_basis_funcs=5, window_size=10)
     >>> sample_points = np.random.randn(100)
     >>> # convolve the basis
     >>> features = ortho_basis.compute_features(sample_points)
@@ -1967,7 +1991,7 @@ class HistoryConv(ConvBasisMixin, HistoryBasis):
     def __init__(
         self,
         window_size: int,
-        label: Optional[str] = "IdentityEval",
+        label: Optional[str] = "HistoryConv",
         conv_kwargs: Optional[dict] = None,
     ):
         ConvBasisMixin.__init__(self, window_size=window_size, conv_kwargs=conv_kwargs)

src/nemos/glm.py

@@ -22,6 +22,7 @@
 from .solvers._compute_defaults import glm_compute_optimal_stepsize_configs
 from .type_casting import jnp_asarray_if, support_pynapple
 from .typing import DESIGN_INPUT_TYPE
+from .utils import format_repr
 
 ModelParams = Tuple[jnp.ndarray, jnp.ndarray]
 
@@ -145,6 +146,12 @@ class GLM(BaseRegressor):
     >>> import nemos as nmo
     >>> # define single neuron GLM model
     >>> model = nmo.glm.GLM()
+    >>> model
+    GLM(
+        observation_model=PoissonObservations(inverse_link_function=exp),
+        regularizer=UnRegularized(),
+        solver_name='GradientDescent'
+    )
     >>> print("Regularizer type: ", type(model.regularizer))
     Regularizer type:  <class 'nemos.regularizer.UnRegularized'>
     >>> print("Observation model: ", type(model.observation_model))
@@ -1116,6 +1123,9 @@ def _get_optimal_solver_params_config(self):
         """Return the functions for computing default step and batch size for the solver."""
         return glm_compute_optimal_stepsize_configs(self)
 
+    def __repr__(self):
+        return format_repr(self, multiline=True)
+
 
 class PopulationGLM(GLM):
     """
@@ -1233,6 +1243,12 @@ class PopulationGLM(GLM):
            [0, 1]], dtype=int32)
     >>> # Create and fit the model
     >>> model = PopulationGLM(feature_mask=feature_mask).fit(X, y)
+    >>> model
+    PopulationGLM(
+        observation_model=PoissonObservations(inverse_link_function=exp),
+        regularizer=UnRegularized(),
+        solver_name='GradientDescent'
+    )
     >>> # Check the fitted coefficients
     >>> print(model.coef_.shape)
     (3, 2)

src/nemos/observation_models.py

@@ -45,6 +45,9 @@ def __init__(self, inverse_link_function: Callable, **kwargs):
         self.inverse_link_function = inverse_link_function
         self.scale = 1.0
 
+    def __repr__(self):
+        return utils.format_repr(self, use_name_keys=["inverse_link_function"])
+
     @property
     def inverse_link_function(self):
         """Getter for the inverse link function for the model."""

src/nemos/regularizer.py

@@ -18,6 +18,7 @@
 from .base_class import Base
 from .proximal_operator import prox_group_lasso
 from .typing import DESIGN_INPUT_TYPE, ProximalOperator
+from .utils import format_repr
 
 __all__ = ["UnRegularized", "Ridge", "Lasso", "GroupLasso"]
 
@@ -93,6 +94,9 @@ def get_proximal_operator(
         """
         pass
 
+    def __repr__(self):
+        return format_repr(self)
+
 
 class UnRegularized(Regularizer):
     """