28 Graph Builder

Introduction

This file formalizes the GraphBuilder class, which implements the separation of graph structure (topology and edge coefficients) from parameters (the vector \boldsymbol{\theta}) for efficient repeated evaluation. In inference workflows such as SVGD ([18]) and MCMC ([19]), the same graph structure is evaluated with hundreds or thousands of different parameter vectors. Parsing the JSON serialization and constructing vertices and edges from scratch for each \boldsymbol{\theta} is wasteful when the topology is fixed.

The GraphBuilder addresses this by parsing the JSON representation once, caching the vertex states, edge connectivity, and coefficient vectors in C++ data structures, and then providing fast build(theta), compute_pmf(theta, times), and compute_moments(theta, nr_moments) methods that construct a concrete graph and evaluate it in a single call. This pattern—serialize once, evaluate many times—provides the performance backbone for the JAX FFI interface formalized in [25].

Prerequisites: [01], [11], [14], [15]

Source files:

src/cpp/parameterized/graph_builder.hpp (class declaration)
src/cpp/parameterized/graph_builder.cpp (implementation)

Definitions

Recall (Definition 2.15). A parameterized edge (v \to z) carries a coefficient vector \mathbf{c} = (c_1, \ldots, c_d) and computes its weight as w(v \to z; \boldsymbol{\theta}) = \sum_j c_j \theta_j (linear mode) or w = \prod_j (c_j \theta_j) (log mode).

Recall (Definition 16.1). Uniformization of a CPH graph with granularity g produces a discrete-time chain enabling PDF computation via the forward algorithm.

Recall (Definition 17.1). The k-th moment \mathbb{E}[\tau^k] is computed via k successive applications of the expected waiting time computation on the graph.

Definition 28.1 (Graph Structure Serialization) Let G = (V, E, W) be a parameterized phase-type graph (Definition 2.15) with p transient vertices, parameter length d, and state length \ell. The serialization of G is a JSON object J with the following fields:

states: a p \times \ell array of integer state vectors, \texttt{states}[i] = \mathbf{s}_{v_i} for i = 0, \ldots, p-1.
edges: a list of triples (i, j, w) for each non-parameterized edge (v_i \to v_j) with fixed weight w.
start_edges: a list of pairs (j, w) for each edge from the starting vertex v_0 with fixed weight w.
param_edges: a list of tuples (i, j, c_1, \ldots, c_d) for each parameterized edge (v_i \to v_j) with coefficient vector (c_1, \ldots, c_d).
start_param_edges: a list of tuples (j, c_1, \ldots, c_d) for each parameterized edge from v_0.
param_length: d, the dimension of \boldsymbol{\theta}.
state_length: \ell, the dimension of state vectors.
n_vertices: p, the number of transient vertices.
weight_mode: "linear" or "log", the weight computation mode.

The serialization separates structural information (topology, coefficients) from parameter values (\boldsymbol{\theta}), enabling the structure to be fixed while \boldsymbol{\theta} varies.

Intuition

The serialization is a snapshot of the graph’s “skeleton”: it records which vertices exist, how they are connected, and what coefficients each parameterized edge carries, but does not commit to specific parameter values. This is analogous to a compiled computation graph in machine learning frameworks, where the graph structure is fixed at compile time and only the input data varies at runtime.

Example

A two-state coalescent model with parameter \theta (inverse population size) serializes as:

{
  "states": [[2], [1]],
  "edges": [],
  "start_edges": [[0, 1.0]],
  "param_edges": [[0, 1, 1.0]],
  "start_param_edges": [],
  "param_length": 1,
  "state_length": 1,
  "n_vertices": 2,
  "weight_mode": "linear"
}

The single parameterized edge from vertex 0 to vertex 1 has coefficient c_1 = 1.0, so w = 1.0 \cdot \theta_1 = \theta_1.

Definition 28.2 (GraphBuilder) A GraphBuilder is a C++ object \mathcal{B} initialized from a serialized graph structure J (Definition 28.1). It caches the parsed representation and provides the operation:

\mathcal{B}.\operatorname{build}(\boldsymbol{\theta}) = G(\boldsymbol{\theta}), \tag{28.1}

which returns a concrete (non-parameterized) graph G(\boldsymbol{\theta}) = (V, E, W_{\boldsymbol{\theta}}) where each edge weight is evaluated at \boldsymbol{\theta}:

W_{\boldsymbol{\theta}}(v_i \to v_j) = \begin{cases} w & \text{if } (v_i \to v_j) \text{ is a fixed edge with weight } w, \\ \sum_{k=1}^{d} c_k \theta_k & \text{if weight mode is linear}, \\ \prod_{k=1}^{d} (c_k \theta_k) & \text{if weight mode is log}. \end{cases} \tag{28.2}

The GraphBuilder additionally provides composite operations \mathcal{B}.\operatorname{compute\_pmf}(\boldsymbol{\theta}, \mathbf{t}) and \mathcal{B}.\operatorname{compute\_moments}(\boldsymbol{\theta}, K) that combine graph construction with evaluation in a single call, avoiding the overhead of exposing the intermediate graph to Python.

Intuition

The GraphBuilder is a factory that stamps out concrete graphs from a template. The JSON parsing (which involves string processing, memory allocation, and validation) happens once at construction; each subsequent build call only performs the lightweight operations of creating vertices, adding edges with computed weights, and returning the graph. For a graph with p vertices and m edges, parsing costs O(|J|) while building costs O(p + m).

Example

In SVGD with 100 particles and 1000 iterations, the GraphBuilder is constructed once from the serialized graph. Each of the 100 \times 1000 = 100{,}000 evaluations calls build(theta) + pdf(t, g), reusing the cached vertex states and edge connectivity. Without caching, each evaluation would re-parse the JSON string.

Definition 28.3 (Batch Evaluation) Given a GraphBuilder \mathcal{B} and a collection of parameter vectors \boldsymbol{\theta}^{(1)}, \ldots, \boldsymbol{\theta}^{(B)} \in \mathbb{R}^d, the batch evaluation is:

\mathcal{B}.\operatorname{batch\_pmf}(\boldsymbol{\Theta}, \mathbf{t}) = \begin{pmatrix} h(\boldsymbol{\theta}^{(1)}, \mathbf{t}) \\ \vdots \\ h(\boldsymbol{\theta}^{(B)}, \mathbf{t}) \end{pmatrix} \in \mathbb{R}^{B \times n}, \tag{28.3}

where \boldsymbol{\Theta} = (\boldsymbol{\theta}^{(1)}, \ldots, \boldsymbol{\theta}^{(B)})^\top \in \mathbb{R}^{B \times d} is the parameter matrix. In the JAX FFI implementation, batch evaluation is triggered by jax.vmap, which adds a batch dimension to the theta buffer. The FFI handler processes each batch element independently, optionally in parallel via OpenMP threads.

Intuition

Batch evaluation is the workhorse of particle-based inference. In SVGD, each particle carries its own \boldsymbol{\theta}^{(b)}, and the log-likelihood and its gradient must be computed for all particles at each iteration. Batch evaluation via vmap + OpenMP parallelizes this across CPU cores, with each core running a separate GraphBuilder::build + forward algorithm sequence. Since each build creates an independent graph, there is no synchronization overhead.

Example

With 100 SVGD particles on an 8-core machine, jax.vmap(compute_pmf_ffi) dispatches 100 evaluations to the FFI handler, which processes them in batches of 8 (one per OpenMP thread), achieving approximately 8\times speedup over sequential evaluation.

Theorems and Proofs

Theorem 28.1 (GraphBuilder Equivalence) Let G = (V, E, W) be a parameterized phase-type graph with serialization J = \operatorname{serialize}(G), and let \mathcal{B} = \operatorname{GraphBuilder}(J). For any valid parameter vector \boldsymbol{\theta} \in \mathbb{R}^d:

The concrete graph G(\boldsymbol{\theta}) = \mathcal{B}.\operatorname{build}(\boldsymbol{\theta}) has the same vertex set V (with the same state vectors), the same edge set E (with the same connectivity), and edge weights satisfying Equation 28.2.
For any time points \mathbf{t} \in \mathbb{R}^n_{\geq 0} and granularity g: \mathcal{B}.\operatorname{compute\_pmf}(\boldsymbol{\theta}, \mathbf{t}) = (f(t_1; G(\boldsymbol{\theta})), \ldots, f(t_n; G(\boldsymbol{\theta}))), \tag{28.4} where f(t; G(\boldsymbol{\theta})) is the PDF computed by Algorithm 16.2 on the concrete graph.
For any number of moments K: \mathcal{B}.\operatorname{compute\_moments}(\boldsymbol{\theta}, K) = (\mathbb{E}[\tau], \mathbb{E}[\tau^2], \ldots, \mathbb{E}[\tau^K]), \tag{28.5} where the moments are computed by Algorithm 17.2 on the concrete graph.

In particular, the results are independent of the serialization/deserialization path: they depend only on the mathematical content of G and \boldsymbol{\theta}.

Proof. We verify each claim.

(1) Graph construction. The build method (lines 118–194 of graph_builder.cpp) iterates over the cached states_ array and creates one vertex per entry, preserving the state vectors from the serialization. It then iterates over edges_, start_edges_, param_edges_, and start_param_edges_, adding edges with weights computed by compute_weight. For fixed edges, the weight is stored directly. For parameterized edges, compute_weight evaluates Equation 28.2: in linear mode, it computes \sum_k c_k \theta_k; in log mode, it computes \exp(\sum_k \log(c_k \theta_k)) = \prod_k (c_k \theta_k). The vertex indices in the serialization correspond to the creation order, and the build method creates vertices in the same order, so edge connectivity is preserved.

(2) PMF computation. The compute_pmf method (lines 309–361) calls build(theta) to obtain G(\boldsymbol{\theta}), then calls g.pdf(t_i, granularity) for each time point (continuous case) or g.dph_pmf(t_i) for the discrete case. These are the same C++ methods that implement Algorithm 16.2 (CPH) and Algorithm 16.1 (DPH) from [14].

(3) Moment computation. The compute_moments method (lines 271–307) calls build(theta) and then compute_moments_impl, which iteratively calls g.expected_waiting_time(rewards) (the C++ implementation of Algorithm 17.2 from [15]) and applies the factorial scaling \mathbb{E}[\tau^k] = k! \cdot r_0^{(k)} where r_0^{(k)} is the k-th iterated expected waiting time at the starting vertex. \square

Algorithms

28.0.0.1 GraphBuilder PMF Computation

Description. This algorithm describes the full pipeline from a serialized graph structure and parameter vector to PDF/PMF values. It is the C++ implementation behind the JAX FFI handler (Algorithm 27.1). The key optimization is that JSON parsing and structure caching happen once (at GraphBuilder construction), while graph instantiation and forward algorithm execution happen per evaluation. For batch processing, the GraphBuilder is reused across all parameter vectors in the batch.

GraphBuilder PMF Computation
1: Let J be a JSON string encoding the graph structure (Definition 26.1)
2: Let theta = (theta_1, ..., theta_d) be the parameter vector
3: Let t = (t_1, ..., t_n) be evaluation points
4: Let discrete be a boolean flag and g the granularity
5:
6: ▷ One-time initialization
7: function GraphBuilder(J)
8:   Parse J into:
9:     states[0..p-1], edges, start_edges,
10:    param_edges, start_param_edges,
11:    d ← param_length, l ← state_length, p ← n_vertices,
12:    mode ← weight_mode
13:  Store all parsed data in builder B
14:  return B
15: end function
16:
17: ▷ Compute edge weight from coefficients and parameters
18: function ComputeWeight(c, theta, mode)
19:   if mode = "linear" then
20:     return Sigma_{k=1}^{d} c_k * theta_k
21:   else if mode = "log" then
22:     s ← 0
23:     for k ← 1 to d do
24:       if c_k * theta_k <= 0 then error("non-positive product")
25:       s ← s + log(c_k * theta_k)
26:     end for
27:     return exp(s)
28:   end if
29: end function
30:
31: ▷ Build concrete graph from cached structure
32: function Build(B, theta)
33:   if length(theta) != d then error("length mismatch")
34:   G ← new Graph(l)
35:   v_start ← G.starting_vertex()
36:   vertices ← empty array of length p
37:   for i ← 0 to p-1 do
38:     if states[i] = (0,...,0) and i = 0 then
39:       vertices[i] ← v_start               ▷ Reuse starting vertex
40:     else
41:       vertices[i] ← G.create_vertex(states[i])
42:     end if
43:   end for
44:   for (from, to, w) in edges do
45:     vertices[from].add_edge(vertices[to], w)
46:   end for
47:   for (to, w) in start_edges do
48:     v_start.add_edge(vertices[to], w)
49:   end for
50:   for (from, to, c_1, ..., c_d) in param_edges do
51:     w ← ComputeWeight((c_1,...,c_d), theta, mode)
52:     vertices[from].add_edge(vertices[to], w)
53:   end for
54:   for (to, c_1, ..., c_d) in start_param_edges do
55:     w ← ComputeWeight((c_1,...,c_d), theta, mode)
56:     v_start.add_edge(vertices[to], w)
57:   end for
58:   return G
59: end function
60:
61: ▷ Compute PMF/PDF values
62: function ComputePMF(B, theta, t, discrete, g)
63:   G ← Build(B, theta)
64:   y ← empty array of length n
65:   for i ← 1 to n do
66:     if discrete then
67:       y_i ← G.dph_pmf(floor(t_i))        ▷ DPH Forward Step
68:     else
69:       y_i ← G.pdf(t_i, g)                ▷ CPH Uniformization
70:     end if
71:   end for
72:   return y
73: end function
74:
75: ▷ Compute moments
76: function ComputeMoments(B, theta, K)
77:   G ← Build(B, theta)
78:   rewards ← empty                         ▷ Standard moments (no rewards)
79:   moments ← empty array of length K
80:   r ← G.expected_waiting_time(rewards)    ▷ Sojourn Cross-Moments
81:   moments[1] ← r[0]                       ▷ E[tau]
82:   for k ← 2 to K do
83:     r ← G.expected_waiting_time(r)        ▷ Iterate with previous result
84:     moments[k] ← k! * r[0]               ▷ E[tau^k] = k! * iterated result
85:   end for
86:   return moments
87: end function

Correspondence table:

Pseudocode variable	Math symbol	Code variable (file:function)
`J`	J (serialized structure)	`structure_json` (`graph_builder.cpp:GraphBuilder`)
`B`	\mathcal{B} (builder object)	`this` (`graph_builder.cpp`)
`theta`	\boldsymbol{\theta}	`theta` / `theta_vec` (`graph_builder.cpp:build`)
`t`	\mathbf{t}	`times` / `times_vec` (`graph_builder.cpp:compute_pmf`)
`d`	d (parameter dimension)	`param_length_` (`graph_builder.hpp`)
`l`	\ell (state dimension)	`state_length_` (`graph_builder.hpp`)
`p`	p (number of vertices)	`n_vertices_` (`graph_builder.hpp`)
`mode`	weight mode	`weight_mode_` (`graph_builder.hpp`)
`states`	—	`states_` (`graph_builder.hpp`)
`edges`	—	`edges_` (`graph_builder.hpp`)
`param_edges`	—	`param_edges_` (`graph_builder.hpp`)
`c`	\mathbf{c} (coefficient vector)	`coefficients` (`graph_builder.hpp:ParameterizedEdge`)
`G`	G(\boldsymbol{\theta})	`g` (`graph_builder.cpp:build`)
`y`	\mathbf{y} = h(\boldsymbol{\theta}, \mathbf{t})	`result_vec` (`graph_builder.cpp:compute_pmf`)
`moments`	(\mathbb{E}[\tau], \ldots, \mathbb{E}[\tau^K])	`result` (`graph_builder.cpp:compute_moments_impl`)
`r`	iterated expected waiting time	`rewards2` (`graph_builder.cpp:compute_moments_impl`)
`K`	number of moments	`nr_moments` (`graph_builder.cpp:compute_moments`)

Complexity. Let p = |V| (vertices), m = |E| (edges), d = |\boldsymbol{\theta}| (parameters), n = |\mathbf{t}| (evaluation points), K_s the number of uniformization steps, and |J| the JSON string length.

GraphBuilder construction: Time O(|J|) for JSON parsing. Space O(p \cdot \ell + m \cdot d) for cached structure.
Build: Time O(p + m) for vertex creation and edge addition. Space O(p + m) for the concrete graph.
ComputePMF: Time O(p + m + n \cdot K_s \cdot m). The O(p + m) is for build; the O(n \cdot K_s \cdot m) is for n forward algorithm runs.
ComputeMoments: Time O(p + m + K \cdot p^2). The O(K \cdot p^2) is for K applications of expected_waiting_time, each involving graph elimination on p vertices.
Batch evaluation (B parameter vectors): Time O(B \cdot (p + m + n \cdot K_s \cdot m)) with parallelism across B via OpenMP.

Algorithm 28.1: Correctness. Follows from Theorem 28.1. The Build function reconstructs G(\boldsymbol{\theta}) from the cached structure with weights evaluated at \boldsymbol{\theta}. The ComputePMF and ComputeMoments functions invoke the same C++ implementations used by Algorithm 16.1 and Algorithm 16.2 ([14]) and Algorithm 17.2 ([15]).

Implementation Notes

Source code mapping:

Algorithm	File	Function	Lines
Algorithm 28.1 (constructor)	`src/cpp/parameterized/graph_builder.cpp`	`GraphBuilder::parse_structure`	L17–L116
Algorithm 28.1 (weight)	`src/cpp/parameterized/graph_builder.cpp`	`GraphBuilder::compute_weight`	L196–L218
Algorithm 28.1 (build)	`src/cpp/parameterized/graph_builder.cpp`	`GraphBuilder::build`	L118–L194
Algorithm 28.1 (PMF)	`src/cpp/parameterized/graph_builder.cpp`	`GraphBuilder::compute_pmf`	L309–L361
Algorithm 28.1 (moments)	`src/cpp/parameterized/graph_builder.cpp`	`GraphBuilder::compute_moments_impl`	L228–L269
Algorithm 28.1 (PMF+moments)	`src/cpp/parameterized/graph_builder.cpp`	`GraphBuilder::compute_pmf_and_moments`	L363–L574

Deviations from pseudocode:

The compute_pmf_and_moments method in the implementation combines PMF and moment computation in a single call, building the graph only once. The pseudocode presents them as separate functions for clarity.
The implementation supports 2D reward matrices for multivariate models (Definition 2.10), computing PDF and moments per feature dimension by iterating over reward columns and applying reward_transform to the base graph. This is not shown in the pseudocode.
The compute_pmf_multivariate method handles sparse observations (zero entries = no observation) by skipping the forward algorithm for those entries.
The GIL release/acquire pattern (py::gil_scoped_release) is a pybind11 implementation detail not reflected in the pseudocode. It copies data from NumPy arrays to C++ vectors before releasing the GIL, then copies results back to NumPy arrays after reacquiring it.

Symbol Index

Symbol	Name	First appearance
J	Serialized graph structure (JSON)	Definition 28.1
\mathcal{B}	GraphBuilder object	Definition 28.2
G(\boldsymbol{\theta})	Concrete graph at parameters \boldsymbol{\theta}	Definition 28.2
W_{\boldsymbol{\theta}}	Weight function evaluated at \boldsymbol{\theta}	Definition 28.2
\boldsymbol{\Theta}	Parameter matrix (B \times d)	Definition 28.3
B	Batch size	Definition 28.3
\ell	State vector dimension	Definition 28.1