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

* implement constructors

* update type hint

* Update nodes.py

* rename new_const to new_value

* Update Var and constructors docstrings

* Update calculator docs

* Deprecate old functions, update docs

* make tfd immediately available in doctests

* speed up doctest

* fix doctests

* Update model overview

* Update import statements on welcome page

* Update Var short description

* Update model docstring

* fix typo

* satisfy pre-commit

* Update model.py

* update indexing docs using new constructors

conftest.py

@@ -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
@@ -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

docs/source/model.rst

@@ -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`.
@@ -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

docs/source/welcome.md

@@ -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.

liesel/goose/optim.py

@@ -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)}
 

liesel/model/legacy.py

@@ -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)
 
@@ -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)

liesel/model/model.py

@@ -963,8 +963,14 @@ 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`, you can simply override the
+        :attr:`.Var.value` attribute. Remeber to call :meth:`.Model.update` afterwards.
 
     Parameters
     ----------
@@ -981,58 +987,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__(