ehgf update step

src/pyhgf/model.py

@@ -34,6 +34,12 @@ class HGF(object):
         network only has one input node.
     model_type :
         The model implemented (can be `"continuous"`, `"binary"` or `"custom"`).
+    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 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.
     n_levels :
         The number of hierarchies in the model, including the input vector. Cannot be
         less than 2.
@@ -58,6 +64,7 @@ def __init__(
         self,
         n_levels: Optional[int] = 2,
         model_type: str = "continuous",
+        update_type: str = "eHGF",
         initial_mu: Dict = {
             "1": 0.0,
             "2": 0.0,
@@ -138,6 +145,7 @@ def __init__(
 
         """
         self.model_type = model_type
+        self.update_type = update_type
         self.verbose = verbose
         self.n_levels = n_levels
         self.edges: Edges
@@ -641,7 +649,7 @@ def add_value_parent(
             "psis_children": tuple(value_coupling for _ in range(len(children_idxs))),
             "psis_parents": None,
             "omega": omega,
-            "rho": rho,
+            "drift": rho,
         }
 
         # add more parameters (optional)
@@ -741,7 +749,7 @@ def add_volatility_parent(
             "psis_children": None,
             "psis_parents": None,
             "omega": omega,
-            "rho": rho,
+            "drift": rho,
         }
 
         # add more parameters (optional)

src/pyhgf/networks.py

@@ -21,6 +21,7 @@
     continuous_input_prediction_error,
     continuous_node_prediction,
     continuous_node_prediction_error,
+    ehgf_continuous_node_prediction_error,
 )
 
 if TYPE_CHECKING:
@@ -333,7 +334,11 @@ def get_update_sequence(
         # --------------------------
 
         # case 1 - default to a continuous node
-        update_fn = continuous_node_prediction_error
+        # choose between the eHGF and standard update step
+        if hgf.update_type == "eHGF":
+            update_fn = ehgf_continuous_node_prediction_error
+        elif hgf.update_type == "standard":
+            update_fn = continuous_node_prediction_error
         prediction_fn = continuous_node_prediction
 
         # case 2 - this is an input node

src/pyhgf/updates/continuous.py

@@ -121,6 +121,131 @@ def continuous_node_prediction_error(
     return attributes
 
 
+@partial(jit, static_argnames=("edges", "node_idx"))
+def ehgf_continuous_node_prediction_error(
+    attributes: Dict, time_step: float, node_idx: int, edges: Edges, **args
+) -> Dict:
+    """eHGF prediction-error step for value and volatility parents of a continuous node.
+
+    This update step uses a different order for the mean and precision as compared to
+    the standard HGF, respectively:
+    1. Update volatility parent(s).
+    2. Update value parent(s).
+
+    If a value/volatility parent has multiple children, all the children will update
+    the parent together, therefor this function should only be called once per group
+    of child nodes. The method :py:meth:`pyhgf.model.HGF.get_update_sequence`
+    ensures that this function is only called once all the children have been
+    updated.
+
+    Then returns the structure of the new parameters.
+
+    Parameters
+    ----------
+    attributes :
+        The attributes of the probabilistic nodes.
+    time_step :
+        The interval between the previous time point and the current time point.
+    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
+    --------
+    update_continuous_input_parents, update_binary_input_parents
+
+    References
+    ----------
+    .. [1] Weber, L. A., Waade, P. T., Legrand, N., Møller, A. H., Stephan, K. E., &
+       Mathys, C. (2023). The generalized Hierarchical Gaussian Filter (Version 1).
+       arXiv. https://doi.org/10.48550/ARXIV.2305.10937
+
+    """
+    # list value and volatility parents
+    value_parents_idxs = edges[node_idx].value_parents
+    volatility_parents_idxs = edges[node_idx].volatility_parents
+
+    # return here if no parents node are provided
+    if (value_parents_idxs is None) and (volatility_parents_idxs is None):
+        return attributes
+
+    ########################
+    # Update value parents #
+    ########################
+    if value_parents_idxs is not None:
+        for value_parent_idx in value_parents_idxs:
+            # if this child is the last one relative to this parent's family, all the
+            # children will update the parent at once, otherwise just pass and wait
+            if edges[value_parent_idx].value_children[-1] == node_idx:
+                # in the eHGF update step, we use the expected precision here
+                # as we haven't computed it yet due to the reverse update order
+                precision_value_parent = attributes[value_parent_idx][
+                    "expected_precision"
+                ]
+
+                # Estimate the mean of the posterior distribution
+                mean_value_parent = prediction_error_mean_value_parent(
+                    attributes, edges, value_parent_idx, precision_value_parent
+                )
+                # Update this parent's parameters
+                attributes[value_parent_idx]["mean"] = mean_value_parent
+
+                # Estimate the precision of the posterior distribution
+                precision_value_parent = prediction_error_precision_value_parent(
+                    attributes, edges, value_parent_idx
+                )
+
+                # Update this parent's parameters
+                attributes[value_parent_idx]["precision"] = precision_value_parent
+
+    #############################
+    # Update volatility parents #
+    #############################
+    if volatility_parents_idxs is not None:
+        for volatility_parent_idx in volatility_parents_idxs:
+            # if this child is the last one relative to this parent's family, all the
+            # children will update the parent at once, otherwise just pass and wait
+            if edges[volatility_parent_idx].volatility_children[-1] == node_idx:
+                # in the eHGF update step, we use the expected precision here
+                # as we haven't computed it yet due to the reverse update order
+                precision_volatility_parent = attributes[volatility_parent_idx][
+                    "expected_precision"
+                ]
+
+                # Estimate the new mean of the volatility parent
+                mean_volatility_parent = prediction_error_mean_volatility_parent(
+                    attributes,
+                    edges,
+                    time_step,
+                    volatility_parent_idx,
+                    precision_volatility_parent,
+                )
+                attributes[volatility_parent_idx]["mean"] = mean_volatility_parent
+
+                # Estimate the new precision of the volatility parent
+                precision_volatility_parent = (
+                    prediction_error_precision_volatility_parent(
+                        attributes, edges, time_step, volatility_parent_idx
+                    )
+                )
+
+                # Update this parent's parameters
+                attributes[volatility_parent_idx][
+                    "precision"
+                ] = precision_volatility_parent
+
+    return attributes
+
+
 @partial(jit, static_argnames=("edges", "node_idx"))
 def continuous_node_prediction(
     attributes: Dict, time_step: float, node_idx: int, edges: Edges, **args