Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Var constructors: new_param, new_obs, new_calc, new_value #201

Merged
merged 23 commits into from
Oct 18, 2024
Merged
Show file tree
Hide file tree
Changes from 20 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
import jax.numpy as jnp
import numpy as np
import pytest
import tensorflow_probability.substrates.jax.distributions as tfd

import liesel.goose as gs
import liesel.model as lsl
Expand All @@ -19,3 +20,4 @@ def add_doctest_imports(doctest_namespace):
doctest_namespace["jnp"] = jnp
doctest_namespace["gs"] = gs
doctest_namespace["lsl"] = lsl
doctest_namespace["tfd"] = tfd
43 changes: 28 additions & 15 deletions docs/source/model.rst
Original file line number Diff line number Diff line change
Expand Up @@ -10,47 +10,52 @@ We usually import ``liesel.model`` as follows::

import liesel.model as lsl

Additionnaly, it often makes sense to import ``jax.numpy`` and tensorflow probability::

import jax.numpy as jnp
import tensorflow_probability.substrates.jax.distributions as tfd


The model building workflow in Liesel consists of two steps:

1. Set up the nodes and variables that make up your model.
2. Set up a :class:`.GraphBuilder`, add your root node(s) to it, and call :meth:`.GraphBuilder.build_model` to build your model graph.
2. Initialize a :class:`.Model` with your root variable(s).

.. note::
This document provides an overview of the most important classes for model building.
You find more guidance on *how* to use them in the respective API documentation
and in the :doc:`tutorials <tutorials_overview>`.


Nodes and Variables
-------------------
Variables: The model building blocks
------------------------------------

The fundamental building blocks of your model graph are given by just four classes.
Each of these building blocks is documented with examples, so make
sure to check them out.
The fundamental building blocks of your model graph are given by just two classes.
Both are documented with examples, so make sure to check them out.

.. autosummary::
:toctree: generated
:recursive:
:nosignatures:

~liesel.model.nodes.Var
~liesel.model.nodes.Value
~liesel.model.nodes.Calc
~liesel.model.nodes.Dist

To set up :class:`.Var` objects, Liesel provides two helper functions:
To set up :class:`.Var` objects, Liesel provides four constructors:

.. autosummary::
:toctree: generated
:recursive:
:nosignatures:

~liesel.model.nodes.obs
~liesel.model.nodes.param
~liesel.model.nodes.Var.new_param
~liesel.model.nodes.Var.new_obs
~liesel.model.nodes.Var.new_calc
~liesel.model.nodes.Var.new_value

Defining :class:`.Var` objects with these functions makes sure that the respective
:attr:`.Var.observerd` and :attr:`.Var.parameter` flags are correctly set. This in
We recommend to always use one of these constructors when initializing a variable.
This makes sure that the respective
:attr:`.Var.observed` and :attr:`.Var.parameter` flags are correctly set. This in
turn ensures that the :attr:`.Var.log_prob` of an *observed* variable will be included
in the :attr:`.Model.log_lik` and the :attr:`.Var.log_prob` of a *parameter* variable
will be included in the :attr:`.Model.log_prior`.
Expand All @@ -59,14 +64,22 @@ will be included in the :attr:`.Model.log_prior`.
Build and plot your model
-------------------------

The most important class here is the :class:`.GraphBuilder`.
The most important class here is the :class:`.Model`.

.. autosummary::
:toctree: generated
:recursive:
:nosignatures:

~liesel.model.model.GraphBuilder
~liesel.model.model.Model
~liesel.model.viz.plot_vars

For advanced users, further interesting functionality can be found here:

.. autosummary::
:toctree: generated
:recursive:
:nosignatures:

~liesel.model.model.GraphBuilder
~liesel.model.viz.plot_nodes
11 changes: 10 additions & 1 deletion docs/source/welcome.md
Original file line number Diff line number Diff line change
Expand Up @@ -28,13 +28,22 @@ $ pip install liesel
If you want to work with the latest development version of Liesel or use PyGraphviz for
prettier plots of the model graphs, see the [README][5] in the main repository.

Now you can get started. We recommend using the following import paths:
Now you can get started. Throughout this documentation, we import Liesel as follows:

```python
import liesel.model as lsl
import liesel.goose as gs
```

We also commonly use the following imports:

```python
import jax
import jax.numpy as jnp
import numpy as np
import tensorflow_probability.substrates.jax.distributions as tfd
```

We provide overviews of the most important building blocks provided by `liesel.model`
and `liesel.goose` in [Model Building (liesel.model)](model_overview) and
[MCMC Sampling (liesel.goose)](goose_overview), respectively.
Expand Down
3 changes: 2 additions & 1 deletion liesel/goose/optim.py
Original file line number Diff line number Diff line change
Expand Up @@ -330,7 +330,8 @@ def optim_flat(
Now, we are ready to run the optimization.
>>> result = gs.optim_flat(model, params=["coef"])
>>> stopper = gs.Stopper(max_iter=1000, patience=10, atol=0.01)
>>> result = gs.optim_flat(model, params=["coef"], stopper=stopper)
>>> {name: jnp.round(value, 2) for name, value in result.position.items()}
{'coef': Array([0.52, 1.29], dtype=float32)}
Expand Down
6 changes: 4 additions & 2 deletions liesel/model/legacy.py
Original file line number Diff line number Diff line change
Expand Up @@ -210,7 +210,8 @@ def Obs(value: Any | Calc, distribution: Dist | None = None, name: str = "") ->
:func:`.obs` : New alias. Sould be used in future code.
"""
warnings.warn(
"Use lsl.obs() instead. This alias will be removed in v0.4.0", FutureWarning
"Use lsl.Var.new_obs() instead. This class will be removed in v0.4.0",
FutureWarning,
)
return obs(value, distribution, name)

Expand Down Expand Up @@ -238,6 +239,7 @@ def Param(value: Any | Calc, distribution: Dist | None = None, name: str = "") -
:func:`.param` : New alias. Sould be used in future code.
"""
warnings.warn(
"Use lsl.param() instead. This alias will be removed in v0.4.0", FutureWarning
"Use lsl.Var.new_param() instead. This class will be removed in v0.4.0",
FutureWarning,
)
return param(value, distribution, name)
63 changes: 20 additions & 43 deletions liesel/model/model.py
Original file line number Diff line number Diff line change
Expand Up @@ -963,8 +963,13 @@ class Model:
A model with a static graph.

.. tip::
While you can create a model directly, it is usually more convenient to use a
:class:`.GraphBuilder` to construct the model.
If you have an existing model and want to make changes to it, you can use the
:meth:`.Model.pop_nodes_and_vars` method to release the nodes and variables
of the model. You can then make changes to them, for example i.e. changing the
distribution of a variable or the inputs of a calculation. Afterwards, you
initialize a *new* model with your changed variables.
If you simply want to change the value of a variable, it is not necessary to
call :meth:`~.Model.pop_nodes_and_vars`.
jobrachem marked this conversation as resolved.
Show resolved Hide resolved

Parameters
----------
Expand All @@ -981,58 +986,30 @@ class Model:

See Also
--------
:class:`.GraphBuilder` : A graph builder, used to set up a model.
.Var.new_obs : Initializes a strong variable that holds observed data.
.Var.new_param : Initializes a strong variable that acts as a model parameter.
.Var.new_calc :
Initializes a weak variable that is a function of other variables.
.Var.new_value : Initializes a strong variable without a distribution.
:class:`.GraphBuilder` :
A graph builder, which can be used to set up and manipulate a model if you need
more control.

Examples
--------

For basic examples on how to set up a model, please refer to the
:class:`.GraphBuilder` documentation.
Here, we set up a basic model based on three variables:

.. rubric:: Modifying an existing model

If you have an existing model and want to make changes to it, you can use the
:meth:`.Model.copy_nodes_and_vars` or the :meth:`.Model.copy_nodes_and_vars`
method to obtain the nodes and variables of the model, make changes to them, and
then create a new model from the modified nodes and variables.

>>> a = lsl.Var(1.0, name="a")
>>> b = lsl.Var(2.0, name="b")
>>> c = lsl.Var(lsl.Calc(lambda x, y: x + y, a, b), name="c")
>>> a = lsl.Var.new_value(1.0, name="a")
>>> b = lsl.Var.new_value(2.0, name="b")
>>> c = lsl.Var.new_calc(lambda x, y: x + y, a, b, name="c")

We now build a model:

>>> model = lsl.GraphBuilder().add(c).build_model()
>>> model = lsl.Model([c])
>>> model
Model(9 nodes, 3 vars)

>>> nodes, vars_ = model.pop_nodes_and_vars()
>>> vars_
{'c': Var(name="c"), 'b': Var(name="b"), 'a': Var(name="a")}

>>> from pprint import pprint # for nicer formatting of the output dicts
>>> pprint(nodes)
{'a_value': Value(name="a_value"),
'a_var_value': VarValue(name="a_var_value"),
'b_value': Value(name="b_value"),
'b_var_value': VarValue(name="b_var_value"),
'c_value': Calc(name="c_value"),
'c_var_value': VarValue(name="c_var_value")}

We can now make changes to the nodes and variables.
Just for show, let's add a distribution to the node ``a``:

>>> import tensorflow_probability.substrates.jax.distributions as tfd
>>> vars_["a"].dist_node = lsl.Dist(tfd.Normal, loc=0.0, scale=1.0)

Now we create a new GraphBuilder and build a new model:

>>> gb = lsl.GraphBuilder()
>>> gb = gb.add(*nodes.values(), *vars_.values())
>>> model = gb.build_model()
>>> model
Model(12 nodes, 3 vars)

"""

def __init__(
Expand Down
Loading