diff --git a/README.md b/README.md
index d89f28e..b58fc30 100644
--- a/README.md
+++ b/README.md
@@ -14,7 +14,7 @@ There are already a few open-source libraries for Self-Organizing Maps in python
 - Ability to easily customize the SOM algorithm (e.g. distance, neighborhood, learning rate and update functions).
 - Capacity to vectorize the computations over many SOMs (e.g. for distributed learning over 2D maps of SOMs).
 
-Thanks to [JAX](https://github.com/google/jax)'s `jit` and `vmap` magic functions, it turned out that performance was also significantly better compared to other frameworks. More precisely, it relies indirectly on JAX via the [Equinox](https://github.com/patrick-kidger/equinox) library that offers an easy-to-use PyTorch-like syntax.
+Thanks to [JAX](https://github.com/google/jax)'s `jit` and `vmap` magic functions, it turned out that performance was also significantly better compared to other frameworks. Under the hood, it relies indirectly on JAX via the [Equinox](https://github.com/patrick-kidger/equinox) library that offers an easy-to-use PyTorch-like syntax.
 
 # Installation
 
@@ -55,10 +55,45 @@ quantization_errors = aux["metrics"]["quantization_error"]
 topographic_errors = aux["metrics"]["topographic_error"]
 ```
 
+You can also define your custom SOM:
+```python
+import somap as smp
+from jaxtyping import Array, Float
+
+class MyCustomSomParams(smp.AbstractSomParams):
+    sigma: float | Float[Array, "..."]
+    alpha: float | Float[Array, "..."]
+
+class MyCustomSom(smp.AbstractSom):
+
+    @staticmethod
+    def generate_algo(p: MyCustomSomParams) -> smp.SomAlgo:
+        return smp.SomAlgo(
+            f_dist=smp.EuclidianDist(),
+            f_nbh=smp.GaussianNbh(sigma=p.sigma),
+            f_lr=smp.ConstantLr(alpha=p.alpha),
+            f_update=smp.SomUpdate(),
+        )
+```
+
+If you need custom distance, neighborhood, learning rate and update functions for your SOM, you can define them by inheriting from `smp.AbstractDist`, `smp.AbstractNbh`, `smp.AbstractLr` and `smp.AbstractUpdate`. See the library source code for how to do it.
+
+
 # Documentation
 
-* [Documentation website](https://mthiboust.github.io/somap/)
-* [Examples from notebooks](notebooks/)
+See: [https://mthiboust.github.io/somap/](https://mthiboust.github.io/somap/)
+
+
+# Next steps
+
+I am currently working on different ways to extend the basic SOM algorithm:
+* **Inputs**: In addition to classic bottom-up driving inputs, a SOM could also receive lateral contextual or top-down modulatory inputs.
+* **Weighted inputs**: Each data point from inputs can be weighted so that fuzzy data is weighted less for the winner selection.
+* **Dynamics**: When receiving continuous inputs in time, past activations can influence the computation of the next step.
+* **Supervized and self-supervized learning**: Top-down inputs and next inputs in time can act as teaching signal for supervized and self-supervized learning.
+* **Multi-agent system***: Each SOM is an agent of a mutli-agent system where thousands of SOMs interact with each other.
+
+Some of these features will land on an other library that depends on ***Somap***.
 
 # Citation
 
@@ -71,5 +106,3 @@ If you found this library to be useful in academic work, then please cite:
   url={https://github.com/mthiboust/somap/},
 }
 ```
-
-
diff --git a/docs/index.md b/docs/index.md
index 1fbaa4a..6a82e48 100644
--- a/docs/index.md
+++ b/docs/index.md
@@ -1,4 +1,4 @@
-# Somap
+# Getting started
 
 ***Somap*** is a flexible, fast and scalable Self-Organizing Maps library in python. It allows you to define and run different flavors of SOMs (Kohonen, DSOM or your custom ones) on square or hexagonal 2D grid, with or without toroidal topology.
 
@@ -8,15 +8,15 @@
     <img src="som_on_mnist_square.png">
 </p>
 
-## Why a new SOM library?
+# Why a new SOM library?
 
 There are already a few open-source libraries for Self-Organizing Maps in python, of which [MiniSom](https://github.com/JustGlowing/minisom) and [SOMPY](https://github.com/sevamoo/SOMPY) seem to be the most popular. I developped ***Somap*** to overcome what I believe to be two shortcomings of existing libraries for my research on bio-inspired AI: 
 - Ability to easily customize the SOM algorithm (e.g. distance, neighborhood, learning rate and update functions).
 - Capacity to vectorize the computations over many SOMs (e.g. for distributed learning over 2D maps of SOMs).
 
-Thanks to [JAX](https://github.com/google/jax)'s `jit` and `vmap` magic functions, it turned out that performance was also significantly better compared to other frameworks. More precisely, it relies indirectly on JAX via the [Equinox](https://github.com/patrick-kidger/equinox) library that offers an easy-to-use PyTorch-like syntax.
+Thanks to [JAX](https://github.com/google/jax)'s `jit` and `vmap` magic functions, it turned out that performance was also significantly better compared to other frameworks. Under the hood, it relies indirectly on JAX via the [Equinox](https://github.com/patrick-kidger/equinox) library that offers an easy-to-use PyTorch-like syntax.
 
-## Installation
+# Installation
 
 Requires Python 3.10+ and a working installation of JAX 0.4.20+. You can follow [these instructions](https://github.com/google/jax#installation) to install JAX with the relevant hardware acceleration support.
 
@@ -26,7 +26,7 @@ Then:
 pip install somap
 ```
 
-## Quick example
+# Quick example
 
 The classic workflow goes as follow:
 ```python
@@ -55,7 +55,47 @@ quantization_errors = aux["metrics"]["quantization_error"]
 topographic_errors = aux["metrics"]["topographic_error"]
 ```
 
-## Citation
+You can also define your custom SOM:
+```python
+import somap as smp
+from jaxtyping import Array, Float
+
+class MyCustomSomParams(smp.AbstractSomParams):
+    sigma: float | Float[Array, "..."]
+    alpha: float | Float[Array, "..."]
+
+class MyCustomSom(smp.AbstractSom):
+
+    @staticmethod
+    def generate_algo(p: MyCustomSomParams) -> smp.SomAlgo:
+        return smp.SomAlgo(
+            f_dist=smp.EuclidianDist(),
+            f_nbh=smp.GaussianNbh(sigma=p.sigma),
+            f_lr=smp.ConstantLr(alpha=p.alpha),
+            f_update=smp.SomUpdate(),
+        )
+```
+
+If you need custom distance, neighborhood, learning rate and update functions for your SOM, you can define them by inheriting from `smp.AbstractDist`, `smp.AbstractNbh`, `smp.AbstractLr` and `smp.AbstractUpdate`. See the library source code for how to do it.
+
+
+# Documentation
+
+See: [https://mthiboust.github.io/somap/](https://mthiboust.github.io/somap/)
+
+
+# Next steps
+
+I am currently working on different ways to extend the basic SOM algorithm:
+* **Inputs**: In addition to classic bottom-up driving inputs, a SOM could also receive lateral contextual or top-down modulatory inputs.
+* **Weighted inputs**: Each data point from inputs can be weighted so that fuzzy data is weighted less for the winner selection.
+* **Dynamics**: When receiving continuous inputs in time, past activations can influence the computation of the next step.
+* **Supervized and self-supervized learning**: Top-down inputs and next inputs in time can act as teaching signal for supervized and self-supervized learning.
+* **Multi-agent system***: Each SOM is an agent of a mutli-agent system where thousands of SOMs interact with each other.
+
+Some of these features will land on an other library that depends on ***Somap***.
+
+# Citation
 
 If you found this library to be useful in academic work, then please cite:
 ```
diff --git a/mkdocs.yml b/mkdocs.yml
index 2edae60..8eb3082 100644
--- a/mkdocs.yml
+++ b/mkdocs.yml
@@ -1,3 +1,5 @@
+# Mainly inspired by https://github.com/patrick-kidger/equinox/blob/main/mkdocs.yml
+
 site_name: Somap
 site_description: Documentation for the Somap self-organizing map python library.
 site_author: Matthieu Thiboust
@@ -42,9 +44,12 @@ markdown_extensions:
       line_spans: __span
       pygments_lang_class: true
   - pymdownx.inlinehilite
-  - pymdownx.snippets
+  - pymdownx.details  # Allowing hidden expandable regions denoted by ???
+  - pymdownx.snippets: # Include one Markdown file into another
+      base_path: docs
   - pymdownx.superfences
-
+  - toc:
+      toc_depth: 4
 
 plugins:
   - search  # default search plugin; needs manually re-enabling when using any other plugins
@@ -53,6 +58,11 @@ plugins:
       default_handler: python
       handlers:
         python:
+          setup_commands:
+              - import pytkdocs_tweaks
+              - pytkdocs_tweaks.main()
+              - import jaxtyping
+              - jaxtyping.set_array_name_format("array")
           paths: [src]
           selection:
               inherited_members: true  # Allow looking up inherited methods
diff --git a/pyproject.toml b/pyproject.toml
index 44d4a56..e802446 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -39,6 +39,7 @@ docs = [
     "mkdocs-material>=9.4.11",
     "mknotebooks>=0.8.0",
     "mkdocstrings[crystal,python]>=0.24.0",
+    "pytkdocs-tweaks>=0.0.7",
 ]
 
 [build-system]