diff --git a/pyhgf/model/network.py b/pyhgf/model/network.py index 3eec09908..082b929a4 100644 --- a/pyhgf/model/network.py +++ b/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 diff --git a/pyhgf/updates/posterior/continuous/__init__.py b/pyhgf/updates/posterior/continuous/__init__.py index fd8740754..e64c9523a 100644 --- a/pyhgf/updates/posterior/continuous/__init__.py +++ b/pyhgf/updates/posterior/continuous/__init__.py @@ -1,5 +1,8 @@ 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", diff --git a/pyhgf/updates/posterior/continuous/continuous_node_posterior_update_unbounded.py b/pyhgf/updates/posterior/continuous/continuous_node_posterior_update_unbounded.py new file mode 100644 index 000000000..5ed7041d6 --- /dev/null +++ b/pyhgf/updates/posterior/continuous/continuous_node_posterior_update_unbounded.py @@ -0,0 +1,66 @@ +# Author: Nicolas Legrand + +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 diff --git a/pyhgf/updates/posterior/continuous/posterior_update_mean_continuous_node_unbounded.py b/pyhgf/updates/posterior/continuous/posterior_update_mean_continuous_node_unbounded.py new file mode 100644 index 000000000..0f4ed643b --- /dev/null +++ b/pyhgf/updates/posterior/continuous/posterior_update_mean_continuous_node_unbounded.py @@ -0,0 +1,112 @@ +# Author: Nicolas Legrand + +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)) diff --git a/pyhgf/updates/posterior/continuous/posterior_update_precision_continuous_node_unbounded.py b/pyhgf/updates/posterior/continuous/posterior_update_precision_continuous_node_unbounded.py new file mode 100644 index 000000000..87222d277 --- /dev/null +++ b/pyhgf/updates/posterior/continuous/posterior_update_precision_continuous_node_unbounded.py @@ -0,0 +1,84 @@ +# Author: Nicolas Legrand + +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)) diff --git a/pyhgf/utils/get_update_sequence.py b/pyhgf/utils/get_update_sequence.py index 1ad304ccc..b385b245a 100644 --- a/pyhgf/utils/get_update_sequence.py +++ b/pyhgf/utils/get_update_sequence.py @@ -9,6 +9,7 @@ from pyhgf.updates.posterior.continuous import ( continuous_node_posterior_update, continuous_node_posterior_update_ehgf, + continuous_node_posterior_update_unbounded, ) from pyhgf.updates.prediction.binary import binary_state_node_prediction from pyhgf.updates.prediction.continuous import continuous_node_prediction @@ -135,7 +136,12 @@ def get_update_sequence( if all([i not in nodes_without_prediction_error for i in all_children]): no_update = False if network.edges[idx].node_type == 2: - if update_type == "eHGF": + if update_type == "unbounded": + if network.edges[idx].volatility_children is not None: + update_fn = continuous_node_posterior_update_unbounded + else: + update_fn = continuous_node_posterior_update + elif update_type == "eHGF": if network.edges[idx].volatility_children is not None: update_fn = continuous_node_posterior_update_ehgf else: