fpgen — Kramers–Moyal / Fokker–Planck / Langevin toolkit

fpgen derives phase-space dynamics from a structured bosonic master equation using exact SymPy algebra. It supports the Wigner, Glauber-Sudarshan P, Husimi Q, and positive-P representations.

Authors: Yu Xue-Hao and Qiao Cong-Feng (University of Chinese Academy of Sciences, UCAS)

The public workflow is

master equation
  -> full Kramers-Moyal series
  -> second-order Fokker-Planck truncation
  -> Ito c-number Langevin equation
  -> closed second-moment dynamics, Jacobian, and symbolic reduction

The repository deliberately stops at symbolic model construction. Numerical fixed-point searches, continuation, bifurcation detection, and GPU execution belong to downstream solvers such as qphase.

Installation

Python 3.10 or later is required.

pip install -e .

For tests, create an fpgen-only uv environment and run pytest:

uv sync --group test
uv run pytest tests -q

Notebook dependencies are separate from the test environment:

uv sync --group notebook

Quick start

import sympy as sp
from fpgen import (
    LindbladChannel,
    MasterEquation,
    boson_modes,
    derive_kramers_moyal,
)

(a,) = boson_modes("a")
omega, kappa = sp.symbols("omega kappa", real=True)

master = MasterEquation(
    modes=[a],
    hamiltonian=omega * a.dag * a,
    channels=[LindbladChannel(a, kappa)],
)

km = derive_kramers_moyal(master, representation="wigner")
fpe = km.truncate(2)
langevin = fpe.to_langevin()

Mode order is explicit and remains fixed throughout the derivation. Operators are noncommutative polynomials supporting addition, multiplication, integer powers, and .dag.

Kramers-Moyal convention

fpgen stores the flux-form coefficients K_mu(z) in

dP/dt = sum_mu (-1)^|mu| / mu! * partial^mu [K_mu(z) P].

The exact normal-form differential operator is first converted to flux form with the multivariate Leibniz rule. Only then does km.truncate(2) discard terms above second order. This order is important: truncating the normal form first can remove lower-order contributions generated by the rearrangement.

Useful attributes include:

km.coefficients          # complete {multi_index: coefficient} mapping
km.order                 # maximum exact KM order
fpe.drift                # first-order coefficient vector
fpe.diffusion            # complex-symmetric augmented FPE diffusion
fpe.discarded            # coefficients removed by truncation
fpe.is_exact             # whether truncation removed nothing

The augmented variable order is (alpha, alpha_bar) for Wigner/P/Q and (alpha, beta) for positive-P. In positive-P, beta is independent of alpha, not its complex conjugate.

Langevin output

The Fokker-Planck equation defines the Ito equation

dot(z) = A(z) + xi,       <xi xi^T> = G_aug.

When the drift admits a canonical factorization, fpgen also reports A(z) = -i H_eff(z) z. Constant drives or ambiguous factorizations leave langevin.hamiltonian unset and provide a diagnostic instead.

langevin.drift
langevin.hamiltonian             # augmented H_eff, or None
langevin.physical_hamiltonian    # alpha-alpha block for Wigner/P/Q
langevin.anomalous_hamiltonian   # alpha-alpha_bar drift block
langevin.diffusion               # G_aug
langevin.normal_correlation      # D = <xi xi^dagger>
langevin.anomalous_correlation   # M = <xi xi^T>
langevin.noise_factor            # canonical B with B B^T = G_aug
blocks = langevin.to_augmented_blocks()
blocks.physical_hamiltonian      # H
blocks.pairing_hamiltonian       # V
blocks.augmented_hamiltonian     # [[H,V],[-V*, -H*]]
blocks.augmented_diffusion       # D_tilde = [[D,M],[M*,D*]]

For Wigner/P/Q, D and M are extracted directly from the FPE diffusion, not reconstructed from the noise factor. The factor B is one valid real-noise realization and is generally not minimal.

The FPE matrix and Hermitian noise correlation use different products. With z = (alpha, alpha*) and the block-exchange matrix J,

G_aug   = <xi_tilde xi_tilde^T>,
D_tilde = <xi_tilde xi_tilde^dagger> = G_aug J.

to_augmented_blocks() validates the conjugate block identities before returning H, V, D, M, H_tilde, and D_tilde.

Second-moment dynamics

The unified entry point derives a full moment system and selects its layout:

dynamics = langevin.to_second_moment_dynamics(
    parameters=(omega, epsilon, kappa),
    layout="auto",       # auto | normal | augmented | independent
    closure="exact",     # exact | factorized_bilinear
)

For Wigner/P/Q, define

R = <alpha alpha^dagger>,       C = <alpha alpha^T>,
R_tilde = <z z^dagger> = [[R,C],[C*,R*]].

The Hermitian augmented equation is

dot(R_tilde) = -i H_tilde R_tilde
               + i R_tilde H_tilde^dagger + D_tilde.

Its two independent blocks are

dot(R) = -i(HR + VC*) + i(RH^dagger + CV^dagger) + D,
dot(C) = -i(HC + CH^T + VR* + RV^T) + M.

The returned object exposes both the matrix and real-vector forms:

dynamics.layout_kind          # normal | augmented | independent
dynamics.normal_moment        # R, if conjugate-partner layout
dynamics.anomalous_moment     # C, for augmented layout
dynamics.moment_matrix        # R, R_tilde, or positive-P S
dynamics.normal_rhs           # dot(R)
dynamics.anomalous_rhs        # dot(C)
dynamics.rhs                  # real vectorized F(x; p)
dynamics.jacobian()           # dF/dx
dynamics.layout_decision      # symbolic evidence for auto selection

Automatic normal fallback

layout="auto" returns the smaller normal state only after symbolic checks prove that dot(R) has no anomalous coordinates and agrees with

dot(R) = -i(HR - RH^dagger) + D.

The fast proof checks V=0 and verifies that every phase-space factor needed by H and D maps exclusively to R. If that proof is unavailable, fpgen constructs the full R/C equations and checks their actual dependencies. layout="normal" raises CovarianceClosureError rather than discarding a required anomalous state.

Closure policy

closure="exact" accepts a state-independent Hamiltonian and diffusion terms whose Ito contribution can be expressed using second moments. It rejects a nonlinear Hamiltonian because its second-moment equation generally contains higher moments.

closure="factorized_bilinear" explicitly applies

alpha_i alpha_j       -> C_ij
alpha_i alpha_j*      -> R_ij
alpha_i* alpha_j      -> R_ji
alpha_i* alpha_j*     -> C_ij*

inside H_tilde and D_tilde before matrix multiplication. This is the nonlinear second-moment model used by the VDP and Kerr examples, but it is a declared factorization, not an exact Ito moment closure. The policy is stored in derivation provenance. Gaussian/Wick closure is never enabled implicitly.

State layouts

The normal real coordinate order remains:

1. diagonal R entries in declared mode order;
2. real parts of upper-triangle R entries in lexicographic index order;
3. imaginary parts in the same order.

The Hermitian augmented layout then appends real and imaginary parts of the upper triangle of the symmetric C, including its diagonal.

For positive-P, alpha and beta are independent. layout="auto" therefore selects a separate complex-symmetric state S=<z z^T> and never imposes the Wigner/P/Q conjugacy constraints.

to_covariance_dynamics() and to_augmented_second_moment_dynamics() remain available as explicit-layout convenience methods; new integrations should use the unified to_second_moment_dynamics() entry point.

Symbolic scalar reduction

CovarianceDynamics can identify stationary equations that are affine in a set of eliminated moments. For an explicitly chosen order parameter q:

search = dynamics.search_linear_reductions(
    retained_dimension=1,
    retained_ids=("r_diag_0",),
    equation_partitions="all",
    partition_limit=None,
    materialization_limit=None,
    return_limit=8,
)
candidate = search.candidates[0]
plan = dynamics.linear_reduce(candidate=candidate)
reduced = plan.materialize(method="fraction_free")

search.manifest()              # coverage, work counts, and truncation reasons
reduced.eliminated_solution
reduced.reduced_residual       # dot(q) = G(q; p)
reduced.numerators             # P(q; p)
reduced.denominators           # Q(q; p), with G = P/Q
reduced.multiplicity_conditions(4)

The eliminated stationary block has the form

A(q; p) y + b(q; p) = 0.

det A != 0 is only the regularity condition under which y can be eliminated. It is not the reduced dynamics. The actual scalar dynamics is dot(q) = G(q; p) = P(q; p) / Q(q; p). On a regular branch, an equilibrium of multiplicity m is characterized by the corresponding derivatives of P, together with the nonvanishing regularity and denominator conditions.

For larger models, callers may keep the exact condensed representation A y + b = 0 and the retained residual instead of expanding a very large rational expression.

The search only describes regular affine-elimination branches. A partition_limit stops work before additional symbolic blocks are built; return_limit only limits the ranked result prefix. These cases are distinct in ReductionSearchResult.coverage and truncation_reasons.

Downstream model interface

The symbolic moment dynamics can be exported as a versioned MomentDynamicsSpec. It records state and parameter order, provenance, sparsity, matrix semantics, physical-domain hints, and a reproducible fingerprint. compile_numpy() returns batch-aware NumPy reference callables for RHS, Jacobians, JVPs, VJPs, parameter derivatives, and complex state-matrix reconstruction.

Moment API 1.0, reduction API 1.0, and model schema 2.0 report moment_layout, explicit normal, anomalous, and independent state-index blocks, matrix_semantics, matrix_symmetry, and physical_domain_hint. Downstream code therefore does not infer matrix semantics from display names.

spec = dynamics.to_model_spec(name="two_mode_vdp")
compiled = spec.compile_numpy()

spec.state_index("r_diag_0")
spec.parameter_index("gamma_a")
spec.supports("state_matrix")
R = compiled.state_matrix(x, p)
F = compiled.rhs(x, p)
J = compiled.jacobian(x, p)

This interface is intended for validation and handoff to numerical projects. It does not make numerical bifurcation solving part of fpgen, and the public notebooks do not run downstream solvers.

Examples

• example.ipynb gives the complete symbolic derivation for a two-mode van der Pol oscillator, ending with the scalar reduced equation and formal high-order equilibrium conditions.
• models.ipynb is a compact gallery covering every maintained model: degenerate parametric oscillator, two-mode van der Pol, pair hopping, two-mode Kerr, three-mode Kerr, and cross-Kerr. Each entry displays the Hamiltonian, diffusion correlations, selected moment layout, Jacobian, retained order parameter, and reduced symbolic dynamics. The DPO entry also displays R_tilde, dot(R), and dot(C).

Notes

• The Langevin convention is Ito; Stratonovich conversion is not implemented.
• Wigner truncation can be approximate when the exact KM series has order greater than two. Inspect fpe.discarded and fpe.is_exact.
• Symbolic reduction describes regular elimination branches. Singular blocks require separate analysis.
• Nonlinear factorized_bilinear moment dynamics is an explicit modeling approximation; use closure="exact" when an exact Ito closure is required.
• Numerical continuation, root finding, certification, and GPU kernels are intentionally outside this repository.

License

Licensed under the MIT License. See LICENSE for details.