Label get param

docs/background/basis/plot_01_1D_basis_function.md

@@ -6,7 +6,7 @@ jupytext:
     format_version: 0.13
     jupytext_version: 1.16.4
 kernelspec:
-  display_name: Python 3
+  display_name: Python 3 (ipykernel)
   language: python
   name: python3
 ---
@@ -42,7 +42,6 @@ warnings.filterwarnings(
 from nemos._documentation_utils._myst_nb_glue import glue_two_step_convolve
 
 glue_two_step_convolve()
-
 ```
 
 (simple_basis_function)=
@@ -71,7 +70,8 @@ order = 4
 n_basis = 10
 
 # Define the 1D basis function object
-bspline = nmo.basis.BSplineEval(n_basis_funcs=n_basis, order=order)
+bspline = nmo.basis.BSplineEval(n_basis_funcs=n_basis, order=order, label="bspline")
+bspline
 ```
 
 We provide the convenience method `evaluate_on_grid` for evaluating the basis on an equi-spaced grid of points that makes it easier to plot and visualize all basis elements.
@@ -85,7 +85,6 @@ plt.plot(x, y, lw=2)
 plt.title("B-Spline Basis")
 ```
 
-
 ## Computing Features
 All bases in the `nemos.basis` module perform a transformation of one or more time series into a set of features. This operation is always carried out by the method  [`compute_features`](nemos.basis._basis.Basis.compute_features). 
 We can group the bases into two categories depending on the type of transformation that [`compute_features`](nemos.basis._basis.Basis.compute_features) applies:
@@ -97,8 +96,8 @@ We can group the bases into two categories depending on the type of transformati
 Let's see how these two categories operate:
 
 ```{code-cell} ipython3
-eval_mode = nmo.basis.BSplineEval(n_basis_funcs=n_basis)
-conv_mode = nmo.basis.BSplineConv(n_basis_funcs=n_basis, window_size=100)
+eval_mode = nmo.basis.BSplineEval(n_basis_funcs=n_basis, label="eval")
+conv_mode = nmo.basis.BSplineConv(n_basis_funcs=n_basis, window_size=100, label="conv")
 
 # define an input
 angles = np.linspace(0, np.pi*4, 201)
@@ -148,7 +147,7 @@ For inputs with more than one dimension, `compute_features` assumes the first ax
 For Eval bases, `compute_features` evaluates the basis and outputs a 2D feature matrix.
 
 ```{code-cell} ipython3
-basis = nmo.basis.RaisedCosineLinearEval(n_basis_funcs=5)
+basis = nmo.basis.RaisedCosineLinearEval(n_basis_funcs=5, label="multidim")
 # generate a 3D array
 inp = np.random.randn(50, 3, 2)
 out = basis.compute_features(inp)
@@ -158,6 +157,12 @@ out.shape
 For each of the $3 \times 2 = 6$ inputs, `n_basis_funcs = 5` features are computed. These are concatenated on the second axis of the feature matrix, for a total of 
 $3 \times 2 \times 5  = 30$ outputs.
 
+This concatenation can be undone by the `split_by_feature` method of basis, which creates a dictionary with keys the labels of the basis and values a reshaped array.
+
+```{code-cell} ipython3
+basis.split_by_feature(out, axis=1)["multidim"].shape
+```
+
 #### Conv Basis
 
 For Conv bases, `compute_features` convolves each input with `n_basis_funcs` kernels and outputs a 2D feature matrix.
@@ -214,7 +219,6 @@ You can specify a range for the support of your basis by setting the `bounds`
 parameter at initialization of Eval bases. 
 Evaluating the basis at any sample outside the bounds will result in a NaN.
 
-
 ```{code-cell} ipython3
 bspline_range = nmo.basis.BSplineEval(n_basis_funcs=n_basis, order=order, bounds=(0.2, 0.8))
 
@@ -226,7 +230,6 @@ print(np.round(bspline_range.compute_features([0.5, 0.1]), 3))
 Let's compare the default behavior of basis (estimating the range from the samples) with
 the fixed range basis.
 
-
 ```{code-cell} ipython3
 samples = np.linspace(0, 1, 200)
 fig, axs = plt.subplots(2,1, sharex=True)
@@ -237,4 +240,3 @@ axs[1].plot(samples, bspline_range.compute_features(samples), color="tomato")
 axs[1].set_title("bounds=[0.2, 0.8]")
 plt.tight_layout()
 ```
-

docs/background/basis/plot_02_ND_basis_function.md

@@ -6,7 +6,7 @@ jupytext:
     format_version: 0.13
     jupytext_version: 1.16.4
 kernelspec:
-  display_name: Python 3
+  display_name: Python 3 (ipykernel)
   language: python
   name: python3
 ---
@@ -123,24 +123,23 @@ $$
 
 Here, we simply add two basis objects, `a_basis` and `b_basis`, together to define the additive basis.
 
-
 ```{code-cell} ipython3
 import matplotlib.pyplot as plt
 import numpy as np
 import nemos as nmo
 
 # Define 1D basis objects
-a_basis = nmo.basis.MSplineEval(n_basis_funcs=15, order=3)
-b_basis = nmo.basis.RaisedCosineLogEval(n_basis_funcs=14)
+a_basis = nmo.basis.MSplineEval(n_basis_funcs=15, order=3, label="a")
+b_basis = nmo.basis.RaisedCosineLogEval(n_basis_funcs=14, label="b")
 
 # Define the 2D additive basis object
 additive_basis = a_basis + b_basis
+additive_basis
 ```
 
 Evaluating the additive basis will require two inputs, one for each coordinate.
 The total number of elements of the additive basis will be the sum of the elements of the 1D basis.
 
-
 ```{code-cell} ipython3
 # Define a trajectory with 1000 time-points representing the recorded trajectory of the animal
 T = 1000
@@ -154,14 +153,14 @@ eval_basis = additive_basis.compute_features(x_coord, y_coord)
 
 print(f"Sum of two 1D splines with {eval_basis.shape[1]} "
       f"basis element and {eval_basis.shape[0]} samples:\n"
-      f"\t- a_basis had {a_basis.n_basis_funcs} elements\n\t- b_basis had {b_basis.n_basis_funcs} elements.")
+      f"\t- a_basis had {additive_basis['a'].n_basis_funcs} elements.\n"
+      f"\t- b_basis had {additive_basis['b'].n_basis_funcs} elements.")
 ```
 
 (plotting-2d-additive-basis-elements)=
 #### Plotting 2D Additive Basis Elements
 Let's select and plot a basis element from each of the basis we added.
 
-
 ```{code-cell} ipython3
 basis_a_element = 5
 basis_b_element = 1
@@ -184,21 +183,18 @@ We can visualize how these elements are extended in 2D by evaluating the additiv
 on a grid of points that spans its domain and plotting the result.
 We use the `evaluate_on_grid` method for this.
 
-
 ```{code-cell} ipython3
 X, Y, Z = additive_basis.evaluate_on_grid(200, 200)
 ```
 
 We can select the indices of the 2D additive basis that corresponds to the 1D original elements.
 
-
 ```{code-cell} ipython3
 basis_elem_idx = [basis_a_element, a_basis.n_basis_funcs + basis_b_element]
 ```
 
 Finally, we can plot the 2D counterparts.
 
-
 ```{code-cell} ipython3
 _, axs = plt.subplots(1, 2, subplot_kw={'aspect': 1})
 
@@ -218,6 +214,19 @@ plt.tight_layout()
 plt.show()
 ```
 
+If we don't want to do the index algebra ourself, we can use the `split_by_feature` method to split `Z` for each additive element of the basis.
+
+```{code-cell} ipython3
+Z_split = additive_basis.split_by_feature(Z, axis=-1)
+print(Z_split["a"].shape, Z_split["b"].shape)
+```
+
+And then index directly the splitted array.
+
+```{code-cell} ipython3
+element_a, element_b = Z_split["a"][basis_a_element], Z_split["b"][basis_b_element]
+```
+
 ### Multiplicative Basis Object
 
 If the aim is to capture interactions between the coordinates, the response function can be modeled as the external
@@ -230,7 +239,6 @@ $$
 In this model, we define the 2D basis function as the product of two 1D basis objects.
 This allows the response to capture non-linear and interaction effects between the x and y coordinates.
 
-
 ```{code-cell} ipython3
 # 2D basis function as the product of the two 1D basis objects
 prod_basis = a_basis * b_basis
@@ -239,7 +247,6 @@ prod_basis = a_basis * b_basis
 Again evaluating the basis will require 2 inputs.
 The number of elements of the product basis will be the product of the elements of the two 1D bases.
 
-
 ```{code-cell} ipython3
 # Evaluate the product basis at the x and y coordinates
 eval_basis = prod_basis.compute_features(x_coord, y_coord)
@@ -255,9 +262,7 @@ print(f"Product of two 1D splines with {eval_basis.shape[1]} "
 Plotting works in the same way as before. To demonstrate that, we select a few pairs of 1D basis elements,
 and we visualize the corresponding product.
 
-
 ```{code-cell} ipython3
-
 X, Y, Z = prod_basis.evaluate_on_grid(200, 200)
 
 # basis element pairs
@@ -299,7 +304,6 @@ A practical example would be characterizing the responses to position
 in a linear maze and the LFP phase angle.
 :::
 
-
 N-Dimensional Basis
 -------------------
 Sometimes it may be useful to model even higher dimensional interactions, for example between the heding direction of
@@ -331,7 +335,6 @@ print(f"Product of three 1D splines results in {prod_basis_3.n_basis_funcs} "
 
 The evaluation of the product of 3 basis is a 4 dimensional tensor; we can visualize slices of it.
 
-
 ```{code-cell} ipython3
 X, Y, W, Z = prod_basis_3.evaluate_on_grid(30, 30, 30)
 
@@ -369,7 +372,6 @@ full domain of the basis.
 Here we demonstrate a shortcut syntax for multiplying bases of the same class.
 This is achieved using the power operator with an integer exponent.
 
-
 ```{code-cell} ipython3
 # First, let's define a basis `power_basis` that is equivalent to `prod_basis_3`,
 # but we use the power syntax this time:

docs/background/plot_00_conceptual_intro.md

@@ -6,7 +6,7 @@ jupytext:
     format_version: 0.13
     jupytext_version: 1.16.4
 kernelspec:
-  display_name: Python 3
+  display_name: Python 3 (ipykernel)
   language: python
   name: python3
 ---
@@ -39,6 +39,7 @@ warnings.filterwarnings(
     category=RuntimeWarning,
 )
 ```
+
 (glm_intro_background)=
 # Generalized Linear Models: An Introduction
 
@@ -154,7 +155,6 @@ to simplify things, we will look at three simple LNP neuron models as
 described above, working through each step of the transform. First, we will
 plot the linear transformation of the input x:
 
-
 ```{code-cell} ipython3
 weights = np.asarray([.5, 4, -4])
 intercepts = np.asarray([.5, -3, -2])
@@ -180,7 +180,6 @@ have to be non-negative! That's what the nonlinearity handles: making sure our
 firing rate is always positive. We can visualize this second stage of the LNP model
 by adding the `plot_nonlinear` keyword to our `lnp_schematic()` plotting function:
 
-
 ```{code-cell} ipython3
 fig = doc_plots.lnp_schematic(input_feature, weights, intercepts,
                               plot_nonlinear=True)
@@ -207,7 +206,6 @@ positive, though note that the y-values have changed drastically.
 Now we're ready to look at the third step of the LNP model, and see what 
 the generated spikes spikes look like!
 
-
 ```{code-cell} ipython3
 # mkdocs_gallery_thumbnail_number = 3
 fig = doc_plots.lnp_schematic(input_feature, weights, intercepts,

docs/background/plot_03_1D_convolution.md

@@ -6,7 +6,7 @@ jupytext:
     format_version: 0.13
     jupytext_version: 1.16.4
 kernelspec:
-  display_name: Python 3
+  display_name: Python 3 (ipykernel)
   language: python
   name: python3
 ---
@@ -46,7 +46,6 @@ warnings.filterwarnings(
 ## Generate synthetic data
 Generate some simulated spike counts.
 
-
 ```{code-cell} ipython3
 import matplotlib.patches as patches
 import matplotlib.pylab as plt
@@ -109,7 +108,6 @@ and anti-causal effects one should use the acausal filters.
 Below we provide the function [`create_convolutional_predictor`](nemos.convolve.create_convolutional_predictor) that runs the convolution in "valid" mode and pads the convolution output
 for the different filter types.
 
-
 ```{code-cell} ipython3
 # pad according to the causal direction of the filter, after squeeze,
 # the dimension is (n_filters, n_samples)
@@ -126,7 +124,6 @@ spk_acausal_conv = nmo.convolve.create_convolutional_predictor(
 
 Plot the results
 
-
 ```{code-cell} ipython3
 # NaN padded area
 rect_causal = patches.Rectangle((0, -2.5), ws, 5, alpha=0.3, color='grey')
@@ -161,7 +158,6 @@ plt.vlines(np.arange(spk.shape[0]), 0, shift_spk, color='k')
 plt.plot(np.arange(spk.shape[0]), spk_acausal_conv)
 plt.ylabel('acausal')
 plt.tight_layout()
-
 ```
 
 ```{code-cell} ipython3
@@ -195,7 +191,6 @@ convolution.
 All the parameters of [`create_convolutional_predictor`](nemos.convolve.create_convolutional_predictor) can be passed to the object directly at initialization. 
 Let's see how we can get the same results through [`Basis`](nemos.basis._basis.Basis).
 
-
 ```{code-cell} ipython3
 # define basis with different predictor causality
 causal_basis = nmo.basis.RaisedCosineLinearConv(

docs/how_to_guide/README.md

@@ -71,7 +71,9 @@ plot_05_transformer_basis.md
 :::{grid-item-card}
 
 <figure>
+<a href="plot_06_sklearn_pipeline_cv_demo.html">
 <img src="../_static/thumbnails/how_to_guide/plot_06_sklearn_pipeline_cv_demo.svg" style="height: 100px", alt="PyTrees."/>
+</a>
 </figure>
 
 ```{toctree}
@@ -85,7 +87,9 @@ plot_06_sklearn_pipeline_cv_demo.md
 :::{grid-item-card}
 
 <figure>
+<a href="plot_07_glm_pytree.html">
 <img src="../_static/thumbnails/how_to_guide/plot_07_glm_pytree.svg" style="height: 100px", alt="PyTrees."/>
+</a>
 </figure>
 
 ```{toctree}
@@ -99,10 +103,11 @@ plot_07_glm_pytree.md
 :::{grid-item-card}
 
 ```{eval-rst}
-
+ 
 .. plot:: scripts/glm_predictors.py plot_categorical_var_design_matrix
    :show-source-link: False
    :height: 100px
+
 ```
 
 ```{toctree}
@@ -133,7 +138,9 @@ custom_predictors.md
 :::{grid-item-card}
 
 <figure>
+<a href="raw_history_feature.html">
 <img src="../_static/glm_population_scheme.svg" style="height: 100px", alt="Coupled GLM."/>
+</a>
 </figure>
 
 ```{toctree}
@@ -144,4 +151,15 @@ raw_history_feature.md
 
 :::
 
+:::{grid-item-card}
+
+
+```{toctree}
+:maxdepth: 2
+
+handling_composite_bases.md
+```
+
+:::
+
 ::::