add uhgf updates

pyhgf/model/network.py

@@ -83,7 +83,7 @@ def input_idxs(self, value):
         self.input_idxs = value
 
     def create_belief_propagation_fn(
-        self, overwrite: bool = True, update_type: str = "eHGF"
+        self, overwrite: bool = True, update_type: str = "unbounded"
     ) -> "Network":
         """Create the belief propagation function.
 
@@ -97,11 +97,16 @@ def create_belief_propagation_fn(
             preexisting values. Otherwise, do not create a new function if the attribute
             `scan_fn` is already defined.
         update_type :
-            The type of update to perform for volatility coupling. Can be `"eHGF"`
-            (defaults) or `"standard"`. The eHGF update step was proposed as an
+            The type of update to perform for volatility coupling. Can be `"unbounded"`
+            (defaults), `"ehgf"` or `"standard"`. The unbounded approximation was
+            recently introduced to avoid negative precisions updates, which greatly
+            improve sampling performance. The eHGF update step was proposed as an
             alternative to the original definition in that it starts by updating the
             mean and then the precision of the parent node, which generally reduces the
-            errors associated with impossible parameter space and improves sampling.
+            occurence of negative precision updates, while not removing them entirely.
+            .. note:
+              The different update steps only apply to nodes having at least one
+              volatility parents. In other cases, the regular HGF updates are applied.
 
         """
         # create the update sequence if it does not already exist

pyhgf/updates/posterior/continuous/__init__.py

@@ -0,0 +1,11 @@
+from .continuous_node_posterior_update import continuous_node_posterior_update
+from .continuous_node_posterior_update_ehgf import continuous_node_posterior_update_ehgf
+from .continuous_node_posterior_update_unbounded import (
+    continuous_node_posterior_update_unbounded,
+)
+
+__all__ = [
+    "continuous_node_posterior_update_ehgf",
+    "continuous_node_posterior_update_unbounded",
+    "continuous_node_posterior_update",
+]

pyhgf/updates/posterior/continuous/continuous_node_posterior_update_unbounded.py

@@ -0,0 +1,66 @@
+# Author: Nicolas Legrand <nicolas.legrand@cas.au.dk>
+
+from functools import partial
+from typing import Dict
+
+from jax import jit
+
+from pyhgf.typing import Edges
+
+from .posterior_update_mean_continuous_node_unbounded import (
+    posterior_update_mean_continuous_node_unbounded,
+)
+from .posterior_update_precision_continuous_node_unbounded import (
+    posterior_update_precision_continuous_node_unbounded,
+)
+
+
+@partial(jit, static_argnames=("edges", "node_idx"))
+def continuous_node_posterior_update_unbounded(
+    attributes: Dict, node_idx: int, edges: Edges, **args
+) -> Dict:
+    """Update the posterior of a continuous node using an unbounded approximation.
+
+    Parameters
+    ----------
+    attributes :
+        The attributes of the probabilistic nodes.
+    node_idx :
+        Pointer to the node that needs to be updated. After continuous updates, the
+        parameters of value and volatility parents (if any) will be different.
+    edges :
+        The edges of the probabilistic nodes as a tuple of
+        :py:class:`pyhgf.typing.Indexes`. The tuple has the same length as node number.
+        For each node, the index list value and volatility parents and children.
+
+    Returns
+    -------
+    attributes :
+        The updated attributes of the probabilistic nodes.
+
+    See Also
+    --------
+    continuous_node_posterior_update_ehgf
+
+    """
+    # update the posterior mean and precision using the eHGF update step
+    # we start with the mean update using the expected precision as an approximation
+    posterior_precision, precision_l1, precision_l2 = (
+        posterior_update_precision_continuous_node_unbounded(
+            attributes,
+            edges,
+            node_idx,
+        )
+    )
+    attributes[node_idx]["precision"] = posterior_precision
+
+    posterior_mean = posterior_update_mean_continuous_node_unbounded(
+        attributes=attributes,
+        edges=edges,
+        node_idx=node_idx,
+        precision_l1=precision_l1,
+        precision_l2=precision_l2,
+    )
+    attributes[node_idx]["mean"] = posterior_mean
+
+    return attributes

pyhgf/updates/posterior/continuous/posterior_update_mean_continuous_node_unbounded.py

@@ -0,0 +1,112 @@
+# Author: Nicolas Legrand <nicolas.legrand@cas.au.dk>
+
+from functools import partial
+from typing import Dict
+
+import jax.numpy as jnp
+from jax import jit
+
+from pyhgf.typing import Edges
+
+
+@partial(jit, static_argnames=("edges", "node_idx"))
+def posterior_update_mean_continuous_node_unbounded(
+    attributes: Dict,
+    edges: Edges,
+    node_idx: int,
+    precision_l1: float,
+    precision_l2: float,
+) -> float:
+    """Posterior update of mean using ubounded update."""
+    volatility_child_idx = edges[node_idx].volatility_children[0]
+    volatility_coupling = attributes[node_idx]["volatility_coupling_children"][0]
+    gamma = attributes[node_idx]["expected_mean"]
+    phi = jnp.log(
+        (1 / attributes[volatility_child_idx]["precision"]) * (2 + jnp.sqrt(3))
+    )
+
+    # first approximation ------------------------------------------------------
+    delta_l1 = (
+        (
+            (1 / attributes[volatility_child_idx]["precision"])
+            + (
+                attributes[volatility_child_idx]["mean"]
+                - attributes[volatility_child_idx]["expected_mean"] ** 2
+            )
+        )
+        / (
+            (1 / attributes[volatility_child_idx]["expected_precision"])
+            + jnp.exp(
+                volatility_coupling * phi
+                + attributes[volatility_child_idx]["tonic_volatility"]
+            )
+        )
+    ) - 1
+    mean_l1 = (
+        attributes[node_idx]["expected_mean"]
+        + (
+            (volatility_coupling * attributes[node_idx]["tonic_volatility"])
+            / (2 * precision_l1)
+        )
+        * delta_l1
+    )
+
+    # second approximation -----------------------------------------------------
+    omega_phi = jnp.exp(
+        volatility_coupling * phi + attributes[node_idx]["tonic_volatility"]
+    ) / (
+        (1 / attributes[volatility_child_idx]["precision"])
+        + jnp.exp(volatility_coupling * phi + attributes[node_idx]["tonic_volatility"])
+    )
+    delta_phi = (
+        (1 / attributes[volatility_child_idx]["precision"])
+        + (
+            attributes[volatility_child_idx]["mean"]
+            - attributes[volatility_child_idx]["expected_mean"]
+        )
+        ** 2
+    ) / (
+        (1 / attributes[volatility_child_idx]["expected_precision"])
+        + jnp.exp(
+            volatility_coupling * phi
+            + attributes[volatility_child_idx]["tonic_volatility"]
+        )
+    ) - 1
+
+    mu_phi = ((2 * precision_l2 - 1) * phi + attributes[node_idx]["expected_mean"]) / (
+        2 * precision_l2
+    )
+
+    mean_l2 = (
+        mu_phi + (volatility_coupling * omega_phi) / (2 * precision_l2) * delta_phi
+    )
+
+    # weigthed interpolation
+    theta_l = jnp.sqrt(
+        1.2
+        * (
+            (1 / attributes[volatility_child_idx]["precision"])
+            + (
+                attributes[volatility_child_idx]["mean"]
+                - attributes[volatility_child_idx]["expected_mean"]
+            )
+            ** 2
+        )
+        / ((1 / attributes[volatility_child_idx]["expected_precision"]) * precision_l1)
+    )
+    phi_l = 8.0
+    theta_r = 0.0
+    phi_r = 1.0
+    mean = (1 - b(gamma, theta_l, phi_l, theta_r, phi_r)) * mean_l1 + b(
+        gamma, theta_l, phi_l, theta_r, phi_r
+    ) * mean_l2
+
+    return mean
+
+
+def s(x, theta, phi):
+    return 1 / (1 + jnp.exp(-phi * (x - theta)))
+
+
+def b(x, theta_l, phi_l, theta_r, phi_r):
+    return s(x, theta_l, phi_l) - (1 - s(x, theta_r, phi_r))

pyhgf/updates/posterior/continuous/posterior_update_precision_continuous_node_unbounded.py

@@ -0,0 +1,84 @@
+# Author: Nicolas Legrand <nicolas.legrand@cas.au.dk>
+
+from functools import partial
+from typing import Dict
+
+import jax.numpy as jnp
+from jax import jit
+
+from pyhgf.typing import Edges
+
+
+@partial(jit, static_argnames=("edges", "node_idx"))
+def posterior_update_precision_continuous_node_unbounded(
+    attributes: Dict, edges: Edges, node_idx: int
+) -> float:
+    """Posterior update of precision using ubounded update."""
+    volatility_child_idx = edges[node_idx].volatility_children[0]
+    volatility_coupling = attributes[node_idx]["volatility_coupling_children"][0]
+    gamma = attributes[node_idx]["expected_mean"]
+
+    # first approximation ------------------------------------------------------
+    precision_l1 = attributes[node_idx][
+        "expected_precision"
+    ] + 0.5 * volatility_coupling**2 * attributes[node_idx]["tonic_volatility"] * (
+        1 - attributes[node_idx]["tonic_volatility"]
+    )
+
+    # second approximation -----------------------------------------------------
+    phi = jnp.log(
+        (1 / attributes[volatility_child_idx]["expected_precision"]) * (2 + jnp.sqrt(3))
+    )
+    omega_phi = jnp.exp(
+        volatility_coupling * phi + attributes[node_idx]["tonic_volatility"]
+    ) / (
+        (1 / attributes[volatility_child_idx]["expected_precision"])
+        + jnp.exp(volatility_coupling * phi + attributes[node_idx]["tonic_volatility"])
+    )
+    delta_phi = (
+        (1 / attributes[volatility_child_idx]["precision"])
+        + (
+            attributes[volatility_child_idx]["mean"]
+            - attributes[volatility_child_idx]["expected_mean"]
+        )
+        ** 2
+    ) / (
+        (1 / attributes[volatility_child_idx]["expected_precision"])
+        + jnp.exp(volatility_coupling * phi + attributes[node_idx]["tonic_volatility"])
+    ) - 1
+
+    precision_l2 = attributes[node_idx][
+        "expected_precision"
+    ] + 0.5 * volatility_coupling**2 * omega_phi * (
+        omega_phi + (2 * omega_phi - 1) * delta_phi
+    )
+
+    # weigthed interpolation
+    theta_l = jnp.sqrt(
+        1.2
+        * (
+            (1 / attributes[volatility_child_idx]["precision"])
+            + (
+                attributes[volatility_child_idx]["mean"]
+                - attributes[volatility_child_idx]["expected_mean"]
+            )
+            ** 2
+        )
+        / ((1 / attributes[volatility_child_idx]["expected_precision"]) * precision_l1)
+    )
+    phi_l = 8.0
+    theta_r = 0.0
+    phi_r = 1.0
+    precision = (1 - b(gamma, theta_l, phi_l, theta_r, phi_r)) * precision_l1 + b(
+        gamma, theta_l, phi_l, theta_r, phi_r
+    ) * precision_l2
+
+    return precision, precision_l1, precision_l2
+
+
+def s(x, theta, phi):
+    return 1 / (1 + jnp.exp(-phi * (x - theta)))
+
+
+def b(x, theta_l, phi_l, theta_r, phi_r):
+    return s(x, theta_l, phi_l) - (1 - s(x, theta_r, phi_r))