Inference basics

Bayesian Inference with SVGD (Stein Variational Gradient Descent)

To estimate model parameters and their confidence intervals, Phasic used uses Stein Variational Gradient Descent (SVGD), which represents the posterior with a set of particles (parameter vectors) and combines ideas from optimization, kernel methods, and functional analysis in a deterministic algorithm that updates these particles iteratively to move them toward the posterior distribution. SVGD exploits the PDF for gradient computation and moments for regularization. It begins with an initial set of particles (typically drawn from the prior or from a simple distribution) and iteratively transport these particles toward the posterior using gradient information. Each particle interacts with all other particles through a kernel function, with the interaction strength depending on the distance between particles. This interaction ensures that particles spread out to cover the posterior distribution, allowing extraction of not only the most likely parameter value (MAP estimate) but also a confidence interval of the estimate.

Such inference is not possible in the standard matrix formulation of phase-type distributions where the large number of likelihood computations required would each be $O(n^3)$. In Phasic, the algorithm of the gaussian elimination on the graph fully exploits the sparseness of state spaces bringing the complexity closer to $n^2$. Because Phasic records and caches the elimination trace, this computation is only done once, removing a dimension of complexity for all subsequent likelihood evaluations. For an expose of these tricks and innovations see Math and Algorithms, Architecture, Caching and Sharing.

from phasic import (
    Graph, with_ipv, GaussPrior, LogGaussPrior, HalfCauchyPrior, DataPrior,
    Adam, ExpStepSize, ExpRegularization, clear_caches, set_log_level
)
import sys
import numpy as np
import jax.numpy as jnp
import pandas as pd
from typing import Optional
import matplotlib.pyplot as plt
from matplotlib.colors import LogNorm
import seaborn as sns
from tqdm.auto import tqdm
import vscodenb

import svgd_tests

sns.set_palette('tab10')
np.random.seed(42)

set_log_level('INFO')
clear_caches()

[INFO] phasic.graph_cache: Cleared 0 cached graphs

  Removed 77 file(s), preserved directory structure

Lets make a coalescent moddel with one parameter, the coalescent rate:

nr_samples = 4

@with_ipv([nr_samples]+[0]*(nr_samples-1))
def coalescent_1param(state):
    transitions = []
    for i in range(state.size):
        for j in range(i, state.size):            
            same = int(i == j)
            if same and state[i] < 2:
                continue
            if not same and (state[i] < 1 or state[j] < 1):
                continue 
            new = state.copy()
            new[i] -= 1
            new[j] -= 1
            new[i+j+1] += 1
            transitions.append([new, [state[i]*(state[j]-same)/(1+same)]])
    return transitions

Construct the graph and set some value of the model parameter:

graph = Graph(coalescent_1param)
graph.plot()

[INFO] phasic.graph_cache: Saved graph to cache: 4e7d3baaa4030c3c... (6 vertices)
[INFO] phasic: Saved graph to cache: 6 vertices

With a large set of observed times to the most recent common ancestor (TMRCA), we can fit our model to find the most likely value of the model parameter. For this turorial we will just set the model parameter to 7 and then sample from the model to get such a data set. This has the advantage that we know what the model parameter is supposed to be.

true_theta = [7]
graph.update_weights(true_theta)
nr_observations = 10000
observed_data = graph.sample(nr_observations)

We can plot a histogram of the observed data along with the theoretical density to see that they match up:

fig, ax = plt.subplots()
ax.hist(observed_data, density=True, bins=50, label='Histogram of samples')
x = np.linspace(0, 1, 100)
ax.plot(x, graph.pdf(x), color='C1', lw=1, label='Theoretical density')
ax.legend() ;

The svgd method provide sensible defaults to the optimization parameters. So in the simplest version, all you have to do to fit parameter values is:

svgd = graph.svgd(observed_data)
svgd.summary()

Parameter  Fixed          MAP        Mean       SD         HPD 95% lo   HPD 95% hi  
0          No             7.038      7.017      0.1039     6.923        7.276       

Particles: 20, Iterations: 1000

A nice property of SVGD is that you can easily evaluate if the optimization converges. We would like to see both the mean and standard deviation across particles to plateau towards the end.

svgd.plot_convergence()

In this case they clearly do not. Fortunately there are several easy remedies for this as shown below.

Learning rate

We can change the learning rate the number of iterations from their default values:

svgd = graph.svgd(observed_data, 
                  learning_rate=0.05,
                  n_iterations=200,
                  )
svgd.plot_convergence()

If the curves are not nice and smooth, you should inspect the parameter values of each particle across iterations:

svgd.plot_trace();

You can also plot the acceleration (the relative rate of change of each particle estimate) for an alternative way to diagnose this behavior:

svgd.plot_trace_accel()

You an see the values jumping abruptly. That happens when the learning rate step is too large compared to the gradient. Since the gradient flattens during the fit, we need to make the learning rate start our large and decay exponentially during training. We can do that by passing “step size schedule” like ExpStepSize to svgd. Its lets you control the sizes of the first as last steps and the exponential transition between them (the tau argument).

step_schedule = ExpStepSize(first_step=0.002, last_step=0.0001, tau=80.0)
step_schedule.plot(nr_iter=500);

svgd = graph.svgd(observed_data, 
                  learning_rate=step_schedule,
                  n_iterations=200,                  
                  )
svgd.plot_convergence()

svgd.plot_trace()

svgd.plot_trace_accel()

The distribution of particle values describe the posterior distribution from which a mean and CI can be computed:

svgd.plot_ci(true_theta=true_theta)

svgd.plot_ci(true_theta=true_theta, ci=0.5)

You can even make an animation of the fit with svgd.animate().

Prior

The prior distribution used in Bayesian inference represents our prior belief, I.e. our knowledge about what range of values the parameter can reasonably take. Phasic comes with two types of priors that are easily defined from basic knowledge about the parameter. The default behaviour of Phasic is to use a Gaussian prior centered on a rough estimate of each parameter using method-of-moments. To see the prior used, you can can plot it like this:

svgd.prior.plot() ;

or reconstruct the GaussPrior from its properties:

GaussPrior(mean=svgd.prior.theta.item(), std=svgd.prior.std.item()).plot() ;

You can define your own Gaussian prior (GaussPrior) or a Half-Cauchy prior (HalfCauchy).

The Half-Cauchy can be defined by the scale parameter or the upper bound of of the CI.

prior = HalfCauchyPrior(ci=0.4, prob=0.8)
prior.plot();

The Gaussian can be defined by either mean and SD or a confidence interval:

# prior = Prior(mean=5, std=2) 
prior = GaussPrior(ci=[2, 5])
prior.plot();

Phasic also has a LogGaussPrior for stating prior belief in terms of “orders of magnitude”:

prior = LogGaussPrior(ci=[1e-7, 1e-2])
prior.plot();

prior = LogGaussPrior(mean=1e-3, std=2) 
prior.plot();

Tip

The GaussPrior and HalfCauchyPrior classes provide the log-probabilities needed for SVGD. If you want to provide your own prior function, it must return such log-probabilities.

Lets add a prior to SVGD inference and rerun it.

step_schedule = ExpStepSize(first_step=0.002, last_step=0.0001, tau=80.0)
step_schedule.plot(nr_iter=500)

prior = GaussPrior(ci=[10, 20])
prior.plot();

prior = GaussPrior(ci=[10, 20])
step_schedule = ExpStepSize(first_step=0.002, last_step=0.0001, tau=80.0)
svgd = graph.svgd(observed_data, 
                  learning_rate=step_schedule,
                  prior=prior,
                  n_iterations=100
                  )

svgd.plot_ci(true_theta=true_theta)

svgd.plot_trace(max_particles=100)


svgd.plot_ci(true_theta=true_theta, ci=0.5)

svgd = graph.svgd(observed_data, 
                  learning_rate=step_schedule,
                  prior=GaussPrior(ci=[1, 5])
                  )
svgd.plot_ci(true_theta=true_theta, ci=0.5)

svgd = graph.svgd(observed_data, 
                  learning_rate=step_schedule,
                  prior=GaussPrior(ci=[10, 20])
                  )
svgd.plot_ci(true_theta=true_theta, ci=0.5)

svgd.plot_trace(true_theta=true_theta, ci=0.5)

svgd.analyze_trace()

CONVERGED (iteration 50/1001)
  Mean stabilized at iteration 50
  Std stabilized at iteration 50

Detected Issues:
  ℹ Converged at 5% of iterations - could reduce n_iterations
Learning Rate: Converged efficiently
  current schedule is good
Particles: High ESS ratio (1) - could reduce particles
  Keep n_particles=20
Iterations: Converged early
  Could reduce to n_iterations=60

Moment regularization

We can also help optimization by applying moment regularization across iterations:

reg_schedule = ExpRegularization(first_reg=10.0, last_reg=0.1, tau=20.0)
reg_schedule.plot(100) ;

svgd = graph.svgd(observed_data, 
                  prior=prior, 
                  learning_rate=step_schedule,
                  regularization=reg_schedule,
                  nr_moments=2
                  )
svgd.plot_ci(true_theta=true_theta)

[WARNING] phasic.svgd: svgd_step: clipped 1/20 particle gradient norm(s) exceeding 10000x the healthy batch median (a loss singularity or a diverging particle was encountered) -- inference for those particles is damped, not exact, for this step.

---------------------------------------------------------------------------
KeyboardInterrupt                         Traceback (most recent call last)
Cell In[31], line 1
----> 1 svgd = graph.svgd(observed_data, 
      2                   prior=prior,
      3                   learning_rate=step_schedule,
      4                   regularization=reg_schedule,

File ~/phasic/.pixi/envs/default/lib/python3.13/site-packages/phasic/__init__.py:6826, in Graph.svgd(self, observed_data, discrete, prior, n_particles, n_iterations, optimizer, learning_rate, bandwidth, theta_init, theta_dim, return_history, seed, verbose, progress, jit, parallel, n_devices, precompile, compilation_config, regularization, nr_moments, positive_params, param_transform, joint_index, rewards, fixed, tied, callback, weight_formula, preconditioner, epoch_starts, daisy_chain_t_eval, daisy_chain_granularity, daisy_chain_probe_theta, daisy_chain_t_eval_tol, final_read, exposure, exposure_param_index, validate_rewards, quiet_assumptions, exact_moment_grad, exact_final_grad, exact_grad)
   6822             # path).
   6823             if _callback_overridden:
   6824                 self._weight_callback = _prev_weight_callback
   6825                 self._weight_mode = _prev_weight_mode
-> 6826             if _formula_overridden:
   6827                 self._restore_weight_formula_state(*_prev_formula_state)

File ~/phasic/.pixi/envs/default/lib/python3.13/site-packages/phasic/svgd.py:6855, in SVGD.optimize(self, rewards, return_history)
   6851                     regularization=reg_value,
   6852                     rewards=rewards
   6853                 )
   6854 
-> 6855             results = run_svgd(
   6856                 log_prob_fn=None,  # Created dynamically per iteration
   6857                 theta_init=self.theta_init,
   6858                 n_steps=self.n_iterations,

File ~/phasic/.pixi/envs/default/lib/python3.13/site-packages/phasic/svgd.py:4457, in run_svgd(log_prob_fn, theta_init, n_steps, learning_rate, kernel, return_history, verbose, progress, compiled_grad, parallel_mode, n_devices, log_prob_fn_factory, regularization_schedule, lr_scale, fixed_mask, fixed_values, optimizer)
   4453         else:
   4454             compiled_grad_to_use = compiled_grad
   4455 
   4456         # Perform SVGD update
-> 4457         particles = svgd_step(particles, log_prob_fn, kernel, current_step_size,
   4458                              compiled_grad=compiled_grad_to_use,
   4459                              parallel_mode=parallel_mode,
   4460                              n_devices=n_devices,

File ~/phasic/.pixi/envs/default/lib/python3.13/site-packages/phasic/svgd.py:4217, in svgd_step(particles, log_prob_fn, kernel, step_size, compiled_grad, parallel_mode, n_devices, fixed_mask, fixed_values, optimizer)
   4213         # Single device vectorization - use vmap only
   4214         if compiled_grad_to_use is not None:
   4215             grad_log_p = vmap(compiled_grad_to_use)(particles_for_grad)
   4216         else:
-> 4217             grad_log_p = vmap(grad(log_prob_for_grad))(particles_for_grad)
   4218     elif actual_parallel_mode == 'none':
   4219         # No parallelization - sequential computation (useful for debugging)
   4220         if compiled_grad_to_use is not None:

    [... skipping hidden 1 frame]

File ~/phasic/.pixi/envs/default/lib/python3.13/site-packages/jax/_src/api.py:1114, in vmap.<locals>.vmap_f(*args, **kwargs)
   1111 try:
   1112   axis_data = batching.AxisData(axis_name, axis_size_, spmd_axis_name,
   1113                                 explicit_mesh_axis)
-> 1114   out_flat = batching.batch(
   1115       flat_fun, axis_data, in_axes_flat,
   1116       lambda: flatten_axes("vmap out_axes", out_tree(), out_axes)
   1117   ).call_wrapped(*args_flat)
   1118 except batching.SpecMatchError as e:
   1119   out_axes_flat = flatten_axes("vmap out_axes", out_tree(), out_axes)

File ~/phasic/.pixi/envs/default/lib/python3.13/site-packages/jax/_src/linear_util.py:212, in WrappedFun.call_wrapped(self, *args, **kwargs)
    210 def call_wrapped(self, *args, **kwargs):
    211   """Calls the transformed function"""
--> 212   return self.f_transformed(*args, **kwargs)

File ~/phasic/.pixi/envs/default/lib/python3.13/site-packages/jax/_src/interpreters/batching.py:644, in _batch_outer(f, axis_data, in_dims, *in_vals)
    642 tag = TraceTag()
    643 with source_info_util.transform_name_stack('vmap'):
--> 644   outs, trace = f(tag, in_dims, *in_vals)
    645 with core.ensure_no_leaks(trace): del trace
    646 return outs

File ~/phasic/.pixi/envs/default/lib/python3.13/site-packages/jax/_src/interpreters/batching.py:660, in _batch_inner(f, axis_data, out_dim_dests, tag, in_dims, *in_vals)
    656   in_tracers = map(partial(to_elt, trace, idx), in_vals, in_dims)
    657 with (core.set_current_trace(trace),
    658       core.extend_axis_env_nd([(axis_data.name, axis_data.size)]),
    659       core.add_spmd_axis_names(axis_data.spmd_name)):
--> 660   outs = f(*in_tracers)
    661   out_dim_dests = out_dim_dests() if callable(out_dim_dests) else out_dim_dests
    662   out_vals = map(partial(from_elt, trace, axis_data.size, axis_data.explicit_mesh_axis),
    663                  range(len(outs)), outs, out_dim_dests)

File ~/phasic/.pixi/envs/default/lib/python3.13/site-packages/jax/_src/interpreters/batching.py:342, in flatten_fun_for_vmap(f, store, in_tree, *args_flat)
    338 @lu.transformation_with_aux2
    339 def flatten_fun_for_vmap(f: Callable,
    340                          store: lu.Store, in_tree: PyTreeDef, *args_flat):
    341   py_args, py_kwargs = tree_unflatten(in_tree, args_flat)
--> 342   ans = f(*py_args, **py_kwargs)
    343   ans, out_tree = tree_flatten(ans, is_leaf=is_vmappable)
    344   store.store(out_tree)

File ~/phasic/.pixi/envs/default/lib/python3.13/site-packages/jax/_src/linear_util.py:429, in _get_result_paths_thunk(_fun, _store, *args, **kwargs)
    427 @transformation_with_aux2
    428 def _get_result_paths_thunk(_fun: Callable, _store: Store, *args, **kwargs):
--> 429   ans = _fun(*args, **kwargs)
    430   result_paths = tuple(f"result{_clean_keystr_arg_names(path)}" for path, _ in generate_key_paths(ans))
    431   if _store:
    432     # In some instances a lu.WrappedFun is called multiple times, e.g.,
    433     # the bwd function in a custom_vjp

    [... skipping hidden 1 frame]

File ~/phasic/.pixi/envs/default/lib/python3.13/site-packages/jax/_src/api.py:405, in grad.<locals>.grad_f(*args, **kwargs)
    402 @wraps(fun, docstr=docstr, argnums=argnums)
    403 @api_boundary
    404 def grad_f(*args, **kwargs):
--> 405   _, g = value_and_grad_f(*args, **kwargs)
    406   return g

    [... skipping hidden 1 frame]

File ~/phasic/.pixi/envs/default/lib/python3.13/site-packages/jax/_src/api.py:478, in value_and_grad.<locals>.value_and_grad_f(*args, **kwargs)
    476   _check_input_dtype_grad(holomorphic, allow_int, leaf)
    477 if not has_aux:
--> 478   ans, vjp_py = _vjp(f_partial, *dyn_args)
    479 else:
    480   ans, vjp_py, aux = _vjp(f_partial, *dyn_args, has_aux=True)

File ~/phasic/.pixi/envs/default/lib/python3.13/site-packages/jax/_src/api.py:2134, in _vjp(fun, has_aux, *primals)
   2132 if not has_aux:
   2133   flat_fun, out_tree = flatten_fun_nokwargs(fun, in_tree)
-> 2134   out_primals, vjp = ad.vjp(flat_fun, primals_flat)
   2135   out_tree = out_tree()
   2136 else:

File ~/phasic/.pixi/envs/default/lib/python3.13/site-packages/jax/_src/interpreters/ad.py:313, in vjp(traceable, primals, has_aux)
    311 def vjp(traceable: lu.WrappedFun, primals, has_aux=False):
    312   if not has_aux:
--> 313     out_primals, pvals, jaxpr, consts = linearize(traceable, *primals)
    314   else:
    315     out_primals, pvals, jaxpr, consts, aux = linearize(traceable, *primals, has_aux=True)

File ~/phasic/.pixi/envs/default/lib/python3.13/site-packages/jax/_src/interpreters/ad.py:287, in linearize(traceable, *primals, **kwargs)
    285 has_aux = kwargs.pop('has_aux', False)
    286 if config.use_direct_linearize.value:
--> 287   return direct_linearize(traceable, primals, kwargs, has_aux=has_aux)
    288 if not has_aux:
    289   jvpfun = jvp(traceable)

File ~/phasic/.pixi/envs/default/lib/python3.13/site-packages/jax/_src/interpreters/ad.py:261, in direct_linearize(traceable, primals, kwargs, has_aux, tag)
    256   aux_primals = [x.primal
    257                  if isinstance(x, LinearizeTracer)
    258                  and x._trace.tag is linearize_trace.tag
    259                  else x for x in aux]
    260 else:
--> 261   ans = traceable.call_wrapped(*tracers)
    262   aux = None
    263 out_primals, out_tangents = unzip2(map(linearize_trace.to_primal_tangent_pair, ans))

File ~/phasic/.pixi/envs/default/lib/python3.13/site-packages/jax/_src/linear_util.py:212, in WrappedFun.call_wrapped(self, *args, **kwargs)
    210 def call_wrapped(self, *args, **kwargs):
    211   """Calls the transformed function"""
--> 212   return self.f_transformed(*args, **kwargs)

File ~/phasic/.pixi/envs/default/lib/python3.13/site-packages/jax/_src/api_util.py:90, in flatten_fun_nokwargs(f, store, in_tree, *args_flat)
     86 @lu.transformation_with_aux2
     87 def flatten_fun_nokwargs(f: Callable, store: lu.Store,
     88                          in_tree: PyTreeDef, *args_flat):
     89   py_args = tree_unflatten(in_tree, args_flat)
---> 90   ans = f(*py_args)
     91   ans, out_tree = tree_flatten(ans)
     92   store.store(out_tree)

File ~/phasic/.pixi/envs/default/lib/python3.13/site-packages/jax/_src/api_util.py:292, in _argnums_partial(_fun, _dyn_argnums, _fixed_args, *dyn_args, **kwargs)
    290 args = [next(fixed_args_).val if x is sentinel else x for x in args]
    291 assert next(fixed_args_, sentinel) is sentinel
--> 292 return _fun(*args, **kwargs)

File ~/phasic/.pixi/envs/default/lib/python3.13/site-packages/jax/_src/linear_util.py:429, in _get_result_paths_thunk(_fun, _store, *args, **kwargs)
    427 @transformation_with_aux2
    428 def _get_result_paths_thunk(_fun: Callable, _store: Store, *args, **kwargs):
--> 429   ans = _fun(*args, **kwargs)
    430   result_paths = tuple(f"result{_clean_keystr_arg_names(path)}" for path, _ in generate_key_paths(ans))
    431   if _store:
    432     # In some instances a lu.WrappedFun is called multiple times, e.g.,
    433     # the bwd function in a custom_vjp

File ~/phasic/.pixi/envs/default/lib/python3.13/site-packages/phasic/svgd.py:6274, in SVGD._log_prob_unified(self, theta, nr_moments, sample_moments, regularization, rewards)
   6270 
   6271         # Evaluate model
   6272         try:
   6273             result = self.model(theta_transformed, self.observed_data, rewards=rewards)
-> 6274         except Exception as e:
   6275             raise ValueError(
   6276                 f"Model evaluation failed. Ensure model has signature model(theta, times, rewards=None). "
   6277                 f"Error: {e}"

    [... skipping hidden 1 frame]

File ~/phasic/.pixi/envs/default/lib/python3.13/site-packages/jax/_src/custom_derivatives.py:747, in custom_vjp.__call__(self, *args, **kwargs)
    743 flat_fwd, out_trees = _flatten_fwd(
    744     fwd_, self.nondiff_argnums, self.symbolic_zeros, debug_fun,
    745     debug_fwd, in_tree, out_type)
    746 flat_bwd = _flatten_bwd(bwd, in_tree, in_avals, out_trees)
--> 747 out_flat = custom_vjp_call_p.bind(flat_fun, flat_fwd, flat_bwd,
    748                                   *args_flat, out_trees=out_trees,
    749                                   symbolic_zeros=self.symbolic_zeros)
    750 _, (out_tree, _, _) = lu.merge_linear_aux(out_type, out_trees)
    751 return tree_unflatten(out_tree, out_flat)

File ~/phasic/.pixi/envs/default/lib/python3.13/site-packages/jax/_src/custom_derivatives.py:999, in CustomVJPCallPrimitive.bind(self, *args, **params)
    998 def bind(self, *args, **params):
--> 999   return self._true_bind(*args, **params)

File ~/phasic/.pixi/envs/default/lib/python3.13/site-packages/jax/_src/core.py:650, in Primitive._true_bind(self, *args, **params)
    648 trace_ctx.set_trace(eval_trace)
    649 try:
--> 650   return self.bind_with_trace(prev_trace, args, params)
    651 finally:
    652   trace_ctx.set_trace(prev_trace)

File ~/phasic/.pixi/envs/default/lib/python3.13/site-packages/jax/_src/custom_derivatives.py:1003, in CustomVJPCallPrimitive.bind_with_trace(self, trace, args, params)
   1001 def bind_with_trace(self, trace, args, params):
   1002   fun, fwd, bwd, tracers = args[0], args[1], args[2], args[3:]
-> 1003   return trace.process_custom_vjp_call(self, fun, fwd, bwd, tracers, **params)

File ~/phasic/.pixi/envs/default/lib/python3.13/site-packages/jax/_src/interpreters/ad.py:955, in LinearizeTrace.process_custom_vjp_call(self, prim, fun, fwd, bwd, tracers, out_trees, symbolic_zeros)
    953 fwd_in_flat = [x for pair in fwd_in for x in pair]   # flatten
    954 with core.set_current_trace(self.parent_trace):
--> 955   res_and_primals_out = fwd.call_wrapped(*fwd_in_flat)
    957 _, res_tree, input_fwds = out_trees()
    958 num_res_out = res_tree.num_leaves - sum(f is not None for f in input_fwds)

File ~/phasic/.pixi/envs/default/lib/python3.13/site-packages/jax/_src/linear_util.py:212, in WrappedFun.call_wrapped(self, *args, **kwargs)
    210 def call_wrapped(self, *args, **kwargs):
    211   """Calls the transformed function"""
--> 212   return self.f_transformed(*args, **kwargs)

File ~/phasic/.pixi/envs/default/lib/python3.13/site-packages/jax/_src/custom_derivatives.py:848, in _flatten_fwd(f, store, nondiff_argnums, symbolic_zeros, debug_primal, debug_fwd, in_tree, maybe_out_type, *args)
    846 if config.mutable_array_checks.value:
    847   _check_for_aliased_refs(f, nondiff_argnums, debug_primal, py_args)
--> 848 pair_out = f(*py_args)
    849 if not isinstance(pair_out, (list, tuple)) or len(pair_out) != 2:
    850   msg = (f"Custom VJP fwd rule {fwd_name} for function {primal_name} "
    851          "must produce a pair (list or tuple of length two) where the first "
    852          "element represents the primal output (equal to those of the "
   (...)    855          "forward pass for use on the backward pass), but "
    856          f"instead of a pair the fwd rule {fwd_name} produced {pair_out}.")

File ~/phasic/.pixi/envs/default/lib/python3.13/site-packages/jax/_src/linear_util.py:429, in _get_result_paths_thunk(_fun, _store, *args, **kwargs)
    427 @transformation_with_aux2
    428 def _get_result_paths_thunk(_fun: Callable, _store: Store, *args, **kwargs):
--> 429   ans = _fun(*args, **kwargs)
    430   result_paths = tuple(f"result{_clean_keystr_arg_names(path)}" for path, _ in generate_key_paths(ans))
    431   if _store:
    432     # In some instances a lu.WrappedFun is called multiple times, e.g.,
    433     # the bwd function in a custom_vjp

File ~/phasic/.pixi/envs/default/lib/python3.13/site-packages/phasic/__init__.py:8424, in Graph.pmf_and_moments_from_graph.<locals>.model_fwd(theta, times, rewards)
   8422         def model_fwd(theta, times, rewards=None):
   8423             # Call the underlying computation, not model (avoid infinite recursion!)
-> 8424             pmf, moments = _compute_pure(theta, times, rewards)
   8425             return (pmf, moments), (theta, times, rewards)

File ~/phasic/.pixi/envs/default/lib/python3.13/site-packages/phasic/__init__.py:8284, in Graph.pmf_and_moments_from_graph.<locals>._compute_pure(theta, times, rewards)
   8280                         rewards_np = None
   8281 
   8282                     return _compute_pmf_and_moments_cached(theta_np, times_np, rewards_np)
   8283 
-> 8284                 result = jax.pure_callback(
   8285                     callback_fn,
   8286                     (pmf_shape, moments_shape),
   8287                     theta, times, rewards_jax,  # Pass all args to allow vmap

File ~/phasic/.pixi/envs/default/lib/python3.13/site-packages/jax/_src/callback.py:388, in pure_callback(callback, result_shape_dtypes, sharding, vectorized, vmap_method, *args, **kwargs)
    385 result_avals = tree_util.tree_map(
    386     lambda x: core.ShapedArray(x.shape, x.dtype), result_shape_dtypes)
    387 flat_result_avals, out_tree = tree_util.tree_flatten(result_avals)
--> 388 out_flat = pure_callback_p.bind(
    389     *flat_args,
    390     callback=_FlatCallback(callback, in_tree),
    391     result_avals=tuple(flat_result_avals),
    392     sharding=sharding,
    393     vmap_method=vmap_method,
    394 )
    395 return tree_util.tree_unflatten(out_tree, out_flat)

File ~/phasic/.pixi/envs/default/lib/python3.13/site-packages/jax/_src/core.py:634, in Primitive.bind(self, *args, **params)
    632 def bind(self, *args, **params):
    633   args = args if self.skip_canonicalization else map(canonicalize_value, args)
--> 634   return self._true_bind(*args, **params)

File ~/phasic/.pixi/envs/default/lib/python3.13/site-packages/jax/_src/core.py:650, in Primitive._true_bind(self, *args, **params)
    648 trace_ctx.set_trace(eval_trace)
    649 try:
--> 650   return self.bind_with_trace(prev_trace, args, params)
    651 finally:
    652   trace_ctx.set_trace(prev_trace)

File ~/phasic/.pixi/envs/default/lib/python3.13/site-packages/jax/_src/core.py:662, in Primitive.bind_with_trace(self, trace, args, params)
    660     with set_current_trace(trace):
    661       return self.to_lojax(*args, **params)  # type: ignore
--> 662   return trace.process_primitive(self, args, params)
    663 trace.process_primitive(self, args, params)  # may raise lojax error
    664 raise Exception(f"couldn't apply typeof to args: {args}")

File ~/phasic/.pixi/envs/default/lib/python3.13/site-packages/jax/_src/interpreters/batching.py:531, in BatchTrace.process_primitive(self, p, tracers, params)
    529 elif p in primitive_batchers:
    530   with core.set_current_trace(self.parent_trace):
--> 531     val_out, dim_out = primitive_batchers[p](vals_in, dims_in, **params)
    532 else:
    533   raise NotImplementedError(f"Batching rule for '{p}' not implemented")

File ~/phasic/.pixi/envs/default/lib/python3.13/site-packages/jax/_src/ffi.py:704, in ffi_batching_rule(prim, args, dims, vmap_method, result_avals, **kwargs)
    700   if kwargs.get("input_layouts") is not None:
    701     kwargs["input_layouts"] = tuple(
    702         None if layout is None else tuple(n + 1 for n in layout) + (0,)
    703         for layout in kwargs["input_layouts"])
--> 704   outvals = prim.bind(
    705     *bcast_args,
    706     vmap_method=vmap_method,
    707     result_avals=batched_result_avals,
    708     **kwargs,
    709   )
    710 elif vmap_method == "sequential" or vmap_method == "sequential_unrolled":
    711   is_batched = [d is not batching.not_mapped for d in dims]

File ~/phasic/.pixi/envs/default/lib/python3.13/site-packages/jax/_src/core.py:634, in Primitive.bind(self, *args, **params)
    632 def bind(self, *args, **params):
    633   args = args if self.skip_canonicalization else map(canonicalize_value, args)
--> 634   return self._true_bind(*args, **params)

File ~/phasic/.pixi/envs/default/lib/python3.13/site-packages/jax/_src/core.py:650, in Primitive._true_bind(self, *args, **params)
    648 trace_ctx.set_trace(eval_trace)
    649 try:
--> 650   return self.bind_with_trace(prev_trace, args, params)
    651 finally:
    652   trace_ctx.set_trace(prev_trace)

File ~/phasic/.pixi/envs/default/lib/python3.13/site-packages/jax/_src/core.py:662, in Primitive.bind_with_trace(self, trace, args, params)
    660     with set_current_trace(trace):
    661       return self.to_lojax(*args, **params)  # type: ignore
--> 662   return trace.process_primitive(self, args, params)
    663 trace.process_primitive(self, args, params)  # may raise lojax error
    664 raise Exception(f"couldn't apply typeof to args: {args}")

File ~/phasic/.pixi/envs/default/lib/python3.13/site-packages/jax/_src/core.py:1189, in EvalTrace.process_primitive(self, primitive, args, params)
   1187 args = map(full_lower, args)
   1188 check_eval_args(args)
-> 1189 return primitive.impl(*args, **params)

File ~/phasic/.pixi/envs/default/lib/python3.13/site-packages/jax/_src/dispatch.py:90, in apply_primitive(prim, *args, **params)
     88 prev = lib.jax_jit.swap_thread_local_state_disable_jit(False)
     89 try:
---> 90   outs = fun(*args)
     91 finally:
     92   lib.jax_jit.swap_thread_local_state_disable_jit(prev)

    [... skipping hidden 1 frame]

File ~/phasic/.pixi/envs/default/lib/python3.13/site-packages/jax/_src/pjit.py:268, in _cpp_pjit.<locals>.cache_miss(*args, **kwargs)
    263 if config.no_tracing.value:
    264   raise RuntimeError(f"re-tracing function {jit_info.fun_sourceinfo} for "
    265                      "`jit`, but 'no_tracing' is set")
    267 (outs, out_flat, out_tree, args_flat, jaxpr,
--> 268  executable, pgle_profiler, const_args) = _python_pjit_helper(
    269      fun, jit_info, *args, **kwargs)
    271 maybe_fastpath_data = _get_fastpath_data(
    272     executable, out_tree, args_flat, out_flat, jaxpr.effects, jaxpr.consts,
    273     jit_info.abstracted_axes, pgle_profiler,
    274     const_args)
    276 return outs, maybe_fastpath_data, _need_to_rebuild_with_fdo(pgle_profiler)

File ~/phasic/.pixi/envs/default/lib/python3.13/site-packages/jax/_src/pjit.py:150, in _python_pjit_helper(fun, jit_info, *args, **kwargs)
    147   out_flat, compiled, profiler, const_args = _pjit_call_impl_python(
    148       *args_flat, **p.params)
    149 else:
--> 150   out_flat = jit_p.bind(*args_flat, **p.params)
    151   compiled = None
    152   profiler = None

File ~/phasic/.pixi/envs/default/lib/python3.13/site-packages/jax/_src/core.py:634, in Primitive.bind(self, *args, **params)
    632 def bind(self, *args, **params):
    633   args = args if self.skip_canonicalization else map(canonicalize_value, args)
--> 634   return self._true_bind(*args, **params)

File ~/phasic/.pixi/envs/default/lib/python3.13/site-packages/jax/_src/core.py:650, in Primitive._true_bind(self, *args, **params)
    648 trace_ctx.set_trace(eval_trace)
    649 try:
--> 650   return self.bind_with_trace(prev_trace, args, params)
    651 finally:
    652   trace_ctx.set_trace(prev_trace)

File ~/phasic/.pixi/envs/default/lib/python3.13/site-packages/jax/_src/core.py:662, in Primitive.bind_with_trace(self, trace, args, params)
    660     with set_current_trace(trace):
    661       return self.to_lojax(*args, **params)  # type: ignore
--> 662   return trace.process_primitive(self, args, params)
    663 trace.process_primitive(self, args, params)  # may raise lojax error
    664 raise Exception(f"couldn't apply typeof to args: {args}")

File ~/phasic/.pixi/envs/default/lib/python3.13/site-packages/jax/_src/core.py:1189, in EvalTrace.process_primitive(self, primitive, args, params)
   1187 args = map(full_lower, args)
   1188 check_eval_args(args)
-> 1189 return primitive.impl(*args, **params)

File ~/phasic/.pixi/envs/default/lib/python3.13/site-packages/jax/_src/pjit.py:1829, in _pjit_call_impl(jaxpr, in_shardings, out_shardings, in_layouts, out_layouts, donated_invars, ctx_mesh, name, keep_unused, inline, compiler_options_kvs, *args)
   1821 donated_argnums = tuple(i for i, d in enumerate(donated_invars) if d)
   1822 cache_key = pxla.JitGlobalCppCacheKeys(
   1823     donate_argnums=donated_argnums, donate_argnames=None,
   1824     device=None, backend=None,
   (...)   1827     in_layouts_treedef=None, in_layouts_leaves=in_layouts,
   1828     out_layouts_treedef=None, out_layouts_leaves=out_layouts)
-> 1829 return xc._xla.pjit(
   1830     name, f, call_impl_cache_miss, [], [], cache_key,
   1831     tree_util.dispatch_registry, pxla.cc_shard_arg,
   1832     _get_cpp_global_cache(cache_key.contains_explicit_attributes))(*args)

File ~/phasic/.pixi/envs/default/lib/python3.13/site-packages/jax/_src/pjit.py:1805, in _pjit_call_impl.<locals>.call_impl_cache_miss(*args_, **kwargs_)
   1801 def call_impl_cache_miss(*args_, **kwargs_):
   1802   # args_ do not include the const args
   1803   # See https://docs.jax.dev/en/latest/internals/constants.html.
   1804   # TODO(necula): remove num_const_args when fixing the C++ path
-> 1805   out_flat, compiled, pgle_profiler, const_args = _pjit_call_impl_python(
   1806       *args, jaxpr=jaxpr, in_shardings=in_shardings,
   1807       out_shardings=out_shardings, in_layouts=in_layouts,
   1808       out_layouts=out_layouts, donated_invars=donated_invars,
   1809       ctx_mesh=ctx_mesh, name=name, keep_unused=keep_unused,
   1810       inline=inline, compiler_options_kvs=compiler_options_kvs)
   1811   fastpath_data = _get_fastpath_data(
   1812       compiled, tree_structure(out_flat), args, out_flat,
   1813       jaxpr.effects, jaxpr.consts, None, pgle_profiler,
   1814       const_args)
   1815   return out_flat, fastpath_data, _need_to_rebuild_with_fdo(pgle_profiler)

File ~/phasic/.pixi/envs/default/lib/python3.13/site-packages/jax/_src/pjit.py:1758, in _pjit_call_impl_python(jaxpr, in_shardings, out_shardings, in_layouts, out_layouts, donated_invars, ctx_mesh, name, keep_unused, inline, compiler_options_kvs, *args)
   1746 # Passing mutable PGLE profile here since it should be extracted by JAXPR to
   1747 # initialize the fdo_profile compile option.
   1748 computation = _resolve_and_lower(
   1749     args, jaxpr=jaxpr, in_shardings=in_shardings,
   1750     out_shardings=out_shardings, in_layouts=in_layouts,
   (...)   1756     compiler_options_kvs=compiler_options_kvs,
   1757 )
-> 1758 compiled = computation.compile()
   1760 # This check is expensive so only do it if enable_checks is on.
   1761 if compiled._auto_spmd_lowering and config.enable_checks.value:

File ~/phasic/.pixi/envs/default/lib/python3.13/site-packages/jax/_src/interpreters/pxla.py:2537, in MeshComputation.compile(self, compiler_options, device_assignment)
   2534 assert isinstance(compilation_device_list, (type(None), xc.DeviceList))
   2536 if self._executable is None or compiler_options_kvs or device_assignment:
-> 2537   executable = UnloadedMeshExecutable.from_hlo(
   2538       self._name, self._hlo, **self.compile_args,
   2539       compiler_options_kvs=compiler_options_kvs,
   2540       device_list=compilation_device_list)
   2541   if not compiler_options_kvs:
   2542     self._executable = executable

File ~/phasic/.pixi/envs/default/lib/python3.13/site-packages/jax/_src/interpreters/pxla.py:3083, in UnloadedMeshExecutable.from_hlo(***failed resolving arguments***)
   3080       break
   3082 util.test_event("pxla_cached_compilation")
-> 3083 xla_executable = _cached_compilation(
   3084     hlo, name, mesh, spmd_lowering,
   3085     tuple_args, auto_spmd_lowering, allow_prop_to_inputs,
   3086     allow_prop_to_outputs, tuple(host_callbacks), backend, device_list,
   3087     pmap_nreps, compiler_options_kvs, pgle_profiler)
   3089 if auto_spmd_lowering:
   3090   assert mesh is not None

File ~/phasic/.pixi/envs/default/lib/python3.13/site-packages/jax/_src/interpreters/pxla.py:2864, in _cached_compilation(computation, name, mesh, spmd_lowering, tuple_args, auto_spmd_lowering, allow_prop_to_inputs, allow_prop_to_outputs, host_callbacks, backend, da, pmap_nreps, compiler_options_kvs, pgle_profiler)
   2856 compile_options = create_compile_options(
   2857     computation, mesh, spmd_lowering, tuple_args, auto_spmd_lowering,
   2858     allow_prop_to_inputs, allow_prop_to_outputs, backend,
   2859     dev, pmap_nreps, compiler_options)
   2861 with dispatch.log_elapsed_time(
   2862     "Finished XLA compilation of {fun_name} in {elapsed_time:.9f} sec",
   2863     fun_name=name, event=dispatch.BACKEND_COMPILE_EVENT):
-> 2864   xla_executable = compiler.compile_or_get_cached(
   2865       backend, computation, dev, compile_options, host_callbacks,
   2866       da, pgle_profiler)
   2867 return xla_executable

File ~/phasic/.pixi/envs/default/lib/python3.13/site-packages/jax/_src/compiler.py:494, in compile_or_get_cached(backend, computation, devices, compile_options, host_callbacks, executable_devices, pgle_profiler)
    492 else:
    493   log_persistent_cache_miss(module_name, cache_key)
--> 494   return _compile_and_write_cache(
    495       backend,
    496       computation,
    497       executable_devices,
    498       compile_options,
    499       host_callbacks,
    500       module_name,
    501       cache_key,
    502   )

File ~/phasic/.pixi/envs/default/lib/python3.13/site-packages/jax/_src/compiler.py:762, in _compile_and_write_cache(backend, computation, executable_devices, compile_options, host_callbacks, module_name, cache_key)
    752 def _compile_and_write_cache(
    753     backend: xc.Client,
    754     computation: ir.Module,
   (...)    759     cache_key: str,
    760 ) -> xc.LoadedExecutable:
    761   start_time = time.monotonic()
--> 762   executable = backend_compile_and_load(
    763       backend, computation, executable_devices, compile_options, host_callbacks
    764   )
    765   compile_time = time.monotonic() - start_time
    766   _cache_write(
    767       cache_key, compile_time, module_name, backend, executable, host_callbacks
    768   )

File ~/phasic/.pixi/envs/default/lib/python3.13/site-packages/jax/_src/profiler.py:364, in annotate_function.<locals>.wrapper(*args, **kwargs)
    361 @wraps(func)
    362 def wrapper(*args, **kwargs):
    363   with TraceAnnotation(name, **decorator_kwargs):
--> 364     return func(*args, **kwargs)

File ~/phasic/.pixi/envs/default/lib/python3.13/site-packages/jax/_src/compiler.py:369, in backend_compile_and_load(backend, module, executable_devices, options, host_callbacks)
    367 else:
    368   if host_callbacks:
--> 369     return backend.compile_and_load(
    370         built_c,
    371         executable_devices=executable_devices,
    372         compile_options=options,
    373         host_callbacks=host_callbacks,
    374     )
    375   # Some backends don't have `host_callbacks` option yet
    376   # TODO(sharadmv): remove this fallback when all backends allow `compile`
    377   # to take in `host_callbacks`
    378   return backend.compile_and_load(
    379       built_c,
    380       executable_devices=executable_devices,
    381       compile_options=options,
    382   )

KeyboardInterrupt:

It does not contribute much in this case that already converges, but for more complex models using regularization can help stabilize the optimization.

Posterior predictive checks

from scipy.stats import ecdf

# ECDF for observed data
res = ecdf(observed_data)
x, ecdf = res.cdf.quantiles, res.cdf.probabilities

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(8,3))
sns.histplot(observed_data,
             kde=False, stat='density', label='Data', ax=ax1)
graph.update_weights(svgd.get_results()['theta_mean'])
ax1.plot(x, graph.pdf(x), lw=1.5, c='C1', label='PDF')
ax1.set_xlabel('Reward (time)')
ax1.set_title('PDF')
ax1.legend()
ax2.plot(x, ecdf, lw=1.5, c='C0', label='ECDF')
ax2.plot(x, graph.cdf(x), lw=1.5, c='C1', label='CDF')
ax2.set_xlabel('Reward (time)')
ax2.set_title('CDF')
ax2.legend() 
ax3 = ax2.twinx()
ax3.plot(x, graph.cdf(x)-ecdf, lw=0.5, c='C2', 
    alpha=0.5,label='CDF-ECDF')
ax3.grid(False)
ax3.set_ylabel('CDF - ECDF', color='C2')
ax3.tick_params(axis='y', colors='C2')
plt.tight_layout()

Adaptive optimizers

Phasic provides three optimizers for SVGD that offer adaptive per-parameter learning rates. These can be particularly useful when gradients have vastly different scales across parameters or when fixed step sizes cause oscillation.

Adam

When using Adam, the learning_rate parameter passed to SVGD is ignored in favor of the optimizer’s learning rate.

Adam (Adaptive Moment Estimation) maintains running estimates of two quantities for each parameter:

First moment ($m$): An exponentially weighted average of past gradients (momentum)
Second moment ($v$): An exponentially weighted average of past squared gradients (gradient variance)

Update rule:

\[ \begin{align} m &\leftarrow \beta_1 \cdot m + (1 - \beta_1) \cdot g \\ v &\leftarrow \beta_2 \cdot v + (1 - \beta_2) \cdot g^2 \\[0.5em] \hat{m} &= \frac{m}{1 - \beta_1^t} \\[0.5em] \hat{v} &= \frac{v}{1 - \beta_2^t} \\[0.5em] \theta &\leftarrow \theta - \alpha \cdot \frac{\hat{m}}{\sqrt{\hat{v}} + \epsilon} \end{align} \]

Each parameter gets its own effective learning rate:

Parameters with consistently large gradients $\rightarrow$ larger $\hat{v}$ $\rightarrow$ smaller effective step
Parameters with consistently small gradients $\rightarrow$ smaller $\hat{v}$ $\rightarrow$ larger effective step

This automatic scaling helps when different parameters naturally have different gradient magnitudes.

When fixed learning rate works well:

Well-behaved optimization landscapes
Parameters with similar gradient scales
When you’ve tuned the learning rate carefully

When Adam helps:

Large datasets causing large gradient magnitudes
Parameters with vastly different scales
“Shark teeth” oscillation patterns in convergence
When you want reasonable results without extensive tuning

Tuning Guidelines:

learning_rate: Start with 0.001-0.01. If convergence is too slow, increase. If unstable, decrease.
beta1: 0.9 works well for most cases. Lower values (0.8) give less momentum.
beta2: 0.999 is standard. Lower values (0.99) adapt faster to gradient changes.
epsilon: Rarely needs adjustment. Increase to 1e-7 if you see numerical issues.

Consider Adam when you observe:

Oscillating convergence (“shark teeth” pattern in loss)
Different parameters converging at different rates
Sensitivity to learning rate choice
Large datasets** causing large gradient magnitudes

Stick with fixed learning rate when:

Convergence is already smooth and fast
You’ve tuned the learning rate well for your problem
You want minimal computational overhead
The optimization landscape is well-behaved

Kingma and Ba (2014) - Adam: A Method for Stochastic Optimization. arXiv:1412.6980.

from phasic import Adam

optimizer = Adam(
    learning_rate=0.001,  # Base learning rate, scaled per-parameter by Adam (alpha)
    beta1=0.9,            # Decay rate for first moment (momentum). Higher = more smoothing
    beta2=0.999,          # Decay rate for second moment. Higher = longer memory of gradient magnitudes
    epsilon=1e-8          # Small constant to prevent division by zero
)

from phasic import Adam

svgd = graph.svgd(
    observed_data,
    prior = prior,
    optimizer=Adam(learning_rate=0.25),
)
svgd.plot_ci(true_theta=true_theta)

It can overcome even a misleading prior:

svgd = graph.svgd(
    observed_data,
    prior = GaussPrior(ci=[1, 3]),
    optimizer=Adam(learning_rate=0.25),
)
svgd.summary()

svgd.plot_convergence() ;

svgd.plot_ci(true_theta=true_theta) ;

svgd.plot_trace() ;

SGDMomentum

SGDMomentum (SGD with momentum) accumulates velocity in directions of persistent gradient descent, helping accelerate convergence and dampen oscillations. Update rule:

\[ \begin{align} v &= m * v + g \\ \theta &= \theta + (l * v) \end{align} \]

where $g$ is gradient $l$ is learning rate, and $m$ and $v$ are the first and second moments.

from phasic.svgd import SGDMomentum

optimizer = SGDMomentum(
    learning_rate=0.01,  # Step size
    momentum=0.9         # Momentum coefficient (0.9 standard, 0.99 high)
)
svgd = graph.svgd(
    observed_data,
    optimizer=optimizer,
)
svgd.summary()

svgd.plot_convergence() ;

svgd.plot_ci(true_theta=true_theta) ;

svgd.plot_trace() ;

RMSprop

RMSprop (Hinton 2012) divides the learning rate by an exponentially decaying average of squared gradients, adapting per-parameter. Update rule:

$$ \[\begin{align} v &= d * v + (1 - d) * g^2 \\ \theta &= \theta + \frac{l * g}{\sqrt{v} + \epsilon} \end{align}\] $$

where $g$ is gradient $l$ is learning rate, $d$ is the decay and $v$ is

from phasic.svgd import RMSprop

optimizer = RMSprop(
    learning_rate=1,  # Base learning rate (default 0.001)
    decay=0.99,       # Decay rate for squared gradient average
    epsilon=1e-8      # Numerical stability
)
svgd = graph.svgd(
    observed_data,
    optimizer=optimizer,
)
svgd.summary()

svgd.plot_convergence() ;

svgd.plot_ci(true_theta=true_theta) ;

svgd.plot_trace() ;

Adagrad

Adagrad (Duchi et al. 2011) accumulates the sum of squared gradients, giving smaller learning rates to parameters with large accumulated gradients. Update rule:

\[ \begin{align} G &= g^2 \\ \theta &= \theta \frac{l * g}{\sqrt{G} + \epsilon} \end{align} \]

The learning rate of Adagrad decays over time as $G$ accumulates. For long runs, RMSprop or Adam may perform better.

from phasic import Adagrad

optimizer = Adagrad(
    learning_rate=2,  # Base learning rate (default 0.001)
    epsilon=1e-8         # Numerical stability
)
svgd = graph.svgd(
    observed_data,
    optimizer=optimizer,
)
svgd.summary()

svgd.plot_convergence() ;

svgd.plot_ci(true_theta=true_theta) ;

svgd.plot_trace() ;

Weight update formula

For models where updating edge weights cannot be expressed as an inner product of parameters and edge coefficients, a callback update function can be passed to graph.update_weights(). For inference with SVGD, however, such non-linear weight updates must be expressed using a “weight formula” string that is compiled and made available to the C backend.

The weight formula string may include t0, t1, … representing the scalars of the theta parameter vector, c0, c1, … representing the coefficient vector of the edge being updated. numeric literals (3, 2.5, 1e-4, .5), arithmetic operators, elementary functions, comparisons, logic, branchless selection. The weight formula producing the default inner product would be 'c0*t0 + c1*t1 + c2*t2'.

The weight formula is added as an attribute to the graph, but can also be passed as a keyword argument to graph.svgd().

graph = Graph(coalescent_1param, theta_dim=1)
graph.weight_formula = "c0*t0"
graph.svgd(observed_data)
# or
# graph.svgd(observed_data, weight_formula="c0*t0")

Here are some example uses:

# Basic arithmetics:
"c0*t0 + c1*t1"            # the default linear inner product (Σ cₖ·tₖ)
"c0*t0 * c1*t1"            # product form (same as weight_mode='log')
"c0*t0 / (1 + c1*t1)"      # ratio / saturating rate
"t0**2"                    # power (right-assoc: 2**3**2 == 2**(3**2) == 512)
"t0**c0"                   # power with a per-edge exponent
"c0 * t0**2 + c1*t1"       # precedence: ** > * > + (no parens needed)
"-c0*t0 + c1"              # unary minus  (note: result must stay ≥ 0)
"2*t0 - 0.5*c0"            # integer/float literals

# Elementary functions: exp log sqrt logistic pow
"exp(c0*t0 + c1*t1)"       # log-link (always positive) — classic GLM rate
"exp(-c0*t0)"              # decaying rate (resistance-style)
"log(c0*t0)"              # requires c0*t0 > 0, else NaN → raises
"sqrt(c0*t0)"             # requires argument ≥ 0
"logistic(c0*t0)"         # numerically-stable 1/(1+exp(-·)) ∈ (0,1)
"pow(t0, c0)"             # same as t0**c0 (2-arg form)
"c2 * logistic(c0*t0) + 1e-3"   # bounded rate with a small floor

# Comparisons: ==  !=  <  >  <=  >= delta (Kronecker)
"delta(c0, c1)"           # 1.0 if c0 == c1 else 0.0  (sugar for c0 == c1)
"delta(c0, 3) * c1*t0"    # this term active only on edges where c0 == 3
"(c0 < c1) * t0"          # gate a term by an edge-structural comparison
"(c0 >= 2) * c1*t1"

# Logic (nonzero = true): and or not
"and(c0 == 3, c1 < 2)"            # 1.0 iff both hold
"or(c0 == 0, c1 == 0) * t0"       # term active if either coeff is zero
"not(c0) * t1"                    # active only where c0 == 0

# Branchless selection: select(cond, a, b) corresponds to cond ? a : b
"select(c0 == 3, c1*t0, c2*t1)"               # pick a branch by edge type
"delta(c0,c1)*t0 + (1 - delta(c0,c1))*t1"     # partition-of-unity (same idea)
"select(and(c0 == 3, c1 < 2), c2*t0, t1)"     # compound condition
"select(c0 > c1, exp(c2*t0), c2*t0)"          # different functional forms per edge

As with the the callback argument to update_weights, the (theta) vector can be shorter than the coefficient vectors. Just pass the param_length keyword argument to the Graph constructor. Allowing for more coefficients is handy if you need to pass more information to the callback function in order to update edge weights.

Two invariants to keep in mind

Conditions cannot depend on the values of theta: The operands of == != < > <= >=, of delta, of and/or/not, and the first argument of select can use constants and c0, c1, … but not t0, t1, … . The value branches of select may use theta freely. For smooth gating on theta, use logistic(...).

"select(t0 > c0, t1, t2)"   # will raise exception because condition depends on theta
"select(c0 > 1, t1, t2)"    # ok, because condition depends only coefficients and constants
"logistic((t0 - c0)*10)*t1 + (1 - logistic((t0 - c0)*10))*t2"  # smooth θ-gate

Weights must be finite and non-negative. The following will raise an exception: log/sqrt of <= 0, pow of a negative base, division by zero, or a negative result.

Examples:

# Coalescent-with-exposure-style: log-link rate with a per-edge offset
"exp(c0*t0 + c1*t1) + c2"

# Migration on/off per edge (c2 is a 0/1 edge flag), with a log-link base rate
"select(c2 == 1, exp(c0*t0), 0.0) + c1*t1"

# Piecewise: one rate law for coalescent edges (c0>0), another for mutation edges
"select(c0 > 0, c0*t0, c1*t1)"

Tip

The weight_formula argument can also be passed to graph.update_weights(). This is useful for testing formulas before running SVGD:

graph.update_weights(theta, weight_formula="c0*t0 + c1*t1")
graph.expectation()

If you need to debug your formula, you can create the a corresponding Python callback returning the formula expressed using the c and t variables:

def weight_update(theta, edge_coefficients):
    t0, t1 = theta
    c0, c1 = edge_coefficients
    return c0*t0 + c1*t1