---
title: "SVGD Inference Procedure"
format:
  html:
    toc: true
    toc-depth: 3
---

# SVGD Inference in phasic

This document provides a comprehensive explanation of Stein Variational Gradient Descent (SVGD) inference in phasic, covering different workflows, configurations, Python/C++ interfaces, methodologies, capabilities, and constraints.

## Overview

SVGD ([Liu & Wang, 2016](https://arxiv.org/abs/1608.04471)) is a deterministic sampling method for Bayesian inference that evolves a set of particles to approximate a target posterior distribution. In phasic, SVGD is used to infer parameters of phase-type distributions from observed data.

Key advantages of SVGD include efficient gradient-based optimization, deterministic particle evolution, and natural parallelization across particles. The phasic implementation supports multi-core and multi-node parallelization, moment-based regularization, and dynamic scheduling strategies.

## Entry Points and Workflows

### High-Level Workflow: Graph.svgd()

The simplest way to run SVGD inference is through the high-level `Graph.svgd()` method at `__init__.py:2451`. This method automatically handles model creation and SVGD configuration.

```python
# Build parameterized graph
graph = Graph(model_callback, nr_samples=5)

# Run SVGD directly
svgd = graph.svgd(
    observed_data=observations,
    theta_dim=2,
    n_particles=100,
    n_iterations=1000,
    learning_rate=0.01,
    regularization=1.0,
    nr_moments=2,
    rewards=None  # Optional: 1D or 2D rewards
)

# Access results
print(svgd.theta_mean, svgd.theta_std)
```

This workflow automatically detects whether to use univariate or multivariate models based on the rewards shape. For 2D rewards with shape `(n_vertices, n_features)`, it creates a multivariate model using `pmf_and_moments_from_graph_multivariate()`. For 1D rewards or no rewards, it uses the standard `pmf_and_moments_from_graph()`.

**Advantages**: Simple API, automatic model selection, handles 1D and 2D rewards seamlessly.

**When to use**: Quick inference tasks, standard workflows, exploratory analysis.

### Object-Oriented Workflow: SVGD Class

For more control over configuration, use the `SVGD` class directly at `svgd.py:1383`.

```python
# Create model explicitly
model = Graph.pmf_and_moments_from_graph(
    graph,
    nr_moments=2,
    discrete=False,
    use_ffi=True  # Enable FFI for multi-core
)

# Create SVGD instance
svgd = SVGD(
    model=model,
    observed_data=observations,
    theta_dim=2,
    n_particles=100,
    n_iterations=1000,
    learning_rate=0.01,
    jit=True,
    parallel='pmap',  # Multi-core parallelization
    n_devices=8,
    regularization=1.0,
    nr_moments=2
)

# Run inference
svgd.optimize(return_history=True)

# Access results
print(svgd.theta_mean, svgd.theta_std)
svgd.plot_posterior()
```

**Advantages**: Explicit model control, detailed configuration options, diagnostic plotting.

**When to use**: Custom configurations, performance tuning, multi-device setups, detailed diagnostics.

### Trace-Based Workflow

For repeated inference with different data or parameter exploration, use trace-based elimination at `trace_elimination.py:1238`.

```python
from phasic.trace_elimination import record_elimination_trace, trace_to_log_likelihood

# Record trace once (O(n³) operation)
trace = record_elimination_trace(graph, param_length=2)

# Create log-likelihood function (exact phase-type PDF)
observed_times = np.array([1.5, 2.3, 0.8, 1.2])
log_lik = trace_to_log_likelihood(
    trace,
    observed_times,
    granularity=100,  # Accuracy control
    use_cpp=True  # 10× faster than Python mode
)

# Use with SVGD
from phasic import SVGD
svgd = SVGD(log_lik, theta_dim=2, n_particles=100)
svgd.fit()
```

The trace-based approach records elimination operations once, then evaluates them efficiently for each parameter vector. This is 5-10× faster than symbolic DAG evaluation for SVGD workloads.

**Two modes**:

- **C++ mode** (`use_cpp=True`): Generates standalone C++ code, compiles to shared library, provides ~0.5ms per evaluation
- **Python mode** (`use_cpp=False`): Uses `instantiate_from_trace()` + `graph.pdf()`, ~5ms per evaluation, useful for debugging

**Advantages**: Fast repeated evaluation, exact phase-type likelihood, suitable for MCMC/HMC.

**When to use**: Multiple SVGD runs, parameter exploration, gradient-based optimization methods.

## Model Creation Interfaces

### FFI Mode: Multi-Core Zero-Copy

The FFI (Foreign Function Interface) mode at `ffi_wrappers.py` provides zero-copy XLA-optimized computation for distributed systems.

```python
model = Graph.pmf_and_moments_from_graph(
    graph,
    nr_moments=2,
    use_ffi=True  # Enable FFI
)
```

The FFI mode uses `jax.ffi.ffi_call` to invoke C++ code with zero-copy data transfer. This enables true multi-core parallelization within a single forward pass, not just across particles.

**Key features**:

- Zero-copy data transfer between Python and C++
- XLA fusion for optimized computation
- Native support for jit, vmap, pmap
- GraphBuilder cached by JSON structure
- Automatic multi-core parallelization

**Performance**: Best for large models (100+ vertices) and multi-core systems.

**Location**: `ffi_wrappers.py::compute_pmf_and_moments_ffi`

### Ctypes Mode: Compiled C++ Fallback

The ctypes mode generates standalone C++ code and compiles it to a shared library.

```python
model = Graph.pmf_and_moments_from_graph(
    graph,
    nr_moments=2,
    use_ffi=False  # Disable FFI, use ctypes
)
```

This mode generates C++ code embedding the graph structure, compiles it once, and loads it via ctypes. The compilation is cached based on graph structure hash.

**Key features**:

- Standalone C++ code generation
- Compilation caching (~/.phasic_cache/)
- No JAX required for compiled code
- Custom VJP for gradients via finite differences

**Performance**: Good for small to medium models, single-core systems.

**Location**: `__init__.py:2895` (pmf_and_moments_from_graph)

### Pybind11 Mode: Direct Bindings

The pybind11 mode uses direct C++ bindings for graph operations.

```python
# Direct graph operations
pdf = graph.pdf(time, granularity=100)
moments = graph.moments(power=2)
sample = graph.sample(nr_samples=1000)
```

These methods are bound directly to C++ implementations via pybind11 at `phasic_pybind.cpp:921`.

**Key features**:

- Direct C++ method calls
- No code generation overhead
- Immediate execution
- Used internally by trace-based approach

**Performance**: Excellent for single evaluations, less efficient for batched operations.

**Location**: `phasic_pybind.cpp::Graph` class bindings

### GraphBuilder: Parameterized Models

The GraphBuilder class provides efficient batch evaluation for parameterized models at `phasic_pybind.cpp`.

```python
from phasic import phasic_pybind as cpp_module
import json

# Create builder from serialized graph
structure_json = json.dumps(graph.serialize())
builder = cpp_module.parameterized.GraphBuilder(structure_json)

# Compute PMF and moments
pmf, moments = builder.compute_pmf_and_moments(
    theta=np.array([1.0, 2.0]),
    times=np.array([1.0, 2.0, 3.0]),
    nr_moments=2,
    discrete=False,
    granularity=100
)
```

The GraphBuilder is cached by the `pmf_and_moments_from_graph()` function to avoid repeated JSON parsing.

**Key features**:

- JSON-based graph structure
- Cached builder instance
- Batch evaluation support
- Used by non-FFI models

**Performance**: Good for models with many evaluations on same structure.

## GraphBuilder Approach: Deep Dive

The GraphBuilder approach at `graph_builder.hpp:29` separates graph topology (structure) from parameter values (theta), enabling efficient repeated evaluation with different parameters.

### JSON-Based Structure Representation

When `Graph.serialize()` is called, the graph structure is converted to JSON format at `__init__.py:1650`:

```json
{
  "states": [[0, 0], [1, 0], [0, 1]],
  "edges": [[0, 1, 1.5], [0, 2, 2.0]],
  "start_edges": [[1, 1.0]],
  "param_edges": [[1, 2, 0.0, 2.5, 0.8]],
  "start_param_edges": [[2, 0.0, 1.0, 0.5]],
  "param_length": 2,
  "state_length": 2,
  "n_vertices": 3
}
```

**Format details**:

- `states`: Vertex states as `[[s₀₀, s₀₁, ...], ...]` for state-indexed vertices
- `edges`: Regular edges as `[from_idx, to_idx, weight]`
- `start_edges`: Edges from starting vertex as `[to_idx, weight]`
- `param_edges`: Parameterized edges as `[from_idx, to_idx, base_weight, coeff₁, coeff₂, ...]`
- `start_param_edges`: Parameterized starting edges as `[to_idx, base_weight, coeff₁, coeff₂, ...]`

### Graph Building Process

The GraphBuilder constructs concrete graphs from structure + parameters at `graph_builder.cpp:106`:

```cpp
Graph GraphBuilder::build(const double* theta, size_t theta_len) {
    // 1. Create empty graph with state dimension
    Graph g(state_length_);

    // 2. Create all vertices from states
    for (int i = 0; i < n_vertices_; i++) {
        vertices.push_back(g.find_or_create_vertex_p(states_[i]));
    }

    // 3. Add regular edges (constant weights)
    for (const auto& edge : edges_) {
        from_v->add_edge(*to_v, edge.weight);
    }

    // 4. Add parameterized edges (computed weights)
    for (const auto& edge : param_edges_) {
        // Compute: weight = base_weight + Σᵢ coeffᵢ * θᵢ
        double weight = edge.base_weight;
        for (int i = 0; i < param_length_; i++) {
            weight += edge.coefficients[i] * theta[i];
        }
        from_v->add_edge(*to_v, weight);
    }

    return g;
}
```

**Key steps**:

1. **Empty graph creation**: Allocates graph with correct state dimension
2. **Vertex creation**: All vertices created from stored states (no callback needed)
3. **Regular edges**: Added with constant weights (no computation)
4. **Parameterized edges**: Weights computed as `base_weight + c · θ` via dot product

**No elimination**: GraphBuilder does not perform elimination. It builds the original graph structure, and elimination happens later when computing PDF/moments.

### Caching Strategy

GraphBuilder caching happens at multiple levels:

**Python-level caching** at `__init__.py:3024`:

```python
# Create GraphBuilder ONCE - captured in model closure
builder = cpp_module.parameterized.GraphBuilder(structure_json)

def _compute_pmf_and_moments_cached(theta_np, times_np, rewards_np=None):
    """Uses cached builder - NO JSON parsing per call."""
    pmf, moments = builder.compute_pmf_and_moments(...)
    return pmf, moments
```

The builder instance is created once and captured in the closure returned by `pmf_and_moments_from_graph()`. This avoids repeated JSON parsing.

**FFI-level caching** (for FFI mode):

FFI handlers maintain thread-local GraphBuilder caches keyed by JSON structure hash. When `compute_pmf_and_moments_ffi()` is called multiple times with the same structure, the builder is reused.

**When builder is created**:

- First call to `pmf_and_moments_from_graph(graph)` for a given graph structure
- First FFI call with a new JSON structure (thread-local cache miss)

**When builder is reused**:

- Subsequent calls to the returned model function with different theta values
- FFI calls with identical JSON structure on same thread

**Cache invalidation**: Builder is NOT invalidated unless the graph structure changes (different vertices, edges, or parameterization).

### Graph Elimination in GraphBuilder Context

GraphBuilder itself does NOT perform elimination. Elimination happens when:

1. **PDF computation** (`graph.pdf()`): Uses forward algorithm (Algorithm 4), which doesn't require elimination
2. **Moment computation** (`graph.moments()`): Requires elimination to convert cyclic → acyclic graph

**Moment computation process** at `graph_builder.cpp:201`:

```cpp
py::array_t<double> GraphBuilder::compute_moments(theta, nr_moments) {
    // 1. Build graph with theta
    Graph g = build(theta, theta.size());

    // 2. Compute moments via reward transformation
    std::vector<double> rewards;  // Empty for standard moments
    std::vector<double> rewards2 = g.expected_waiting_time(rewards);

    moments[0] = rewards2[0];  // E[T]

    // 3. Iteratively compute higher moments: E[T^k]
    for (int k = 2; k <= nr_moments; k++) {
        for (int j = 0; j < rewards3.size(); j++) {
            rewards3[j] = rewards2[j] * std::pow(rewards2[j], k-1);
        }
        rewards2 = g.expected_waiting_time(rewards3);
        moments[k-1] = factorial(k) * rewards2[0];
    }

    return moments;
}
```

The `g.expected_waiting_time()` method performs graph elimination internally using Algorithm 3 (Gaussian elimination on graph structure).

**Algorithm 3 (Graph Elimination)** operates in-place on the graph:

```
For each vertex i in topological order:
    For each parent p of i:
        For each child c of i:
            If p → c edge exists:
                Update weight: w(p→c) += w(p→i) * w(i→c)
            Else:
                Create bypass edge: p → c with weight w(p→i) * w(i→c)
        Remove edge p → i
        Renormalize edges from p to sum to 1
```

**Complexity**: O(n³) worst case for dense graphs, O(n²) for sparse graphs.

**No caching of eliminated graphs**: Each `compute_moments()` call performs elimination independently. The GraphBuilder caches the structure, not the eliminated result.

### Performance Characteristics

**Strengths**:

- Fast graph construction: O(E) where E = number of edges
- No JSON parsing overhead after first call (cached builder)
- Suitable for SVGD: build graph once per particle per iteration
- Memory efficient: stores structure once, not per theta

**Weaknesses**:

- Elimination not cached: O(n³) per `compute_moments()` call
- Not optimal for repeated evaluation with same theta
- Graph construction overhead (though small)

**Best for**: SVGD inference where each particle has different theta, and each iteration uses different theta values.

## Trace Approach: Deep Dive

The trace approach at `trace_elimination.py:301` records elimination operations once, then replays them efficiently for any parameter values.

### Trace Recording Process

Trace recording performs graph elimination while recording all arithmetic operations in a linear sequence.

**Phase 1: Compute Vertex Rates** at `trace_elimination.py:426`:

```python
for i, v in enumerate(vertices_list):
    # Get all edges (regular + parameterized)
    edges = v.edges()
    param_edges = v.parameterized_edges()

    # Build weight expressions
    weight_indices = []

    # Regular edges: CONST operations
    for edge in edges:
        weight_idx = builder.add_const(edge.weight())
        weight_indices.append(weight_idx)

    # Parameterized edges: DOT operations
    for param_edge in param_edges:
        coeffs = param_edge.edge_state(param_length)
        base_weight = param_edge.weight()

        # weight = base + c₁θ₁ + c₂θ₂ + ... + cₙθₙ
        dot_idx = builder.add_dot(coeffs)
        if base_weight != 0.0:
            base_idx = builder.add_const(base_weight)
            weight_idx = builder.add_add(base_idx, dot_idx)
        else:
            weight_idx = dot_idx

        weight_indices.append(weight_idx)

    # rate = 1 / sum(weights)
    sum_idx = builder.add_sum(weight_indices)
    vertex_rates[i] = builder.add_inv(sum_idx)
```

**Phase 2: Convert Edges to Probabilities** at `trace_elimination.py:488`:

```python
for i, v in enumerate(vertices_list):
    for edge in v.edges():
        weight = edge.weight()
        weight_idx = builder.add_const(weight)

        # prob = weight * rate
        prob_idx = builder.add_mul(weight_idx, vertex_rates[i])

        edge_probs[i].append(prob_idx)
        vertex_targets[i].append(to_idx)
```

**Phase 3: Elimination Loop** at `trace_elimination.py:554`:

```python
for i in range(n_vertices):
    for parent_idx in parents[i]:
        parent_to_i_prob = edge_probs[parent_idx][parent_to_i_edge_idx]

        # For each child of i
        for child_idx in vertex_targets[i]:
            i_to_child_prob = edge_probs[i][child_edge_idx]

            # Bypass: parent → child via i
            bypass_prob = builder.add_mul(parent_to_i_prob, i_to_child_prob)

            if (parent_idx, child_idx) in edge_map:
                # Update existing edge
                old_prob = edge_probs[parent_idx][parent_to_child_edge_idx]
                new_prob = builder.add_add(old_prob, bypass_prob)
                edge_probs[parent_idx][parent_to_child_edge_idx] = new_prob
            else:
                # Create new edge
                edge_probs[parent_idx].append(bypass_prob)
                vertex_targets[parent_idx].append(child_idx)

        # Remove edge parent → i
        edge_probs[parent_idx][parent_to_i_edge_idx] = -1

        # Renormalize parent's edges
        total_idx = builder.add_sum([edge_probs[parent_idx][j] for j in valid_indices])
        for j in valid_indices:
            old_prob = edge_probs[parent_idx][j]
            new_prob = builder.add_div(old_prob, total_idx)
            edge_probs[parent_idx][j] = new_prob
```

This implements Algorithm 3 (Gaussian elimination) while recording operation indices instead of computing values.

### Trace Structure

The EliminationTrace dataclass at `trace_elimination.py:106` contains:

```python
@dataclass
class EliminationTrace:
    operations: List[Operation]              # Linear sequence of operations
    vertex_rates: np.ndarray                 # Operation indices for rates
    edge_probs: List[List[int]]              # Operation indices for edge probs
    vertex_targets: List[List[int]]          # Target vertex indices
    states: np.ndarray                       # Vertex states
    starting_vertex_idx: int                 # Starting vertex index
    n_vertices: int                          # Number of vertices
    state_length: int                        # State vector dimension
    param_length: int                        # Number of parameters
    is_discrete: bool                        # Discrete vs continuous
    metadata: Dict[str, Any]                 # Statistics, phase info
```

**Operation types** at `trace_elimination.py:50`:

```python
class OpType(Enum):
    CONST = "const"    # Constant value
    PARAM = "param"    # Parameter θᵢ
    DOT = "dot"        # c₁θ₁ + c₂θ₂ + ... + cₙθₙ
    ADD = "add"        # a + b
    MUL = "mul"        # a * b
    DIV = "div"        # a / b
    INV = "inv"        # 1 / a
    SUM = "sum"        # sum([a, b, c, ...])
```

Each operation references earlier operations by index, forming a DAG (directed acyclic graph) of computations.

### Trace Evaluation

Evaluation replays operations with concrete parameter values at `trace_elimination.py:689`:

```python
def evaluate_trace(trace: EliminationTrace, params: np.ndarray):
    # Allocate value array
    n_ops = len(trace.operations)
    values = np.zeros(n_ops, dtype=np.float64)

    # Execute operations in order
    for i, op in enumerate(trace.operations):
        if op.op_type == OpType.CONST:
            values[i] = op.const_value

        elif op.op_type == OpType.PARAM:
            values[i] = params[op.param_idx]

        elif op.op_type == OpType.DOT:
            values[i] = np.dot(op.coefficients, params)

        elif op.op_type == OpType.ADD:
            values[i] = values[op.operands[0]] + values[op.operands[1]]

        elif op.op_type == OpType.MUL:
            values[i] = values[op.operands[0]] * values[op.operands[1]]

        elif op.op_type == OpType.DIV:
            values[i] = values[op.operands[0]] / values[op.operands[1]]

        elif op.op_type == OpType.INV:
            values[i] = 1.0 / values[op.operands[0]]

        elif op.op_type == OpType.SUM:
            values[i] = sum(values[idx] for idx in op.operands)

    # Extract results
    vertex_rates = values[trace.vertex_rates]
    edge_probs = [values[np.array(ops)] for ops in trace.edge_probs]

    return {
        'vertex_rates': vertex_rates,
        'edge_probs': edge_probs,
        'vertex_targets': trace.vertex_targets,
        'states': trace.states,
        'starting_vertex_idx': trace.starting_vertex_idx
    }
```

**Complexity**: O(n_ops) where n_ops = number of operations. Typically O(n²) for n vertices after elimination.

### JAX Integration

JAX-compatible evaluation uses `jax.numpy` for automatic differentiation at `trace_elimination.py:1082`:

```python
def evaluate_trace_jax(trace: EliminationTrace, params):
    import jax.numpy as jnp

    values = jnp.zeros(n_ops, dtype=jnp.float64)

    for i, op in enumerate(trace.operations):
        if op.op_type == OpType.CONST:
            values = values.at[i].set(op.const_value)
        elif op.op_type == OpType.DOT:
            values = values.at[i].set(jnp.dot(op.coefficients, params))
        # ... other operations using JAX array syntax

    return results
```

This enables `jax.jit`, `jax.grad`, `jax.vmap`, and `jax.pmap` transformations on the trace evaluation function.

### Trace Caching

Traces are cached at Python level to avoid repeated recording.

**Cache location**: `~/.phasic_cache/traces/`

**Cache key** at `trace_cache.py:228`:

```python
hash_dict = {
    "state_length": graph_data.get("state_length"),
    "n_vertices": len(graph_data.get("vertices", [])),
    "edges": []  # Edge structure (parameterized status + connectivity)
}

# Compute SHA256 hash of structure
json_str = json.dumps(hash_dict, sort_keys=True)
hash_hex = hashlib.sha256(json_str.encode()).hexdigest()

cache_file = cache_dir / f"{hash_hex}.json"
```

**Automatic caching** at `trace_elimination.py:676`:

```python
# Save trace to cache after recording
try:
    from .trace_cache import save_trace_to_cache_python
    save_trace_to_cache_python(graph, trace)
except Exception:
    pass  # Silently ignore cache errors
```

**Cache management utilities**:

- `clear_trace_cache()`: Remove all cached traces
- `get_trace_cache_stats()`: Get cache size and file count
- `list_cached_traces()`: List all cached traces with metadata
- `cleanup_old_traces(max_size_mb, max_age_days)`: LRU eviction

### Instantiation from Trace

The `instantiate_from_trace()` function at `trace_elimination.py:1004` creates a concrete graph from evaluated trace:

```python
def instantiate_from_trace(trace: EliminationTrace, params: np.ndarray):
    # 1. Evaluate trace to get concrete rates and probabilities
    result = evaluate_trace(trace, params)

    # 2. Create new graph
    graph = Graph(trace.state_length)

    # 3. Create all vertices
    state_to_vertex = {}
    for i in range(trace.n_vertices):
        state = tuple(trace.states[i].tolist())
        v = graph.find_or_create_vertex(trace.states[i].tolist())
        state_to_vertex[state] = v

    # 4. Add edges with concrete weights
    for i in range(trace.n_vertices):
        inv_rate = result['vertex_rates'][i]

        for j, to_idx in enumerate(result['vertex_targets'][i]):
            prob = result['edge_probs'][i][j]

            # Convert probability back to weight
            # weight = prob / rate = prob * inv_rate
            weight = prob / inv_rate

            from_vertex.add_edge(to_vertex, weight)

    return graph
```

The instantiated graph can then be used for PDF/PMF computation via `graph.pdf()` or `graph.dph_pmf()`.

### Retracing Conditions

A new trace must be recorded when:

1. **Graph structure changes**: Different vertices, edges, or connectivity
2. **Parameterization changes**: Different parameter count or edge parameterization
3. **Cache miss**: No cached trace exists for the graph structure hash

A trace can be reused when:

1. **Same structure**: Identical vertices, edges, and parameterization
2. **Different parameter values**: Only theta values change, not structure
3. **Cache hit**: Cached trace exists with matching hash

**Cache hit example**:

```python
# First call: records trace
graph1 = Graph(model, nr_samples=5)
trace1 = record_elimination_trace(graph1, param_length=2)

# Second call with same structure: uses cached trace
graph2 = Graph(model, nr_samples=5)
trace2 = record_elimination_trace(graph2, param_length=2)

# trace2 loaded from cache, not recomputed
```

### C++ Code Generation Mode

For `trace_to_log_likelihood()` with `use_cpp=True`, the trace is embedded in standalone C++ code at `trace_elimination.py:1354`:

```python
cpp_code = _generate_cpp_from_trace(trace, observed_data, granularity)

# Compile to shared library (cached by trace hash)
lib_path = _compile_trace_library(cpp_code, trace_hash)

# Wrap for JAX
log_likelihood = _wrap_trace_log_likelihood_for_jax(lib_path, trace.param_length)
```

This generates C++ code with static arrays containing the trace operations, compiles it once, and loads via ctypes. The compiled library is cached indefinitely.

**Performance**: ~0.5ms per evaluation vs ~5ms for Python mode (10× speedup).

### Performance Characteristics

**Strengths**:

- Elimination done once: O(n³) recording, O(n²) evaluation
- Fast repeated evaluation: 5-10× faster than symbolic DAG
- JAX compatible: jit, grad, vmap, pmap all work
- Memory efficient: linear operation sequence
- Scales to 100K+ vertices

**Weaknesses**:

- Initial recording overhead: O(n³) for first call
- Cache storage: ~KB per trace
- Not suitable for single evaluation (use GraphBuilder)

**Best for**: Repeated evaluation with different theta (SVGD, MCMC, HMC), parameter exploration, gradient-based optimization.

## Comparison: GraphBuilder vs Trace

| Aspect | GraphBuilder | Trace |
|--------|--------------|-------|
| **Elimination** | Every call (O(n³)) | Once during recording (O(n³)) |
| **Evaluation** | Build graph (O(E)) + eliminate (O(n³)) | Replay operations (O(n²)) |
| **Caching** | Structure only (JSON) | Full elimination result |
| **Memory** | Low (structure) | Medium (operation sequence) |
| **First call** | Fast (no recording) | Slow (recording + cache save) |
| **Nth call (N>1)** | Same as first | 5-10× faster than first |
| **SVGD (1000 evals)** | 1000 × (build + eliminate) | 1 × record + 1000 × replay |
| **Best for** | Few evaluations per structure | Many evaluations per structure |
| **JAX support** | Via FFI or ctypes | Native JAX arrays |
| **Code generation** | No | Yes (C++ mode) |

**Decision guide**:

- **Use GraphBuilder** if: <10 evaluations per graph structure, or structure changes frequently
- **Use Trace** if: >10 evaluations per graph structure, or gradient-based optimization (SVGD, HMC)

For SVGD with 1000 iterations and 100 particles (100K evaluations), Trace is 5-10× faster overall despite initial recording overhead.

## Configuration Options

### JIT Compilation

Control JIT compilation via the `jit` parameter.

```python
# Enable JIT (default from config)
svgd = SVGD(model, data, theta_dim=2, jit=True)

# Disable JIT (debugging)
svgd = SVGD(model, data, theta_dim=2, jit=False)

# Auto-detect from config
svgd = SVGD(model, data, theta_dim=2, jit=None)
```

JIT compilation provides significant speedup but adds initial compilation overhead. Compiled functions are cached in memory and on disk.

**Cache locations**:

- Memory: `SVGD._compiled_cache` (class-level dictionary)
- Disk: `~/.phasic_cache/compiled_svgd_*.pkl`

**Cache keys**: Model ID, theta shape, data shape, regularization parameters, rewards.

### Parallelization Strategies

Choose parallelization strategy via the `parallel` parameter at `svgd.py:1596`.

```python
# Single device vectorization (default for 1 device)
svgd = SVGD(model, data, theta_dim=2, parallel='vmap')

# Multi-device parallelization (default for >1 device)
svgd = SVGD(model, data, theta_dim=2, parallel='pmap', n_devices=8)

# No parallelization (sequential, for debugging)
svgd = SVGD(model, data, theta_dim=2, parallel='none')

# Auto-select based on available devices
svgd = SVGD(model, data, theta_dim=2, parallel=None)
```

**vmap** (`svgd_step:1180`): Vectorizes gradient computation across particles on a single device using `jax.vmap`. Best for single-device systems or when devices are limited.

**pmap** (`svgd_step:1158`): Parallelizes gradient computation across multiple devices using `jax.pmap`. Requires `n_particles % n_devices == 0`. Best for multi-core CPUs or multi-GPU systems.

**none** (`svgd_step:1186`): Sequential computation with Python loop. Useful for debugging gradient issues or when parallelization causes problems.

**Auto-selection logic**: Uses pmap if `len(jax.devices()) > 1`, otherwise vmap.

### Regularization

Configure moment-based regularization to guide inference toward matching sample moments.

```python
# Constant regularization
svgd = SVGD(
    model, data,
    theta_dim=2,
    regularization=1.0,  # Fixed λ
    nr_moments=2
)

# Exponential decay schedule
from phasic import ExpRegularization
reg_schedule = ExpRegularization(
    first_reg=5.0,  # Strong initial regularization
    last_reg=0.1,   # Weak final regularization
    tau=500.0       # Decay time constant
)
svgd = SVGD(model, data, theta_dim=2, regularization=reg_schedule, nr_moments=2)

# Exponential CDF schedule (bidirectional)
from phasic import ExponentialCDFRegularization
reg_schedule = ExponentialCDFRegularization(
    first_reg=0.0,  # No initial regularization
    last_reg=1.0,   # Strong final regularization
    tau=500.0
)
svgd = SVGD(model, data, theta_dim=2, regularization=reg_schedule, nr_moments=2)
```

The regularization term adds a penalty to match model moments to sample moments:

```
log p(theta | data) = log p(data | theta) + log p(theta) - λ * ||E[T^k|theta] - sample_moments||²
```

**Sample moments** (`svgd.py:1348`): Computed from observed data at SVGD initialization: `[E[X], E[X²], ..., E[X^k]]`.

**Model moments**: Computed via reward transformation in C++ for each particle.

**Aggregation for multivariate**: For 2D moments with shape `(n_features, nr_moments)`, the penalty is computed by averaging across features: `mean(moments, axis=0)` at `svgd.py:2097`.

**Gradient precompilation**: Disabled when using regularization schedules to allow dynamic regularization values per iteration at `svgd.py:2426`.

### Step Size Schedules

Control learning rate dynamics via step size schedules.

```python
# Constant step size (default)
svgd = SVGD(model, data, theta_dim=2, learning_rate=0.01)

# Exponential decay
from phasic import ExpStepSize
step_schedule = ExpStepSize(
    first_step=0.1,   # Initial step size
    last_step=0.01,   # Final step size
    tau=500.0         # Decay time constant
)
svgd = SVGD(model, data, theta_dim=2, learning_rate=step_schedule)

# Adaptive step size based on particle spread
from phasic import AdaptiveStepSize
step_schedule = AdaptiveStepSize(
    base_step=0.01,
    kl_target=0.1,
    adjust_rate=0.1
)
svgd = SVGD(model, data, theta_dim=2, learning_rate=step_schedule)
```

Step size schedules help prevent divergence with large datasets and enable exploration-exploitation trade-offs.

**Auto-scaling** (`svgd.py:1715`): Learning rate is automatically scaled by `1.0 / max(1.0, n_observations / 1000.0)` to prevent gradient explosion with large datasets.

### Parameter Transformations

Constrain parameters to specific domains via transformations.

```python
# Positive parameters (default for phase-type models)
svgd = SVGD(model, data, theta_dim=2, positive_params=True)
# Uses softplus: θ_constrained = log(1 + exp(θ_unconstrained))

# Unconstrained parameters
svgd = SVGD(model, data, theta_dim=2, positive_params=False)

# Custom transformation
import jax.nn as jnn
svgd = SVGD(
    model, data, theta_dim=2,
    positive_params=False,  # Must be False
    param_transform=lambda theta: jnn.sigmoid(theta)  # Map to [0, 1]
)
```

SVGD operates in unconstrained space but applies transformations before model evaluation. Gradients are computed via the chain rule automatically.

**Default**: `positive_params=True` at `svgd.py:1749` since phase-type parameters (rates) must be positive.

## Multivariate Models

For multivariate phase-type distributions where each feature dimension has its own reward vector.

```python
# Create multivariate model
model = Graph.pmf_and_moments_from_graph_multivariate(
    graph,
    nr_moments=2,
    discrete=False
)

# Setup multivariate data
n_times = 100
n_features = 3
observed_data = jnp.array([...])  # Shape: (100, 3)

# 2D rewards: one reward vector per feature
rewards_2d = jnp.array([...])  # Shape: (n_vertices, 3)

# Run SVGD
svgd = graph.svgd(
    observed_data=observed_data,
    theta_dim=2,
    rewards=rewards_2d
)
```

Each feature dimension is computed independently with its reward vector using `jax.lax.scan` to keep all computation in compiled code at `__init__.py:2640`.

**Output shapes**:

- PMF: `(n_times, n_features)` for 2D rewards, `(n_times,)` for 1D
- Moments: `(n_features, nr_moments)` for 2D rewards, `(nr_moments,)` for 1D

**Log-likelihood**: Sum over all observation elements: `Σᵢⱼ log(PMF[i,j])`.

**Moment aggregation**: Mean across features for regularization: `mean(moments, axis=0)`.

## SVGD Algorithm Details

### Kernel Computation

SVGD uses an RBF kernel with automatic bandwidth selection via the median heuristic at `svgd.py:1008`.

```
K[i,j] = exp(-||θᵢ - θⱼ||² / (2h²))
∇K[i,j] = -K[i,j] * (θᵢ - θⱼ) / h²
```

Bandwidth selection at `svgd.py:752`:

```
h = median(pairwise_distances) / log(n_particles + 1)
```

The kernel matrix and its gradient are computed via vectorized operations at `svgd.py:1026` with no Python loops.

### Update Rule

The SVGD update at `svgd.py:1084` combines likelihood gradients and kernel repulsion:

```
φ = (K @ ∇log p) / n + Σⱼ ∇K / n
θ ← θ + ε * φ
```

- **Positive term**: `K @ ∇log p` attracts particles toward high posterior density
- **Negative term**: `Σⱼ ∇K` repels particles to maintain diversity

This is computed via efficient einsum operations at `svgd.py:1113` without explicit loops.

### Prior

Default prior is standard normal on unconstrained parameters at `svgd.py:1955`:

```
log p(θ) = -0.5 * Σᵢ θᵢ²
```

Custom priors can be provided via the `prior` parameter.

## Capabilities

1. **Bayesian Parameter Inference**: Full posterior approximation via particle-based SVGD
2. **Moment-Based Regularization**: Guide inference to match empirical moments
3. **Multi-Core Parallelization**: pmap for multi-device, vmap for single-device vectorization
4. **Multi-Node Distribution**: SLURM cluster support via `initialize_distributed()`
5. **Dynamic Schedules**: Exponential decay, adaptive, and custom schedules for learning rate and regularization
6. **Parameter Constraints**: Automatic positive constraints via softplus, custom transformations
7. **Multivariate Models**: 2D rewards for multivariate phase-type distributions
8. **Exact Phase-Type Likelihood**: Forward algorithm (Algorithm 4) for accurate PDF computation
9. **Trace-Based Optimization**: Record once, evaluate many times (5-10× speedup)
10. **Compilation Caching**: Memory and disk caching for JIT-compiled functions

## Constraints and Limitations

1. **JAX Required**: Most features require JAX to be installed and configured
2. **Model Signature**: Model must return `(pmf, moments)` tuple for regularized SVGD
3. **Particle Divisibility**: For pmap, `n_particles` must be divisible by `n_devices`
4. **Schedule Precompilation**: Gradient precompilation disabled when using dynamic regularization schedules
5. **Reward Vectors**: Not yet supported with exact phase-type likelihood in trace mode
6. **Memory Requirements**: Large particle counts and high-dimensional parameter spaces require significant RAM
7. **Compilation Time**: Initial JIT compilation can take several minutes for large models
8. **pmap Mesh Conflicts**: JAX 0.8+ requires explicit device mesh to avoid conflicts in nested pmap calls
9. **Moment Count Mismatch**: Model must provide at least `nr_moments` moments if using regularization
10. **Transform Conflicts**: Cannot use both `positive_params=True` and custom `param_transform`

## Performance Considerations

### FFI vs Ctypes vs Pybind11

- **FFI**: Best for multi-core systems, zero-copy, XLA fusion. Requires JAX. Best performance for large models.
- **Ctypes**: Good for single-core, standalone C++ code. No JAX required for compiled code. Good for deployment.
- **Pybind11**: Best for single evaluations, direct C++ calls. Used internally by trace mode.

### JIT Compilation

- **First run**: Slow due to compilation overhead (minutes for large models)
- **Subsequent runs**: Fast due to cached compiled functions (memory + disk)
- **Cache hits**: ~100× speedup for cached gradients
- **Disable for debugging**: Set `jit=False` to avoid compilation issues

### Parallelization

- **vmap**: 5-10× speedup on single device vs sequential
- **pmap**: 3-5× speedup per device (diminishing returns with more devices)
- **Overhead**: pmap has communication overhead, use vmap if <4 devices

### Regularization Schedules

- **Constant**: Gradient can be precompiled, ~10× faster
- **Dynamic**: Gradient computed per iteration, slower but more flexible
- **Trade-off**: Use constant if possible, schedules when exploration/exploitation balance is critical

## References

- [Røikjer, Hobolth & Munch (2022)](https://doi.org/10.1007/s11222-022-10155-6) - Phase-Type Distributions for Statistics and Computing
- [Liu & Wang (2016)](https://arxiv.org/abs/1608.04471) - Stein Variational Gradient Descent

---

*Generated with Claude Code*
