diff --git a/docs/_images/sweep.png b/docs/_images/sweep.png
new file mode 100644
index 0000000..8869107
Binary files /dev/null and b/docs/_images/sweep.png differ
diff --git a/docs/_sources/tutorial.rst.txt b/docs/_sources/tutorial.rst.txt
index b7804b5..1797887 100644
--- a/docs/_sources/tutorial.rst.txt
+++ b/docs/_sources/tutorial.rst.txt
@@ -3,5 +3,5 @@ Examples
 
 .. automodule:: examples.00_intro
 
-
+.. automodule:: examples.01_sweep
 
diff --git a/docs/genindex.html b/docs/genindex.html
index a14c825..8b53cf6 100644
--- a/docs/genindex.html
+++ b/docs/genindex.html
@@ -131,15 +131,22 @@ <h2 id="E">E</h2>
 </li>
       <li><a href="modules.html#vbjax.neural_mass.BOLDTheta.epsilon">epsilon (vbjax.neural_mass.BOLDTheta attribute)</a>
 </li>
-  </ul></td>
-  <td style="width: 33%; vertical-align: top;"><ul>
       <li><a href="modules.html#vbjax.neural_mass.MPRTheta.eta">eta (vbjax.neural_mass.MPRTheta attribute)</a>
 </li>
+  </ul></td>
+  <td style="width: 33%; vertical-align: top;"><ul>
       <li>
     examples.00_intro
 
       <ul>
         <li><a href="tutorial.html#module-examples.00_intro">module</a>
+</li>
+      </ul></li>
+      <li>
+    examples.01_sweep
+
+      <ul>
+        <li><a href="tutorial.html#module-examples.01_sweep">module</a>
 </li>
       </ul></li>
   </ul></td>
@@ -277,6 +284,8 @@ <h2 id="M">M</h2>
 
       <ul>
         <li><a href="tutorial.html#module-examples.00_intro">examples.00_intro</a>
+</li>
+        <li><a href="tutorial.html#module-examples.01_sweep">examples.01_sweep</a>
 </li>
         <li><a href="modules.html#module-vbjax.connectome">vbjax.connectome</a>
 </li>
diff --git a/docs/images/sweep.png b/docs/images/sweep.png
new file mode 100644
index 0000000..8869107
Binary files /dev/null and b/docs/images/sweep.png differ
diff --git a/docs/modules.html b/docs/modules.html
index 7cd994d..84c5c65 100644
--- a/docs/modules.html
+++ b/docs/modules.html
@@ -14,7 +14,7 @@
     <script src="_static/sphinx_highlight.js?v=dc90522c"></script>
     <link rel="index" title="Index" href="genindex.html" />
     <link rel="search" title="Search" href="search.html" />
-    <link rel="prev" title="Tutorial" href="tutorial.html" /> 
+    <link rel="prev" title="Examples" href="tutorial.html" /> 
   </head><body>
     <div class="related" role="navigation" aria-label="related navigation">
       <h3>Navigation</h3>
@@ -26,7 +26,7 @@ <h3>Navigation</h3>
           <a href="py-modindex.html" title="Python Module Index"
              >modules</a> |</li>
         <li class="right" >
-          <a href="tutorial.html" title="Tutorial"
+          <a href="tutorial.html" title="Examples"
              accesskey="P">previous</a> |</li>
         <li class="nav-item nav-item-0"><a href="index.html">vbjax v0.0.10 documentation</a> &#187;</li>
         <li class="nav-item nav-item-this"><a href="">neural_mass</a></li> 
@@ -923,7 +923,7 @@ <h3><a href="index.html">Table of Contents</a></h3>
   <div>
     <h4>Previous topic</h4>
     <p class="topless"><a href="tutorial.html"
-                          title="previous chapter">Tutorial</a></p>
+                          title="previous chapter">Examples</a></p>
   </div>
   <div role="note" aria-label="source link">
     <h3>This Page</h3>
@@ -956,7 +956,7 @@ <h3>Navigation</h3>
           <a href="py-modindex.html" title="Python Module Index"
              >modules</a> |</li>
         <li class="right" >
-          <a href="tutorial.html" title="Tutorial"
+          <a href="tutorial.html" title="Examples"
              >previous</a> |</li>
         <li class="nav-item nav-item-0"><a href="index.html">vbjax v0.0.10 documentation</a> &#187;</li>
         <li class="nav-item nav-item-this"><a href="">neural_mass</a></li> 
diff --git a/docs/objects.inv b/docs/objects.inv
index d6e1284..218fa2e 100644
Binary files a/docs/objects.inv and b/docs/objects.inv differ
diff --git a/docs/py-modindex.html b/docs/py-modindex.html
index 91623fe..d8aa38d 100644
--- a/docs/py-modindex.html
+++ b/docs/py-modindex.html
@@ -59,6 +59,11 @@ <h1>Python Module Index</h1>
        <td>&#160;&#160;&#160;
        <a href="tutorial.html#module-examples.00_intro"><code class="xref">examples.00_intro</code></a></td><td>
        <em></em></td></tr>
+     <tr class="cg-1">
+       <td></td>
+       <td>&#160;&#160;&#160;
+       <a href="tutorial.html#module-examples.01_sweep"><code class="xref">examples.01_sweep</code></a></td><td>
+       <em></em></td></tr>
      <tr class="pcap"><td></td><td>&#160;</td><td></td></tr>
      <tr class="cap" id="cap-v"><td></td><td>
        <strong>v</strong></td><td></td></tr>
diff --git a/docs/searchindex.js b/docs/searchindex.js
index a420dee..059f2af 100644
--- a/docs/searchindex.js
+++ b/docs/searchindex.js
@@ -1 +1 @@
-Search.setIndex({"docnames": ["index", "modules", "tutorial"], "filenames": ["index.rst", "modules.rst", "tutorial.rst"], "titles": ["Welcome to vbjax\u2019s documentation!", "neural_mass", "Examples"], "terms": {"i": [0, 1, 2], "jax": [0, 1, 2], "base": [0, 1], "packag": [0, 1], "work": 0, "virtual": 0, "brain": 0, "style": 0, "model": [0, 2], "pip": 0, "can": [0, 1], "us": [0, 1, 2], "sourc": [0, 1], "git": 0, "clone": 0, "http": 0, "github": 0, "com": 0, "ins": 0, "amu": 0, "cd": 0, "dev": 0, "The": [0, 1, 2], "primari": 0, "addit": [0, 1], "depend": 0, "googl": 0, "which": [0, 2], "itself": 0, "onli": 0, "numpi": [0, 1, 2], "scipi": [0, 1], "opt": [0, 1], "einsum": 0, "so": [0, 1], "should": 0, "safe": 0, "add": [0, 1], "your": 0, "exist": 0, "project": 0, "cuda": 0, "If": [0, 1], "have": 0, "enabl": 0, "gpu": 0, "requisit": 0, "like": 0, "upgrad": 0, "cuda11_pip": 0, "f": 0, "storag": 0, "googleapi": 0, "releas": 0, "jax_cuda_releas": 0, "html": 0, "m1": 0, "m2": 0, "On": 0, "newer": 0, "appl": 0, "machin": 0, "w": [0, 1], "support": [0, 1], "experiment": 0, "just": 0, "two": 0, "extra": [0, 1], "ml": 0, "dtype": [0, 1], "0": [0, 1, 2], "2": [0, 1, 2], "metal": 0, "about": 0, "third": 0, "test": 0, "fail": 0, "due": 0, "absenc": 0, "certain": 0, "oper": [0, 1], "n": [0, 1], "dim": 0, "scatter": [0, 1], "gather": [0, 1], "fft": 0, "mai": 0, "faster": 0, "becaus": 0, "cpu": 0, "alreadi": 0, "excel": 0, "memori": [0, 1], "bandwidth": 0, "latenc": 0, "hide": 0, "BUT": 0, "softwar": 0, "stack": 0, "version": 0, "make": [0, 1], "align": 0, "star": 0, "look": 0, "child": 0, "plai": 0, "contain": [0, 1], "imag": [0, 2], "ar": [0, 2], "avail": 0, "auto": 0, "built": 0, "action": 0, "workflow": 0, "docker": 0, "yml": 0, "run": 0, "rm": 0, "ghcr": 0, "io": 0, "main": 0, "python3": [0, 1], "c": [0, 1, 2], "import": [0, 1, 2], "print": 0, "__version__": 0, "nvidia": 0, "runtim": 0, "all": [0, 1], "enough": 0, "discov": 0, "index": [0, 1], "search": 0, "page": 0, "neural_mass": 0, "boldtheta": [0, 1], "bveptheta": [0, 1], "jrstate": [0, 1], "jrtheta": [0, 1], "mprstate": [0, 1], "mprtheta": [0, 1], "bold_dfun": [0, 1], "bvep_dfun": [0, 1], "compute_bold_theta": [0, 1], "jr_dfun": [0, 1], "mpr_dfun": [0, 1, 2], "monitor": 0, "compute_sarvas_gain": [0, 1], "make_bold": [0, 1], "make_cov": [0, 1], "make_fc": [0, 1], "make_fft": [0, 1], "make_gain": [0, 1], "make_offlin": [0, 1], "make_timeavg": [0, 1], "loop": [0, 2], "heun_step": [0, 1], "make_dd": [0, 1], "make_od": [0, 1], "make_sdd": [0, 1], "make_sd": [0, 1, 2], "connectom": 0, "make_conn_latent_mvnorm": [0, 1], "layer": 0, "make_dense_lay": [0, 1], "spars": 0, "csr_to_jax_bcoo": [0, 1], "make_sg_spmv": [0, 1], "make_spmv": [0, 1], "util": 0, "to_jax": [0, 1], "to_np": [0, 1], "shtlc": 0, "grid_pairwise_dist": [0, 1], "kernel_conv_prep": [0, 1], "kernel_diff": [0, 1], "kernel_dist_origin": [0, 1], "kernel_estimate_shtn": [0, 1], "kernel_gaussian": [0, 1], "kernel_laplac": [0, 1], "kernel_mexican_hat": [0, 1], "kernel_sh_norm": [0, 1], "make_grid": [0, 1], "make_grid_shtn": [0, 1], "make_lm": [0, 1], "make_shtdiff": [0, 1], "make_shtdiff_np": [0, 1], "randn_relax": [0, 1], "class": 1, "vbjax": [1, 2], "tau_": 1, "tau_f": 1, "tau_o": 1, "alpha": 1, "te": 1, "v0": 1, "e0": 1, "epsilon": 1, "nu_0": 1, "r_0": 1, "recip_tau_": 1, "recip_tau_f": 1, "recip_tau_o": 1, "recip_alpha": 1, "recip_e0": 1, "k1": 1, "k2": 1, "k3": 1, "tupl": 1, "alia": 1, "field": 1, "number": 1, "3": 1, "6": 1, "7": 1, "15": 1, "16": 1, "17": 1, "8": 1, "9": 1, "13": 1, "14": 1, "11": 1, "12": 1, "10": 1, "1": [1, 2], "4": 1, "5": 1, "tau0": 1, "i1": 1, "x0": 1, "y0": 1, "y1": 1, "y2": 1, "y3": 1, "y4": 1, "y5": 1, "A": [1, 2], "b": 1, "nu_max": 1, "r": 1, "j": 1, "a_1": 1, "a_2": 1, "a_3": 1, "a_4": 1, "mu": 1, "v": 1, "tau": 1, "delta": 1, "eta": 1, "cr": 1, "cv": 1, "sfvq": 1, "x": [1, 2], "p": [1, 2], "y": 1, "65": 1, "41": 1, "98": 1, "32": [1, 2], "04": 1, "40": 1, "25": 1, "q": 1, "o": [1, 2], "att": 1, "d": 1, "dc": 1, "arrai": 1, "shape": 1, "dt": [1, 2], "bold": 1, "fmri": 1, "period": 1, "gain": 1, "none": 1, "matrix": 1, "suitabl": 1, "seeg": 1, "eeg": 1, "meg": 1, "step_fn": 1, "sample_fn": 1, "arg": 1, "comput": 1, "sampl": [1, 2], "an": 1, "offlin": 1, "batch": 1, "fashion": 1, "time": [1, 2], "averag": 1, "function": [1, 2], "build": 1, "step": [1, 2], "dfun": [1, 2], "heun": 1, "scheme": 1, "state": [1, 2], "right": 1, "hand": 1, "side": 1, "forc": 1, "term": [1, 2], "nh": 1, "unrol": 1, "invok": 1, "gfun": [1, 2], "integr": 1, "autonom": 1, "ordinari": 1, "differenti": [1, 2], "equat": [1, 2], "od": 1, "paramet": [1, 2], "float": 1, "form": 1, "deriv": [1, 2], "return": [1, 2], "t": 1, "take": [1, 2], "one": 1, "accord": 1, "iter": 1, "call": 1, "note": 1, "In": 1, "both": 1, "case": 1, "compat": 1, "set": [1, 2], "provid": 1, "either": 1, "some": 1, "pytre": 1, "structur": 1, "vb": [1, 2], "np": [1, 2], "_": [1, 2], "lambda": 1, "r_": 1, "125": 1, "0625": 1, "float32": 1, "weak_typ": 1, "true": [1, 2], "zero_delai": 1, "fals": [1, 2], "stochast": [1, 2], "delai": 1, "sde": 1, "int": 1, "maximum": 1, "xt": 1, "drift": 1, "coeffici": 1, "diffus": 1, "numer": 1, "valu": 1, "thi": [1, 2], "constant": 1, "linear": 1, "x_t": 1, "z_t": 1, "each": [1, 2], "start": [1, 2], "doe": 1, "normal": 1, "distribut": 1, "nois": [1, 2], "must": 1, "user": 1, "histori": 1, "buffer": 1, "pass": 1, "corrector": 1, "stage": 1, "method": 1, "predictor": 1, "perform": 1, "reason": 1, "unless": 1, "good": 1, "compromis": 1, "zero": 1, "sdde": 1, "ones": 1, "22": 1, "z": [1, 2], "randn": [1, 2], "5093468": 1, "30794007": 1, "07600437": 1, "03876263": 1, "sc": 1, "nc": 1, "return_ful": 1, "latent": 1, "multivari": 1, "over": 1, "nconn": 1, "given": [1, 2], "parcel": 1, "option": 1, "compon": 1, "space": 1, "bool": 1, "whether": 1, "inform": 1, "svd": 1, "u_mean": 1, "array_lik": 1, "mean": 1, "u_cov": 1, "covari": 1, "xfm": 1, "map": 1, "vector": [1, 2], "full": 1, "u": 1, "left": 1, "singular": 1, "correspond": 1, "embed": 1, "": 1, "vt": 1, "nconf": 1, "confus": 1, "induc": 1, "dimension": 1, "reduct": 1, "in_dim": 1, "latent_dim": 1, "out_dim": 1, "init_scl": 1, "extra_in": 1, "act_fn": 1, "pjitfunct": 1, "leaky_relu": 1, "kei": 1, "42": 1, "uint32": 1, "dens": 1, "neural": 1, "network": [1, 2], "size": 1, "csr_matrix": 1, "convert": 1, "csr": 1, "format": [1, 2], "coo": 1, "use_pmap": 1, "shard": 1, "positionalshard": 1, "spmv": 1, "kernel": 1, "gener": [1, 2], "is_symmetr": 1, "use_scipi": 1, "closur": 1, "multipl": 1, "default": 1, "symmetr": 1, "implement": 1, "spase": 1, "multipli": 1, "gradient": 1, "type": 1, "ndarrai": 1, "move": 1, "via": 1, "dlpack": 1, "sht": 1, "local": 1, "coupl": [1, 2], "theta": 1, "phi": 1, "pairwis": 1, "distanc": 1, "grid": 1, "intens": 1, "larg": 1, "k": 1, "prepar": 1, "evalu": 1, "convolut": 1, "l": 1, "shtn": 1, "order": 1, "origin": 1, "estim": 1, "effect": 1, "object": 1, "spatial": 1, "spectral": 1, "gaussian": 1, "laplacian": 1, "mexican": 1, "hat": 1, "nlat": 1, "nlon": 1, "creat": [1, 2], "latitud": 1, "longitud": 1, "lmax": 1, "0004": 1, "return_l": 1, "construct": 1, "diff": 1, "modul": 1, "from": 1, "tng": 1, "anaconda3": 1, "env": 1, "lib": 1, "site": 1, "__init__": 1, "py": [1, 2], "plain": 1, "For": 1, "random": 1, "captur": 1, "exampl": 0, "simpl": 2, "montbrio": 2, "node": 2, "few": 2, "00_intro": 2, "makedir": 2, "exist_ok": 2, "show": 2, "how": 2, "librari": 2, "simul": 2, "defin": 2, "argument": 2, "def": 2, "03": 2, "sum": 2, "axi": 2, "interv": 2, "initi": 2, "condit": 2, "01": 2, "dimes": 2, "first": 2, "remain": 2, "500": 2, "mpr_default_theta": 2, "plot_stat": 2, "plot": 2, "name": 2, "output": 2, "file": 2, "rv": 2, "jpg": 2, "example1": 2}, "objects": {"examples": [[2, 0, 0, "-", "00_intro"]], "vbjax": [[1, 0, 0, "-", "connectome"], [1, 0, 0, "-", "layers"], [1, 0, 0, "-", "loops"], [1, 0, 0, "-", "monitor"], [1, 0, 0, "-", "neural_mass"], [1, 0, 0, "-", "shtlc"], [1, 0, 0, "-", "sparse"], [1, 0, 0, "-", "util"]], "vbjax.connectome": [[1, 1, 1, "", "make_conn_latent_mvnorm"]], "vbjax.layers": [[1, 1, 1, "", "make_dense_layers"]], "vbjax.loops": [[1, 1, 1, "", "heun_step"], [1, 1, 1, "", "make_dde"], [1, 1, 1, "", "make_ode"], [1, 1, 1, "", "make_sdde"], [1, 1, 1, "", "make_sde"]], "vbjax.monitor": [[1, 1, 1, "", "compute_sarvas_gain"], [1, 1, 1, "", "make_bold"], [1, 1, 1, "", "make_cov"], [1, 1, 1, "", "make_fc"], [1, 1, 1, "", "make_fft"], [1, 1, 1, "", "make_gain"], [1, 1, 1, "", "make_offline"], [1, 1, 1, "", "make_timeavg"]], "vbjax.neural_mass": [[1, 2, 1, "", "BOLDTheta"], [1, 2, 1, "", "BVEPTheta"], [1, 2, 1, "", "JRState"], [1, 2, 1, "", "JRTheta"], [1, 2, 1, "", "MPRState"], [1, 2, 1, "", "MPRTheta"], [1, 1, 1, "", "bold_dfun"], [1, 1, 1, "", "bvep_dfun"], [1, 1, 1, "", "compute_bold_theta"], [1, 1, 1, "", "jr_dfun"], [1, 1, 1, "", "mpr_dfun"]], "vbjax.neural_mass.BOLDTheta": [[1, 3, 1, "", "alpha"], [1, 3, 1, "", "e0"], [1, 3, 1, "", "epsilon"], [1, 3, 1, "", "k1"], [1, 3, 1, "", "k2"], [1, 3, 1, "", "k3"], [1, 3, 1, "", "nu_0"], [1, 3, 1, "", "r_0"], [1, 3, 1, "", "recip_alpha"], [1, 3, 1, "", "recip_e0"], [1, 3, 1, "", "recip_tau_f"], [1, 3, 1, "", "recip_tau_o"], [1, 3, 1, "", "recip_tau_s"], [1, 3, 1, "", "tau_f"], [1, 3, 1, "", "tau_o"], [1, 3, 1, "", "tau_s"], [1, 3, 1, "", "te"], [1, 3, 1, "", "v0"]], "vbjax.neural_mass.BVEPTheta": [[1, 3, 1, "", "I1"], [1, 3, 1, "", "tau0"], [1, 3, 1, "", "x0"]], "vbjax.neural_mass.JRState": [[1, 3, 1, "", "y0"], [1, 3, 1, "", "y1"], [1, 3, 1, "", "y2"], [1, 3, 1, "", "y3"], [1, 3, 1, "", "y4"], [1, 3, 1, "", "y5"]], "vbjax.neural_mass.JRTheta": [[1, 3, 1, "", "A"], [1, 3, 1, "", "B"], [1, 3, 1, "", "I"], [1, 3, 1, "", "J"], [1, 3, 1, "", "a"], [1, 3, 1, "", "a_1"], [1, 3, 1, "", "a_2"], [1, 3, 1, "", "a_3"], [1, 3, 1, "", "a_4"], [1, 3, 1, "", "b"], [1, 3, 1, "", "mu"], [1, 3, 1, "", "nu_max"], [1, 3, 1, "", "r"], [1, 3, 1, "", "v0"]], "vbjax.neural_mass.MPRState": [[1, 3, 1, "", "V"], [1, 3, 1, "", "r"]], "vbjax.neural_mass.MPRTheta": [[1, 3, 1, "", "Delta"], [1, 3, 1, "", "I"], [1, 3, 1, "", "J"], [1, 3, 1, "", "cr"], [1, 3, 1, "", "cv"], [1, 3, 1, "", "eta"], [1, 3, 1, "", "tau"]], "vbjax.shtlc": [[1, 1, 1, "", "grid_pairwise_distance"], [1, 1, 1, "", "kernel_conv_prep"], [1, 1, 1, "", "kernel_diff"], [1, 1, 1, "", "kernel_dist_origin"], [1, 1, 1, "", "kernel_estimate_shtns"], [1, 1, 1, "", "kernel_gaussian"], [1, 1, 1, "", "kernel_laplace"], [1, 1, 1, "", "kernel_mexican_hat"], [1, 1, 1, "", "kernel_sh_normalized"], [1, 1, 1, "", "make_grid"], [1, 1, 1, "", "make_grid_shtns"], [1, 1, 1, "", "make_lm"], [1, 1, 1, "", "make_shtdiff"], [1, 1, 1, "", "make_shtdiff_np"], [1, 1, 1, "", "randn_relaxed"]], "vbjax.sparse": [[1, 1, 1, "", "csr_to_jax_bcoo"], [1, 1, 1, "", "make_sg_spmv"], [1, 1, 1, "", "make_spmv"]], "vbjax.util": [[1, 1, 1, "", "to_jax"], [1, 1, 1, "", "to_np"]]}, "objtypes": {"0": "py:module", "1": "py:function", "2": "py:class", "3": "py:attribute"}, "objnames": {"0": ["py", "module", "Python module"], "1": ["py", "function", "Python function"], "2": ["py", "class", "Python class"], "3": ["py", "attribute", "Python attribute"]}, "titleterms": {"welcom": 0, "vbjax": 0, "": 0, "document": 0, "introduct": 0, "instal": 0, "gee": 0, "pee": 0, "you": 0, "tutori": 0, "content": 0, "modul": 0, "neural_mass": 1, "monitor": 1, "loop": 1, "connectom": 1, "layer": 1, "spars": 1, "util": 1, "shtlc": 1, "exampl": 2}, "envversion": {"sphinx.domains.c": 3, "sphinx.domains.changeset": 1, "sphinx.domains.citation": 1, "sphinx.domains.cpp": 9, "sphinx.domains.index": 1, "sphinx.domains.javascript": 3, "sphinx.domains.math": 2, "sphinx.domains.python": 4, "sphinx.domains.rst": 2, "sphinx.domains.std": 2, "sphinx.ext.todo": 2, "sphinx.ext.viewcode": 1, "sphinx": 60}, "alltitles": {"neural_mass": [[1, "module-vbjax.neural_mass"]], "monitor": [[1, "module-vbjax.monitor"]], "loops": [[1, "module-vbjax.loops"]], "connectome": [[1, "module-vbjax.connectome"]], "layers": [[1, "module-vbjax.layers"]], "sparse": [[1, "module-vbjax.sparse"]], "util": [[1, "module-vbjax.util"]], "shtlc": [[1, "module-vbjax.shtlc"]], "Welcome to vbjax\u2019s documentation!": [[0, "welcome-to-vbjax-s-documentation"]], "Introduction": [[0, "introduction"]], "Installation": [[0, "installation"]], "gee pee you": [[0, "gee-pee-you"]], "Tutorial": [[0, "tutorial"]], "Contents:": [[0, null], [0, null]], "Modules": [[0, "modules"]], "Examples": [[2, "module-examples.00_intro"]]}, "indexentries": {"examples.00_intro": [[2, "module-examples.00_intro"]], "module": [[2, "module-examples.00_intro"]]}})
\ No newline at end of file
+Search.setIndex({"docnames": ["index", "modules", "tutorial"], "filenames": ["index.rst", "modules.rst", "tutorial.rst"], "titles": ["Welcome to vbjax\u2019s documentation!", "neural_mass", "Examples"], "terms": {"i": [0, 1, 2], "jax": [0, 1, 2], "base": [0, 1], "packag": [0, 1], "work": 0, "virtual": 0, "brain": 0, "style": 0, "model": [0, 2], "pip": 0, "can": [0, 1], "us": [0, 1, 2], "sourc": [0, 1], "git": 0, "clone": 0, "http": 0, "github": 0, "com": 0, "ins": 0, "amu": 0, "cd": 0, "dev": 0, "The": [0, 1, 2], "primari": 0, "addit": [0, 1], "depend": 0, "googl": 0, "which": [0, 2], "itself": 0, "onli": 0, "numpi": [0, 1, 2], "scipi": [0, 1], "opt": [0, 1], "einsum": 0, "so": [0, 1], "should": 0, "safe": 0, "add": [0, 1], "your": 0, "exist": 0, "project": 0, "cuda": 0, "If": [0, 1, 2], "have": 0, "enabl": 0, "gpu": [0, 2], "requisit": 0, "like": 0, "upgrad": 0, "cuda11_pip": 0, "f": [0, 2], "storag": 0, "googleapi": 0, "releas": 0, "jax_cuda_releas": 0, "html": 0, "m1": 0, "m2": 0, "On": 0, "newer": 0, "appl": 0, "machin": 0, "w": [0, 1, 2], "support": [0, 1], "experiment": 0, "just": [0, 2], "two": 0, "extra": [0, 1], "ml": 0, "dtype": [0, 1], "0": [0, 1, 2], "2": [0, 1, 2], "metal": 0, "about": 0, "third": 0, "test": 0, "fail": 0, "due": 0, "absenc": 0, "certain": 0, "oper": [0, 1], "n": [0, 1], "dim": 0, "scatter": [0, 1], "gather": [0, 1], "fft": 0, "mai": 0, "faster": 0, "becaus": 0, "cpu": [0, 2], "alreadi": 0, "excel": 0, "memori": [0, 1], "bandwidth": 0, "latenc": 0, "hide": 0, "BUT": 0, "softwar": 0, "stack": 0, "version": 0, "make": [0, 1], "align": 0, "star": 0, "look": 0, "child": 0, "plai": 0, "contain": [0, 1], "imag": [0, 2], "ar": [0, 2], "avail": 0, "auto": 0, "built": 0, "action": 0, "workflow": 0, "docker": 0, "yml": 0, "run": [0, 2], "rm": 0, "ghcr": 0, "io": 0, "main": 0, "python3": [0, 1], "c": [0, 1, 2], "import": [0, 1, 2], "print": [0, 2], "__version__": 0, "nvidia": 0, "runtim": 0, "all": [0, 1], "enough": [0, 2], "discov": 0, "exampl": 0, "index": [0, 1], "search": 0, "page": 0, "neural_mass": 0, "boldtheta": [0, 1], "bveptheta": [0, 1], "jrstate": [0, 1], "jrtheta": [0, 1], "mprstate": [0, 1], "mprtheta": [0, 1], "bold_dfun": [0, 1], "bvep_dfun": [0, 1], "compute_bold_theta": [0, 1], "jr_dfun": [0, 1], "mpr_dfun": [0, 1, 2], "monitor": 0, "compute_sarvas_gain": [0, 1], "make_bold": [0, 1], "make_cov": [0, 1], "make_fc": [0, 1], "make_fft": [0, 1], "make_gain": [0, 1], "make_offlin": [0, 1], "make_timeavg": [0, 1], "loop": [0, 2], "heun_step": [0, 1], "make_dd": [0, 1], "make_od": [0, 1], "make_sdd": [0, 1], "make_sd": [0, 1, 2], "connectom": 0, "make_conn_latent_mvnorm": [0, 1], "layer": 0, "make_dense_lay": [0, 1], "spars": 0, "csr_to_jax_bcoo": [0, 1], "make_sg_spmv": [0, 1], "make_spmv": [0, 1], "util": 0, "to_jax": [0, 1], "to_np": [0, 1], "shtlc": 0, "grid_pairwise_dist": [0, 1], "kernel_conv_prep": [0, 1], "kernel_diff": [0, 1], "kernel_dist_origin": [0, 1], "kernel_estimate_shtn": [0, 1], "kernel_gaussian": [0, 1], "kernel_laplac": [0, 1], "kernel_mexican_hat": [0, 1], "kernel_sh_norm": [0, 1], "make_grid": [0, 1], "make_grid_shtn": [0, 1], "make_lm": [0, 1], "make_shtdiff": [0, 1], "make_shtdiff_np": [0, 1], "randn_relax": [0, 1], "class": 1, "vbjax": [1, 2], "tau_": 1, "tau_f": 1, "tau_o": 1, "alpha": 1, "te": 1, "v0": 1, "e0": 1, "epsilon": 1, "nu_0": 1, "r_0": 1, "recip_tau_": 1, "recip_tau_f": 1, "recip_tau_o": 1, "recip_alpha": 1, "recip_e0": 1, "k1": 1, "k2": 1, "k3": 1, "tupl": 1, "alia": 1, "field": 1, "number": [1, 2], "3": [1, 2], "6": [1, 2], "7": [1, 2], "15": 1, "16": 1, "17": 1, "8": [1, 2], "9": [1, 2], "13": 1, "14": 1, "11": 1, "12": 1, "10": 1, "1": [1, 2], "4": [1, 2], "5": 1, "tau0": 1, "i1": 1, "x0": 1, "y0": 1, "y1": 1, "y2": 1, "y3": 1, "y4": 1, "y5": 1, "A": [1, 2], "b": 1, "nu_max": 1, "r": [1, 2], "j": 1, "a_1": 1, "a_2": 1, "a_3": 1, "a_4": 1, "mu": 1, "v": [1, 2], "tau": 1, "delta": 1, "eta": [1, 2], "cr": 1, "cv": 1, "sfvq": 1, "x": [1, 2], "p": [1, 2], "y": 1, "65": 1, "41": 1, "98": 1, "32": [1, 2], "04": 1, "40": 1, "25": 1, "q": 1, "o": [1, 2], "att": 1, "d": 1, "dc": 1, "arrai": 1, "shape": [1, 2], "dt": [1, 2], "bold": 1, "fmri": 1, "period": 1, "gain": 1, "none": 1, "matrix": 1, "suitabl": 1, "seeg": 1, "eeg": 1, "meg": 1, "step_fn": 1, "sample_fn": 1, "arg": 1, "comput": 1, "sampl": [1, 2], "an": 1, "offlin": 1, "batch": [1, 2], "fashion": 1, "time": [1, 2], "averag": 1, "function": [1, 2], "build": 1, "step": [1, 2], "dfun": [1, 2], "heun": 1, "scheme": 1, "state": [1, 2], "right": 1, "hand": 1, "side": 1, "forc": 1, "term": [1, 2], "nh": 1, "unrol": 1, "invok": 1, "gfun": [1, 2], "integr": 1, "autonom": 1, "ordinari": 1, "differenti": [1, 2], "equat": [1, 2], "od": 1, "paramet": [1, 2], "float": 1, "form": 1, "deriv": [1, 2], "return": [1, 2], "t": [1, 2], "take": [1, 2], "one": 1, "accord": 1, "iter": 1, "call": 1, "note": 1, "In": 1, "both": 1, "case": 1, "compat": 1, "set": [1, 2], "provid": 1, "either": 1, "some": 1, "pytre": 1, "structur": 1, "vb": [1, 2], "np": [1, 2], "_": [1, 2], "lambda": 1, "r_": 1, "125": 1, "0625": 1, "float32": 1, "weak_typ": 1, "true": [1, 2], "zero_delai": 1, "fals": [1, 2], "stochast": [1, 2], "delai": 1, "sde": 1, "int": 1, "maximum": 1, "xt": 1, "drift": 1, "coeffici": 1, "diffus": 1, "numer": 1, "valu": [1, 2], "thi": [1, 2], "constant": 1, "linear": 1, "x_t": 1, "z_t": 1, "each": [1, 2], "start": [1, 2], "doe": 1, "normal": 1, "distribut": 1, "nois": [1, 2], "must": 1, "user": 1, "histori": 1, "buffer": 1, "pass": 1, "corrector": 1, "stage": 1, "method": 1, "predictor": 1, "perform": 1, "reason": 1, "unless": 1, "good": 1, "compromis": 1, "zero": 1, "sdde": 1, "ones": [1, 2], "22": 1, "z": [1, 2], "randn": [1, 2], "5093468": 1, "30794007": 1, "07600437": 1, "03876263": 1, "sc": 1, "nc": 1, "return_ful": 1, "latent": 1, "multivari": 1, "over": [1, 2], "nconn": 1, "given": [1, 2], "parcel": 1, "option": 1, "compon": 1, "space": 1, "bool": 1, "whether": 1, "inform": 1, "svd": 1, "u_mean": 1, "array_lik": 1, "mean": 1, "u_cov": 1, "covari": 1, "xfm": 1, "map": 1, "vector": [1, 2], "full": 1, "u": 1, "left": 1, "singular": 1, "correspond": 1, "embed": 1, "": [1, 2], "vt": 1, "nconf": 1, "confus": 1, "induc": 1, "dimension": 1, "reduct": 1, "in_dim": 1, "latent_dim": 1, "out_dim": 1, "init_scl": 1, "extra_in": 1, "act_fn": 1, "pjitfunct": 1, "leaky_relu": 1, "kei": 1, "42": 1, "uint32": 1, "dens": 1, "neural": 1, "network": [1, 2], "size": [1, 2], "csr_matrix": 1, "convert": 1, "csr": 1, "format": [1, 2], "coo": 1, "use_pmap": 1, "shard": 1, "positionalshard": 1, "spmv": 1, "kernel": 1, "gener": [1, 2], "is_symmetr": 1, "use_scipi": 1, "closur": 1, "multipl": 1, "default": 1, "symmetr": 1, "implement": 1, "spase": 1, "multipli": 1, "gradient": 1, "type": 1, "ndarrai": 1, "move": 1, "via": 1, "dlpack": 1, "sht": 1, "local": 1, "coupl": [1, 2], "theta": 1, "phi": 1, "pairwis": 1, "distanc": 1, "grid": [1, 2], "intens": 1, "larg": 1, "k": [1, 2], "prepar": [1, 2], "evalu": 1, "convolut": 1, "l": 1, "shtn": 1, "order": 1, "origin": 1, "estim": 1, "effect": 1, "object": 1, "spatial": 1, "spectral": 1, "gaussian": 1, "laplacian": 1, "mexican": 1, "hat": 1, "nlat": 1, "nlon": 1, "creat": [1, 2], "latitud": 1, "longitud": 1, "lmax": 1, "0004": 1, "return_l": 1, "construct": 1, "diff": 1, "modul": 1, "from": 1, "tng": 1, "anaconda3": 1, "env": 1, "lib": 1, "site": 1, "__init__": 1, "py": [1, 2], "plain": 1, "For": 1, "random": 1, "captur": 1, "simpl": 2, "montbrio": 2, "node": 2, "few": 2, "00_intro": 2, "makedir": 2, "exist_ok": 2, "show": 2, "how": 2, "librari": 2, "simul": 2, "defin": 2, "argument": 2, "def": 2, "03": 2, "sum": 2, "axi": 2, "interv": 2, "initi": 2, "condit": 2, "01": 2, "dimes": 2, "first": 2, "remain": 2, "500": 2, "mpr_default_theta": 2, "plot_stat": 2, "plot": 2, "name": 2, "output": 2, "file": 2, "rv": 2, "jpg": 2, "example1": 2, "consid": 2, "sweep": 2, "Then": 2, "engin": 2, "parallel": 2, "core": 2, "pmap": 2, "otherwis": 2, "vmap": 2, "we": 2, "rest": 2, "result": 2, "01_sweep": 2, "pylab": 2, "pl": 2, "mpr_p": 2, "sigma": 2, "par": 2, "sig": 2, "explor": 2, "_replac": 2, "mpr": 2, "rv0": 2, "sim": 2, "std": 2, "400": 2, "eval": 2, "metric": 2, "done": 2, "n_node": 2, "using_cpu": 2, "local_devic": 2, "platform": 2, "run_batch": 2, "in_ax": 2, "els": 2, "ng": 2, "sig_i": 2, "enumer": 2, "logarithm": 2, "scale": 2, "log_k": 2, "mgrid": 2, "1j": 2, "reshap": 2, "big": 2, "c_": 2, "exp": 2, "ravel": 2, "copi": 2, "expect": 2, "chunk": 2, "transpos": 2, "now": 2, "block_until_readi": 2, "append": 2, "toc": 2, "elaps": 2, "tic": 2, "1f": 2, "figur": 2, "figsiz": 2, "zip": 2, "subplot": 2, "imshow": 2, "vmin": 2, "vmax": 2}, "objects": {"examples": [[2, 0, 0, "-", "00_intro"], [2, 0, 0, "-", "01_sweep"]], "vbjax": [[1, 0, 0, "-", "connectome"], [1, 0, 0, "-", "layers"], [1, 0, 0, "-", "loops"], [1, 0, 0, "-", "monitor"], [1, 0, 0, "-", "neural_mass"], [1, 0, 0, "-", "shtlc"], [1, 0, 0, "-", "sparse"], [1, 0, 0, "-", "util"]], "vbjax.connectome": [[1, 1, 1, "", "make_conn_latent_mvnorm"]], "vbjax.layers": [[1, 1, 1, "", "make_dense_layers"]], "vbjax.loops": [[1, 1, 1, "", "heun_step"], [1, 1, 1, "", "make_dde"], [1, 1, 1, "", "make_ode"], [1, 1, 1, "", "make_sdde"], [1, 1, 1, "", "make_sde"]], "vbjax.monitor": [[1, 1, 1, "", "compute_sarvas_gain"], [1, 1, 1, "", "make_bold"], [1, 1, 1, "", "make_cov"], [1, 1, 1, "", "make_fc"], [1, 1, 1, "", "make_fft"], [1, 1, 1, "", "make_gain"], [1, 1, 1, "", "make_offline"], [1, 1, 1, "", "make_timeavg"]], "vbjax.neural_mass": [[1, 2, 1, "", "BOLDTheta"], [1, 2, 1, "", "BVEPTheta"], [1, 2, 1, "", "JRState"], [1, 2, 1, "", "JRTheta"], [1, 2, 1, "", "MPRState"], [1, 2, 1, "", "MPRTheta"], [1, 1, 1, "", "bold_dfun"], [1, 1, 1, "", "bvep_dfun"], [1, 1, 1, "", "compute_bold_theta"], [1, 1, 1, "", "jr_dfun"], [1, 1, 1, "", "mpr_dfun"]], "vbjax.neural_mass.BOLDTheta": [[1, 3, 1, "", "alpha"], [1, 3, 1, "", "e0"], [1, 3, 1, "", "epsilon"], [1, 3, 1, "", "k1"], [1, 3, 1, "", "k2"], [1, 3, 1, "", "k3"], [1, 3, 1, "", "nu_0"], [1, 3, 1, "", "r_0"], [1, 3, 1, "", "recip_alpha"], [1, 3, 1, "", "recip_e0"], [1, 3, 1, "", "recip_tau_f"], [1, 3, 1, "", "recip_tau_o"], [1, 3, 1, "", "recip_tau_s"], [1, 3, 1, "", "tau_f"], [1, 3, 1, "", "tau_o"], [1, 3, 1, "", "tau_s"], [1, 3, 1, "", "te"], [1, 3, 1, "", "v0"]], "vbjax.neural_mass.BVEPTheta": [[1, 3, 1, "", "I1"], [1, 3, 1, "", "tau0"], [1, 3, 1, "", "x0"]], "vbjax.neural_mass.JRState": [[1, 3, 1, "", "y0"], [1, 3, 1, "", "y1"], [1, 3, 1, "", "y2"], [1, 3, 1, "", "y3"], [1, 3, 1, "", "y4"], [1, 3, 1, "", "y5"]], "vbjax.neural_mass.JRTheta": [[1, 3, 1, "", "A"], [1, 3, 1, "", "B"], [1, 3, 1, "", "I"], [1, 3, 1, "", "J"], [1, 3, 1, "", "a"], [1, 3, 1, "", "a_1"], [1, 3, 1, "", "a_2"], [1, 3, 1, "", "a_3"], [1, 3, 1, "", "a_4"], [1, 3, 1, "", "b"], [1, 3, 1, "", "mu"], [1, 3, 1, "", "nu_max"], [1, 3, 1, "", "r"], [1, 3, 1, "", "v0"]], "vbjax.neural_mass.MPRState": [[1, 3, 1, "", "V"], [1, 3, 1, "", "r"]], "vbjax.neural_mass.MPRTheta": [[1, 3, 1, "", "Delta"], [1, 3, 1, "", "I"], [1, 3, 1, "", "J"], [1, 3, 1, "", "cr"], [1, 3, 1, "", "cv"], [1, 3, 1, "", "eta"], [1, 3, 1, "", "tau"]], "vbjax.shtlc": [[1, 1, 1, "", "grid_pairwise_distance"], [1, 1, 1, "", "kernel_conv_prep"], [1, 1, 1, "", "kernel_diff"], [1, 1, 1, "", "kernel_dist_origin"], [1, 1, 1, "", "kernel_estimate_shtns"], [1, 1, 1, "", "kernel_gaussian"], [1, 1, 1, "", "kernel_laplace"], [1, 1, 1, "", "kernel_mexican_hat"], [1, 1, 1, "", "kernel_sh_normalized"], [1, 1, 1, "", "make_grid"], [1, 1, 1, "", "make_grid_shtns"], [1, 1, 1, "", "make_lm"], [1, 1, 1, "", "make_shtdiff"], [1, 1, 1, "", "make_shtdiff_np"], [1, 1, 1, "", "randn_relaxed"]], "vbjax.sparse": [[1, 1, 1, "", "csr_to_jax_bcoo"], [1, 1, 1, "", "make_sg_spmv"], [1, 1, 1, "", "make_spmv"]], "vbjax.util": [[1, 1, 1, "", "to_jax"], [1, 1, 1, "", "to_np"]]}, "objtypes": {"0": "py:module", "1": "py:function", "2": "py:class", "3": "py:attribute"}, "objnames": {"0": ["py", "module", "Python module"], "1": ["py", "function", "Python function"], "2": ["py", "class", "Python class"], "3": ["py", "attribute", "Python attribute"]}, "titleterms": {"welcom": 0, "vbjax": 0, "": 0, "document": 0, "introduct": 0, "instal": 0, "gee": 0, "pee": 0, "you": 0, "tutori": 0, "content": 0, "modul": 0, "neural_mass": 1, "monitor": 1, "loop": 1, "connectom": 1, "layer": 1, "spars": 1, "util": 1, "shtlc": 1, "exampl": 2}, "envversion": {"sphinx.domains.c": 3, "sphinx.domains.changeset": 1, "sphinx.domains.citation": 1, "sphinx.domains.cpp": 9, "sphinx.domains.index": 1, "sphinx.domains.javascript": 3, "sphinx.domains.math": 2, "sphinx.domains.python": 4, "sphinx.domains.rst": 2, "sphinx.domains.std": 2, "sphinx.ext.todo": 2, "sphinx.ext.viewcode": 1, "sphinx": 60}, "alltitles": {"Welcome to vbjax\u2019s documentation!": [[0, "welcome-to-vbjax-s-documentation"]], "Introduction": [[0, "introduction"]], "Installation": [[0, "installation"]], "gee pee you": [[0, "gee-pee-you"]], "Tutorial": [[0, "tutorial"]], "Contents:": [[0, null], [0, null]], "Modules": [[0, "modules"]], "neural_mass": [[1, "module-vbjax.neural_mass"]], "monitor": [[1, "module-vbjax.monitor"]], "loops": [[1, "module-vbjax.loops"]], "connectome": [[1, "module-vbjax.connectome"]], "layers": [[1, "module-vbjax.layers"]], "sparse": [[1, "module-vbjax.sparse"]], "util": [[1, "module-vbjax.util"]], "shtlc": [[1, "module-vbjax.shtlc"]], "Examples": [[2, "module-examples.00_intro"]]}, "indexentries": {"examples.00_intro": [[2, "module-examples.00_intro"]], "examples.01_sweep": [[2, "module-examples.01_sweep"]], "module": [[2, "module-examples.00_intro"], [2, "module-examples.01_sweep"]]}})
\ No newline at end of file
diff --git a/docs/source/tutorial.rst b/docs/source/tutorial.rst
index b7804b5..1797887 100644
--- a/docs/source/tutorial.rst
+++ b/docs/source/tutorial.rst
@@ -3,5 +3,5 @@ Examples
 
 .. automodule:: examples.00_intro
 
-
+.. automodule:: examples.01_sweep
 
diff --git a/docs/tutorial.html b/docs/tutorial.html
index ed0458c..6b3e38e 100644
--- a/docs/tutorial.html
+++ b/docs/tutorial.html
@@ -89,7 +89,105 @@ <h3>Navigation</h3>
 </pre></div>
 </div>
 <figure class="align-default">
-<a class="reference internal image-reference" href="_images/example1.jpg"><img alt="_images/example1.jpg" src="_images/example1.jpg" style="width: 320.0px; height: 240.0px;" /></a>
+<a class="reference internal image-reference" href="_images/example1.jpg"><img alt="_images/example1.jpg" src="_images/example1.jpg" style="width: 480.0px; height: 360.0px;" /></a>
+</figure>
+<p id="module-examples.01_sweep"><strong>Example 2</strong>: Consider a network of coupled Montbrio model nodes and sweep
+over the parameters of the model.</p>
+<p>Starting with a few imports</p>
+<div class="literal-block-wrapper docutils container" id="id2">
+<div class="code-block-caption"><span class="caption-text">../../examples/01_sweep.py</span><a class="headerlink" href="#id2" title="Link to this code">¶</a></div>
+<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="kn">import</span> <span class="nn">os</span>
+<span class="kn">import</span> <span class="nn">time</span>
+<span class="kn">import</span> <span class="nn">vbjax</span> <span class="k">as</span> <span class="nn">vb</span>
+<span class="kn">import</span> <span class="nn">pylab</span> <span class="k">as</span> <span class="nn">pl</span>  
+<span class="kn">import</span> <span class="nn">jax</span><span class="o">,</span> <span class="nn">jax.numpy</span> <span class="k">as</span> <span class="nn">np</span>
+<span class="n">os</span><span class="o">.</span><span class="n">makedirs</span><span class="p">(</span><span class="s1">&#39;images&#39;</span><span class="p">,</span> <span class="n">exist_ok</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
+</pre></div>
+</div>
+</div>
+<p>This example shows how to use the <cite>vbjax</cite> library to simulate a network of
+Montbrio model nodes. The network is defined by the function <cite>network</cite> which
+takes as arguments the state of the network and the parameters of the model.</p>
+<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="k">def</span> <span class="nf">network</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="n">p</span><span class="p">):</span>
+    <span class="n">r</span><span class="p">,</span> <span class="n">v</span> <span class="o">=</span> <span class="n">x</span>
+    <span class="n">k</span><span class="p">,</span> <span class="n">_</span><span class="p">,</span> <span class="n">mpr_p</span> <span class="o">=</span> <span class="n">p</span>
+    <span class="n">c</span> <span class="o">=</span> <span class="n">k</span><span class="o">*</span><span class="n">r</span><span class="o">.</span><span class="n">sum</span><span class="p">(),</span> <span class="n">k</span><span class="o">*</span><span class="n">v</span><span class="o">.</span><span class="n">sum</span><span class="p">()</span>
+    <span class="k">return</span> <span class="n">vb</span><span class="o">.</span><span class="n">mpr_dfun</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="n">c</span><span class="p">,</span> <span class="n">mpr_p</span><span class="p">)</span>
+</pre></div>
+</div>
+<p>The function noise is used to generate the noise term of the stochastic</p>
+<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="k">def</span> <span class="nf">noise</span><span class="p">(</span><span class="n">_</span><span class="p">,</span> <span class="n">p</span><span class="p">):</span>
+    <span class="n">_</span><span class="p">,</span> <span class="n">sigma</span><span class="p">,</span> <span class="n">_</span> <span class="o">=</span> <span class="n">p</span>
+    <span class="k">return</span> <span class="n">sigma</span>
+</pre></div>
+</div>
+<p>The function run is used to simulate the network for a given set of parameters.</p>
+<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="k">def</span> <span class="nf">run</span><span class="p">(</span><span class="n">pars</span><span class="p">,</span> <span class="n">mpr_p</span><span class="o">=</span><span class="n">vb</span><span class="o">.</span><span class="n">mpr_default_theta</span><span class="p">):</span>
+    <span class="n">k</span><span class="p">,</span> <span class="n">sig</span><span class="p">,</span> <span class="n">eta</span> <span class="o">=</span> <span class="n">pars</span>                      <span class="c1"># explored pars</span>
+    <span class="n">p</span> <span class="o">=</span> <span class="n">k</span><span class="p">,</span> <span class="n">sig</span><span class="p">,</span> <span class="n">mpr_p</span><span class="o">.</span><span class="n">_replace</span><span class="p">(</span><span class="n">eta</span><span class="o">=</span><span class="n">eta</span><span class="p">)</span>     <span class="c1"># set mpr</span>
+    <span class="n">xs</span> <span class="o">=</span> <span class="n">loop</span><span class="p">(</span><span class="n">rv0</span><span class="p">,</span> <span class="n">zs</span><span class="p">,</span> <span class="n">p</span><span class="p">)</span>                   <span class="c1"># run sim</span>
+    <span class="n">std</span> <span class="o">=</span> <span class="n">xs</span><span class="p">[</span><span class="mi">400</span><span class="p">:,</span> <span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">std</span><span class="p">()</span>                 <span class="c1"># eval metric</span>
+    <span class="k">return</span> <span class="n">std</span>                              <span class="c1"># done</span>
+</pre></div>
+</div>
+<p>Then prepare the simulation</p>
+<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="n">n_nodes</span> <span class="o">=</span> <span class="mi">8</span>
+</pre></div>
+</div>
+<p>define the number of nodes in the network.</p>
+<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="n">using_cpu</span> <span class="o">=</span> <span class="n">jax</span><span class="o">.</span><span class="n">local_devices</span><span class="p">()[</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">platform</span> <span class="o">==</span> <span class="s1">&#39;cpu&#39;</span>
+<span class="k">if</span> <span class="n">using_cpu</span><span class="p">:</span>
+    <span class="n">run_batches</span> <span class="o">=</span> <span class="n">jax</span><span class="o">.</span><span class="n">pmap</span><span class="p">(</span><span class="n">jax</span><span class="o">.</span><span class="n">vmap</span><span class="p">(</span><span class="n">run</span><span class="p">,</span> <span class="n">in_axes</span><span class="o">=</span><span class="mi">1</span><span class="p">),</span> <span class="n">in_axes</span><span class="o">=</span><span class="mi">0</span><span class="p">)</span>
+<span class="k">else</span><span class="p">:</span>
+    <span class="n">run_batches</span> <span class="o">=</span> <span class="n">jax</span><span class="o">.</span><span class="n">vmap</span><span class="p">(</span><span class="n">run</span><span class="p">,</span> <span class="n">in_axes</span><span class="o">=</span><span class="mi">1</span><span class="p">)</span>
+</pre></div>
+</div>
+<p>defines the engine to be used for the simulation. If the simulation is run on
+the CPU, then the simulation is parallelized over the cores of the CPU using <cite>jax.pmap</cite>.
+Otherwise, the simulation is parallelized over the GPU using <cite>jax.vmap</cite>.</p>
+<p>then we prepare the network and the noise samples</p>
+<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="c1"># prepare network</span>
+<span class="n">_</span><span class="p">,</span> <span class="n">loop</span> <span class="o">=</span> <span class="n">vb</span><span class="o">.</span><span class="n">make_sde</span><span class="p">(</span><span class="mf">0.01</span><span class="p">,</span> <span class="n">network</span><span class="p">,</span> <span class="n">noise</span><span class="p">)</span>
+<span class="n">rv0</span> <span class="o">=</span> <span class="n">vb</span><span class="o">.</span><span class="n">randn</span><span class="p">(</span><span class="mi">2</span><span class="p">,</span> <span class="n">n_nodes</span><span class="p">)</span>
+</pre></div>
+</div>
+<p>The rest run the simulation on set of parameters and plot the results.</p>
+<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="c1"># sweep sigma but just a few values are enough</span>
+<span class="n">sigmas</span> <span class="o">=</span> <span class="p">[</span><span class="mf">0.0</span><span class="p">,</span> <span class="mf">0.2</span><span class="p">,</span> <span class="mf">0.3</span><span class="p">,</span> <span class="mf">0.4</span><span class="p">]</span>
+<span class="n">results</span> <span class="o">=</span> <span class="p">[]</span>
+<span class="n">ng</span> <span class="o">=</span> <span class="n">vb</span><span class="o">.</span><span class="n">cores</span><span class="o">*</span><span class="mi">4</span> <span class="k">if</span> <span class="n">using_cpu</span> <span class="k">else</span> <span class="mi">32</span>
+
+<span class="k">for</span> <span class="n">i</span><span class="p">,</span> <span class="n">sig_i</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="n">sigmas</span><span class="p">):</span>
+    <span class="c1"># create grid of k (on logarithmic scale) and eta</span>
+    <span class="n">log_ks</span><span class="p">,</span> <span class="n">etas</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">mgrid</span><span class="p">[</span><span class="o">-</span><span class="mf">9.0</span><span class="p">:</span><span class="o">-</span><span class="mf">2.0</span><span class="p">:</span><span class="mi">1</span><span class="n">j</span><span class="o">*</span><span class="n">ng</span><span class="p">,</span> <span class="o">-</span><span class="mf">4.0</span><span class="p">:</span><span class="o">-</span><span class="mf">6.0</span><span class="p">:</span><span class="mi">1</span><span class="n">j</span><span class="o">*</span><span class="n">ng</span><span class="p">]</span>
+
+    <span class="c1"># reshape grid to big batch of values</span>
+    <span class="n">pars</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">c_</span><span class="p">[</span>
+        <span class="n">np</span><span class="o">.</span><span class="n">exp</span><span class="p">(</span><span class="n">log_ks</span><span class="o">.</span><span class="n">ravel</span><span class="p">()),</span>
+        <span class="n">np</span><span class="o">.</span><span class="n">ones</span><span class="p">(</span><span class="n">log_ks</span><span class="o">.</span><span class="n">size</span><span class="p">)</span><span class="o">*</span><span class="n">sig_i</span><span class="p">,</span>
+        <span class="n">etas</span><span class="o">.</span><span class="n">ravel</span><span class="p">()]</span><span class="o">.</span><span class="n">T</span><span class="o">.</span><span class="n">copy</span><span class="p">()</span>
+
+    <span class="c1"># cpu w/ pmap expects a chunk for each core</span>
+    <span class="k">if</span> <span class="n">using_cpu</span><span class="p">:</span>
+        <span class="n">pars</span> <span class="o">=</span> <span class="n">pars</span><span class="o">.</span><span class="n">reshape</span><span class="p">((</span><span class="mi">3</span><span class="p">,</span> <span class="n">vb</span><span class="o">.</span><span class="n">cores</span><span class="p">,</span> <span class="o">-</span><span class="mi">1</span><span class="p">))</span><span class="o">.</span><span class="n">transpose</span><span class="p">((</span><span class="mi">1</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">2</span><span class="p">))</span>
+
+    <span class="c1"># now run</span>
+    <span class="n">result</span> <span class="o">=</span> <span class="n">run_batches</span><span class="p">(</span><span class="n">pars</span><span class="p">)</span><span class="o">.</span><span class="n">block_until_ready</span><span class="p">()</span>
+    <span class="n">results</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">result</span><span class="p">)</span>
+
+<span class="n">toc</span> <span class="o">=</span> <span class="n">time</span><span class="o">.</span><span class="n">time</span><span class="p">()</span>
+<span class="nb">print</span><span class="p">(</span><span class="sa">f</span><span class="s1">&#39;elapsed time for sweep </span><span class="si">{</span><span class="n">toc</span><span class="w"> </span><span class="o">-</span><span class="w"> </span><span class="n">tic</span><span class="si">:</span><span class="s1">0.1f</span><span class="si">}</span><span class="s1"> s&#39;</span><span class="p">)</span>
+
+
+<span class="n">pl</span><span class="o">.</span><span class="n">figure</span><span class="p">(</span><span class="n">figsize</span><span class="o">=</span><span class="p">(</span><span class="mi">8</span><span class="p">,</span> <span class="mi">2</span><span class="p">))</span>
+<span class="n">pl</span><span class="o">.</span><span class="n">figure</span><span class="p">(</span><span class="n">figsize</span><span class="o">=</span><span class="p">(</span><span class="mi">8</span><span class="p">,</span><span class="mi">2</span><span class="p">))</span>
+<span class="k">for</span> <span class="n">i</span><span class="p">,</span> <span class="p">(</span><span class="n">sig_i</span><span class="p">,</span> <span class="n">result</span><span class="p">)</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="nb">zip</span><span class="p">(</span><span class="n">sigmas</span><span class="p">,</span> <span class="n">results</span><span class="p">)):</span>
+    <span class="n">pl</span><span class="o">.</span><span class="n">subplot</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="mi">4</span><span class="p">,</span> <span class="n">i</span> <span class="o">+</span> <span class="mi">1</span><span class="p">)</span>
+    <span class="n">pl</span><span class="o">.</span><span class="n">imshow</span><span class="p">(</span><span class="n">result</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="n">log_ks</span><span class="o">.</span><span class="n">shape</span><span class="p">),</span> <span class="n">vmin</span><span class="o">=</span><span class="mf">0.2</span><span class="p">,</span> <span class="n">vmax</span><span class="o">=</span><span class="mf">0.7</span><span class="p">)</span>
+</pre></div>
+</div>
+<figure class="align-default">
+<a class="reference internal image-reference" href="_images/sweep.png"><img alt="_images/sweep.png" src="_images/sweep.png" style="width: 800.0px; height: 200.0px;" /></a>
 </figure>
 </section>
 
diff --git a/examples/00_intro.py b/examples/00_intro.py
index b9bef4d..ddb1534 100644
--- a/examples/00_intro.py
+++ b/examples/00_intro.py
@@ -50,7 +50,7 @@
     :lines: 13
 
 .. figure:: ../../examples/images/example1.jpg
-    :scale: 50 %
+    :scale: 75 %
 
 
 """
diff --git a/examples/01_sweep.py b/examples/01_sweep.py
new file mode 100644
index 0000000..447e816
--- /dev/null
+++ b/examples/01_sweep.py
@@ -0,0 +1,147 @@
+""" 
+**Example 2**: Consider a network of coupled Montbrio model nodes and sweep
+over the parameters of the model.
+
+Starting with a few imports
+
+.. literalinclude:: ../../examples/01_sweep.py
+    :start-after: example-st\u0061rt
+    :lines: 1-6
+    :caption: 
+
+This example shows how to use the `vbjax` library to simulate a network of
+Montbrio model nodes. The network is defined by the function `network` which
+takes as arguments the state of the network and the parameters of the model.
+
+.. literalinclude:: ../../examples/01_sweep.py
+    :start-after: example-st\u0061rt
+    :lines: 9-13
+
+The function noise is used to generate the noise term of the stochastic
+
+.. literalinclude:: ../../examples/01_sweep.py
+    :start-after: example-st\u0061rt
+    :lines: 16-18
+
+The function run is used to simulate the network for a given set of parameters.
+
+.. literalinclude:: ../../examples/01_sweep.py
+    :start-after: example-st\u0061rt
+    :lines: 21-26
+
+Then prepare the simulation
+
+.. literalinclude:: ../../examples/01_sweep.py
+    :start-after: example-st\u0061rt
+    :lines: 30
+
+define the number of nodes in the network.
+
+.. literalinclude:: ../../examples/01_sweep.py
+    :start-after: example-st\u0061rt
+    :lines: 31-35
+
+defines the engine to be used for the simulation. If the simulation is run on
+the CPU, then the simulation is parallelized over the cores of the CPU using `jax.pmap`.
+Otherwise, the simulation is parallelized over the GPU using `jax.vmap`.
+
+then we prepare the network and the noise samples
+
+.. literalinclude:: ../../examples/01_sweep.py
+    :start-after: example-st\u0061rt
+    :lines: 37-39
+
+The rest run the simulation on set of parameters and plot the results.
+
+.. literalinclude:: ../../examples/01_sweep.py
+    :start-after: example-st\u0061rt
+    :lines: 44-75
+
+.. figure:: ../../examples/images/sweep.png
+    :scale: 100 %
+
+
+"""
+# example-start
+import os
+import time
+import vbjax as vb
+import pylab as pl  
+import jax, jax.numpy as np
+os.makedirs('images', exist_ok=True)
+
+
+def network(x, p):
+    r, v = x
+    k, _, mpr_p = p
+    c = k*r.sum(), k*v.sum()
+    return vb.mpr_dfun(x, c, mpr_p)
+
+
+def noise(_, p):
+    _, sigma, _ = p
+    return sigma
+
+
+def run(pars, mpr_p=vb.mpr_default_theta):
+    k, sig, eta = pars                      # explored pars
+    p = k, sig, mpr_p._replace(eta=eta)     # set mpr
+    xs = loop(rv0, zs, p)                   # run sim
+    std = xs[400:, 0].std()                 # eval metric
+    return std                              # done
+
+
+# prepare simulation
+n_nodes = 8
+using_cpu = jax.local_devices()[0].platform == 'cpu'
+if using_cpu:
+    run_batches = jax.pmap(jax.vmap(run, in_axes=1), in_axes=0)
+else:
+    run_batches = jax.vmap(run, in_axes=1)
+
+# prepare network
+_, loop = vb.make_sde(0.01, network, noise)
+rv0 = vb.randn(2, n_nodes)
+zs = vb.randn(1000, *rv0.shape)
+
+tic = time.time()
+
+# sweep sigma but just a few values are enough
+sigmas = [0.0, 0.2, 0.3, 0.4]
+results = []
+ng = vb.cores*4 if using_cpu else 32
+
+for i, sig_i in enumerate(sigmas):
+    # create grid of k (on logarithmic scale) and eta
+    log_ks, etas = np.mgrid[-9.0:-2.0:1j*ng, -4.0:-6.0:1j*ng]
+
+    # reshape grid to big batch of values
+    pars = np.c_[
+        np.exp(log_ks.ravel()),
+        np.ones(log_ks.size)*sig_i,
+        etas.ravel()].T.copy()
+
+    # cpu w/ pmap expects a chunk for each core
+    if using_cpu:
+        pars = pars.reshape((3, vb.cores, -1)).transpose((1, 0, 2))
+
+    # now run
+    result = run_batches(pars).block_until_ready()
+    results.append(result)
+
+toc = time.time()
+print(f'elapsed time for sweep {toc - tic:0.1f} s')
+
+
+pl.figure(figsize=(8, 2))
+pl.figure(figsize=(8,2))
+for i, (sig_i, result) in enumerate(zip(sigmas, results)):
+    pl.subplot(1, 4, i + 1)
+    pl.imshow(result.reshape(log_ks.shape), vmin=0.2, vmax=0.7)
+    pl.ylabel('k') if i==0 else (), pl.xlabel('eta')
+    pl.title(f'sig = {sig_i:0.1f}')
+
+pl.tight_layout()
+pl.savefig('images/sweep.png')
+
+# example-end
\ No newline at end of file
