Glm class restructure

docs/developers_notes/03-observation_models.md

@@ -24,7 +24,7 @@ For subclasses derived from `Observations` to function correctly, they must impl
 
 ### Public Methods
 
-- **pseudo_r2**: Method for computing the pseudo-$R^2$ of the model based on the residual deviance. There is no consensus definition for the pseudo-$R^2$, what we used here is the definition by Choen at al. 2003[^2]. 
+- **pseudo_r2**: Method for computing the pseudo-$R^2$ of the model based on the residual deviance. There is no consensus definition for the pseudo-$R^2$, what we used here is the definition by Cohen at al. 2003[^2]. 
 
 
 ### Auxiliary Methods

docs/developers_notes/04-regularizer.md

@@ -40,7 +40,7 @@ Additionally, the class provides auxiliary methods for checking that the solver
 
 ### Auxiliary Methods
 
-- **`_check_solver`**: This method ensures that the provided solver name is in the list of allowed optimizers for the specific `Regularizer` object. This is crucial for maintaining consistency and correctness in the solver's operation.
+- **`_check_solver`**: This method ensures that the provided solver name is in the list of allowed solvers for the specific `Regularizer` object. This is crucial for maintaining consistency and correctness in the solver's operation.
 
 - **`_check_solver_kwargs`**: This method checks if the provided keyword arguments are valid for the specified solver. This helps in catching and preventing potential errors in solver configuration.
 
@@ -50,7 +50,7 @@ The `UnRegularized` class extends the base `Regularizer` class and is designed s
 
 ### Attributes
 
-- **`allowed_solvers`**: A list of string identifiers for the optimization algorithms that can be used with this solver class. The optimization methods listed here are specifically suitable for unregularized optimization problems.
+- **`allowed_solvers`**: A list of string identifiers for the optimization solvers that can be used with this regularizer class. The optimization methods listed here are specifically suitable for unregularized optimization problems.
 
 ### Methods
 
@@ -72,7 +72,7 @@ The `Ridge` class extends the `Regularizer` class to handle optimization problem
 
 ### Attributes
 
-- **`allowed_solvers`**: A list containing string identifiers of optimization algorithms compatible with Ridge regularization.
+- **`allowed_solvers`**: A list containing string identifiers of optimization solvers compatible with Ridge regularization.
 
 - **`regularizer_strength`**: A floating-point value determining the strength of the Ridge regularization. Higher values correspond to stronger regularization which tends to drive the model parameters towards zero.
 
@@ -97,7 +97,7 @@ optim_results = runner(init_params, exog_vars, endog_vars)
 `ProxGradientRegularizer` class extends the `Regularizer` class to utilize the Proximal Gradient method for optimization. It leverages the `jaxopt` library's Proximal Gradient optimizer, introducing the functionality of a proximal operator.
 
 ### Attributes:
-- **`allowed_solvers`**: A list containing string identifiers of optimization algorithms compatible with this solver, specifically the "ProximalGradient".
+- **`allowed_solvers`**: A list containing string identifiers of optimization solvers compatible with this solver, specifically the "ProximalGradient".
 
 ### Methods:
 - **`__init__`**: The constructor method for the `ProxGradientRegularizer` class. It accepts the name of the solver algorithm (`solver_name`), an optional dictionary of additional keyword arguments (`solver_kwargs`) for the solver, the regularization strength (`regularizer_strength`), and an optional mask array (`mask`).
@@ -149,7 +149,7 @@ When developing a functional (i.e., concrete) `Regularizer` class:
 - **Must** inherit from `Regularizer` or one of its derivatives.
 - **Must** implement the `instantiate_solver` method to tailor the solver instantiation based on the provided loss function.
 - For any Proximal Gradient method, **must** include a `get_prox_operator` method to define the proximal operator.
-- **Must** possess an `allowed_solvers` attribute to list the optimizer names that are permissible to be used with this solver.
+- **Must** possess an `allowed_solvers` attribute to list the solver names that are permissible to be used with this regularizer.
 - **May** embed additional attributes and methods such as `mask` and `_check_mask` if required by the specific Solver subclass for handling special optimization scenarios.
 - **May** include a `regularizer_strength` attribute to control the strength of the regularization in scenarios where regularization is applicable.
 - **May** rely on a custom solver implementation for specific optimization problems, but the implementation **must** adhere to the `jaxopt` API.

docs/examples/plot_glm_demo.py

@@ -24,21 +24,21 @@
  data.
 
 """
-import json
 
 import jax
 import matplotlib.pyplot as plt
 import numpy as np
-import sklearn.model_selection as sklearn_model_selection
 from matplotlib.patches import Rectangle
+from sklearn import model_selection
 
 import nemos as nmo
+from nemos import simulation
 
-# Enable float64 precision (optional)
+# enable float64 precision (optional)
 jax.config.update("jax_enable_x64", True)
 
 np.random.seed(111)
-# Random design tensor. Shape (n_time_points, n_neurons, n_features).
+# random design tensor. Shape (n_time_points, n_neurons, n_features).
 X = 0.5*np.random.normal(size=(100, 1, 5))
 
 # log-rates & weights, shape (n_neurons, ) and (n_neurons, n_features) respectively.
@@ -78,12 +78,12 @@
 # ### Model Configuration
 # One could visualize the model hyperparameters by calling `get_params` method.
 
-# Get the glm model parameters only
+# get the glm model parameters only
 print("\nGLM model parameters:")
 for key, value in model.get_params(deep=False).items():
     print(f"\t- {key}: {value}")
 
-# Get the glm model parameters, including the all the
+# get the glm model parameters, including the all the
 # attributes
 print("\nNested parameters:")
 for key, value in model.get_params(deep=True).items():
@@ -127,11 +127,11 @@
 
 # %%
 # !!! warning
-#     Each `Regularizer` has an associated attribute `Regularizer.allowed_optimizers`
+#     Each `Regularizer` has an associated attribute `Regularizer.allowed_solvers`
 #     which lists the optimizers that are suited for each optimization problem.
 #     For example, a `Ridge` is differentiable and can be fit with `GradientDescent`
 #     , `BFGS`, etc., while a `Lasso` should use the `ProximalGradient` method instead.
-#     If the provided `solver_name` is not listed in the `allowed_optimizers` this will raise an
+#     If the provided `solver_name` is not listed in the `allowed_solvers` this will raise an
 #     exception.
 
 # %%
@@ -141,7 +141,7 @@
 # Additionally one may provide an initial parameter guess.
 # The same exact syntax works for any configuration.
 
-# Fit a ridge regression Poisson GLM
+# fit a ridge regression Poisson GLM
 model = nmo.glm.GLM()
 model.set_params(regularizer__regularizer_strength=0.1)
 model.fit(X, spikes)
@@ -162,7 +162,9 @@
 # **Ridge**
 
 parameter_grid = {"regularizer__regularizer_strength": np.logspace(-1.5, 1.5, 6)}
-cls = sklearn_model_selection.GridSearchCV(model, parameter_grid, cv=5)
+# in practice, you should use more folds than 2, but for the purposes of this
+# demo, 2 is sufficient.
+cls = model_selection.GridSearchCV(model, parameter_grid, cv=2)
 cls.fit(X, spikes)
 
 print("Ridge results        ")
@@ -176,7 +178,7 @@
 # **Lasso**
 
 model.set_params(regularizer=nmo.regularizer.Lasso())
-cls = sklearn_model_selection.GridSearchCV(model, parameter_grid, cv=5)
+cls = model_selection.GridSearchCV(model, parameter_grid, cv=2)
 cls.fit(X, spikes)
 
 print("Lasso results        ")
@@ -194,7 +196,7 @@
 
 regularizer = nmo.regularizer.GroupLasso("ProximalGradient", mask=mask)
 model.set_params(regularizer=regularizer)
-cls = sklearn_model_selection.GridSearchCV(model, parameter_grid, cv=5)
+cls = model_selection.GridSearchCV(model, parameter_grid, cv=2)
 cls.fit(X, spikes)
 
 print("\nGroup Lasso results")
@@ -223,48 +225,112 @@
 # %%
 # ## Recurrently Coupled GLM
 # Defining a recurrent model follows the same syntax. In this example
-# we will simulate two coupled neurons. and we will inject a transient
+# we will simulate two coupled neurons, and we will inject a transient
 # input driving the rate of one of the neurons.
-#
-# For brevity, we will import the model parameters instead of generating
-# them on the fly.
 
-# load parameters
-with open("coupled_neurons_params.json", "r") as fh:
-    config_dict = json.load(fh)
 
-# basis weights & intercept for the GLM (both coupling and feedforward)
-# (the last coefficient is the weight of the feedforward input)
-basis_coeff = np.asarray(config_dict["coef_"])[:, :-1]
+# Neural population parameters
+n_neurons = 2
+coupling_filter_duration = 100
+
+# %%
+# We can now to define coupling filters that we will use to simulate
+# the pairwise interactions between the neurons. We will model the
+# filters as a difference of two Gamma probability density function.
+# The negative component will capture inhibitory effects such as the
+# refractory period of a neuron, while the positive component will
+# describe excitation.
+
+np.random.seed(101)
+
+# Gamma parameter for the inhibitory component of the fi;ter
+inhib_a = 1
+inhib_b = 1
+
+# Gamma parameters for the excitatory component of the filter
+excit_a = np.random.uniform(1.1, 5, size=(n_neurons, n_neurons))
+excit_b = np.random.uniform(1.1, 5, size=(n_neurons, n_neurons))
+
+# define 2x2 coupling filters of the specific with create_temporal_filter
+coupling_filter_bank = np.zeros((coupling_filter_duration, n_neurons, n_neurons))
+for unit_i in range(n_neurons):
+    for unit_j in range(n_neurons):
+        coupling_filter_bank[:, unit_i, unit_j] = nmo.simulation.difference_of_gammas(
+            coupling_filter_duration,
+            inhib_a=inhib_a,
+            excit_a=excit_a[unit_i, unit_j],
+            inhib_b=inhib_b,
+            excit_b=excit_b[unit_i, unit_j],
+        )
+
+# shrink the filters for simulation stability
+coupling_filter_bank *= 0.8
 
-# Mask the weights so that only the first neuron receives the imput
-basis_coeff[:, 40:] = np.abs(basis_coeff[:, 40:]) * np.array([[1.], [0.]])
+# %%
+# If we represent our filters in terms of basis functions, we can simulate our network by
+# directly calling the `simulate` method of the `nmo.glm.GLMRecurrent` class.
 
-intercept = np.asarray(config_dict["intercept_"])
+# define a basis function
+n_basis_funcs = 20
+basis = nmo.basis.RaisedCosineBasisLog(n_basis_funcs)
 
-# basis function, inputs and initial spikes
-coupling_basis = jax.numpy.asarray(config_dict["coupling_basis"])
-feedforward_input = jax.numpy.asarray(config_dict["feedforward_input"])
-init_spikes = jax.numpy.asarray(config_dict["init_spikes"])
+# approximate the coupling filters in terms of the basis function
+_, coupling_basis = basis.evaluate_on_grid(coupling_filter_bank.shape[0])
+coupling_coeff = simulation.regress_filter(coupling_filter_bank, coupling_basis)
+intercept = -4 * np.ones(n_neurons)
 
 # %%
-# We can explore visualize the coupling filters and the input.
+# We can check that our approximation worked by plotting the original filters
+# and the basis expansion
 
 # plot coupling functions
 n_basis_coupling = coupling_basis.shape[1]
-fig, axs = plt.subplots(2,2)
+fig, axs = plt.subplots(n_neurons, n_neurons)
 plt.suptitle("Coupling filters")
-for unit_i in range(2):
-    for unit_j in range(2):
-        axs[unit_i,unit_j].set_title(f"unit {unit_j} -> unit {unit_i}")
-        coeff = basis_coeff[unit_i, unit_j * n_basis_coupling: (unit_j + 1) * n_basis_coupling]
-        axs[unit_i, unit_j].plot(np.dot(coupling_basis, coeff))
+for unit_i in range(n_neurons):
+    for unit_j in range(n_neurons):
+        axs[unit_i, unit_j].set_title(f"unit {unit_j} -> unit {unit_i}")
+        coeff = coupling_coeff[unit_i, unit_j]
+        axs[unit_i, unit_j].plot(coupling_filter_bank[:, unit_i, unit_j], label="gamma difference")
+        axs[unit_i, unit_j].plot(np.dot(coupling_basis, coeff), ls="--", color="k", label="basis function")
+axs[0, 0].legend()
 plt.tight_layout()
 
-fig, axs = plt.subplots(1,1)
-plt.title("Feedforward inputs")
-plt.plot(feedforward_input[:, 0])
+# %%
+# Define a squared stimulus current for the first neuron, and no stimulus for
+# the second neuron
+
+# define a squared current parameters
+simulation_duration = 1000
+stimulus_onset = 200
+stimulus_offset = 500
+stimulus_intensity = 1.5
+
+# create the input tensor of shape (n_samples, n_neurons, n_dimension_stimuli)
+feedforward_input = np.zeros((simulation_duration, n_neurons, 1))
+# inject square input to the first neuron only
+feedforward_input[stimulus_onset: stimulus_offset, 0] = stimulus_intensity
+
+# plot the input
+fig, axs = plt.subplots(1,2)
+plt.suptitle("Feedforward inputs")
+axs[0].set_title("Input to neuron 0")
+axs[0].plot(feedforward_input[:, 0])
+
+axs[1].set_title("Input to neuron 1")
+axs[1].plot(feedforward_input[:, 1])
+axs[1].set_ylim(axs[0].get_ylim())
+
+
+# the input for the simulation will be the dot product
+# of input_coeff with the feedforward_input
+input_coeff = np.ones((n_neurons, 1))
+
+# stack the coefficients in a single matrix
+basis_coeff = np.hstack((coupling_coeff.reshape(n_neurons, -1), input_coeff))
 
+# initialize the spikes for the recurrent simulation
+init_spikes = np.zeros((coupling_filter_duration, n_neurons))
 
 # %%
 # We can now simulate spikes by calling the `simulate_recurrent` method.
@@ -277,7 +343,7 @@
 # call simulate, with both the recurrent coupling
 # and the input
 spikes, rates = model.simulate_recurrent(
-    random_key,
+    jax.random.PRNGKey(123),
     feedforward_input=feedforward_input,
     coupling_basis_matrix=coupling_basis,
     init_y=init_spikes
@@ -298,8 +364,8 @@
 p0, = plt.plot(rates[:, 0])
 p1, = plt.plot(rates[:, 1])
 
-plt.vlines(np.where(spikes[:, 0])[0], 0.00, 0.01, color=p0.get_color(), label="neu 0")
-plt.vlines(np.where(spikes[:, 1])[0], -0.01, 0.00, color=p1.get_color(), label="neu 1")
+plt.vlines(np.where(spikes[:, 0])[0], 0.00, 0.01, color=p0.get_color(), label="rate neuron 0")
+plt.vlines(np.where(spikes[:, 1])[0], -0.01, 0.00, color=p1.get_color(), label="rate neuron 1")
 plt.plot(np.exp(basis_coeff[0, -1] * feedforward_input[:, 0, 0] + intercept[0]), color='k', lw=0.8, label="stimulus")
 ax.add_patch(patch)
 plt.ylim(-0.011, .13)

src/nemos/__init__.py

@@ -7,5 +7,6 @@
     observation_models,
     regularizer,
     sample_points,
+    simulation,
     utils,
 )

src/nemos/basis.py

@@ -629,6 +629,8 @@ def evaluate(self, sample_pts: NDArray) -> NDArray:
         The evaluation is performed by looping over each element and using `splev`
         from SciPy to compute the basis values.
         """
+        (sample_pts,) = self._check_evaluate_input(sample_pts)
+
         # add knots
         knot_locs = self._generate_knots(sample_pts, 0.0, 1.0)
 
@@ -711,6 +713,8 @@ def evaluate(self, sample_pts: NDArray) -> NDArray:
         SciPy to compute the basis values.
 
         """
+        (sample_pts,) = self._check_evaluate_input(sample_pts)
+
         knot_locs = self._generate_knots(sample_pts, 0.0, 1.0, is_cyclic=True)
 
         # for cyclic, do not repeat knots