mimosa-ml 0.1.0a0__tar.gz

mimosa_ml-0.1.0a0/.github/workflows/publish.yml

@@ -0,0 +1,32 @@
+name: Publish Python 🐍📦 to PyPI
+
+on:
+  push:
+    tags:
+      - "v*"  # déclenche quand on push un tag qui commence par v
+
+jobs:
+  build-and-publish:
+    runs-on: ubuntu-latest
+    environment: pypi
+    permissions:
+      id-token: write
+
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.10"
+
+      - name: Install build dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install build twine
+
+      - name: Build package
+        run: python -m build
+
+      - name: Publish to PyPI
+        uses: pypa/gh-action-pypi-publish@release/v1

mimosa_ml-0.1.0a0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Simon Lejoly
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

mimosa_ml-0.1.0a0/PKG-INFO

@@ -0,0 +1,240 @@
+Metadata-Version: 2.4
+Name: mimosa-ml
+Version: 0.1.0a0
+Summary: MIMOSA: Multi-Input Multi-Output Sample Analysis - A fully-featured multi-task Gaussian process framework for analysing functional data.
+Project-URL: Homepage, https://github.com/SimLej18/mimosa
+Author-email: Simon Lejoly <simon.lejoly@unamur.be>
+License-Expression: MIT
+License-File: LICENSE
+Classifier: Development Status :: 3 - Alpha
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Requires-Python: >=3.10
+Requires-Dist: beartype
+Requires-Dist: equinox
+Requires-Dist: jax
+Requires-Dist: jaxtyping
+Requires-Dist: kernax-ml
+Requires-Dist: matplotlib
+Requires-Dist: numpy
+Requires-Dist: optimistix
+Requires-Dist: polars
+Provides-Extra: dev
+Requires-Dist: pytest; extra == 'dev'
+Description-Content-Type: text/markdown
+
+# MIMOSA: Multi-Input Multi-Output Sample Analysis 
+
+A fully-featured *multi-task Gaussian process framework* for analysing functional data.
+It is based on the [MagmaClust framework](https://jmlr.org/papers/v24/20-1321.html), with many enhancements for efficiency and modularity.
+
+Its core features currently include:
+* Joint learning across multiple samples (Multi-Task GPs) with unaligned sampling grids
+* Clustering of samples
+* Multi-dimensional inputs and outputs
+* Probabilistic predictions with uncertainty quantification
+* A multitude of training configuration to model complex relationships between tasks, clusters and inputs
+* Deep kernax integration for complex kernels/means functions
+* Support for training and predicting on GPUs and TPUs
+* Full Jax/Equinox compatibility for vmap/grad/jit
+
+Towards v1.0.0, we plan to add:
+* Inter-output correlation learning
+* Adaptation to non-gaussian data and classification tasks
+* Sparse learning and approximations for scalability
+* Visualisation tools for data and model outputs
+* Clever initialisation for hyper-parameters of the models
+
+> **⚠️ Project Status**:  MIMOSA is currently in early development. 
+> While current features are functional in most cases, bugs can happen.
+> The API and features are still subject to change.
+> If you use the package for research/production, we encourage you to version-lock your mimosa dependency.
+
+---
+## Installation
+
+You can install mimosa through pip:
+
+```bash
+pip install mimosa-ml
+```
+
+You can then import components of the package as:
+
+```python
+from mimosa.hyperpost import Hyperpost  # Import example
+```
+
+N.b.: notice the difference between `mimosa-ml` (the PyPI ID of the package) and `mimosa` (the internal package name).
+
+Python compatibility: mimosa is compatible with Python 3.10+. 
+We encourage using a more recent version to be compatible with latest releases of Jax and Equinox.
+
+---
+## Philosophy
+
+We basically took Gaussian processes, and combined them with a ton of ✨*fancy components*✨
+* ➡️ Multi-task learning based on [Magma](https://jmlr.org/papers/v24/20-1321.html)
+* ➡️ Clustering as a mixture of Magma GPs
+* ➡️ Multi-dimensional inputs thanks to efficient kernels from [Kernax](https://github.com/SimLej18/kernax-ml)
+* ➡️ Multi-dimensional uncorrelated outputs by simply broadcasting the algorithm to another dimension
+* ➡️ Multi-output correlation discovery by learning a [convolution process]([TODO]) with specific kernel
+
+
+We use **JAX** as a back-end for every computation! 
+`jit` compilation saves us a lot of time, 
+`vmap` allows us to efficiently batch over all dimensions of the data and parameters,
+`grad` spares us the hastle of implementing gradients for every component of the model, 
+and XLA support allows us to train and predict on GPUs and TPUs for scalability.
+
+A GP library is nothing without a fast and modular **Kernel library**. 
+MIMOSA relies on the efficiency-focused [Kernax package](https://github.com/SimLej18/kernax-ml) for every kernel use 
+throughout the algorithms.
+
+---
+## Examples
+
+Basic examples can be found in the `examples` directory. A quick example using the top-level API from Model can be found here:
+
+```python
+import jax.random as jr
+from jax import vmap
+from kernax import ZeroMean, VarianceKernel, SEKernel, WhiteNoiseKernel
+
+from mimosa.data_structures import Dimensions, Parameters, ModelConfig
+from mimosa.synthetic import generate_data
+from mimosa.models import BasicModel
+from mimosa.grid import UnionGrid
+from mimosa.plot import plot_dataset, plot_clusters, plot_single_task_prediction
+from mimosa.sampling import sample_gp
+
+key = jr.PRNGKey(0)
+
+# 20 tasks, 2 clusters, 1D inputs/outputs, 25 points/task
+dims = Dimensions(T=20, K=2, I=1, O=1, F=1, N=25, G=50)
+config = ModelConfig(isotopic_tasks=False)
+
+# Ground-truth parameters used to generate the synthetic dataset
+init_params = Parameters(
+	cluster_mean=ZeroMean(),
+	cluster_kernel=VarianceKernel(1.0) * SEKernel(length_scale=1.0),
+	task_kernel=VarianceKernel(0.1) * SEKernel(length_scale=2.0),
+	noise_kernel=WhiteNoiseKernel(noise=0.01),
+)
+
+key, subkey = jr.split(key)
+dataset, grid, hyperprior, mixture, parameters, cluster_means, tasks = generate_data(
+	subkey, dims, init_params, config, input_range=(-2.5, 2.5))
+
+# Fit a model on the generated dataset, starting from the same parameters
+fit_grid = UnionGrid()(dataset.inputs)
+model = BasicModel(jr.PRNGKey(1), n_clusters=dims.K)
+fitted_params, fitted_mixture = model.fit(dataset, fit_grid, mixture.proportions, parameters, n_iter=50)
+
+# Predict the posterior distribution of every task, in every cluster
+predictions = model.predict(dataset, fit_grid, fitted_mixture, fitted_params)
+hyperposterior = model.hyperpost(dataset, fit_grid, fitted_mixture, fitted_params)
+
+# Plot the fitted mean-processes over the dataset
+fig, ax = plot_dataset(dataset, mixture=mixture, alpha=.3)
+fig, ax = plot_clusters(fit_grid, hyperposterior=hyperposterior, fig=fig, ax=ax, legend=False)
+```
+
+![Fitted mean-processes over the dataset](docs/images/mean_processes.png)
+
+```python
+# Draw samples from task 0's predictive distribution and plot them
+t_id, o_id = 0, 0
+k_id = int(fitted_mixture.assignments[t_id])  # task's dominant cluster
+prediction = predictions[t_id, k_id, o_id]
+
+key, sample_key = jr.split(key)
+sample_keys = jr.split(sample_key, 64)
+samples = vmap(lambda k: sample_gp(k, prediction.mean, prediction.covariance))(sample_keys)
+
+fig, ax = plot_single_task_prediction(
+	dataset, fit_grid, hyperposterior, fitted_mixture, t_id, o_id, samples=samples)
+```
+
+![Prediction samples for task 0](docs/images/task0_prediction_samples.png)
+
+MIMOSA is designed to have multiple layer of customisation. The "Model" API covers the most general use case of the
+algorithm, but you should consider building your own training loops and combining/developping your own sub-modules to 
+create tailored models and pipelines.
+
+---
+## About configurations
+
+MIMOSA has many training configurations to make it as precise as possible when modelling your specific data.
+These configurations are stored in a `ModelConfig` dataclass. 
+Here is a quick overview of what each configuration option does. 
+
+* `shared_task_hps` (default `True`): whether task kernel/noise hyper-parameters are shared across all tasks, or learnt independently per task.
+* `shared_cluster_hps` (default `True`): whether cluster mean/kernel hyper-parameters are shared across all clusters (mean-processes), or learnt independently per cluster.
+* `shared_output_hps` (default `True`): whether hyper-parameters are shared across output dimensions, or learnt independently per output.
+* `shared_features_hps` (default `True`): whether hyper-parameters are shared across features, or learnt independently per feature.
+* `cluster_specific_task_hps` (default `True`): whether the task/noise kernels themselves vary depending on the cluster a task belongs to, or are shared across clusters.
+* `isotopic_tasks` (default `True`): whether all tasks are observed at the same input locations, or each task has its own set of inputs.
+* `isotopic_features` (default `True`): whether all features are observed at the same input locations, or each feature has its own set of inputs.
+
+These configurations are independent of your choices of parameters (mean functions/kernels), leading to a large
+number of ways to model your data.
+
+The best way to get a feeling for which setup might be the best for your specific needs is to play
+around with `generate_data` and varying configurations and parameters until you get something that looks like
+your specific dataset.
+
+---
+## Troubleshooting
+
+As with any GP-based framework, mimosa is subject to numerical stability issues.
+
+We highly recommend to use Jax' 64bit precision:
+
+```python
+import jax
+jax.config.update("jax_enable_x64", True)  # Should be called before any other jax/kernax/mimosa import
+# ...other imports
+```
+
+If you still have troubles (NaN in hyperpost/likelihoods, optimisation does not converge, odd-looking mean 
+processes/predictions/samples, ...), consider increasing the `jitter` used in the specific component or
+directly at the Model level:
+
+```python
+model = BasicModel(prng_key=key, n_clusters=K, jitter=jnp.asarray(1e-5))
+```
+
+By default, Mimosa's jitter is fixed at 1e-8 for every component of the framework.
+
+---
+## Authors and citation
+
+MIMOSA is primarily developed by the *Magma Task Force*, composed of:
+
+* [Arthur Leroy](https://arthur-leroy.netlify.app/), researcher at Paris Saclay and INRAe (FR), main author of the original tool [MagmaClustR](https://arthurleroy.github.io/MagmaClustR/) and coordinator of the Task Force.
+* [Simon Lejoly](https://researchportal.unamur.be/fr/persons/slejoly/), PhD student at UNamur (BE), main developper of the package and author of the [Kernax package](https://github.com/SimLej18/kernax-ml)
+* [Alexia Grenouillat]([TODO]), PhD student at INSA Toulouse (FR), developper of multi-feature correlation discovery.
+* [Térence Viellard]([TODO]), PhD student at INRAe (FR), working on sparse approximations and scalability.
+
+If you use the package for your research, please consider citing the following paper:
+
+[TODO]
+
+If your research field could provide an interesting example of using MIMOSA, please consider publishing a toy example in
+the documentation [TODO].
+
+---
+## Contributing
+
+At this stage of development, we are not yet accepting contributions to the codebase. However, if you are interested 
+in contributing to the project, please feel free to reach out to us via email or through our GitHub repository.
+
+---
+## License
+
+This work is distributed under the [MIT License](https://opensource.org/licenses/MIT). See `LICENSE` for more information.
+
+---
+
+

mimosa_ml-0.1.0a0/README.md

@@ -0,0 +1,215 @@
+# MIMOSA: Multi-Input Multi-Output Sample Analysis 
+
+A fully-featured *multi-task Gaussian process framework* for analysing functional data.
+It is based on the [MagmaClust framework](https://jmlr.org/papers/v24/20-1321.html), with many enhancements for efficiency and modularity.
+
+Its core features currently include:
+* Joint learning across multiple samples (Multi-Task GPs) with unaligned sampling grids
+* Clustering of samples
+* Multi-dimensional inputs and outputs
+* Probabilistic predictions with uncertainty quantification
+* A multitude of training configuration to model complex relationships between tasks, clusters and inputs
+* Deep kernax integration for complex kernels/means functions
+* Support for training and predicting on GPUs and TPUs
+* Full Jax/Equinox compatibility for vmap/grad/jit
+
+Towards v1.0.0, we plan to add:
+* Inter-output correlation learning
+* Adaptation to non-gaussian data and classification tasks
+* Sparse learning and approximations for scalability
+* Visualisation tools for data and model outputs
+* Clever initialisation for hyper-parameters of the models
+
+> **⚠️ Project Status**:  MIMOSA is currently in early development. 
+> While current features are functional in most cases, bugs can happen.
+> The API and features are still subject to change.
+> If you use the package for research/production, we encourage you to version-lock your mimosa dependency.
+
+---
+## Installation
+
+You can install mimosa through pip:
+
+```bash
+pip install mimosa-ml
+```
+
+You can then import components of the package as:
+
+```python
+from mimosa.hyperpost import Hyperpost  # Import example
+```
+
+N.b.: notice the difference between `mimosa-ml` (the PyPI ID of the package) and `mimosa` (the internal package name).
+
+Python compatibility: mimosa is compatible with Python 3.10+. 
+We encourage using a more recent version to be compatible with latest releases of Jax and Equinox.
+
+---
+## Philosophy
+
+We basically took Gaussian processes, and combined them with a ton of ✨*fancy components*✨
+* ➡️ Multi-task learning based on [Magma](https://jmlr.org/papers/v24/20-1321.html)
+* ➡️ Clustering as a mixture of Magma GPs
+* ➡️ Multi-dimensional inputs thanks to efficient kernels from [Kernax](https://github.com/SimLej18/kernax-ml)
+* ➡️ Multi-dimensional uncorrelated outputs by simply broadcasting the algorithm to another dimension
+* ➡️ Multi-output correlation discovery by learning a [convolution process]([TODO]) with specific kernel
+
+
+We use **JAX** as a back-end for every computation! 
+`jit` compilation saves us a lot of time, 
+`vmap` allows us to efficiently batch over all dimensions of the data and parameters,
+`grad` spares us the hastle of implementing gradients for every component of the model, 
+and XLA support allows us to train and predict on GPUs and TPUs for scalability.
+
+A GP library is nothing without a fast and modular **Kernel library**. 
+MIMOSA relies on the efficiency-focused [Kernax package](https://github.com/SimLej18/kernax-ml) for every kernel use 
+throughout the algorithms.
+
+---
+## Examples
+
+Basic examples can be found in the `examples` directory. A quick example using the top-level API from Model can be found here:
+
+```python
+import jax.random as jr
+from jax import vmap
+from kernax import ZeroMean, VarianceKernel, SEKernel, WhiteNoiseKernel
+
+from mimosa.data_structures import Dimensions, Parameters, ModelConfig
+from mimosa.synthetic import generate_data
+from mimosa.models import BasicModel
+from mimosa.grid import UnionGrid
+from mimosa.plot import plot_dataset, plot_clusters, plot_single_task_prediction
+from mimosa.sampling import sample_gp
+
+key = jr.PRNGKey(0)
+
+# 20 tasks, 2 clusters, 1D inputs/outputs, 25 points/task
+dims = Dimensions(T=20, K=2, I=1, O=1, F=1, N=25, G=50)
+config = ModelConfig(isotopic_tasks=False)
+
+# Ground-truth parameters used to generate the synthetic dataset
+init_params = Parameters(
+	cluster_mean=ZeroMean(),
+	cluster_kernel=VarianceKernel(1.0) * SEKernel(length_scale=1.0),
+	task_kernel=VarianceKernel(0.1) * SEKernel(length_scale=2.0),
+	noise_kernel=WhiteNoiseKernel(noise=0.01),
+)
+
+key, subkey = jr.split(key)
+dataset, grid, hyperprior, mixture, parameters, cluster_means, tasks = generate_data(
+	subkey, dims, init_params, config, input_range=(-2.5, 2.5))
+
+# Fit a model on the generated dataset, starting from the same parameters
+fit_grid = UnionGrid()(dataset.inputs)
+model = BasicModel(jr.PRNGKey(1), n_clusters=dims.K)
+fitted_params, fitted_mixture = model.fit(dataset, fit_grid, mixture.proportions, parameters, n_iter=50)
+
+# Predict the posterior distribution of every task, in every cluster
+predictions = model.predict(dataset, fit_grid, fitted_mixture, fitted_params)
+hyperposterior = model.hyperpost(dataset, fit_grid, fitted_mixture, fitted_params)
+
+# Plot the fitted mean-processes over the dataset
+fig, ax = plot_dataset(dataset, mixture=mixture, alpha=.3)
+fig, ax = plot_clusters(fit_grid, hyperposterior=hyperposterior, fig=fig, ax=ax, legend=False)
+```
+
+![Fitted mean-processes over the dataset](docs/images/mean_processes.png)
+
+```python
+# Draw samples from task 0's predictive distribution and plot them
+t_id, o_id = 0, 0
+k_id = int(fitted_mixture.assignments[t_id])  # task's dominant cluster
+prediction = predictions[t_id, k_id, o_id]
+
+key, sample_key = jr.split(key)
+sample_keys = jr.split(sample_key, 64)
+samples = vmap(lambda k: sample_gp(k, prediction.mean, prediction.covariance))(sample_keys)
+
+fig, ax = plot_single_task_prediction(
+	dataset, fit_grid, hyperposterior, fitted_mixture, t_id, o_id, samples=samples)
+```
+
+![Prediction samples for task 0](docs/images/task0_prediction_samples.png)
+
+MIMOSA is designed to have multiple layer of customisation. The "Model" API covers the most general use case of the
+algorithm, but you should consider building your own training loops and combining/developping your own sub-modules to 
+create tailored models and pipelines.
+
+---
+## About configurations
+
+MIMOSA has many training configurations to make it as precise as possible when modelling your specific data.
+These configurations are stored in a `ModelConfig` dataclass. 
+Here is a quick overview of what each configuration option does. 
+
+* `shared_task_hps` (default `True`): whether task kernel/noise hyper-parameters are shared across all tasks, or learnt independently per task.
+* `shared_cluster_hps` (default `True`): whether cluster mean/kernel hyper-parameters are shared across all clusters (mean-processes), or learnt independently per cluster.
+* `shared_output_hps` (default `True`): whether hyper-parameters are shared across output dimensions, or learnt independently per output.
+* `shared_features_hps` (default `True`): whether hyper-parameters are shared across features, or learnt independently per feature.
+* `cluster_specific_task_hps` (default `True`): whether the task/noise kernels themselves vary depending on the cluster a task belongs to, or are shared across clusters.
+* `isotopic_tasks` (default `True`): whether all tasks are observed at the same input locations, or each task has its own set of inputs.
+* `isotopic_features` (default `True`): whether all features are observed at the same input locations, or each feature has its own set of inputs.
+
+These configurations are independent of your choices of parameters (mean functions/kernels), leading to a large
+number of ways to model your data.
+
+The best way to get a feeling for which setup might be the best for your specific needs is to play
+around with `generate_data` and varying configurations and parameters until you get something that looks like
+your specific dataset.
+
+---
+## Troubleshooting
+
+As with any GP-based framework, mimosa is subject to numerical stability issues.
+
+We highly recommend to use Jax' 64bit precision:
+
+```python
+import jax
+jax.config.update("jax_enable_x64", True)  # Should be called before any other jax/kernax/mimosa import
+# ...other imports
+```
+
+If you still have troubles (NaN in hyperpost/likelihoods, optimisation does not converge, odd-looking mean 
+processes/predictions/samples, ...), consider increasing the `jitter` used in the specific component or
+directly at the Model level:
+
+```python
+model = BasicModel(prng_key=key, n_clusters=K, jitter=jnp.asarray(1e-5))
+```
+
+By default, Mimosa's jitter is fixed at 1e-8 for every component of the framework.
+
+---
+## Authors and citation
+
+MIMOSA is primarily developed by the *Magma Task Force*, composed of:
+
+* [Arthur Leroy](https://arthur-leroy.netlify.app/), researcher at Paris Saclay and INRAe (FR), main author of the original tool [MagmaClustR](https://arthurleroy.github.io/MagmaClustR/) and coordinator of the Task Force.
+* [Simon Lejoly](https://researchportal.unamur.be/fr/persons/slejoly/), PhD student at UNamur (BE), main developper of the package and author of the [Kernax package](https://github.com/SimLej18/kernax-ml)
+* [Alexia Grenouillat]([TODO]), PhD student at INSA Toulouse (FR), developper of multi-feature correlation discovery.
+* [Térence Viellard]([TODO]), PhD student at INRAe (FR), working on sparse approximations and scalability.
+
+If you use the package for your research, please consider citing the following paper:
+
+[TODO]
+
+If your research field could provide an interesting example of using MIMOSA, please consider publishing a toy example in
+the documentation [TODO].
+
+---
+## Contributing
+
+At this stage of development, we are not yet accepting contributions to the codebase. However, if you are interested 
+in contributing to the project, please feel free to reach out to us via email or through our GitHub repository.
+
+---
+## License
+
+This work is distributed under the [MIT License](https://opensource.org/licenses/MIT). See `LICENSE` for more information.
+
+---
+
+

mimosa_ml-0.1.0a0/examples/basic_example.py

@@ -0,0 +1,140 @@
+"""
+Basic usage example of mimosa-ml.
+
+Walks through the full pipeline: configure dimensions -> generate a synthetic dataset and remove
+some points at random -> plot the raw dataset -> fit a BasicModel -> plot the fitted clusters ->
+predict a task's output -> plot the prediction (mean + confidence interval) and samples from it.
+
+Meant to be run cell-by-cell (e.g. in PyCharm/VSCode's "#%%" notebook mode), or as a plain script.
+"""
+import jax
+
+jax.config.update("jax_enable_x64", True)
+jax.config.update("jax_disable_jit", False)
+import jax.random as jr
+import jax.numpy as jnp
+from jax import vmap
+import matplotlib.pyplot as plt
+
+from kernax import ZeroMean, VarianceKernel, SEKernel, WhiteNoiseKernel
+
+from mimosa import (
+	Dimensions, ModelConfig, DataRemovalConfig, Parameters,
+	BasicModel, generate_data, RandomDataRemover, save_csv, load_csv, build_parameters,
+)
+from mimosa.grid import UnionGrid
+from mimosa.plot import plot_dataset, plot_clusters, plot_single_task_prediction
+from mimosa.sampling import sample_gp
+
+key = jr.PRNGKey(42)
+plt.rcParams['figure.dpi']=300
+
+#%% 1. Configuration
+# Dimensions: T tasks, K clusters, I/O input/output dims, N points observed
+# per task, G points in the full grid.
+dims = Dimensions(T=32, K=2, I=1, O=2, F=1, N=50, G=150)
+
+# ModelConfig controls which hyperparameters are shared (across tasks/clusters/outputs/features) and
+# whether tasks/features share input locations.
+model_config = ModelConfig(
+	shared_task_hps=True,
+	shared_cluster_hps=True,
+	shared_output_hps=True,
+	cluster_specific_task_hps=False,
+	isotopic_tasks=False,
+)
+
+# How many points to remove at random per task, to simulate missing data.
+removal_config = DataRemovalConfig(max_missing=5, random_missing_count=True, same_missing_across_outputs=False)
+
+#%% 2. Generative parameters
+# These are the "true" parameters used to synthesise the toy dataset below. Swap any kernel/mean for
+# another kernax one (e.g. MaternKernel, PeriodicKernel, LinearMean, ...) to change the shape of the
+# data generated.
+true_params = Parameters(
+	cluster_mean=ZeroMean(),
+	cluster_kernel=VarianceKernel(5.0) * SEKernel(length_scale=.5),
+	task_kernel=VarianceKernel(1.0) * SEKernel(length_scale=.4),
+	noise_kernel=WhiteNoiseKernel(noise=.05),
+)
+
+#%% 3. Generate synthetic data, then remove points at random
+key, gen_key, removal_key = jr.split(key, 3)
+
+dataset, grid, hyperprior, true_mixture, true_params, cluster_means, tasks = generate_data(
+	gen_key, dims, true_params, model_config, input_range=(-2.5, 2.5)
+)
+dataset = RandomDataRemover()(removal_key, dataset, removal_config)
+
+#%% 3bis. Alternatively, you can load a dataset from a local file through load_csv.
+# Check load_csv's doc or open the csv file to see the expected file format.
+save_csv("./dummy.csv", dataset)
+dataset = load_csv("./dummy.csv")
+
+#%% 4. Plot the raw dataset (coloured by each task's true cluster)
+fig, ax = plot_dataset(dataset, mixture=true_mixture, figsize=(8*dims.O, 6))
+fig.suptitle("Synthetic dataset (colored by true cluster)")
+plt.show()
+
+#%% 5. Instantiate the model
+# n_clusters can differ from the true K above (the model doesn't know it); jitter is the numerical
+# stabiliser added before Cholesky factorizations, only increase it if you hit factorization errors.
+key, model_key = jr.split(key)
+model = BasicModel(prng_key=model_key, n_clusters=dims.K)
+
+# Starting guess for the parameters to fit. In practice these would be a rough, uninformed guess
+# rather than the true generative ones — feel free to try different starting kernels/values here.
+init_params = Parameters(
+		cluster_mean=ZeroMean(),
+		cluster_kernel=VarianceKernel(2.0) * SEKernel(length_scale=1.),
+		task_kernel=VarianceKernel(.5) * SEKernel(length_scale=.5),
+		noise_kernel=WhiteNoiseKernel(noise=0.5))
+
+# build_parameters batches the base kernels/mean below to match model_config's sharing structure
+# (same helper generate_data uses internally), so their shapes line up with what model.fit expects.
+init_params = build_parameters(init_params, dims, model_config)
+
+# proportions of the dataset in each cluster, a priori
+mixture_proportions = jnp.repeat(1 / dims.K, dims.K)  # fixed, equal weight per cluster
+
+#%% 6. Fit
+# Grid construction (union of every task's input points) isn't jit-compatible, so it's built once
+# here by the caller, outside of fit/predict, rather than owned by the model — see
+# mimosa.grid.GridBuilder. Swap UnionGrid for another GridBuilder to change how the grid is built.
+fitted_grid = UnionGrid()(dataset.inputs)
+
+fitted_params, fitted_mixture = model.fit(dataset, fitted_grid, mixture_proportions, init_params, n_iter=50)
+
+#%% 7. Plot the fitted clusters (mean-processes)
+hyperposterior = model.hyperpost(dataset, fitted_grid, fitted_mixture, fitted_params, jitter=model.jitter)
+
+fig, ax = plot_dataset(dataset, mixture=true_mixture, figsize=(8*dims.O, 6), alpha=.1)
+fig, ax = plot_clusters(fitted_grid, hyperposterior=hyperposterior, figsize=(8*dims.O, 6), fig=fig, ax=ax)
+fig.suptitle("Fitted clusters (mean-processes) on the dataset")
+plt.show()
+
+#%% 8. Predict
+predictions = model.predict(dataset, fitted_grid, fitted_mixture, fitted_params)  # MultivariateNormal, batched (T, K, O, G)
+
+t_id, o_id = 0, 0
+k_id = int(fitted_mixture.assignments[t_id])  # task's dominant cluster
+prediction = predictions[t_id, k_id, o_id]
+
+#%% 9. Plot the prediction: observed points, cluster means, and predictive mean + confidence interval
+fig, ax = plot_single_task_prediction(
+	dataset, fitted_grid, hyperposterior, fitted_mixture, t_id, o_id, prediction=prediction, figsize=(8*dims.O, 6)
+)
+fig.suptitle(f"Prediction — task {t_id}, output {o_id}")
+plt.show()
+
+#%% 10. Draw samples from the prediction and plot them alongside it
+key, sample_key = jr.split(key)
+n_samples = 64
+sample_keys = jr.split(sample_key, n_samples)
+samples = vmap(lambda k: sample_gp(k, prediction.mean, prediction.covariance))(sample_keys)  # (S, G)
+
+fig, ax = plot_single_task_prediction(
+	dataset, fitted_grid, hyperposterior, fitted_mixture, t_id, o_id, samples=samples, figsize=(8*dims.O, 6)
+)
+fig.suptitle(f"Prediction samples — task {t_id}, output {o_id}")
+plt.show()