fixed the shapes

flatironinstitute · Nov 8, 2023 · 4ea0276 · 4ea0276
1 parent a2a5cb7
commit 4ea0276
Showing 1 changed file with 7 additions and 7 deletions.
diff --git a/src/nemos/glm.py b/src/nemos/glm.py
@@ -140,7 +140,7 @@ def predict(self, X: Union[NDArray, jnp.ndarray]) -> jnp.ndarray:
         Returns
         -------
         :
-            The predicted rates with shape (n_neurons, n_time_bins).
+            The predicted rates with shape (n_time_bins, n_neurons).
 
         Raises
         ------
@@ -395,15 +395,15 @@ def simulate(
         feedforward_input :
             External input matrix to the model, representing factors like convolved currents,
             light intensities, etc. When not provided, the simulation is done with coupling-only.
-            Expected shape: (n_timesteps, n_neurons, n_basis_input).
+            Expected shape: (n_time_bins, n_neurons, n_basis_input).
 
         Returns
         -------
         simulated_activity :
             Simulated activity (spike counts for PoissonGLMs) for each neuron over time.
-            Shape: (n_neurons, n_timesteps).
+            Shape: (n_time_bins, n_neurons).
         firing_rates :
-            Simulated rates for each neuron over time. Shape, (n_neurons, n_timesteps).
+            Simulated rates for each neuron over time. Shape, (n_neurons, n_time_bins).
 
         Raises
         ------
@@ -495,7 +495,7 @@ def simulate_recurrent(
         feedforward_input :
             External input matrix to the model, representing factors like convolved currents,
             light intensities, etc. When not provided, the simulation is done with coupling-only.
-            Expected shape: (n_timesteps, n_neurons, n_basis_input).
+            Expected shape: (n_time_bins, n_neurons, n_basis_input).
         init_y :
             Initial observation (spike counts for PoissonGLM) matrix that kickstarts the simulation.
             Expected shape: (window_size, n_neurons).
@@ -507,9 +507,9 @@ def simulate_recurrent(
         -------
         simulated_activity :
             Simulated activity (spike counts for PoissonGLMs) for each neuron over time.
-            Shape, (n_neurons, n_timesteps).
+            Shape, (n_time_bins, n_neurons).
         firing_rates :
-            Simulated rates for each neuron over time. Shape, (n_neurons, n_timesteps).
+            Simulated rates for each neuron over time. Shape, (n_time_bins, n_neurons,).
 
         Raises
         ------