More docs

Author: jdehning

docs/api/comp_model.rst

@@ -1,7 +1,7 @@
 Compartmental Model
 ===================
 
-.. module:: icomo.comp_model
+.. module:: icomo
 
 Defining a compartmental model
 -------------------------------
@@ -10,3 +10,10 @@ Defining a compartmental model
 .. autoclass:: CompModel
     :exclude-members: y, dy, edges, __weakref__
     :members:
+
+Helper functions
+----------------
+
+.. autofunction:: erlang_kernel
+
+.. autofunction:: delayed_copy_kernel

docs/api/diffrax_wrapper.rst

@@ -1,8 +1,16 @@
 Diffrax Wrapper
 ===============
 
-.. module:: icomo.diffrax_wrapper
+.. module:: icomo
+
+These are helper functions that wrap some function from the diffrax module like the
+:meth:`diffrax.diffeqsolve` function. This is only done for convenience, to define
+some default values for the parameters of the function. Don't hesitate to directly the
+functions from the diffrax module if you need more flexibility.
+
 
 .. autofunction:: diffeqsolve
 
 .. autofunction:: interpolate_func
+
+

docs/api/experimental.rst

@@ -0,0 +1,13 @@
+Experimental Inference tools
+==================
+
+The following tools are experimental and may change in future releases. They help
+with the Bayesian inference of the model parameters using the `PyMC <https://www.pymc
+.io/>`_ library.
+
+.. module:: icomo.experimental
+
+.. automodule:: icomo.experimental
+    :imported-members:
+    :members:
+    :show-inheritance:

docs/api/jax2pytensor.rst

@@ -0,0 +1,7 @@
+Transform Jax to Pytensor
+=========================
+
+.. module:: icomo
+
+.. autofunction:: icomo.jax2pytensor
+

docs/api/tree_tools.rst

@@ -0,0 +1,12 @@
+Tree tools
+==========
+
+Some tools for working with python trees, useful for working with nested data structures.
+
+.. module:: icomo
+
+.. autofunction:: walk_tree
+
+.. autofunction:: nested_indexing
+
+

docs/index.md

@@ -15,8 +15,13 @@ Getting started <example>
 :caption: API
 
 api/comp_model
+api/jax2pytensor
 api/diffrax_wrapper
-api/references
+api/tree_tools
+api/experimental
+```
+
+```{toctree}
 contributing
 ```
 

icomo/__init__.py

@@ -5,15 +5,14 @@
 from importlib import metadata as importlib_metadata
 
 from . import comp_model
-from .comp_model import (
-    CompModel,
-    erlang_kernel,
-)
+from .comp_model import CompModel, erlang_kernel, delayed_copy_kernel
 
 from .diffrax_wrapper import diffeqsolve, interpolate_func
 
 from .jax2pytensor import jax2pytensor
 
+from .tree_tools import nested_indexing, walk_tree
+
 from . import experimental