Graph

phasic.Graph(arg, ipv=None, graph_cache=None, **kwargs)

Attributes

Name	Description
cache_trace	Whether this graph caches the elimination trace.
dyn_ordering	Whether dynamic minimum-degree elimination ordering is enabled.
hierarchical	Deprecated: use cache_trace instead.
trace_valid	Whether the cached trace is valid (not dirty).
weight_callback	Custom callback for computing edge weights from theta and coefficients.
weight_mode	Weight computation mode for parameterized edges.

Methods

Name	Description
absorbing_state_rewards	Get a reward vector that is 1 for states with edges to absorbing states.
accumulated_occupancy	Compute expected occupancy (visits or time) for each state.
accumulated_visiting_time	Compute expected time spent in each state (continuous only).
accumulated_visits	Compute expected number of visits to each state (discrete only).
add_epoch	Add an epoch boundary, returning a new graph with epoch transition edges.
as_matrices	Convert the graph to its matrix representation.
backward_probabilities	Compute P(reach target | start at v) for each vertex v.
cdf	Compute cumulative distribution function.
clear_from_cache	Delete this graph’s entries from the on-disk caches.
clone	Create a deep copy of this graph.
compute_trace	Compute elimination trace with optional hierarchical caching.
copy	Returns a deep copy of the graph.
covariance	Compute covariance matrix for multivariate phase-type distributions.
daisy_chain_joint_probs	JAX-traceable model: joint-probs at the t-states after a daisy chain.
discretize	Create a discretized copy of this graph.
distribution_context	Create a distribution context for efficient repeated sampling.
eliminate_to_dag	Perform symbolic graph elimination to create a reusable DAG structure.
epoch_context	Build a single-epoch context for reading cumulative t-state probabilities.
expectation	Compute expected value (first moment) of the phase-type distribution.
expected_sojourn_time	Compute expected sojourn time (residence time) in each state.
expected_waiting_time	Compute expected waiting time until absorption.
extend	Extend the graph by continuing to visit unvisited vertices using a callback.
find_or_create_vertex	Find or create a vertex with the given state.
from_matrices	Construct a Graph from matrix representation.
from_serialized	Reconstruct Graph from serialize() output.
joint_prob_graph	Build the joint-probability graph from this parameterized graph.
joint_stop_prob_graph	Build the joint stop-probability graph for daisy-chained inference.
laplace_transform	Create a Laplace-transformed graph.
mcmc	Run MCMC Metropolis-Hastings inference for Bayesian parameter estimation.
method_of_moments	Find parameter estimates by matching model moments to sample moments.
moments	Compute k-th moment of the phase-type distribution.
moments_from_graph	Convert a parameterized Graph to a JAX-compatible function that computes moments.
normalize	Normalize edge weights to make the graph a proper probability distribution.
pdf	Compute probability density/mass function using forward algorithm.
plot	Plot a graph using graphviz.
plot_scc_decomp	Visualise the SCC decomposition of this graph as a
pmf_and_moments_from_graph	Convert a parameterized Graph to a function that computes both PMF/PDF and moments.
pmf_and_moments_from_graph_multivariate	Create a multivariate phase-type model that handles 2D observations and rewards.
pmf_from_cpp	Load a phase-type model from a user’s C++ file and return a JAX-compatible function.
pmf_from_graph	Convert a Python-built Graph to a JAX-compatible function with full gradient support.
pmf_from_graph_joint_index	Create a JAX-compatible model for joint index distributions.
pmf_from_graph_parameterized	Convert a parameterized Python graph builder to a JAX-compatible function.
prewarm_cache	Populate the on-disk Stage A2 cache for this graph. If parameterized,
probability_matching	Find parameter estimates by matching model probabilities to empirical probabilities.
profile	Profile this graph and recommend parallel_elimination,
pull_cache	Download the published compute artifact for this graph if available.
push_cache	Publish this graph’s compute artifact to the phasic-traces registry.
reward_transform	Apply reward transformation to create a new graph with modified rewards.
reward_transform_discrete	Apply reward transformation for discrete-time distributions.
reward_visit_probability	Probability of visiting any rewarded vertex before absorption.
sample	Generate random samples from the phase-type distribution.
sample_path	Sample complete path(s) through the Markov chain.
sample_path_conditioned	Sample path(s) conditioned on reaching a target terminal state.
serialize	Serialize graph to array representation for efficient computation.
state_probability	Compute probability of being in each state at a given time.
stop_probability	Compute probability of being in each state at a given time.
svgd	Run Stein Variational Gradient Descent (SVGD) inference for Bayesian parameter estimation.
update_ipv	Set the initial probability vector after construction.
update_weights	Update parameterized edge weights with given parameters.
variance	Compute variance of the phase-type distribution.

absorbing_state_rewards

phasic.Graph.absorbing_state_rewards()

Get a reward vector that is 1 for states with edges to absorbing states.

This reward vector is used for computing the Laplace transform value from a Laplace-transformed graph.

Returns

: np.ndarray

Reward vector of length n_vertices. Element i is 1.0 if vertex i has an edge to an absorbing state, 0.0 otherwise.

Examples

>>> graph = Graph(coalescent_callback)
>>> rewards = graph.absorbing_state_rewards()
>>> laplace_graph = graph.laplace_transform(0.5)
>>> L_0_5 = laplace_graph.expectation(rewards=rewards)

See Also

laplace_transform : Create a Laplace-transformed graph

accumulated_occupancy

phasic.Graph.accumulated_occupancy(*args, **kwargs)

Compute expected occupancy (visits or time) for each state.

Parameters

*args : tuple = ()

Additional positional arguments passed to C++ implementation.

****kwargs** : dict = {}

Additional keyword arguments passed to C++ implementation.

Returns

: np.ndarray

Expected visits (discrete) or time (continuous) in each state.

Notes

This is a convenience method that dispatches to either: - accumulated_visits() for discrete distributions - accumulated_visiting_time() for continuous distributions

accumulated_visiting_time

phasic.Graph.accumulated_visiting_time(*args, **kwargs)

Compute expected time spent in each state (continuous only).

Parameters

*args : tuple = ()

Additional positional arguments passed to C++ implementation.

****kwargs** : dict = {}

Additional keyword arguments passed to C++ implementation.

Returns

: np.ndarray

Expected time spent in each state until absorption.

Notes

Only available for continuous-time phase-type distributions. For discrete distributions, use accumulated_visits() instead.

accumulated_visits

phasic.Graph.accumulated_visits(*args, **kwargs)

Compute expected number of visits to each state (discrete only).

Parameters

*args : tuple = ()

Additional positional arguments passed to C++ implementation.

****kwargs** : dict = {}

Additional keyword arguments passed to C++ implementation.

Returns

: np.ndarray

Expected number of visits to each state until absorption.

Raises

: ValueError

If graph is not discrete.

Notes

Only available for discrete-time phase-type (DPH) distributions. The graph must be discretized via discretize() before calling this method.

add_epoch

phasic.Graph.add_epoch(time, callback=None, **kwargs)

Add an epoch boundary, returning a new graph with epoch transition edges.

Computes transition rates from stop_probability(time) / accumulated_occupancy(time) and wires up sister vertices in the next epoch. The first call also adds an epoch index to the state vector.

Coefficient layout

After at least one add_epoch call, the per-edge coefficient vector follows a uniform interleaved layout. With base coefficient count B (i.e. the length of each coefficient list returned by the original callback), epoch k owns slots [k*(B+1) .. (k+1)*(B+1) - 1]: the first B are that epoch’s dynamics coefficients, and the last is the transition rate from epoch k-1 into epoch k. Epoch 0’s transition slot is a dummy — its coefficient is zero on every edge, so any value supplied at that index has no effect.

For example, with B = 1 and N add_epoch calls (i.e. N + 1 epochs) the coefficient vector accepted by update_weights is::

[r_0, t_0_dummy, r_1, t_{0->1}, r_2, t_{1->2}, ..., r_N, t_{(N-1)->N}]

which lets a list of per-epoch rates and times be passed directly via zip::

graph.update_weights([v for pair in zip(rates, times) for v in pair])

param_length() grows from B (no epochs) to 2*(B+1) after the first add_epoch and by B + 1 per subsequent call.

Parameters

time : float

Time at which the epoch boundary occurs.

callback : callable = None

Callback for building out new-epoch vertices. If None, uses the stored callback from graph construction.

****kwargs** : Any = {}

Additional keyword arguments merged with stored callback kwargs.

Returns

: Graph

New graph with epoch transitions wired up.

as_matrices

phasic.Graph.as_matrices()

Convert the graph to its matrix representation.

Returns a NamedTuple containing the traditional phase-type distribution matrices and associated information.

Returns

: MatrixRepresentation

NamedTuple with the following attributes: - states: np.ndarray of shape (n_states, state_dim), dtype=int32 State vector for each vertex - sim: np.ndarray of shape (n_states, n_states), dtype=float64 Sub-intensity matrix - ipv: np.ndarray of shape (n_states,), dtype=float64 Initial probability vector - indices: np.ndarray of shape (n_states,), dtype=int32 1-based indices for vertices (for use with vertex_at())

Examples

>>> g = Graph(1)
>>> start = g.starting_vertex()
>>> v1 = g.find_or_create_vertex([1])
>>> v2 = g.find_or_create_vertex([2])
>>> start.add_edge(v1, 1.0)
>>> v1.add_edge(v2, 2.0)
>>> g.normalize()
>>>
>>> matrices = g.as_matrices()
>>> print(matrices.sim)  # Sub-intensity matrix (attribute access)
>>> print(matrices.ipv)  # Initial probability vector
>>> # Can also use index access like a tuple
>>> states, sim, ipv, indices = matrices

backward_probabilities

phasic.Graph.backward_probabilities(target_vertices)

Compute P(reach target | start at v) for each vertex v.

For each vertex, computes the probability of eventually reaching one of the target terminal states. Uses backward induction over the graph structure.

Parameters

target_vertices : list of int

Indices of target terminal vertices.

Returns

: np.ndarray

Array of length vertices_length() with backward probability for each vertex. Values are between 0 and 1.

cdf

phasic.Graph.cdf(time, **kwargs)

Compute cumulative distribution function.

Parameters

time : float or ArrayLike

Time point(s) at which to evaluate the CDF.

****kwargs** : dict = {}

Additional keyword arguments passed to C++ implementation.

Returns

: float or np.ndarray

CDF value(s) P(T ≤ t) at the specified time point(s).

Notes

For continuous distributions: F(t) = P(T ≤ t) = 1 - α · exp(S·t) · 1 For discrete distributions: F(n) = P(N ≤ n) = sum of PMF up to n

clear_from_cache

phasic.Graph.clear_from_cache(
    graph_cache=True,
    parameterized_reward_compute=True,
)

Delete this graph’s entries from the on-disk caches.

Parameters

graph_cache : bool = True

If True, remove the serialised Graph entry from ~/.phasic_cache/graphs/<callback_hash>.json. Requires that this graph was built from a callback (so the callback + construction kwargs are available to recompute the same hash the constructor used). Manually constructed graphs have no callback-hash key and will raise.

parameterized_reward_compute : bool = True

If True, remove this graph’s Stage A2 symbolic elimination entry from ~/.phasic_cache/parameterized_reward_compute/<content_hash>.bin. Per-SCC entries (scc_<synth_hash>.bin) are not touched because they are content-addressed by SCC subgraph topology and may be shared with other graphs that have the same substructure.

Returns

: dict[str, int]

{'graph_cache': n, 'parameterized_reward_compute': m} where each value is the number of files actually removed (0 if the entry was not present, or if the flag was False).

Raises

: ValueError

If neither flag is True.

: RuntimeError

If graph_cache=True was requested but no callback is stored on this instance (e.g. graph built via Graph(state_length) then populated manually).

Notes

Missing files are not an error — they just mean the graph was never cached (or has already been cleared). The only conditions that raise are user errors (no flags, no callback key available).

For bulk operations across many graphs prefer :func:phasic.cache.clear_param_compute_cache and :func:phasic.clear_all_graph_caches.

Examples

>>> g = Graph(model, indexer=indexer, theta_dim=1)
>>> g.update_weights([1.0])
>>> _ = g.expectation()  # populates parameterized_reward_compute
>>> g.clear_from_cache(graph_cache=False) # keep graph cache
{'graph_cache': 0, 'parameterized_reward_compute': 1}

clone

phasic.Graph.clone()

Create a deep copy of this graph.

The cloned graph preserves the cache_trace setting but starts with a fresh (invalid) trace cache.

Returns

: Graph

A new Graph instance with the same structure.

compute_trace

phasic.Graph.compute_trace(
    param_length=None,
    hierarchical=True,
    min_size=50,
    parallel='auto',
    verbose=False,
    force=False,
)

Compute elimination trace with optional hierarchical caching.

.. deprecated:: The Python EliminationTrace machinery is no longer wired to the public moments() / expectation() / variance() entry points (those route directly to the C++ implementation, which uses the Stage A0-cached parameterized_reward_compute_graph). compute_trace is preserved for callers that still drive the trace pipeline directly, but emits a DeprecationWarning and will be removed in a future release.

When Graph was created with cache_trace=True, the trace is cached on the instance for use by moments(), expectation(), etc. In this mode, the operation is NON-DESTRUCTIVE (graph is preserved via cloning).

When Graph was created with cache_trace=False (default), the operation is DESTRUCTIVE and will empty the graph during trace recording.

Parameters

param_length : int = None

Number of parameters (auto-detect if None)

hierarchical : bool = True

If True, use hierarchical SCC-based caching (recommended). If False, use direct trace recording without caching. Caching provides 10-100x speedup on repeated calls.

min_size : int = 50

Minimum vertices to subdivide (only used if hierarchical=True)

parallel : str = 'auto'

Parallelization: ‘auto’, ‘vmap’, ‘pmap’, or ‘sequential’

verbose : bool = False

If True, show progress bars for major computation stages

force : bool = False

If True, recompute trace even if a cached trace exists and is valid. Only applicable when Graph was created with cache_trace=True.

Returns

: EliminationTrace

Elimination trace (from cache or computed)

Notes

Caching Mode (Graph created with cache_trace=True): - Trace is cached on the instance and reused by moments(), expectation(), etc. - Graph is cloned before trace recording, preserving the original - Subsequent calls return cached trace unless force=True or graph was modified

Non-Caching Mode (default): - The graph elimination algorithm is DESTRUCTIVE - vertices are eliminated - Use disk caching (hierarchical=True parameter) to avoid re-recording

Examples

>>> # Cache mode - non-destructive, cached on instance
>>> g = Graph(model, nr_samples=5, cache_trace=True)
>>> g.normalize()
>>> trace = g.compute_trace()  # Graph preserved, trace cached
>>> g.update_weights([1.0, 2.0])
>>> mean = g.expectation()  # Uses cached trace
>>>
>>> # Standard mode with disk caching
>>> g = Graph(model, nr_samples=5)
>>> trace = g.compute_trace()  # Graph emptied, trace returned

copy

phasic.Graph.copy()

Returns a deep copy of the graph.

Creates an independent copy with all vertices, edges, and metadata. Modifications to the copy will not affect the original graph.

Returns

: Graph

Deep copy of the graph

covariance

phasic.Graph.covariance(rewards1=[], rewards2=[], **kwargs)

Compute covariance matrix for multivariate phase-type distributions.

Parameters

rewards1 : list of float or ndarray = []

The first set of rewards, which should be applied to the phase-type distribution. Must have length equal to vertices_length().

rewards2 : list of float or ndarray = []

The second set of rewards, which should be applied to the phase-type distribution. Must have length equal to vertices_length().

****kwargs** : dict = {}

Additional keyword arguments passed to C++ implementation.

Returns

: np.ndarray

Covariance matrix for the multivariate distribution.

Notes

This method is for multivariate phase-type distributions with multiple marginals. For univariate distributions, use variance().

daisy_chain_joint_probs

phasic.Graph.daisy_chain_joint_probs(
    epoch_thetas,
    epoch_dts,
    initial_ipv,
    t_eval=None,
    fixed_indices=None,
    granularity=0,
)

JAX-traceable model: joint-probs at the t-states after a daisy chain.

Daisies through len(epoch_dts) epoch transitions (update_ipv → update_weights → stop_probability(dt)), then in the final epoch sets the propagated IPV and the final-epoch theta on the JSP graph and reads joint_stop_probabilities(t_eval) at the t-state vertex positions. The number of epochs is len(epoch_thetas), which must equal len(epoch_dts) + 1: each epoch i in 0..n_epochs-2 has a transition of duration epoch_dts[i]; the final epoch n_epochs-1 has no transition out.

The joint probability of an outcome is the relative probability of paths ending in (and being trapped in) that t-state. Because the JSP graph’s t-aux loops trap mass, stop_probability(t) at a t-vertex monotonically approaches the joint probability as t → ∞. t_eval should be large enough that the chain has had time to absorb most mass at the t-states. For slow-mutation models t_eval may need to be quite large; the default scales with sum(epoch_dts) to provide a conservative starting point. For an adaptive t_eval chosen via a residual-mass probe, use Graph.svgd(..., daisy_chain_t_eval='auto') or call Graph._probe_daisy_t_eval(...) directly.

Parameters

epoch_thetas : (array - like, shape(n_epochs, theta_dim))

Per-epoch parameter vectors. JAX-traced.

epoch_dts : (array - like, shape(n_epochs - 1))

Durations of the first n_epochs - 1 epochs. Static Python sequence (not JAX-traced).

initial_ipv : (array - like, shape(n_ipv))

IPV for epoch 0, in the JSP graph’s IPV layout.

t_eval : float = None

Time at which the final-epoch joint stop-probabilities are read. Defaults to max(sum(epoch_dts) * 4, 10.0).

fixed_indices : sequence of int = None

Flat-theta indices held fixed by SVGD. Forwarded to the custom_vjp backward pass to skip finite-difference gradient evaluations on those slots.

granularity : int = 0

Uniformization granularity passed to the underlying stop_probability call in the C++ FFI handler. 0 (default) auto-picks a safe value from the graph’s max rate and the requested time. Larger values give finer discretisation (slower, more accurate); smaller positive values trade accuracy for speed. The same granularity is used for every epoch’s stop_probability call.

Returns

: (jax.Array, shape(n_t_vertices))

Joint-probs at the t-states, in the order of self._t_vertex_indices.

discretize

phasic.Graph.discretize(rate, skip_existing=False, **kwargs)

Create a discretized copy of this graph.

Returns a new graph with auxiliary vertices added for each transient state. The original graph is not modified.

Parameters

rate : float or callable

Discretization rate. If callable, receives state array and **kwargs, must return a scalar rate (for non-parameterized graphs) or a coefficient vector (for parameterized graphs).

skip_existing : bool = False

If True, skip vertices that already have auxiliary vertices.

Returns

: Graph

New discretized graph with .rewards attribute containing the reward vector (1 for auxiliary vertices, 0 otherwise).

distribution_context

phasic.Graph.distribution_context(*args, **kwargs)

Create a distribution context for efficient repeated sampling.

Parameters

*args : tuple = ()

Additional positional arguments passed to C++ implementation.

****kwargs** : dict = {}

Additional keyword arguments passed to C++ implementation.

Returns

: object

Distribution context object that can be used for efficient sampling.

Notes

The distribution context precomputes data structures needed for sampling, making repeated sample() calls much faster than sampling directly from the graph.

eliminate_to_dag

phasic.Graph.eliminate_to_dag()

Perform symbolic graph elimination to create a reusable DAG structure.

This method performs the O(n³) graph elimination algorithm ONCE and returns a symbolic DAG where edges contain expression trees instead of concrete values. The DAG can then be instantiated with different parameters in O(n) time each.

This is the key optimization for SVGD and other inference methods that require evaluating the same graph structure with many different parameter vectors.

Returns

: SymbolicDAG

Symbolic DAG that can be instantiated with parameters

Raises

: RuntimeError

If the graph is not parameterized or elimination fails

Examples

>>> # Create parameterized graph
>>> g = Graph(1)
>>> v_a = g.create_vertex([0])
>>> v_b = g.create_vertex([1])
>>> v_c = g.create_vertex([2])
>>> v_a.add_edge_parameterized(v_b, 0.0, [1.0, 0.0, 0.0])
>>> v_b.add_edge_parameterized(v_c, 0.0, [0.0, 1.0, 0.0])

>>> # Eliminate to symbolic DAG (once)
>>> dag = g.eliminate_to_dag()
>>> print(dag)  # SymbolicDAG(vertices=3, params=3, acyclic=True)

>>> # Fast instantiation for SVGD (100-1000× faster!)
>>> for theta in particle_swarm:
...     g_concrete = dag.instantiate(theta)
...     log_prob = -g_concrete.expectation()  # Fast!

Performance

• Elimination: O(n³) - performed once
• Instantiation: O(n) - performed per particle
• Expected speedup for SVGD: 100-1000×

See Also

SymbolicDAG : The returned symbolic DAG class SymbolicDAG.instantiate : Create concrete graph from parameters

epoch_context

phasic.Graph.epoch_context()

Build a single-epoch context for reading cumulative t-state probabilities.

Allowed only on continuous joint-prob graphs (built via joint_prob_graph(..., discrete=False)). Builds the joint-stop-probability (JSP) graph internally, projects this graph’s starting-vertex IPV into the JSP graph’s IPV layout, inherits the current theta (whatever was last passed to this graph’s update_weights), and returns an :class:EpochContext ready for cumulative_probs calls.

If you have not called update_weights on this graph yet, you must call ctx.update_weights(theta) on the returned context before reading probabilities.

For the multi-epoch / daisy-chain reads used by SVGD, see :meth:daisy_chain_joint_probs.

Returns

: EpochContext

A context object whose cumulative_probs(t) returns P(absorbed via t-state by time t) for each t-state. Converges to :meth:joint_prob_table (built with the same settings and discrete=True) as t → ∞.

Raises

: ValueError

If this graph is not a joint-prob graph, or is the discrete variant (use :meth:joint_prob_table for asymptotes directly).

expectation

phasic.Graph.expectation(rewards=[], **kwargs)

Compute expected value (first moment) of the phase-type distribution.

Parameters

rewards : ArrayLike = []

Reward vector for reward-transformed expectation. If not provided, computes E[T] where T is time until absorption.

****kwargs** : dict = {}

Additional keyword arguments passed to C++ implementation.

Returns

: float

Expected value E[T] or reward-transformed expectation.

Notes

For parameterized graphs, this method uses trace-based computation which requires O(n) memory instead of O(n²) for the matrix-based approach. Set cache_trace=True when creating the graph to cache the trace for faster repeated evaluations.

This is equivalent to moments(1, rewards).

expected_sojourn_time

phasic.Graph.expected_sojourn_time(*args, **kwargs)

Compute expected sojourn time (residence time) in each state.

Parameters

*args : tuple = ()

Positional arguments passed to C++ implementation.

****kwargs** : dict = {}

Keyword arguments passed to C++ implementation.

Returns

: ArrayLike

Expected sojourn time for each vertex.

Notes

This method wraps the C++ implementation of expected_sojourn_time. Returns the expected accumulated reward for the starting vertex propagated through all paths. This is what expectation() uses internally. The difference from expected_residence_time is subtle - expected_residence_time decomposes this into per-vertex contributions.

Available for both continuous and discrete phase-type distributions.

expected_waiting_time

phasic.Graph.expected_waiting_time(*args, **kwargs)

Compute expected waiting time until absorption.

Parameters

*args : tuple = ()

Positional arguments passed to C++ implementation.

****kwargs** : dict = {}

Keyword arguments passed to C++ implementation.

Returns

: float

Expected waiting time until absorption.

Notes

This method wraps the C++ implementation of expected_waiting_time. Only available for continuous-time phase-type distributions.

extend

phasic.Graph.extend(callback=None, vertex_index=None, **kwargs)

Extend the graph by continuing to visit unvisited vertices using a callback.

After manually adding vertices to the graph (e.g., via find_or_create_vertex), this method continues the callback-based graph building process, visiting all newly added unvisited vertices.

Parameters

callback : callable = None

Callback function to use for extending the graph. If None, uses the callback from the original Graph construction. The callback should accept a state array and return: - For non-parameterized: list of (state, weight) tuples - For parameterized: list of (state, [coefficients]) tuples

****kwargs** : Any = {}

Additional keyword arguments to pass to the callback function. If None, uses the kwargs from the original Graph construction.

Raises

: RuntimeError

If no callback is provided and the graph was not constructed with a callback.

Examples

>>> # Build initial graph
>>> graph = Graph(my_callback, nr_samples=5)
>>>
>>> # Manually add new vertex
>>> special_vertex = graph.find_or_create_vertex([100, 200])
>>> graph.starting_vertex().add_edge(special_vertex, [1.5])
>>>
>>> # Continue with callback to explore new vertices
>>> graph.extend()  # Uses original callback
>>>
>>> # Or use different callback
>>> graph.extend(callback=my_other_callback, param=value)

find_or_create_vertex

phasic.Graph.find_or_create_vertex(state)

Find or create a vertex with the given state.

This method wraps the C++ implementation to track trace invalidation.

from_matrices

phasic.Graph.from_matrices(ipv, sim, states=None)

Construct a Graph from matrix representation.

Parameters

ipv : np.ndarray

Initial probability vector, shape (n_states,)

sim : np.ndarray

Sub-intensity matrix, shape (n_states, n_states)

states : np.ndarray = None

State vectors, shape (n_states, state_dim), dtype=int32 If None, uses default states [0], [1], [2], …

Returns

: Graph

The reconstructed phase-type distribution graph

Examples

>>> ipv = np.array([0.6, 0.4])
>>> sim = np.array([[-2.0, 1.0], [0.0, -3.0]])
>>> g = Graph.from_matrices(ipv, sim)
>>> pdf = g.pdf(1.0)

from_serialized

phasic.Graph.from_serialized(data)

Reconstruct Graph from serialize() output.

This method enables distributed trace recording by allowing graphs to be serialized to JSON, sent across the network via JAX pmap, and reconstructed on worker processes.

Parameters

data : dict[str, Any]

Dictionary returned by Graph.serialize() containing: - states: np.ndarray of shape (n_vertices, state_length) - edges: np.ndarray of shape (n_edges, 3) with [from_idx, to_idx, weight] - start_edges: np.ndarray of shape (n_start_edges, 2) with [to_idx, weight] - param_edges: np.ndarray of shape (n_param_edges, 2+param_length) Format (v0.22.0+): [from_idx, to_idx, coeff1, coeff2, …] - start_param_edges: np.ndarray of shape (n_start_param_edges, 1+param_length) Format (v0.22.0+): [to_idx, coeff1, coeff2, …] - param_length: int - state_length: int - n_vertices: int

Returns

: Graph

Reconstructed graph with identical structure to the original

Raises

: ValueError

If data is missing required fields, has wrong shapes, or malformed arrays

: RuntimeError

If graph reconstruction fails (e.g., edges to non-existent vertices)

Examples

>>> import json
>>> import numpy as np
>>> from phasic import Graph
>>>
>>> # Create and serialize graph
>>> g = Graph(1)
>>> v0 = g.starting_vertex()
>>> v1 = g.find_or_create_vertex([1])
>>> v0.add_edge_parameterized(v1, 0.0, [2.0])
>>> serialized = g.serialize(theta_dim=1)
>>>
>>> # Convert to JSON (for network transmission)
>>> json_dict = {k: v.tolist() if isinstance(v, np.ndarray) else v
...              for k, v in serialized.items()}
>>> json_str = json.dumps(json_dict)
>>>
>>> # Reconstruct on worker (e.g., different machine)
>>> received_dict = json.loads(json_str)
>>> received_dict['states'] = np.array(received_dict['states'], dtype=np.int32)
>>> received_dict['edges'] = np.array(received_dict['edges'], dtype=np.float64)
>>> received_dict['start_edges'] = np.array(received_dict['start_edges'], dtype=np.float64)
>>> received_dict['param_edges'] = np.array(received_dict['param_edges'], dtype=np.float64)
>>> received_dict['start_param_edges'] = np.array(received_dict['start_param_edges'], dtype=np.float64)
>>> g_reconstructed = Graph.from_serialized(received_dict)
>>> assert g_reconstructed.vertices_length() == g.vertices_length()

joint_prob_graph

phasic.Graph.joint_prob_graph(
    base_graph_indexer=None,
    reward_only=None,
    reward_rates_callback=None,
    mutation_rate=1.0,
    reward_limit=None,
    tot_reward_limit=np.inf,
    discrete=True,
)

Build the joint-probability graph from this parameterized graph.

Optimised reimplementation of :meth:_joint_prob_graph that produces a bit-identical graph (same vertices, same insertion order, same edges/weights, same serialization) while building much faster.

The speedup comes from hoisting the per-base-state-index → reward-index mapping out of the per-vertex construction loop: that mapping is a pure function of the indexers and was previously recomputed (via millions of index_to_props / props_to_index calls) once per joint vertex.

When a custom reward_rates_callback is supplied the fast path is bypassed and the callback is invoked exactly as in :meth:_joint_prob_graph, so behaviour is unchanged for custom callbacks.

joint_stop_prob_graph

phasic.Graph.joint_stop_prob_graph()

Build the joint stop-probability graph for daisy-chained inference.

Promotes the notebook helper (in docs/pages/tutorial/time_inhom_joint_prob.ipynb) into the library. The transformation: for each “t-vertex” in the source joint-prob graph (a vertex that has a transition to an absorbing vertex), wire a “trapping aux loop” — an aux vertex with state [0,...,0] and bidirectional unit-weight parameterised edges to/from the t-vertex. Mass that reaches a t-vertex shuttles between t-vertex and aux forever, so stop_probability(t)[t_vertex] + stop_probability(t)[aux] equals the cumulative joint absorption mass at that t-state by time t.

Initial-probability-vector edges (starting vertex outgoing) are added at construction time to every non-aux, non-trash, non- absorbing vertex with weight 0. The user calls update_ipv before each epoch to set them; the daisy-chain loop calls it between epochs to propagate surviving mass forward. The IPV vector layout matches joint_stop_probabilities output (vertex order, skipping aux vertices and the trash pair, expanded back to full vertex space — see _collapse_t_aux).

Returns

: Graph

New graph with attributes _joint_stop_prob_graph = True, _t_vertex_indices (sorted list of new-graph t-vertex indices), _t_aux_map (dict mapping new-graph t-vertex index → new-graph aux index), and the propagated _joint_prob_base_graph_indexer / _rewarded_props / _cache_trace from the source.

Raises

: ValueError

If self is not a joint-prob graph (no _joint_prob_base_graph_indexer) or has no parameterised edges (param_length() == 0).

laplace_transform

phasic.Graph.laplace_transform(theta)

Create a Laplace-transformed graph.

Returns a new graph where each transient state has an additional edge to the absorbing state with weight theta (or theta added to existing absorbing edge weight).

To compute the Laplace transform value L(theta) = E[exp(-theta * T)], call expectation() on the result with a reward vector that is 1 for states that had edges to absorbing in the original graph and 0 otherwise.

Parameters

theta : float

The Laplace transform parameter.

Returns

: Graph

A new graph representing the Laplace-transformed distribution.

Examples

>>> import numpy as np
>>> graph = Graph(coalescent_callback)
>>> # Get reward vector: 1 for states with absorbing edges, 0 otherwise
>>> rewards = graph.absorbing_state_rewards()
>>> laplace_graph = graph.laplace_transform(0.5)
>>> laplace_value = laplace_graph.expectation(rewards=rewards)

Notes

The Laplace transform of a phase-type distribution PH(alpha, S) is:

L(theta) = E[exp(-theta * T)] = alpha * (theta*I - S)^(-1) * s

where s is the exit rate vector indicating states with direct transitions to absorption.

For continuous distributions only. For discrete distributions, use the z-transform instead.

See Also

absorbing_state_rewards : Get the reward vector for Laplace transform computation

mcmc

phasic.Graph.mcmc(
    observed_data,
    discrete=None,
    prior=None,
    n_samples=10000,
    n_chains=4,
    burn_in=1000,
    thin=1,
    proposal_scale=None,
    theta_init=None,
    theta_dim=None,
    seed=None,
    verbose=False,
    progress=True,
    jit=None,
    positive_params=True,
    param_transform=None,
    rewards=None,
    fixed=None,
)

Run MCMC Metropolis-Hastings inference for Bayesian parameter estimation.

Samples from the posterior distribution p(theta | data) using random-walk Metropolis-Hastings with multiple independent chains.

Parameters

observed_data : array or SparseObservations

Observed data. A plain 1-D array for univariate models. For multivariate models use SparseObservations.

discrete : bool = None

If True, computes discrete PMF. If False, continuous PDF. If None, inferred from the graph.

prior : callable or list of Prior objects = None

Log prior function. If None, constructs DataPrior.

n_samples : int = 10000

Number of posterior samples per chain (after burn-in).

n_chains : int = 4

Number of independent chains.

burn_in : int = 1000

Number of initial samples to discard.

thin : int = 1

Keep every thin-th sample.

proposal_scale : float or array = None

Proposal standard deviation. If None, defaults to 0.1.

theta_init : array = None

Initial parameter values. Shape (n_chains, theta_dim) or (theta_dim,).

theta_dim : int = None

Dimension of parameter vector. Inferred from graph if None.

seed : int = None

Random seed.

verbose : bool = False

Print progress information.

progress : bool = True

Display progress bar.

jit : bool or None = None

JIT-compile log-probability function.

positive_params : bool = True

Constrain parameters to positive domain via softplus.

param_transform : callable = None

Custom parameter transformation.

rewards : array = None

Reward vector/matrix for multivariate distributions.

fixed : list or array = None

Fixed parameters: [(index, value), …] or binary mask.

Returns

: MCMC

MCMC object with results accessible via get_results() and summary().

Examples

>>> g = Graph(...)  # parameterized graph
>>> data = np.random.exponential(0.5, size=200)
>>> mcmc = g.mcmc(data, n_samples=5000, n_chains=4)
>>> mcmc.summary()
>>> print(mcmc.get_results()['theta_mean'])

method_of_moments

phasic.Graph.method_of_moments(
    observed_data,
    nr_moments=None,
    rewards=None,
    fixed=None,
    theta_dim=None,
    theta_init=None,
    std_multiplier=2.0,
    discrete=None,
    verbose=True,
    weighted=True,
)

Find parameter estimates by matching model moments to sample moments.

Solves the nonlinear least-squares problem::

minimize ||moments_fn(theta) - sample_moments||^2
subject to  theta > 0

The returned MoMResult.prior list can be passed directly to :meth:Graph.svgd as the prior argument, providing data-informed priors centred on the method-of-moments estimates.

Parameters

observed_data : array or SparseObservations

Observed data. A plain 1-D array is accepted for univariate models (must not contain NaN). For multivariate models use SparseObservations (see dense_to_sparse()), in which case rewards must also be provided.

nr_moments : int = None

Number of moments to match per feature dimension. If None (default), automatically chosen based on weighted: when weighted=True, adaptive selection prunes high-order moments by condition number; when weighted=False, uses the heuristic max(2 * n_free_params, 4). Still auto-increased if a user-specified value gives fewer equations than free parameters.

rewards : np.ndarray = None

Reward vectors. None for standard moments, 1-D for a single reward vector, 2-D (n_features, n_vertices) for multivariate.

fixed : list = None

List of (index, value) tuples pinning specific parameters.

theta_dim : int = None

Number of model parameters. Inferred from the graph when None.

theta_init : np.ndarray = None

Initial guess for the free parameters (excluding fixed ones). If None a coordinate-wise grid search is used.

std_multiplier : float = 2.0

Factor applied to the asymptotic standard error to obtain the prior standard deviation: prior_std = std_multiplier * se.

discrete : bool = None

True for discrete (PMF) models, False for continuous (PDF). If None, inferred from self.is_discrete.

verbose : bool = True

Print progress information.

weighted : bool = True

If True (default), use two-step efficient GMM with optimal weighting matrix W = Σ⁻¹. This down-weights high-variance moments and generally produces tighter standard errors. If False, use unweighted least squares (legacy behavior).

Returns

: MoMResult

Dataclass with theta, std, prior, success, etc.

Examples

>>> g = Graph(...)  # parameterized graph
>>> data = np.random.exponential(0.5, size=200)
>>> mom = g.method_of_moments(data)
>>> print(mom.theta)            # parameter estimate
>>> svgd = g.svgd(data, prior=mom.prior)  # use as SVGD prior

moments

phasic.Graph.moments(power, rewards=[], discrete=False, **kwargs)

Compute k-th moment of the phase-type distribution.

Parameters

power : int

Moment order (1 for first moment, 2 for second moment, etc.).

rewards : ArrayLike = []

Reward vector for reward-transformed moments. If not provided, uses unit rewards (standard moments).

discrete : bool = False

If True, compute discrete-time moments (DPH distribution). Requires that the graph was discretized via discretize().

****kwargs** : dict = {}

Additional keyword arguments passed to C++ implementation.

Returns

: float

The k-th moment E[T^k] where T is the time until absorption.

Notes

If the graph is parameterized, this method uses trace-based computation for 5-10x speedup on repeated evaluations and O(n) memory usage. Otherwise falls back to direct C++ graph elimination.

For higher moments (k > 2), numerical stability may become an issue for complex distributions.

moments_from_graph

phasic.Graph.moments_from_graph(graph, nr_moments=2, use_ffi=False)

Convert a parameterized Graph to a JAX-compatible function that computes moments.

This method creates a function that computes the first nr_moments moments of the phase-type distribution: [E[T], E[T^2], …, E[T^nr_moments]].

Moments are computed using the existing C++ graph.moments(power) method for efficiency.

Parameters

graph : Graph

Parameterized graph built using the Python API with parameterized edges. Must have edges created with add_edge_parameterized().

nr_moments : int = 2

Number of moments to compute. For example: - 1: Returns [E[T]] (mean only) - 2: Returns [E[T], E[T^2]] (mean and second moment) - 3: Returns [E[T], E[T^2], E[T^3]]

use_ffi : bool = False

If True, uses Foreign Function Interface approach.

Returns

: callable

JAX-compatible function with signature: moments_fn(theta) -> jnp.array(nr_moments,) Returns array of moments: [E[T], E[T^2], …, E[T^k]]

Examples

>>> # Create parameterized coalescent model
>>> def coalescent(state, nr_samples=2):
...     if len(state) == 0:
...         return [(np.array([nr_samples]), 1.0, [1.0])]
...     if state[0] > 1:
...         n = state[0]
...         rate = n * (n - 1) / 2
...         return [(np.array([n-1]), 0.0, [rate])]
...     return []
>>>
>>> graph = Graph(coalescent, nr_samples=3)
>>> moments_fn = Graph.moments_from_graph(graph, nr_moments=2)
>>>
>>> # Compute moments for given theta
>>> theta = jnp.array([0.5])
>>> moments = moments_fn(theta)  # [E[T], E[T^2]]
>>> print(f"Mean: {moments[0]}, Second moment: {moments[1]}")
>>>
>>> # Variance can be computed as: Var[T] = E[T^2] - E[T]^2
>>> variance = moments[1] - moments[0]**2

Notes

• Requires graph to have parameterized edges (created with parameterized=True)
• Moments are raw moments, not central moments
• For variance, compute: Var[T] = E[T^2] - E[T]^2
• For standard deviation: std[T] = sqrt(Var[T])

normalize

phasic.Graph.normalize(*args, **kwargs)

Normalize edge weights to make the graph a proper probability distribution.

Parameters

*args : tuple = ()

Additional positional arguments passed to C++ implementation.

****kwargs** : dict = {}

Additional keyword arguments passed to C++ implementation.

Returns

: float

Scaling factor applied to normalize the weights.

Notes

This method modifies the graph in-place and invalidates any cached trace.

For continuous distributions: Normalizes transition rates. For discrete distributions: Normalizes transition probabilities to sum to 1.

The normalization ensures that the initial probability vector and transition matrix satisfy the requirements for a valid phase-type distribution.

pdf

phasic.Graph.pdf(time, granularity=0, **kwargs)

Compute probability density/mass function using forward algorithm.

Parameters

time : float or ArrayLike

Time point(s) at which to evaluate the PDF/PMF.

granularity : int = 0

Granularity for uniformization (default: 0, auto-detected as 2*max_rate). Higher values improve accuracy but increase computation time.

****kwargs** : dict = {}

Additional keyword arguments passed to C++ implementation.

Returns

: float or np.ndarray

PDF/PMF value(s) at the specified time point(s).

Notes

This method uses the forward algorithm (Algorithm 4) via uniformization to compute the exact phase-type PDF/PMF, not an approximation.

For continuous distributions: f(t) = α · exp(S·t) · s* For discrete distributions: p(n) = probability of absorption at jump n

plot

phasic.Graph.plot(
    graph,
    filename=None,
    wrap=True,
    label_fmt=None,
    rate_fmt=None,
    subgraphfun=None,
    by_state=None,
    by_index=None,
    max_nodes=100,
    dark=None,
    constraint=True,
    ranksep=1,
    nodesep=1,
    rankdir='LR',
    size=(7, 7),
    fontsize=12,
    rainbow=True,
    penwidth=1,
    taillabel=False,
    seed=1,
    graph_attr={},
    node_attr={},
    edge_attr={},
)

Plot a graph using graphviz.

Parameters

graph : Graph

The phasic graph object to visualize.

filename : str | None = None

If provided, save the graph to this file. The file extension determines the output format (e.g., 'graph.pdf').

wrap : bool | int = True

Whether to wrap vertex labels, and if so, the maximum number of characters per line. By default True.

label_fmt : Callable[…, str] | None = None

Callable for format node labels:

rate_fmt : Callable[float] | None = None

Callable for format edge labels:

subgraphfun : Callable[…, str] | None = None

Deprecated. Use by_state instead. Callback function defining subgraph clusters by state.

by_state : Callable[…, str] | None = None

Callback function defining subgraph clusters. Takes a state as input and returns a string used as the subgraph label.

by_index : Callable[[int], str] | None = None

Callback function defining subgraph clusters. Takes a vertex index as input and returns a string used as the subgraph label.

max_nodes : int = 100

Maximum number of vertices to plot, by default 100.

dark : bool | None = None

Whether to use dark mode for the graph. Detected automatically from the VS Code theme if vscodenb is available.

constraint : bool = True

Graphviz constraint attribute, by default True.

ranksep : float = 1

Graphviz ranksep attribute, by default 1.

nodesep : float = 1

Graphviz nodesep attribute, by default 1.

rankdir : str = 'LR'

Graphviz rankdir attribute, by default "LR".

size : tuple[int, int] = (7, 7)

Graphviz size as (width, height), by default (7, 7).

fontsize : int = 12

Graphviz fontsize attribute, by default 12.

rainbow : bool = True

Whether to color edges with random colors, by default True.

penwidth : float = 1

Graphviz penwidth attribute, by default 1.

taillabel : bool = False

Use taillabel instead of xlabel, by default False

seed : int = 1

Random seed for graph layout, by default 1.

graph_attr : dict = {}

graphviz graph attributes to override defaults.

node_attr : dict = {}

graphviz node attributes to override defaults.

edge_attr : dict = {}

graphviz edge attributes to override defaults.

Returns

: graphviz.Digraph | None

Graphviz Digraph object for display in Jupyter notebooks, or None if the graph exceeds max_nodes.

plot_scc_decomp

phasic.Graph.plot_scc_decomp(
    figsize=None,
    cmap='viridis',
    show_indices=True,
    annotate_sizes=True,
    title=False,
    ax=None,
)

Visualise the SCC decomposition of this graph as a level-wise treemap.

Rows correspond to the levels of the SCC condensation. The source-side (start vertex, where the chain enters) is drawn at the top of the figure; the sink-side (absorbing state) is at the bottom — time flows downward.

Within a row, each tile is one SCC, with width proportional to the SCC’s vertex count, drawn at a common absolute scale shared across all rows. So a narrow row really has fewer total vertices than a wide one. SCCs at the same level are eliminated independently when parallel_elimination=True is enabled, so wide rows signal good parallelism potential and narrow rows are elimination bottlenecks.

Level labels on the left margin look like L7 (16) — the level number followed by the count of parallel SCCs at that level. Note that the C-side composer processes levels in the opposite of the plot’s vertical order (sink-first, bottom-up), but that detail does not affect interpretation: parallelism is per-row in either direction.

This is a structural visualisation only — it does not depend on any runtime telemetry. To assess actual cache hit/miss behaviour after a compose, use phasic.cache.scc_compose_stats().

Parameters

figsize : tuple of float = None

Matplotlib figure size in inches. Ignored if ax is provided.

cmap : str = 'viridis'

Matplotlib colormap name. Tiles are coloured by SCC index to make adjacent SCCs visually distinct.

show_indices : bool = True

Print the SCC index inside each tile when the tile is wide enough.

annotate_sizes : bool = True

Print the vertex count alongside the index.

title : bool = False

Whether to show a title. Default is False. If True adds a one-line summary of the decomposition (number of SCCs, levels, widest row).

ax : matplotlib.axes.Axes or None = None

Existing axes to draw into. If None, a new figure is created.

Returns

: matplotlib.axes.Axes

The axes the treemap was drawn into.

Examples

>>> import phasic
>>> g = phasic.Graph(my_callback)
>>> ax = g.plot_scc_decomp()
>>> ax.figure.savefig('scc.pdf')

See Also

scc_decomposition : underlying SCC structure phasic.distributed_scc.compute_scc_levels : level grouping used by this plot

pmf_and_moments_from_graph

phasic.Graph.pmf_and_moments_from_graph(
    graph,
    nr_moments=2,
    discrete=False,
    use_ffi=False,
    theta_dim=None,
)

Convert a parameterized Graph to a function that computes both PMF/PDF and moments.

This is more efficient than calling pmf_from_graph() and moments_from_graph() separately because it builds the graph once and computes both quantities.

Parameters

graph : Graph

Parameterized graph built using the Python API with parameterized edges.

nr_moments : int = 2

Number of moments to compute

discrete : bool = False

If True, computes discrete PMF. If False, computes continuous PDF.

use_ffi : bool = False

If True, uses Foreign Function Interface approach.

theta_dim : int = None

Number of parameters for parameterized edges. If not provided, will be auto-detected by probing edge states. Providing this explicitly avoids potential issues with auto-detection reading garbage memory.

Returns

: callable

JAX-compatible function with signature: model(theta, times, rewards=None) -> (pmf_values, moments) Where: - theta: Parameter vector - times: Time points or jump counts - rewards: Optional reward vector (one per vertex). If None, computes standard moments. - pmf_values: jnp.array(len(times),) - PMF/PDF values at each time - moments: jnp.array(nr_moments,) - [E[T], E[T^2], …] or [E[R·T], E[R·T^2], …]

Examples

>>> # Create parameterized model
>>> graph = Graph(coalescent, nr_samples=3)
>>> model = Graph.pmf_and_moments_from_graph(graph, nr_moments=2)
>>>
>>> # Compute both PMF and moments
>>> theta = jnp.array([0.5])
>>> times = jnp.array([1.0, 2.0, 3.0])
>>> pmf_vals, moments = model(theta, times)
>>>
>>> print(f"PMF at times: {pmf_vals}")
>>> print(f"Moments: {moments}")  # [E[T], E[T^2]]
>>>
>>> # Compute reward-transformed moments
>>> rewards = jnp.array([1.0, 2.0, 0.5, 1.5])  # One per vertex
>>> pmf_vals, reward_moments = model(theta, times, rewards=rewards)
>>> print(f"Reward moments: {reward_moments}")  # [E[R·T], E[R·T^2]]
>>>
>>> # Use in SVGD with moment regularization
>>> svgd = SVGD(model, observed_pmf, theta_dim=1)
>>> svgd.fit_regularized(observed_times=data, nr_moments=2, regularization=1.0)

Notes

• More efficient than separate calls to pmf_from_graph() and moments_from_graph()
• Required for using moment-based regularization in SVGD.fit_regularized()
• The moments are always computed from the same graph used for PMF/PDF
• Import phasic before importing jax / creating jax arrays. phasic enables jax_enable_x64 on import; the C FFI requires float64 buffers. jax arrays created before import phasic are float32 and trip the FFI check “Wrong buffer dtype: expected F64 but got F32”. If you hit this, restart the kernel and import phasic first (or recreate the arrays).

pmf_and_moments_from_graph_multivariate

phasic.Graph.pmf_and_moments_from_graph_multivariate(
    graph,
    nr_moments=2,
    discrete=False,
    use_ffi=False,
    theta_dim=None,
)

Create a multivariate phase-type model that handles 2D observations and rewards.

This wrapper enables computing joint likelihoods for multivariate phase-type distributions where each feature dimension has its own reward vector defining the marginal distribution.

Parameters

graph : Graph

Parameterized graph built using the Python API with parameterized edges.

nr_moments : int = 2

Number of moments to compute per feature dimension

discrete : bool = False

If True, computes discrete PMF. If False, computes continuous PDF.

use_ffi : bool = False

If True, uses Foreign Function Interface approach.

theta_dim : int = None

Number of parameters for parameterized edges.

Returns

: callable

JAX-compatible function with signature: model(theta, times, rewards=None) -> (pmf_values, moments) Where: - theta: Parameter vector (theta_dim,) - times: Time points - can be 1D (n_times,) or 2D (n_times, n_features) - rewards: Reward vectors - can be: * None: Standard moments * 1D (n_vertices,): Single reward vector (backward compatible) * 2D (n_vertices, n_features): Multivariate - one reward per feature - pmf_values: PMF/PDF values - shape matches input: * 1D rewards → 1D output (n_times,) * 2D rewards → 2D output (n_times, n_features) - moments: Moment values: * 1D rewards → 1D output (nr_moments,) * 2D rewards → 2D output (n_features, nr_moments)

Examples

>>> # Create parameterized model
>>> graph = Graph(coalescent, nr_samples=3)
>>> model = Graph.pmf_and_moments_from_graph_multivariate(graph, nr_moments=2)
>>>
>>> # 1D case (backward compatible)
>>> theta = jnp.array([0.5])
>>> times = jnp.array([1.0, 2.0, 3.0])
>>> rewards_1d = jnp.array([1.0, 2.0, 0.5, 1.5])  # (n_vertices,)
>>> pmf_vals, moments = model(theta, times, rewards_1d)
>>> print(pmf_vals.shape)  # (3,)
>>> print(moments.shape)   # (2,)
>>>
>>> # 2D case (multivariate)
>>> rewards_2d = jnp.array([[1.0, 2.0, 0.5, 1.5],   # Feature 1 reward vector
...                          [0.5, 1.0, 2.0, 0.8]])  # Feature 2 reward vector
...                                                  # (n_features, n_vertices)
>>> times_2d = jnp.array([[1.0, 1.5],
...                        [2.0, 2.5],
...                        [3.0, 3.5]])  # (n_times, n_features)
>>> pmf_vals, moments = model(theta, times_2d, rewards_2d)
>>> print(pmf_vals.shape)  # (3, 2)
>>> print(moments.shape)   # (2, 2)
>>>
>>> # Use in SVGD with 2D observations
>>> observed_data = jnp.array([[1.5, 2.1], [0.8, 1.2], [2.3, 3.1]])
>>> svgd = SVGD(model, observed_data, theta_dim=1, rewards=rewards_2d)
>>> results = svgd.optimize()

Notes

• For 2D rewards, each feature dimension is computed independently using the corresponding row of the rewards matrix (rewards[j, :] for feature j)
• Log-likelihood is computed as sum over all observation elements
• Backward compatible: 1D rewards behave exactly as pmf_and_moments_from_graph()

pmf_from_cpp

phasic.Graph.pmf_from_cpp(cpp_file, discrete=False)

Load a phase-type model from a user’s C++ file and return a JAX-compatible function.

The C++ file should include ‘user_model.h’ and implement:

phasic::Graph build_model(const double* theta, int n_params) { // Build and return Graph instance }

For efficient repeated evaluations with the same parameters without JAX overhead, use load_cpp_builder() instead to get a builder function that creates Graph objects.

Parameters

cpp_file : str or pathlib.Path

Path to the user’s C++ file

discrete : bool = False

If True, uses discrete phase-type distribution (DPH) computation. If False, uses continuous phase-type distribution (PDF).

Raises

: ImportError

If JAX is not installed. Install with: pip install jax jaxlib

: FileNotFoundError

If the specified C++ file does not exist

Returns

: callable

JAX-compatible function (theta, times) -> pmf_values that supports JIT, grad, vmap, etc.

Examples

JAX-compatible approach (default - for SVGD, gradients, optimization):

>>> model = Graph.pmf_from_cpp("my_model.cpp")
>>> theta = jnp.array([1.0, 2.0])
>>> times = jnp.linspace(0, 10, 100)
>>> pmf = model(theta, times)
>>> gradient = jax.grad(lambda p: jnp.sum(model(p, times)))(theta)

Discrete phase-type distribution:

>>> model = Graph.pmf_from_cpp("my_model.cpp", discrete=True)
>>> theta = jnp.array([1.0, 2.0])
>>> jumps = jnp.array([1, 2, 3, 4, 5])
>>> dph_pmf = model(theta, jumps)

For direct C++ access without JAX (faster for repeated evaluations):

>>> builder = load_cpp_builder("my_model.cpp")
>>> graph = builder(np.array([1.0, 2.0]))  # Build graph once
>>> pdf1 = graph.pdf(1.0)  # Use many times
>>> pdf2 = graph.pdf(2.0)  # No rebuild needed

pmf_from_graph

phasic.Graph.pmf_from_graph(
    graph,
    discrete=False,
    use_cache=True,
    theta_dim=None,
)

Convert a Python-built Graph to a JAX-compatible function with full gradient support.

This method automatically detects if the graph has parameterized edges (edges with state vectors) and generates optimized C++ code to enable full JAX transformations including gradients, vmap, and jit compilation.

For direct C++ access without JAX wrapping, use the Graph object’s methods directly: graph.pdf(time), graph.dph_pmf(jump), graph.moments(power), etc.

Raises

: ImportError

If JAX is not installed. Install with: pip install jax jaxlib

Parameters

graph : Graph

Graph built using the Python API. Can have regular edges or parameterized edges.

discrete : bool = False

If True, uses discrete phase-type distribution (DPH) computation. If False, uses continuous phase-type distribution (PDF).

use_cache : bool = True

If True, uses symbolic DAG cache to avoid re-computing expensive symbolic elimination for graphs with the same structure. Default: True Set to False to disable caching (useful for testing).

Returns

: callable

If graph has parameterized edges: JAX-compatible function (theta, times) -> pmf_values Supports JIT, grad, vmap, etc. If graph has no parameterized edges: JAX-compatible function (times) -> pmf_values Supports JIT (backward compatible signature)

Examples

Non-parameterized graph (regular edges only)

>>> g = Graph(1)
>>> start = g.starting_vertex()
>>> v0 = g.find_or_create_vertex([0])
>>> v1 = g.find_or_create_vertex([1])
>>> start.add_edge(v0, 1.0)
>>> v0.add_edge(v1, 2.0)  # fixed weight
>>>
>>> model = Graph.pmf_from_graph(g)
>>> times = jnp.linspace(0, 5, 50)
>>> pdf = model(times)  # No theta needed

Parameterized graph (with edge states for gradient support)

>>> g = Graph(1)
>>> start = g.starting_vertex()
>>> v0 = g.find_or_create_vertex([0])
>>> v1 = g.find_or_create_vertex([1])
>>> start.add_edge(v0, 1.0)
>>> v0.add_edge_parameterized(v1, 0.0, [2.0, 0.5])  # weight = 2.0*theta[0] + 0.5*theta[1]
>>>
>>> model = Graph.pmf_from_graph(g)
>>> theta = jnp.array([1.0, 3.0])
>>> pdf = model(theta, times)  # weight becomes 2.0*1.0 + 0.5*3.0 = 3.5
>>>
>>> # Full JAX support for parameterized graphs
>>> grad_fn = jax.grad(lambda t: jnp.sum(model(t, times)))
>>> gradient = grad_fn(theta)  # Gradients work!

For direct C++ access (no JAX overhead), use graph methods:

>>> pdf_value = g.pdf(1.5)  # Direct C++ call
>>> pmf_value = g.dph_pmf(3)  # Direct C++ call

With symbolic DAG caching (default)

>>> model = Graph.pmf_from_graph(g, use_cache=True)  # First call: computes and caches
>>> model2 = Graph.pmf_from_graph(g, use_cache=True)  # Subsequent: instant from cache!

pmf_from_graph_joint_index

phasic.Graph.pmf_from_graph_joint_index(
    graph,
    theta_dim=None,
    fixed_mask=None,
    exclude_vertices=None,
    observed_indices=None,
)

Create a JAX-compatible model for joint index distributions.

In joint index mode, likelihood is computed from exact expected sojourn times rather than PDF/PMF values. The observed_data should contain vertex indices (integers) instead of time values.

This is used for joint index distributions in population genetics models where the observed data represents which states (vertices) were visited.

Uses the fast expected_sojourn_time() method which computes exact sojourn times for all states in a single pass through the elimination trace.

Parameters

graph : Graph

Parameterized graph built using the Python API with parameterized edges.

theta_dim : int = None

Number of parameters for parameterized edges. If not provided, will be auto-detected from the graph.

fixed_mask : jnp.ndarray = None

Binary mask indicating which parameters are fixed (1=fixed, 0=learnable). If provided, gradients for fixed dimensions will be zero in the custom VJP. This is used to skip finite difference computation for fixed parameters.

exclude_vertices : list of int = None

Vertex indices to exclude from the normalization constant. Use this for ascertainment bias correction — e.g., exclude the zero-mutation terminal state when conditioning on observing at least one mutation. The excluded vertices are removed from the denominator, so the model returns P(s | s not in excluded set).

observed_indices : array-like of int = None

When supplied, enables baked observation dedup: the model will be built around the unique vertex indices in observed_indices (typically far fewer than the number of observations, with many repeats). The FFI sojourn-times call sees k = n_unique instead of k = n_obs (the FFI consumer’s inner loop scales linearly in k), and the returned per-observation array is reconstructed via a scatter through the inverse-index mapping. The model also wraps its forward in a custom_vmap rule so that under vmap(grad(loss))(particles) the per-particle dispatch fuses into a single batched FFI call. When observed_indices is None (default), the legacy path is used and vertex_indices is read from the model’s runtime argument as before.

Returns

: callable

JAX-compatible function with signature: model(theta, vertex_indices, rewards=None) -> (sojourn_times, dummy_moments) Where: - theta: Parameter vector - vertex_indices: Array of vertex indices (integers). Ignored when observed_indices was supplied at construction time; the baked indices are used instead. - rewards: Ignored (must be None for joint_index mode) - sojourn_times: Expected sojourn times for the specified vertices - dummy_moments: Zeros array (moments not supported in joint_index mode)

Notes

• Uses expected_sojourn_time() for fast exact computation
• Much faster than iterating accumulated_visiting_time() until convergence
• Moment regularization is not supported (regularization must be 0)
• Reward transformation is not supported (rewards must be None)
• When observed_indices is supplied, dedup typically gives a 10-100x wall-clock win for SVGD on joint-prob graphs with many repeated observations. See Graph.svgd joint_index path.

pmf_from_graph_parameterized

phasic.Graph.pmf_from_graph_parameterized(graph_builder, discrete=False)

Convert a parameterized Python graph builder to a JAX-compatible function.

This allows users to define parameterized models where the graph structure or edge weights depend on parameters.

Parameters

graph_builder : callable

Function (theta) -> Graph that builds a graph with given parameters

discrete : bool = False

If True, uses discrete phase-type distribution (DPH) computation. If False, uses continuous phase-type distribution (PDF).

Returns

: callable

JAX-compatible function (theta, times) -> pdf_values that supports JIT, grad, vmap, etc.

Examples

>>> def build_exponential(rate):
...     g = Graph(1)
...     start = g.starting_vertex()
...     v0 = g.find_or_create_vertex([0])
...     v1 = g.find_or_create_vertex([1])
...     start.add_edge(v0, 1.0)
...     v0.add_edge(v1, float(rate))
...     return g
>>>
>>> model = Graph.pmf_from_graph_parameterized(build_exponential)
>>> theta = jnp.array([1.5])
>>> times = jnp.linspace(0, 5, 50)
>>> pdf = model(theta, times)

prewarm_cache

phasic.Graph.prewarm_cache()

Populate the on-disk Stage A2 cache for this graph. If parameterized, update_weights must be called beforehand.

Triggers the C-side symbolic elimination (ptd_precompute_reward_compute_graph) and writes its result to ~/.phasic_cache/parameterized_reward_compute/<content_hash>.bin. Subsequent processes (notebook restart, SLURM workers, fresh CLI runs) that build the same graph will load the cached elimination instead of redoing the O(n^3) work.

If phasic.configure(parallel_elimination=True) is active and the graph is parameterised, the hierarchical SCC composer takes over: it writes one scc_<synth_hash>.bin entry per SCC large enough to cache (PHASIC_MIN_SCC_SIZE_TO_CACHE) and skips the monolithic parent <content_hash>.bin. Future runs that want to benefit from the cache must also have parallel_elimination=True active — the parent file and SCC files are populated by different code paths and not interchangeable. Run prewarm_cache once under each configuration you plan to consume from.

Notes

Cost is one full elimination plus one waiting-time read; the cache write is a side effect of the elimination.

Examples

Pre-warm on the head node before farming out to SLURM workers:

>>> g = Graph(model, indexer=indexer, theta_dim=2)
>>> g.prewarm_cache()
>>> # ~/.phasic_cache/parameterized_reward_compute/<hash>.bin now exists
>>> # Workers that rebuild the same graph will load it instantly

Pre-warm SCC cache too (cyclic graph):

>>> from phasic import configure
>>> with configure(parallel_elimination=True):
...     g.prewarm_cache(theta=[1.0, 0.5])
...     # scc_*.bin entries are now also on disk

probability_matching

phasic.Graph.probability_matching(
    observed_data,
    fixed=None,
    theta_dim=None,
    theta_init=None,
    std_multiplier=2.0,
    verbose=True,
)

Find parameter estimates by matching model probabilities to empirical probabilities.

For joint probability graphs (created via :meth:joint_prob_graph), observations are feature-count tuples and the model outputs a probability table. This method minimises the squared difference between model and empirical probabilities.

The returned ProbMatchResult.prior list can be passed directly to :meth:Graph.svgd as the prior argument.

Parameters

observed_data : list of tuples

Observed feature-count tuples, e.g. [(0, 1, 0), (1, 0, 0), ...]. Each tuple must match a row in the joint probability table.

fixed : list = None

List of (index, value) tuples pinning specific parameters.

theta_dim : int = None

Number of model parameters. Inferred from the graph when None.

theta_init : np.ndarray = None

Initial guess for the free parameters (excluding fixed ones). If None a coordinate-wise grid search is used.

std_multiplier : float = 2.0

Factor applied to the asymptotic standard error to obtain the prior standard deviation: prior_std = std_multiplier * se.

verbose : bool = True

Print progress information.

Returns

: ProbMatchResult

Dataclass with theta, std, prior, empirical_probs, model_probs, unique_indices, etc.

Raises

: ValueError

If the graph is not a joint probability graph.

Examples

>>> jg = graph.joint_prob_graph(indexer, tot_reward_limit=2)
>>> jg.update_weights([1.0, 0.5])
>>> obs = [tuple(row) for row in jg.joint_prob_table().iloc[:, :-1].values]
>>> result = jg.probability_matching(obs, fixed=[(1, 0.5)])
>>> svgd = jg.svgd(obs, prior=result.prior, fixed=[(1, 0.5)])

profile

phasic.Graph.profile(theta=None, probe_dyn='auto')

Profile this graph and recommend parallel_elimination, dyn_ordering, and the evaluation path (forward-PDF vs joint/sojourn).

Thin wrapper around :func:phasic.profile_graph; see that function for details. Returns a :class:~phasic.profile.GraphProfile.

print(graph.profile())

pull_cache

phasic.Graph.pull_cache(force=False)

Download the published compute artifact for this graph if available.

Looks up this graph’s content hash in the munch-group/phasic-traces registry. On a hit, downloads the parent .bin (and any per-SCC files) to ~/.phasic_cache/parameterized_reward_compute/ so the next call to :meth:expectation / :meth:pdf / :meth:moments reuses the published elimination instead of recomputing.

No git, gh, or daemon required for the consumer side — fetches are plain HTTPS to raw.githubusercontent.com.

Parameters

force : bool = False

If True, re-download even when a local cache file already exists.

Returns

: bool

True if a fresh download succeeded or a local cache file was already present. False if no registry entry matches this graph’s content hash.

Raises

: phasic.exceptions.PTDBackendError

On network failure or SHA-256 mismatch.

: phasic.exceptions.PTDFormatError

If the artifact’s format_revision exceeds what this phasic build supports.

Examples

>>> g = phasic.Graph(my_callback, ipv=[5])
>>> if g.pull_cache():
...     print('reusing published elimination')
>>> e = g.expectation()

push_cache

phasic.Graph.push_cache(
    id,
    description,
    domain=None,
    model_type=None,
    tags=None,
    license='MIT',
    author=None,
    registry_repo='munch-group/phasic-traces',
    dry_run=False,
    overwrite_branch=False,
)

Publish this graph’s compute artifact to the phasic-traces registry.

Populates the C-side elimination cache if necessary (calls :meth:expectation), saves the parent .bin and any per-SCC files, clones the registry repo to a temporary directory, splices in a new entry, pushes a feature branch, and opens a pull request via gh.

Parameters

id : str

Human-readable identifier for the registry entry (e.g. 'coal_n5_theta1'). Must be unique within the registry.

description : str

One-line free-form description of the model.

domain : str = None

Filtering keys stored in the entry’s metadata.

model_type : str = None

Filtering keys stored in the entry’s metadata.

tags : list[str] = None

Tags stored in the entry’s metadata.

license : str = 'MIT'

SPDX license identifier.

author : str = None

Override; default is "Name <email>" from git config user.{name,email}.

registry_repo : str = 'munch-group/phasic-traces'

GitHub owner/name of the registry repository.

dry_run : bool = False

If True, build the artifacts and return a JSON string of the would-be entry without cloning or pushing.

overwrite_branch : bool = False

If a stale phasic-publish/<id> branch already exists on the remote (e.g. from an earlier failed push) the call refuses unless this is True. With True, the push uses --force-with-lease: it overwrites the stale branch but still refuses if a third party has pushed concurrently.

Returns

: str

URL of the opened PR (or, if dry_run=True, a JSON string).

Raises

: phasic.exceptions.PTDBackendError

If gh is missing or unauthenticated, the entry id already exists in registry.json, the feature branch already exists on the remote (without overwrite_branch=True), or git/gh fails.

Notes

Running gh auth login once in a terminal sets up authentication for all future push_cache() calls. push_cache does not prompt inside the notebook.

Examples

>>> g = phasic.Graph(my_callback, ipv=[5])
>>> g.expectation()  # populate C-side cache (optional; push_cache does this)
>>> entry_json = g.push_cache(
...     id='my_model_v1',
...     description='Kingman coalescent for n=5',
...     domain='population-genetics',
...     model_type='coalescent',
...     dry_run=True,
... )

reward_transform

phasic.Graph.reward_transform(rewards, *, validate_rewards=True)

Apply reward transformation to create a new graph with modified rewards.

Parameters

rewards : np.ndarray

Reward vector of length n_vertices. Each element specifies the reward associated with visiting the corresponding vertex.

validate_rewards : bool = True

If True, run the coverage check and emit a UserWarning when some trajectories accumulate zero reward — i.e. the transformed distribution has a point mass at :math:r = 0. The transformation always proceeds regardless: an (atom + continuous) mixture is a legitimate phase-type distribution, useful for Laplace transforms, conditional expectations, and zero-inflated likelihood inference via Graph.svgd. Set False to silence the warning. Shape errors are always raised regardless of this flag.

Returns

: Graph

New graph with reward-transformed distribution.

Notes

Reward transformation is used to compute moments and expectations for different reward structures. The transformation modifies the graph to compute E[∑ rewards[i] * time_in_state_i] instead of E[T].

For continuous distributions: Uses continuous reward transformation. For discrete distributions: Uses discrete reward transformation.

The returned graph is a new Graph object (not modified in-place).

See Also

moments : Compute moments with optional reward vector expectation : Compute expectation with optional reward vector

reward_transform_discrete

phasic.Graph.reward_transform_discrete(rewards)

Apply reward transformation for discrete-time distributions.

Parameters

rewards : np.ndarray

Reward vector of length n_vertices.

Returns

: Graph

New graph with discrete reward transformation applied.

Notes

This method is specific to discrete-time phase-type (DPH) distributions. For automatic dispatch, use reward_transform() instead.

See Also

reward_transform : General reward transformation (dispatches to this for discrete graphs)

reward_visit_probability

phasic.Graph.reward_visit_probability(rewards, theta=None)

Probability of visiting any rewarded vertex before absorption.

For a 1D reward vector rewards and (parameterized) graph at parameter theta, returns the scalar probability that a trajectory absorbed from the starting vertex visits at least one vertex with reward > 0 before being absorbed. This is the phase-type quantity

p(theta) = sum_{v in starts} alpha_v * h_v(theta)

where h_v(theta) = P(reach a rewarded vertex | start at v) and alpha is the IPV. h is computed via Graph.backward_probabilities using the rewarded vertices as the target set; the result depends only on graph topology, IPV, and theta — NOT on observed reward times.

Parameters

rewards : (array - like, shape(n_vertices))

Per-vertex reward; only the nonzero pattern is used.

theta : array - like = None

Parameter vector. If None, uses the graph’s current parameter setting. If a JAX tracer is passed, the JAX-differentiable FFI path is used and the return is a jax.Array. Otherwise the result is a Python float.

Returns

: (float or jax.Array, shape())

Probability p(theta) in [0, 1].

Notes

When p < 1, some absorbing trajectories accumulate zero cumulative reward. The induced reward distribution is then a mixture (point mass at 0 + continuous part for r > 0); use Graph.svgd(..., rewards=...) to fit it with a zero-inflated likelihood automatically.

sample

phasic.Graph.sample(n, *, validate_rewards=True, **kwargs)

Generate random samples from the phase-type distribution.

Parameters

n : int

Number of samples to generate.

validate_rewards : bool = True

If True and rewards= is provided in kwargs, run the coverage check (BFS for a trajectory that skips every rewarded vertex). A coverage failure means the reward-transformed distribution has a point mass at :math:r = 0; Graph.sample will then return some zero-valued samples and emits a UserWarning so the user knows the mixture shape. Shape errors are always raised regardless of this flag — wrong-length rewards are bugs, not modelling choices.

****kwargs** : dict = {}

Additional keyword arguments passed to C++ implementation, notably rewards= (a per-vertex reward vector).

Returns

: np.ndarray

Array of n samples from the distribution.

Notes

Sampling is done by simulating the underlying Markov chain until absorption. For more efficient repeated sampling, first create a distribution context using distribution_context().

sample_path

phasic.Graph.sample_path(n=1)

Sample complete path(s) through the Markov chain.

Simulates the underlying Markov chain from the starting vertex until absorption, recording every vertex visited and the cumulative time at which each vertex was entered.

Parameters

n : int = 1

Number of paths to sample.

Returns

: dict or list of dict

If n=1, returns a single dict with keys: - ‘vertex_indices’: np.ndarray of vertex indices visited - ‘entry_times’: np.ndarray of cumulative entry times If n>1, returns a list of such dicts.

Notes

The first entry is always the starting vertex with entry_time=0. The last entry is the absorbing vertex. Sojourn times can be computed as the differences between consecutive entry times.

Examples

>>> g = Graph(...)  # build graph
>>> path = g.sample_path()
>>> path['vertex_indices']  # array of visited vertex indices
>>> path['entry_times']     # cumulative times
>>> sojourn_times = np.diff(path['entry_times'])

sample_path_conditioned

phasic.Graph.sample_path_conditioned(target_vertices, n=1)

Sample path(s) conditioned on reaching a target terminal state.

Uses guided forward sampling: at each step, the next state is chosen proportional to edge_weight * h(next_state) where h is the backward probability of reaching the target. This ensures every sampled path ends at one of the target vertices.

Parameters

target_vertices : list of int

Indices of target terminal vertices.

n : int = 1

Number of paths to sample.

Returns

: dict or list of dict

Same format as sample_path().

serialize

phasic.Graph.serialize(theta_dim=None)

Serialize graph to array representation for efficient computation.

Parameters

theta_dim : int = None

Number of parameters for parameterized edges. If not provided, will be auto-detected by probing edge states. Providing this explicitly avoids potential issues with auto-detection reading garbage memory.

Returns

: dict

Dictionary containing: - ‘states’: Array of vertex states (n_vertices, state_dim) - ‘edges’: Array of regular edges [from_idx, to_idx, weight] (n_edges, 3) - ‘start_edges’: Array of starting vertex regular edges [to_idx, weight] (n_start_edges, 2) - ‘param_edges’: Array of parameterized edges [from_idx, to_idx, x1, x2, …] (n_param_edges, theta_dim+2) - ‘start_param_edges’: Array of starting vertex parameterized edges [to_idx, x1, x2, …] (n_start_param_edges, theta_dim+1) NOTE: start_param_edges should be empty (starting edges are not parameterized) - ‘param_length’: Length of parameter vector (0 if no parameterized edges) - ‘state_length’: Integer state dimension - ‘n_vertices’: Number of vertices

state_probability

phasic.Graph.state_probability(time, **kwargs)

Compute probability of being in each state at a given time.

Parameters

time : float or int

Time point (continuous) or jump number (discrete) at which to evaluate state probabilities.

****kwargs** : dict = {}

Additional keyword arguments passed to C++ implementation.

Returns

: np.ndarray

Probability of being in each state at the specified time.

Notes

For continuous distributions: probability vector at time t For discrete distributions: probability vector after n jumps Computed via matrix exponentiation or uniformization.

stop_probability

phasic.Graph.stop_probability(time, **kwargs)

Compute probability of being in each state at a given time.

Parameters

time : float or int

Time point (continuous) or jump number (discrete) at which to evaluate state probabilities.

****kwargs** : dict = {}

Additional keyword arguments passed to C++ implementation.

Returns

: np.ndarray

Probability of being in each state at the specified time.

Notes

For continuous distributions: probability vector at time t For discrete distributions: probability vector after n jumps Computed via matrix exponentiation or uniformization.

svgd

phasic.Graph.svgd(
    observed_data,
    discrete=None,
    prior=None,
    n_particles=None,
    n_iterations=1000,
    optimizer=None,
    learning_rate=None,
    bandwidth='median_per_dim',
    theta_init=None,
    theta_dim=None,
    return_history=True,
    seed=None,
    verbose=False,
    progress=True,
    jit=None,
    parallel=None,
    n_devices=None,
    precompile=True,
    compilation_config=None,
    regularization=0.0,
    nr_moments=2,
    positive_params=True,
    param_transform=None,
    joint_index=False,
    rewards=None,
    fixed=None,
    tied=None,
    callback=None,
    preconditioner='auto',
    epoch_starts=None,
    daisy_chain_t_eval=None,
    daisy_chain_granularity=0,
    daisy_chain_probe_theta=None,
    daisy_chain_t_eval_tol=0.001,
    exposure=None,
    exposure_param_index=None,
    validate_rewards=True,
)

Run Stein Variational Gradient Descent (SVGD) inference for Bayesian parameter estimation.

SVGD finds the posterior distribution p(theta | data) by optimizing a set of particles to approximate the posterior. This method works with parameterized models created by pmf_from_graph() or pmf_from_cpp() where the model signature is model(theta, times).

Parameters

observed_data : array or SparseObservations

Observed data. A plain 1-D array is accepted for univariate models (must not contain NaN). For multivariate models use SparseObservations (see dense_to_sparse()), in which case rewards must also be provided.

discrete : bool = None

If True, computes discrete PMF. If False, computes continuous PDF. If undefined it is inferred from the graph.is_discrete attribute.

prior : callable or list of Prior objects = None

Log prior function for parameters. Can be: - Single callable: prior(theta) -> scalar, applied to entire theta vector - List of Prior objects: One prior per parameter dimension. Use None for fixed parameters: prior=[GaussPrior(ci=[0,1]), None, GaussPrior(ci=[0,1])] - DataPrior instance: Data-informed prior estimated from the observed data. If None (default), constructs DataPrior(graph, observed_data, sd=5) which uses method-of-moments (or probability matching for joint probability graphs) to estimate prior means from the data, with a wide spread (5x the asymptotic standard error). Falls back to standard normal if DataPrior construction fails. Note: this data-driven default differs from the lower-level SVGD class, whose prior=None is a plain standard normal (-0.5 * sum(theta**2)). Pass an explicit prior= to either API for identical behaviour. With fixed parameters: When using a list of priors with the fixed parameter, you must provide None at indices corresponding to fixed parameters. This is validated at initialization. Example: prior=[GaussPrior(ci=[0,1]), None, GaussPrior(ci=[0,1])], fixed=[(1, 0.5)] # theta[1] fixed, prior[1] must be None

n_particles : int or None = None

Number of SVGD particles. None resolves to 20 * theta_dim. More particles = better posterior approximation but slower.

n_iterations : int = 1000

Number of SVGD optimization steps

optimizer : object = None

Learning rate optimizer instance from phasic.optimizers. Default is Adam when learning_rate=None and regularization=0. Options include Adamelia, Adam, SGDMomentum, RMSprop, Adagrad. When an optimizer is used, the learning_rate parameter is ignored (the optimizer has its own learning rate).

learning_rate : float or None = None

SVGD step size. If None (default), uses Adame optimizer with adaptive learning rates. If a float is provided, uses fixed learning rate approach. Larger values = faster convergence but may be unstable.

bandwidth : str, float, or np.ndarray = 'median_per_dim'

Kernel bandwidth selection method: - ‘median_per_dim’: Per-dimension median heuristic (default). Uses a separate bandwidth per parameter dimension for an anisotropic kernel. - ‘median’: Scalar median heuristic (isotropic kernel) - float: Fixed scalar bandwidth value - np.ndarray: Fixed per-dimension bandwidth vector

theta_init : ArrayLike = None

Initial particle positions (n_particles, theta_dim). If None, initializes randomly from standard normal.

theta_dim : int = None

Dimension of theta parameter vector. Can be: - Set at graph construction: Graph(callback, theta_dim=2) - Overridden here for SVGD inference (if graph was modified/augmented) If None, inferred from the graph’s parameterized edge structure via param_length(). Only required if theta_init is None and the graph has no parameterized edges. The value specified here overrides any theta_dim set during graph construction, which is useful if you’ve modified the graph structure (e.g., via extend()).

return_history : bool = True

If True, return particle positions throughout optimization

seed : int = None

Random seed for reproducibility

verbose : bool = False

Print progress information

jit : bool or None = None

Enable JIT compilation. If None, uses value from phasic.get_config().jit. JIT compilation provides significant speedup but adds initial compilation overhead.

parallel : str or None = None

Parallelization strategy: - ‘vmap’: Vectorize across particles (single device) - ‘pmap’: Parallelize across devices (uses multiple CPUs/GPUs) - ‘none’: No parallelization (sequential, useful for debugging) - None: Auto-select (pmap if multiple devices, vmap otherwise)

n_devices : int or None = None

Number of devices to use for pmap. Only used when parallel=‘pmap’. If None, uses all available devices.

precompile : bool = True

(Deprecated: use jit parameter instead) Precompile model and gradient functions for faster execution. First run will take longer but subsequent iterations will be much faster.

compilation_config : CompilationConfig, dict, str, or pathlib.Path = None

JAX compilation optimization configuration. Can be: - CompilationConfig object from phasic.CompilationConfig - dict with CompilationConfig parameters - str/Path to JSON config file - None (uses default balanced configuration)

positive_params : bool = True

If True, applies softplus transformation to ensure all parameters are positive. Recommended for phase-type models where parameters represent rates. SVGD operates in unconstrained space, but model receives positive parameters.

param_transform : callable = None

Custom parameter transformation function: transform(theta_unconstrained) -> theta_constrained. If provided, SVGD optimizes in unconstrained space and applies this transformation before calling the model. Cannot be used together with positive_params. Example: lambda theta: jnp.concatenate([jnp.exp(theta[:1]), jax.nn.softplus(theta[1:])])

joint_index : bool = False

If True, use joint index mode where observed_data contains vertex indices (integers) instead of time values. In this mode, likelihood is computed from converged accumulated_visits() values rather than PDF/PMF values. This is used for joint index distributions in population genetics models. When joint_index=True: - observed_data should contain vertex indices (integers) - Forces discrete=True behavior - Moment regularization is not supported (raises NotImplementedError if regularization > 0)

rewards : ArrayLike = None

Reward vectors for computing reward-transformed likelihoods. Can be: - None: Standard phase-type likelihood (default) - 1D array (n_vertices,): Single reward vector for univariate models - 2D array (n_vertices, n_features): Multivariate rewards - one reward vector per feature dimension. Requires use of pmf_and_moments_from_graph_multivariate() model. For multivariate models, observed_data should also be 2D (n_times, n_features).

fixed : list of (index, value) tuples or 1D array = None

Pin selected parameters at known constants so SVGD only optimises the learnable dimensions. Equivalent to a point-mass prior at value on the pinned slots, with those positions removed from the kernel and the gradient computation. Two accepted forms: (a) Index/value tuples (recommended) — [(idx, value), ...]. Each tuple pins parameter idx at value:: fixed=[(1, 0.01)] # theta[1] = 0.01, rest learned fixed=[(0, 2.5), (2, 0.1)] # theta[0]=2.5, theta[2]=0.1 (b) Binary mask (legacy) — a 1D array of length theta_dim where 1 pins the slot at the value 1.0 and 0 leaves it learnable. Use form (a) whenever the fix value is not 1.0:: fixed=[0, 1] # theta[1] pinned at 1.0 When to use. Fix a parameter when its value is known from prior data, when you want to test sensitivity to a single dimension while holding the rest at MLE, or when a slot is structurally unidentifiable from the observed data alone (e.g. a global rate scale that the model absorbs into another parameter). Interaction with prior. If prior is a per-slot list, the entries at fixed indices must be None; mismatches raise at construction time. A scalar prior callable is fine and is auto-masked at the fixed slots. Daisy-chain semantics (epoch_starts=[...]). The flattened theta has shape (n_epochs * param_length,), but fixed entries still use the local per-epoch index ([0, param_length)) and are broadcast across all epochs. To pin a parameter at different values per epoch, pass a list/array of length n_epochs as the value:: fixed=[(1, 1.0)] # local_idx=1 pinned at 1.0 in EVERY epoch fixed=[(1, [1.0, 2.5])] # local_idx=1: 1.0 in epoch 0, 2.5 in epoch 1 fixed=[(0, 5.0), (1, [2.0, 8.0])] # mix scalar and per-epoch values Combination with tied. Compatible, with one rule: a given (local_idx, epoch) slot may be either fixed or a member of a tied group, not both. Overlaps raise at construction time (rule R20).

tied : list of (local_idx, [epoch_a, epoch_b, …]) tuples = None

Daisy-chain only. Tie a parameter slot across two or more epochs so SVGD treats them as a single learnable value. Within each entry the first epoch is the master (the slot SVGD actually optimises); every subsequent epoch is a slave whose value is replaced with the master’s on every forward evaluation, and whose gradient is routed back into the master. Examples:: tied=[(0, [0, 1])] # local_idx 0 shared across epochs 0 and 1 tied=[(1, [0, 2, 3])] # local_idx 1 shared across epochs 0, 2, 3 tied=[(0, [0, 1]), (1, [1, 2])] # two independent ties When to use. Use tied when a population parameter (e.g. mutation rate per base, baseline hazard) is biologically constant across a subset of epochs while other parameters change. Tying reduces dimensionality, improves identifiability, and makes posteriors tighter without silently fusing epochs whose other parameters should differ. Requirements. - Requires epoch_starts=... (rule R16); a tie within a single epoch is meaningless and rejected with rule R17. - local_idx is the per-epoch index in [0, param_length) (rule R18); the epoch indices are in [0, n_epochs). - Each epoch index may appear at most once per tied entry (rule R19). - Each (local_idx, epoch) slot may belong to at most one tied group, and may not also be fixed (rule R20). Combination with fixed. Compatible as long as no slot is claimed by both — see the rule R20 note in fixed above. To pin a parameter at the same value in every epoch use fixed=[(local_idx, value)] (broadcast). To pin a parameter at different values per epoch use the per-epoch list form of fixed. Use tied only when the shared value should be learned rather than known a priori. Combination with exposure. Compatible. The exposed parameter (exposure_param_index) may itself be tied across epochs; the per-observation exposure scaling multiplies the tied (shared) value in every epoch where the slot appears as master or slave. Not available on the direct SVGD(model=...) path — tying is a pre-model-construction concern owned by Graph.svgd.

preconditioner : str, preconditioner instance, or None = 'auto'

Preconditioning method for multi-scale parameters: - ‘auto’ or ‘jacobian’: Moment Jacobian preconditioning (default, recommended). Uses column norms of the moment Jacobian matrix for scaling. Simpler and more robust than Fisher preconditioning. - ‘fisher’: Fisher diagonal preconditioning. Uses empirical Fisher information matrix diagonal. Can be unstable when PMF values are small. - None or ‘none’: No preconditioning (original behavior) - MomentJacobianPreconditioner or FisherPreconditioner instance: Custom preconditioner

epoch_starts : array-like of float = None

Enables daisy-chain (time-inhomogeneous) inference. epoch_starts[0] == 0; subsequent entries are the start times of additional epochs. n_epochs = len(epoch_starts). Each epoch fits its own param_length parameters, so the flattened theta has length n_epochs * param_length. Requires a continuous-time joint-prob graph (discrete=False). See fixed and tied for how those kwargs interpret indices under daisy-chain.

daisy_chain_t_eval : float, str, or None = None

Time at which the final-epoch joint stop-probabilities are read off the JSP graph’s t-vertices. Only used when epoch_starts is set. - Numeric: used as-is. - None (default): falls back to max(sum(epoch_dts) * 4, 10.0). - 'auto': an adaptive probe (Graph._probe_daisy_t_eval) walks the daisy chain at daisy_chain_probe_theta and grows t_eval until the residual non-t-vertex transient mass falls below daisy_chain_t_eval_tol. Typically picks a much smaller value than the conservative default and so cuts SVGD wall time significantly.

daisy_chain_granularity : int = 0

Uniformization granularity passed to stop_probability inside the daisy-chain FFI handler. 0 = auto (the underlying C++ picks a safe value from the graph’s max rate × time). Larger values give finer discretisation; smaller positive values trade accuracy for speed.

daisy_chain_probe_theta : array - like = None

Theta used by the daisy_chain_t_eval='auto' probe. Shape (param_length,) (broadcast across all epochs) or (n_epochs, param_length) (per-epoch). Defaults to ones.

daisy_chain_t_eval_tol : float = 1e-3

Residual-mass tolerance used by the daisy_chain_t_eval='auto' probe.

exposure : float, array-like, or None = None

Per-observation exposure :math:\alpha_i — a known multiplicative scaling on a rate-typed component of :math:\boldsymbol{\theta}. For observation i the model is evaluated at :math:\boldsymbol{\theta}^{(i)} where :math:\theta^{(i)}_j = \theta_j for :math:j \neq k and :math:\theta^{(i)}_k = \theta_k \cdot \alpha_i, with :math:k = exposure_param_index. The exposed rate parameter and :math:\alpha_i jointly determine each observation’s expected event count (or hazard, or PMF, depending on the model). This is the GLM “exposure” / “offset” construct: it linearises the relationship between a rate parameter and an observation-specific outcome that scales with a known quantity. Concrete instances: - Coalescent-with-mutation: :math:\alpha_i = segment length :math:L_i in bases; :math:\theta_k is the per-base mutation rate. - Survival / failure-time: :math:\alpha_i = time-at-risk for unit :math:i; :math:\theta_k is the hazard rate. - Spatial Poisson: :math:\alpha_i = area or volume of region :math:i; :math:\theta_k is the intensity per unit area. Forms: - None (default): no exposure correction; existing behaviour. - scalar: same :math:\alpha applied to every observation. - 1D array of length n_observations: per-observation :math:\alpha. For dense 2D observed_data of shape (n_observations, n_features) the same :math:\alpha_i is shared across all features of observation :math:i. Requires exposure_param_index to be set. Not supported for SparseObservations (raises NotImplementedError).

exposure_param_index : int or None = None

Index :math:k of the rate-typed parameter in :math:\boldsymbol{\theta} that exposure scales. Required when exposure is set. Must be in [0, param_length). Under daisy-chain (epoch_starts=[…]), :math:\boldsymbol{\theta} has flat layout (n_epochs * param_length,); exposure_param_index remains the local per-epoch index and is broadcast across every epoch internally.

Returns

: dict

Inference results containing: - ‘particles’: Final posterior samples (n_particles, theta_dim) - ‘theta_mean’: Posterior mean estimate - ‘theta_std’: Posterior standard deviation - ‘history’: Particle evolution over iterations (if return_history=True)

Raises

: ImportError

If JAX is not installed

: ValueError

If model is not parameterized or theta_dim cannot be inferred

Examples

>>> import jax.numpy as jnp
>>> from phasic import Graph
>>>
>>> # Build parameterized coalescent model
>>> def coalescent_callback(state, nr_samples=3):
...     if len(state) == 0:
...         return [(np.array([nr_samples]), 1.0, [1.0])]
...     if state[0] > 1:
...         n = state[0]
...         rate = n * (n - 1) / 2
...         return [(np.array([n - 1]), 0.0, [rate])]
...     return []
>>>
>>> g = Graph.from_callback_parameterized(coalescent_callback, nr_samples=4)
>>> model = Graph.pmf_from_graph(g, discrete=False)
>>>
>>> # Generate synthetic observed data
>>> true_theta = jnp.array([2.0])
>>> times = jnp.linspace(0.1, 3.0, 15)
>>> observed_pdf = model(true_theta, times)
>>>
>>> # Run SVGD inference
>>> results = Graph.svgd(
...     observed_data=observed_pdf,
...     theta_dim=1,
...     n_particles=30,
...     n_iterations=500,
...     learning_rate=0.01
... )
>>>
>>> print(f"True theta: {true_theta}")
>>> print(f"Posterior mean: {results['theta_mean']}")
>>> print(f"Posterior std: {results['theta_std']}")

Notes

• SVGD requires a parameterized model. Non-parameterized models (signature: model(times)) cannot be used for inference as there are no parameters to estimate.
• The likelihood is computed as sum(log(model(theta, observed_data)))
• For better results, ensure observed_data has sufficient information about the parameters
• Learning rate and number of iterations may need tuning for different problems

update_ipv

phasic.Graph.update_ipv(ipv)

Set the initial probability vector after construction.

Parameters

ipv : (array - like, shape(n_ipv_edges))

New weights for starting-vertex edges, in construction order. Length must equal the number of starting-vertex edges in the graph.

Notes

IPV is a property of the model, not an inference parameter. Use this method to:

• set the IPV after constructing the graph if your callback does not bake one in;
• re-run the same graph against a different initial distribution without rebuilding it;
• propagate epoch state in a daisy chain (handled internally by the daisy-chain machinery between epochs).

The symbolic compute graph cache (Stage A0) survives this call, so subsequent forward computations (expectation, pdf, compute_pmf, …) reuse the cached elimination.

For users on the Python EliminationTrace path (cache_trace=True), IPV remains baked in at trace-record time — that path is not affected by update_ipv.

update_weights

phasic.Graph.update_weights(theta, callback=None, log=False)

Update parameterized edge weights with given parameters.

This method wraps the C++ implementation to cache theta for use with trace-based computation.

Parameters

theta : ArrayLike

Parameter vector to set edge weights.

callback : callable = None

Custom callback for weight computation.

log : bool = False

If True, use log-space computation.

Notes

This updates edge weights on the live Graph object (C++ side). It does NOT invalidate the cached trace since it only changes parameter values, not graph structure.

For the JAX/FFI/SVGD pipeline, use graph.weight_mode instead. That property controls how pmf_from_graph() and pmf_and_moments_from_graph() compute edge weights from theta.

variance

phasic.Graph.variance(rewards=[], **kwargs)

Compute variance of the phase-type distribution.

Parameters

rewards : ArrayLike = []

Reward vector for reward-transformed variance. If not provided, computes Var(T) where T is time until absorption.

****kwargs** : dict = {}

Additional keyword arguments passed to C++ implementation.

Returns

: float

Variance Var(T) or reward-transformed variance.

Notes

For parameterized graphs, this method uses trace-based computation which requires O(n) memory instead of O(n²) for the matrix-based approach. Set cache_trace=True when creating the graph to cache the trace for faster repeated evaluations.

Computed as Var(T) = E[T^2] - E[T]^2 using moments.