Solvation API¶

Configuration¶

Configuration dataclasses for MC solvation simulations.

Provides frozen (immutable) configuration for molecule, QPE parameters, and solvation simulation settings with the three-mode Hamiltonian framework (hf_corrected / fixed / dynamic).

class q2m3.solvation.config.MoleculeConfig(name, symbols, coords, charge, active_electrons=2, active_orbitals=2, basis='sto-3g')[source]¶

Bases: object

Molecular system configuration for solvation simulations.

Parameters:

name (str)
symbols (list[str])
coords (list[list[float]])
charge (int)
active_electrons (int)
active_orbitals (int)
basis (str)

name¶

Human-readable identifier

Type:: str

symbols¶

Atomic symbols (e.g., [“H”, “H”])

Type:: list[str]

coords¶

Atomic coordinates in Angstrom, shape (n_atoms, 3)

Type:: list[list[float]]

charge¶

Total molecular charge (REQUIRED — no default)

Type:: int

active_electrons¶

Active electrons for QPE active space

Type:: int

active_orbitals¶

Active orbitals for QPE active space

Type:: int

basis¶

Basis set name

Type:: str

property coords_array: ndarray¶: Coordinates as numpy array of shape (n_atoms, 3).

validate()[source]¶

Validate configuration consistency.

Return type:: None

class q2m3.solvation.config.QPEConfig(n_estimation_wires=4, n_trotter_steps=10, n_shots=0, device_seed=None, target_resolution=0.003, energy_range=0.2, qpe_interval=10)[source]¶

Bases: object

QPE (Quantum Phase Estimation) parameters.

Parameters:

n_estimation_wires (int)
n_trotter_steps (int)
n_shots (int)
device_seed (int | None)
target_resolution (float)
energy_range (float)
qpe_interval (int)

n_estimation_wires¶

Qubits for phase estimation register

Type:: int

n_trotter_steps¶

Trotter steps for time evolution

Type:: int

n_shots¶

Measurement shots (0 = analytical, >0 = shots)

Type:: int

device_seed¶

Optional seed for shot-based device sampling reproducibility

Type:: int | None

target_resolution¶

Target energy resolution in Hartree

Type:: float

energy_range¶

Energy range for shifted QPE in Hartree

Type:: float

qpe_interval¶

MC step interval between QPE evaluations (hf_corrected mode)

Type:: int

class q2m3.solvation.config.SolvationConfig(molecule, qpe_config=<factory>, hamiltonian_mode='dynamic', embedding_mode='diagonal', n_waters=10, n_mc_steps=500, temperature=300.0, translation_step=0.3, rotation_step=0.2618, initial_water_distance=4.0, random_seed=42, verbose=True, ir_cache_dir=None, ir_cache_enabled=True, ir_cache_force_recompile=False)[source]¶

Bases: object

Complete MC solvation simulation configuration.

The hamiltonian_mode field selects the energy evaluation strategy:

“hf_corrected”: E_HF + delta_corr_pol (fast, approximate)
“fixed”: Compile-once vacuum Hamiltonian (medium)
“dynamic”: Per-step MM-embedded Hamiltonian (slow, most complete)

The embedding_mode field selects the MM embedding operator boundary:

“diagonal”: Identity/Z coefficient update compatibility mode
“full_oneelectron”: Fixed-MO full one-electron operator support, available only with hamiltonian_mode=”fixed”

Parameters:

molecule (MoleculeConfig)
qpe_config (QPEConfig)
hamiltonian_mode (Literal['hf_corrected', 'fixed', 'dynamic'])
embedding_mode (Literal['diagonal', 'full_oneelectron'])
n_waters (int)
n_mc_steps (int)
temperature (float)
translation_step (float)
rotation_step (float)
initial_water_distance (float)
random_seed (int)
verbose (bool)
ir_cache_dir (str | None)
ir_cache_enabled (bool)
ir_cache_force_recompile (bool)

molecule¶

Molecular system configuration

Type:: q2m3.solvation.config.MoleculeConfig

qpe_config¶

QPE parameters

Type:: q2m3.solvation.config.QPEConfig

hamiltonian_mode¶

Energy evaluation strategy

Type:: Literal[‘hf_corrected’, ‘fixed’, ‘dynamic’]

embedding_mode¶

MM embedding mode for QPE Hamiltonian construction

Type:: Literal[‘diagonal’, ‘full_oneelectron’]

n_waters¶

Number of TIP3P water molecules

Type:: int

n_mc_steps¶

Total MC sampling steps

Type:: int

temperature¶

Simulation temperature in Kelvin

Type:: float

translation_step¶

Max translation per MC move in Angstrom

Type:: float

rotation_step¶

Max rotation per MC move in radians

Type:: float

initial_water_distance¶

Initial water ring radius in Angstrom

Type:: float

random_seed¶

Random seed for reproducibility

Type:: int

verbose¶

Print progress information

Type:: bool

ir_cache_dir¶

Directory for Catalyst LLVM IR cache, or None for the default cache path

Type:: str | None

ir_cache_enabled¶

Enable reuse of cached Catalyst LLVM IR

Type:: bool

ir_cache_force_recompile¶

Ignore cached IR and force recompilation

Type:: bool

property n_qpe_evaluations: int¶: Compute expected QPE evaluation count based on mode.

property kt: float¶: Thermal energy kT in Hartree.

validate()[source]¶

Validate configuration consistency.

Returns self to support chaining: config.validate()

Return type:: SolvationConfig

Orchestrator¶

MC Solvation Orchestrator.

End-to-end workflow for QPE-driven Monte Carlo solvation simulations. Supports three Hamiltonian modes:

hf_corrected: E_HF + E_MM for acceptance, interval-based QPE diagnostics
fixed: Compile-once vacuum Hamiltonian, QPE every step
dynamic: Per-step MM-embedded Hamiltonian, QPE every step

Usage:

from q2m3.solvation import SolvationConfig, MoleculeConfig, run_solvation

config = SolvationConfig(
    molecule=MoleculeConfig(name="H2", symbols=["H","H"],
                            coords=[[0,0,0],[0,0,0.74]], charge=0),
    hamiltonian_mode="fixed",
    n_waters=10,
    n_mc_steps=100,
)
result = run_solvation(config, show_plots=False)

q2m3.solvation.orchestrator.run_solvation(config, show_plots=True)[source]¶

Execute complete MC solvation workflow.

Workflow: 1. Validate config 2. Initialize solvent ring 3. Compute vacuum HF energy 4. Build QPE circuit (build_qpe_circuit -> QPECircuitBundle) 5. Precompute vacuum cache 6. Create step callback (mode-dependent) 7. Compute initial energy (first step_callback call -> triggers @qjit) 8. Create MC loop (create_mc_loop) 9. Execute MC loop 10. Compute Mulliken charges (diagnostic) 11. Print statistics + plot (if requested) 12. Return result dict

Parameters:

config (SolvationConfig) – Complete solvation simulation configuration
show_plots (bool) – Whether to display energy trajectory plots

Returns:

dict with Layer 1 (MCResult) + Layer 2 (orchestrator) fields

Return type:

dict[str, Any]

q2m3.solvation.orchestrator.replay_quantum_trajectory(config, trajectory_solvent_states)[source]¶

Replay fixed or dynamic quantum evaluations on a saved solvent trajectory.

Parameters:

config (SolvationConfig)
trajectory_solvent_states (ndarray)

Return type:

dict[str, Any]

Circuit Builder¶

QPE circuit builder for MC solvation simulations.

Builds parameterized @qjit-compiled QPE circuits that accept Hamiltonian coefficients as runtime arguments. Compile once, call with different coefficients each MC step — avoids 67s recompilation per step.

Key Catalyst constraints (POC-validated): - Use X gates instead of qml.BasisState (Catalyst @qjit + ctrl() bug) - TrotterProduct requires check_hermitian=False (JAX tracer incompatible) - MAX_TROTTER_STEPS_RUNTIME caps Trotter steps to avoid MLIR OOM - MSB-first qubit ordering (matching PennyLane QFT convention)

class q2m3.solvation.circuit_builder.QPECircuitBundle(compiled_circuit, base_coeffs, ops, base_time, op_index_map, energy_shift, n_estimation_wires, n_system_qubits, active_orbitals, n_trotter_steps, measurement_mode, is_fixed_circuit=False, embedding_mode='diagonal')[source]¶

Bases: object

All artifacts from QPE circuit compilation.

Parameters:

compiled_circuit (Callable)
base_coeffs (ndarray)
ops (list)
base_time (float)
op_index_map (dict)
energy_shift (float)
n_estimation_wires (int)
n_system_qubits (int)
active_orbitals (int)
n_trotter_steps (int)
measurement_mode (str)
is_fixed_circuit (bool)
embedding_mode (str)

compiled_circuit¶

@qjit function: () → result (fixed) or (coeffs_arr) → result (dynamic)

Type:: collections.abc.Callable

base_coeffs¶

Initial vacuum + shift coefficients

Type:: numpy.ndarray

ops¶

PennyLane operators (compile-time constants)

Type:: list

base_time¶

Evolution time for phase-to-energy conversion

Type:: float

op_index_map¶

{“identity_idx”: int, “z_wire_idx”: {wire: idx}}

Type:: dict

energy_shift¶

E_HF shift applied to Hamiltonian

Type:: float

n_estimation_wires¶

Number of QPE estimation qubits

Type:: int

n_system_qubits¶

Number of system qubits (= 2 * active_orbitals)

Type:: int

active_orbitals¶

Number of active spatial orbitals

Type:: int

n_trotter_steps¶

Actual Trotter steps (may be capped)

Type:: int

measurement_mode¶

“probs” or “shots”

Type:: str

is_fixed_circuit¶

True = zero-arg circuit (fixed mode); False = takes coeffs_arr

Type:: bool

embedding_mode¶

MM embedding mode used to build the fixed operator support

Type:: str

q2m3.solvation.circuit_builder.build_qpe_circuit(config, qm_coords, hf_energy, *, _keep_intermediate=False)[source]¶

Build parameterized QPE circuit accepting runtime coefficients.

The circuit builder constructs the base Hamiltonian, decomposes it into coefficients and PennyLane operators, applies the HF energy shift, derives shifted QPE parameters, caps runtime Trotter depth when needed, and returns a QPECircuitBundle. In diagonal mode the base Hamiltonian is vacuum operator support; in full_oneelectron mode it is a fixed MM operator.

Parameters:

config (SolvationConfig) – Solvation configuration
qm_coords (ndarray) – QM region coordinates in Angstrom, shape (n_atoms, 3)
hf_energy (float) – HF energy for energy shift and base_time calculation
_keep_intermediate (bool)

Returns:

QPECircuitBundle with compiled circuit and metadata

Return type:

QPECircuitBundle

q2m3.solvation.circuit_builder.decompose_hamiltonian(H)[source]¶

Extract (coeffs, ops) from PennyLane Sum/Hamiltonian.

PennyLane 0.44+ returns Sum of SProd from qml.qchem.molecular_hamiltonian. We need separate coefficients and operators for qml.dot(coeffs, ops).

Parameters:: H – PennyLane Hamiltonian (Sum of SProd operators)
Returns:: Tuple of (coefficients, operators) where coefficients are floats and operators are PennyLane operator instances.
Return type:: tuple[list[float], list]

q2m3.solvation.circuit_builder.build_operator_index_map(ops, n_system_qubits, coeffs)[source]¶

Scan ops to find Identity and single-Z(wire) indices for MM coefficient updates.

If any Z(wire) for wire in range(n_system_qubits) is missing, appends a coeff=0.0 placeholder to ensure all spin orbitals can receive MM corrections.

Parameters:

ops (list) – List of PennyLane operators (from decompose_hamiltonian)
n_system_qubits (int) – Number of system qubits (= 2 * active_orbitals)
coeffs (list[float]) – Corresponding coefficient list (may be extended)

Returns:

index_map: {“identity_idx”: int, “z_wire_idx”: {wire: idx, …}}
extended_coeffs: coefficients with any missing Z terms appended
extended_ops: operators with any missing Z terms appended

Return type:

Tuple of (index_map, extended_coeffs, extended_ops) where

Energy Callbacks¶

Energy computation module for MC solvation simulations.

Provides vacuum caching, three-mode coefficient callbacks, and step callback factories for QPE-driven Monte Carlo solvation energy evaluation.

Three Hamiltonian modes:

hf_corrected: HF energy for MC acceptance, interval-based QPE diagnostics
fixed: Compile-once vacuum coefficients, QPE every step
dynamic: Per-step MM-embedded Hamiltonian, QPE every step

All step callbacks are pure Python (not @qjit); they invoke pre-compiled @qjit circuits internally (ADR-003).

class q2m3.solvation.energy.StepResult(e_qpe, e_mm_sol_sol, e_hf_ref, callback_time, qpe_time)[source]¶

Bases: object

Result from one MC step’s energy computation.

Parameters:

e_qpe (float)
e_mm_sol_sol (float)
e_hf_ref (float)
callback_time (float)
qpe_time (float)

q2m3.solvation.energy.precompute_vacuum_cache(config)[source]¶

Run vacuum SCF once and cache results for MC loop.

Parameters:: config (SolvationConfig) – Solvation configuration with molecule definition
Returns:: h_core_vac, mo_coeff, mo_coeff_active, energy_nuc_vac, e_vacuum, active_idx
Return type:: Dictionary with cached vacuum SCF data

q2m3.solvation.energy.create_coefficient_callback(config, circuit_bundle, vacuum_cache)[source]¶

Factory: create coefficient update callback for the configured mode.

For “hf_corrected” / “fixed”: returns base_coeffs unchanged. For “dynamic”: runs solvated SCF, computes delta, patches coefficients.

Parameters:

config (SolvationConfig) – Solvation configuration
circuit_bundle (QPECircuitBundle) – QPE circuit bundle with base_coeffs and op_index_map
vacuum_cache (dict) – Pre-computed vacuum data

Returns:

(solvent_states, qm_coords_flat) -> np.ndarray[n_terms]

Return type:

Callable

q2m3.solvation.energy.create_step_callback(circuit_bundle, config, vacuum_cache)[source]¶

Factory: create per-step energy computation callback (fixed/dynamic modes).

Internal pipeline: 1. coefficient_callback → new_coeffs 2. compiled_circuit(jnp.array(new_coeffs)) → measurement_result 3. phase_extraction(measurement_result) → e_qpe 4. compute_mm_energy(solvents) → e_mm_sol_sol 5. Return StepResult

Parameters:

circuit_bundle (QPECircuitBundle) – QPE circuit bundle with compiled circuit
config (SolvationConfig) – Solvation configuration
vacuum_cache (dict) – Pre-computed vacuum data

Returns:

(solvent_states, qm_coords_flat) -> StepResult

Return type:

Callable

q2m3.solvation.energy.create_hf_corrected_step_callback(config, vacuum_cache, qm_coords, e_vacuum, circuit_bundle=None)[source]¶

Factory: hf_corrected mode step callback with deferred compilation.

Every step: PySCF QMMM → e_hf_ref; compute_mm_energy → e_mm. Every qpe_interval steps: additionally run QPE circuit for diagnostics.

If circuit_bundle is None, QPE circuit is lazily built on the first QPE step (deferred compilation — avoids startup compilation cost).

Non-QPE steps return StepResult(e_qpe=NaN, qpe_time=0.0).

Parameters:

config (SolvationConfig) – Solvation configuration
vacuum_cache (dict) – Pre-computed vacuum data
qm_coords (np.ndarray) – QM coordinates array for circuit building
e_vacuum (float) – Vacuum HF energy for energy shift
circuit_bundle (QPECircuitBundle | None) – Optional pre-built circuit bundle (None = deferred)

Returns:

(solvent_states, qm_coords_flat) -> StepResult

Return type:

Callable

q2m3.solvation.energy.compute_hf_energy_vacuum(molecule)[source]¶

Compute vacuum HF energy for the QM region.

Parameters:: molecule (MoleculeConfig) – Molecular system configuration
Returns:: Vacuum HF energy in Hartree
Return type:: float

q2m3.solvation.energy.compute_hf_energy_solvated(molecule, solvent_molecules)[source]¶

Compute solvated HF energy with MM electrostatic embedding.

Parameters:

molecule (MoleculeConfig) – QM region molecular configuration
solvent_molecules (Sequence[SolventMolecule]) – List of solvent molecules (MM region)

Returns:

Solvated HF energy in Hartree

Return type:

float

q2m3.solvation.energy.compute_mulliken_charges(molecule, solvent_states=None)[source]¶

Compute Mulliken charges for vacuum or solvated geometry.

Parameters:

molecule (MoleculeConfig) – QM region molecular configuration
solvent_states (ndarray | None) – Optional solvent state array (n_mol, 6); None for vacuum

Returns:

mulliken_charge}, e.g. {“O0”: -0.45, “H1”: +0.22}

Return type:

Dictionary {atom_label

Monte Carlo Loop¶

Monte Carlo loop for QPE-driven solvation simulations.

Implements a pure-Python Metropolis MC loop (ADR-003: no @qjit) that delegates energy evaluation to pre-compiled step callbacks.

Three Hamiltonian modes determine acceptance energy:

hf_corrected: E_HF_ref + E_MM(sol-sol) (fast, approximate)
fixed / dynamic: E_QPE + E_MM(sol-sol) (full quantum)

class q2m3.solvation.mc_loop.MCResult(initial_energy, final_energy, best_energy, best_qpe_energy, acceptance_rate, avg_energy, quantum_energies, hf_energies, callback_times, quantum_times, final_solvent_states, best_solvent_states, best_qpe_solvent_states, trajectory_solvent_states, n_quantum_evaluations, n_accepted)[source]¶

Bases: object

Complete MC simulation result.

Parameters:

initial_energy (float)
final_energy (float)
best_energy (float)
best_qpe_energy (float)
acceptance_rate (float)
avg_energy (float)
quantum_energies (ndarray)
hf_energies (ndarray)
callback_times (ndarray)
quantum_times (ndarray)
final_solvent_states (ndarray)
best_solvent_states (ndarray)
best_qpe_solvent_states (ndarray)
trajectory_solvent_states (ndarray)
n_quantum_evaluations (int)
n_accepted (int)

q2m3.solvation.mc_loop.create_mc_loop(config, step_callback)[source]¶

Factory: create QPE-driven MC loop (pure Python).

Returns a closure that runs the MC simulation with Metropolis acceptance.

Parameters:

config (SolvationConfig) – Solvation simulation configuration.
step_callback (Callable[[ndarray, ndarray], StepResult]) – Energy evaluator for each MC trial. Signature: (solvent_states, qm_coords) -> StepResult

Returns:

mc_loop(initial_solvents, qm_coords, seed, initial_energy) -> MCResult

Return type:

Callable[[…], MCResult]

Phase Extraction¶

Phase extraction utilities for QPE measurement results.

Pure NumPy implementation — no PennyLane, Catalyst, or JAX dependencies.

Two extraction modes:

Analytical (probs): probability-weighted expected bin, preserving sub-bin sensitivity (~0.1 mHa) compared to argmax (~12 mHa for 4-bit QPE).
Shots-based (samples): decode each bitstring to a bin index, then compute energy statistics across shots.

Core formula (shared):

phase = expected_bin / n_bins energy = -2π · phase / base_time + energy_shift

q2m3.solvation.phase_extraction.extract_energy_from_probs(probs, base_time, energy_shift, n_estimation_wires)[source]¶

Extract energy from analytical probability distribution.

Uses probability-weighted expected bin index for sub-bin sensitivity.

Parameters:

probs (ndarray) – Probability array of shape (2^n_estimation_wires,).
base_time (float) – Base evolution time for phase-to-energy conversion.
energy_shift (float) – Energy offset applied to the Hamiltonian.
n_estimation_wires (int) – Number of QPE estimation qubits.

Returns:

Estimated energy in Hartree.

Return type:

float

q2m3.solvation.phase_extraction.extract_energy_from_shots(samples, base_time, energy_shift, n_estimation_wires, return_statistics=False)[source]¶

Extract energy from shots-based measurement samples.

Each row of samples is a bitstring (MSB-first) that encodes a bin index. All per-shot energies are computed, then either the mean or a statistics dict is returned.

Parameters:

samples (ndarray) – Binary array of shape (n_shots, n_estimation_wires).
base_time (float) – Base evolution time for phase-to-energy conversion.
energy_shift (float) – Energy offset applied to the Hamiltonian.
n_estimation_wires (int) – Number of QPE estimation qubits.
return_statistics (bool) – If True, return dict with mean, std, sem, n_shots.

Returns:

Mean energy (float) or statistics dict.

Raises:

ValueError – If samples is empty.

Return type:

float | dict

Solvent Models¶

Explicit Solvent Models for QM/MM Simulations

This module provides data-driven solvent model definitions for use with PySCF’s qmmm module. It focuses on: 1. Solvent geometry and partial charges (for pyscf.qmmm.mm_charge) 2. Rigid body coordinate transformations 3. Classical MM energy for solvent-solvent interactions (MC acceptance)

Design Philosophy:

Data-driven: Solvent models defined as dataclasses, easy to extend
PySCF-compatible: Output coordinates in Angstrom for qmmm.mm_charge
Minimal: Only implement what PySCF doesn’t provide directly

Usage with PySCF:

from pyscf import qmmm

model = TIP3P_WATER molecules = [SolventMolecule(model, pos, angles) for …] mm_coords, mm_charges = get_mm_embedding_data(molecules)

# Convert to Bohr for PySCF mm_coords_bohr = mm_coords * ANGSTROM_TO_BOHR mf = qmmm.mm_charge(mf, mm_coords_bohr, mm_charges)

class q2m3.solvation.solvent.SolventAtom(symbol, local_coord, charge, lj_sigma=0.0, lj_epsilon=0.0)[source]¶

Bases: object

Single atom in a solvent model.

Parameters:

symbol (str)
local_coord (tuple[float, float, float])
charge (float)
lj_sigma (float)
lj_epsilon (float)

symbol¶

Atomic symbol (e.g., “O”, “H”)

Type:: str

local_coord¶

Position in molecule-fixed frame (Angstrom)

Type:: tuple[float, float, float]

charge¶

Partial charge (elementary charge units)

Type:: float

lj_sigma¶

Lennard-Jones sigma parameter (Angstrom), 0 if no LJ

Type:: float

lj_epsilon¶

Lennard-Jones epsilon parameter (kcal/mol), 0 if no LJ

Type:: float

class q2m3.solvation.solvent.SolventModel(name, atoms, reference='')[source]¶

Bases: object

Solvent model definition containing geometry, charges, and LJ parameters.

This is an immutable data structure that defines a solvent type. Actual solvent molecule instances are represented by SolventMolecule.

Parameters:

name (str)
atoms (tuple[SolventAtom, ...])
reference (str)

name¶

Model identifier (e.g., “TIP3P”, “SPC/E”)

Type:: str

atoms¶

Tuple of SolventAtom defining the molecule structure

Type:: tuple[q2m3.solvation.solvent.SolventAtom, …]

reference¶

Literature reference for the model parameters

Type:: str

property local_coords: ndarray¶: Local coordinates as (n_atoms, 3) array.

get_lj_params(atom_idx)[source]¶

Get (sigma, epsilon) for atom at given index.

Parameters:: atom_idx (int)
Return type:: tuple[float, float]

class q2m3.solvation.solvent.SolventMolecule(model, position=<factory>, euler_angles=<factory>)[source]¶

Bases: object

A solvent molecule instance with position and orientation.

Parameters:

model (SolventModel)
position (ndarray)
euler_angles (ndarray)

model¶

The solvent model (defines geometry and parameters)

Type:: q2m3.solvation.solvent.SolventModel

position¶

Center of mass position in Angstrom

Type:: numpy.ndarray

euler_angles¶

Orientation as (roll, pitch, yaw) in radians

Type:: numpy.ndarray

get_atom_coords()[source]¶

Get global coordinates of all atoms.

Returns:: Array of shape (n_atoms, 3) in Angstrom
Return type:: ndarray

get_charges()[source]¶

Get partial charges for all atoms.

Return type:: ndarray

to_state_vector()[source]¶

Convert to state vector [x, y, z, roll, pitch, yaw].

Return type:: ndarray

classmethod from_state_vector(model, state)[source]¶

Create molecule from state vector.

Parameters:

model (SolventModel)
state (ndarray)

Return type:

SolventMolecule

q2m3.solvation.solvent.get_mm_embedding_data(molecules)[source]¶

Extract MM embedding data for pyscf.qmmm.mm_charge().

Parameters:

molecules (Sequence[SolventMolecule]) – Sequence of SolventMolecule instances

Returns:

coords: Atom coordinates in Angstrom, shape (n_total_atoms, 3)
charges: Partial charges, shape (n_total_atoms,)

Return type:

Tuple of

Usage with PySCF:: coords, charges = get_mm_embedding_data(molecules) coords_bohr = coords * ANGSTROM_TO_BOHR mf = qmmm.mm_charge(mf, coords_bohr, charges)

q2m3.solvation.solvent.get_mm_embedding_data_bohr(molecules)[source]¶

Get MM embedding data with coordinates in Bohr (ready for PySCF).

This is a convenience function that converts coordinates to Bohr, which is the unit expected by pyscf.qmmm.mm_charge().

Returns:: Tuple of (coords_bohr, charges)
Parameters:: molecules (Sequence[SolventMolecule])
Return type:: tuple[ndarray, ndarray]

q2m3.solvation.solvent.compute_mm_energy(molecules)[source]¶

Compute classical MM energy between solvent molecules.

This calculates the Lennard-Jones and Coulomb interactions between all pairs of solvent molecules. Used for the MC acceptance criterion.

Note: This is a simple pairwise implementation. For production use with large systems, consider using OpenMM or other optimized libraries.

Parameters:: molecules (Sequence[SolventMolecule]) – Sequence of SolventMolecule instances
Returns:: Total MM energy in Hartree. Exact intermolecular atom overlap is treated as a hard-core invalid configuration and returns +inf.
Return type:: float

q2m3.solvation.solvent.initialize_solvent_ring(model, n_molecules, center, radius, random_seed=42)[source]¶

Initialize solvent molecules in a ring arrangement.

Molecules are placed equidistant in a ring in the x-y plane, with random initial orientations.

Parameters:

model (SolventModel) – Solvent model to use
n_molecules (int) – Number of molecules
center (ndarray) – Center of the ring (typically solute center)
radius (float) – Ring radius in Angstrom
random_seed (int) – Random seed for reproducible orientations

Returns:

List of SolventMolecule instances

Return type:

list[SolventMolecule]

q2m3.solvation.solvent.molecules_to_state_array(molecules)[source]¶

Convert list of molecules to state array of shape (n_mol, 6).

Parameters:: molecules (Sequence[SolventMolecule])
Return type:: ndarray

q2m3.solvation.solvent.state_array_to_molecules(model, states)[source]¶

Convert state array to list of molecules.

Parameters:

model (SolventModel)
states (ndarray)

Return type:

list[SolventMolecule]

IR Cache¶

Catalyst LLVM IR cache for QPE circuit compilation.

Implements two-phase compilation caching:: Phase A (cache miss): Subprocess full @qjit(keep_intermediate=True) -> save LLVM IR Phase B (cache hit): replace_ir + jit_compile -> skip MLIR pipeline (~5x faster)

Cache key = circuit topology parameters only. Runtime Hamiltonian coefficients do not affect LLVM IR structure, so IR compiled for n_waters=5 is valid for n_waters=20. Fixed full-one-electron embedding can change operator support, so embedding_mode is part of the key.

q2m3.solvation.ir_cache.default_cache_dir()[source]¶

Find project root (containing pyproject.toml) and return cache dir.

Return type:: Path

q2m3.solvation.ir_cache.compute_cache_key(config)[source]¶

Generate deterministic cache key from circuit topology parameters.

The key encodes parameters that affect LLVM IR structure: - Molecule identity (name, basis, active space) - QPE circuit parameters (estimation wires, Trotter steps, shots) - Hamiltonian circuit style and embedding operator support

Parameters that only affect runtime coefficient values (n_waters, temperature, coordinates) are NOT included.

Parameters:: config (SolvationConfig)
Return type:: str

q2m3.solvation.ir_cache.cache_path_for_config(config)[source]¶

Return the full path for the cached IR file.

Parameters:: config (SolvationConfig)
Return type:: Path

q2m3.solvation.ir_cache.is_cache_available(config)[source]¶

Check if a valid cache file exists for this config.

Parameters:: config (SolvationConfig)
Return type:: bool

q2m3.solvation.ir_cache.resolve_compiled_circuit(config, bundle)[source]¶

Cache-aware QPE circuit compilation resolution.

On cache hit: replace_ir + jit_compile (seconds) On cache miss: subprocess full compile -> save IR -> replace_ir + jit_compile

The bundle’s compiled_circuit is modified in-place by replace_ir.

Parameters:

config (SolvationConfig) – Solvation configuration with cache settings
bundle (QPECircuitBundle) – QPECircuitBundle from build_qpe_circuit()

Returns:

(bundle, cache_stats) — bundle is the same object (replace_ir modifies in-place)

Return type:

tuple[QPECircuitBundle, dict[str, Any]]

Analysis¶

Solvation analysis functions for δ_corr-pol and related statistics.

Pure NumPy module — no Catalyst or JAX dependency. All public functions are pure functions: take NumPy arrays, return frozen dataclasses.

class q2m3.solvation.analysis.DeltaCorrPolResult(mean_ha, std_ha, sem_ha, t_statistic, n_samples, is_significant, per_step_delta)[source]¶

Bases: object

δ_corr-pol analysis result.

Note: frozen=True prevents field reassignment but does NOT prevent mutation of mutable fields like np.ndarray (e.g., obj.per_step_delta[0] = 999 is still possible). This is an accepted trade-off: frozen provides intent-level immutability and hashability for simple fields, while ndarray copy-on-read is not enforced to avoid performance overhead on large arrays. Consumers should treat all fields as read-only.

Parameters:

mean_ha (float)
std_ha (float)
sem_ha (float)
t_statistic (float)
n_samples (int)
is_significant (bool)
per_step_delta (ndarray)

class q2m3.solvation.analysis.EnergyPhaseResult(early_mean, early_std, late_mean, late_std, n_per_phase)[source]¶

Bases: object

Energy trend phase analysis result (early vs late phase).

Parameters:

early_mean (float)
early_std (float)
late_mean (float)
late_std (float)
n_per_phase (int)

class q2m3.solvation.analysis.QPEHFConsistencyResult(mean_offset_mha, pearson_correlation, early_offset_mha, late_offset_mha, offset_drift_mha, n_samples)[source]¶

Bases: object

QPE-HF trajectory consistency metrics.

Parameters:

mean_offset_mha (float)
pearson_correlation (float)
early_offset_mha (float)
late_offset_mha (float)
offset_drift_mha (float)
n_samples (int)

class q2m3.solvation.analysis.EquilibrationResult(frac_decreasing, n_windows, qpe_drift_mha, is_monotonic)[source]¶

Bases: object

Equilibration diagnostic result (windowed monotonicity).

Parameters:

frac_decreasing (float)
n_windows (int)
qpe_drift_mha (float)
is_monotonic (bool)

class q2m3.solvation.analysis.ModeComparisonResult(delta_corr_pol, energy_phases, qpe_hf_consistency, equilibration, energy_distributions, trotter_bias_mha)[source]¶

Bases: object

Full three-mode comparison analysis result.

Parameters:

delta_corr_pol (DeltaCorrPolResult)
energy_phases (dict[str, EnergyPhaseResult])
qpe_hf_consistency (QPEHFConsistencyResult | None)
equilibration (EquilibrationResult | None)
energy_distributions (dict[str, dict])
trotter_bias_mha (float)

q2m3.solvation.analysis.compute_delta_corr_pol(quantum_energies_fixed, quantum_energies_dynamic)[source]¶

Compute per-step δ_corr-pol from paired fixed/dynamic QPE energies.

δ_corr-pol(R) = E_QPE(dynamic, R) - E_QPE(fixed, R): ≈ E_corr(solvated, R) - E_corr(vacuum)

Parameters:

quantum_energies_fixed (ndarray) – QPE energies from fixed Hamiltonian mode.
quantum_energies_dynamic (ndarray) – QPE energies from dynamic Hamiltonian mode. Must be the same length as quantum_energies_fixed.

Returns:

DeltaCorrPolResult with paired-sample statistics.

Return type:

DeltaCorrPolResult

q2m3.solvation.analysis.analyze_energy_phases(quantum_energies, n_phases=3)[source]¶

Analyze early/late energy trends across MC sampling phases.

Splits the energy trajectory into n_phases equal segments and compares the first (early) and last (late) segments.

Parameters:

quantum_energies (ndarray) – Energy trajectory from QPE-driven MC simulation.
n_phases (int) – Number of phases to divide the trajectory into.

Returns:

EnergyPhaseResult with mean and std for early and late phases.

Return type:

EnergyPhaseResult

q2m3.solvation.analysis.compute_qpe_hf_consistency(quantum_energies, hf_energies)[source]¶

Compute QPE-HF trajectory consistency metrics.

Measures how well the QPE and HF energy trajectories track each other. A positive, high Pearson correlation indicates consistent trends.

Parameters:

quantum_energies (ndarray) – QPE energy trajectory.
hf_energies (ndarray) – HF energy trajectory (same length as quantum_energies).

Returns:

QPEHFConsistencyResult with offset and correlation statistics.

Return type:

QPEHFConsistencyResult

q2m3.solvation.analysis.detect_equilibration(quantum_energies, min_samples=10)[source]¶

Windowed monotonicity diagnostic for MC equilibration.

Splits energies into 5 equal windows and checks if window means are monotonically decreasing (suggesting non-equilibration).

Parameters:

quantum_energies (ndarray) – QPE energy trajectory.
min_samples (int) – Minimum number of samples required. Returns None if len(quantum_energies) < min_samples.

Returns:

EquilibrationResult, or None if insufficient samples.

Return type:

EquilibrationResult | None

q2m3.solvation.analysis.run_mode_comparison(result_fixed, result_dynamic, result_hf_corrected=None, e_vacuum=0.0)[source]¶

Orchestrate full three-mode comparison analysis.

Combines lower-level analysis functions to produce a complete ModeComparisonResult from run_solvation() return dicts.

Parameters:

result_fixed (dict) – run_solvation() result dict for fixed Hamiltonian mode.
result_dynamic (dict) – run_solvation() result dict for dynamic Hamiltonian mode.
result_hf_corrected (dict | None) – Optional result dict for hf_corrected mode.
e_vacuum (float) – Vacuum reference energy (Ha) for Trotter bias calculation.

Returns:

ModeComparisonResult with all analysis fields populated.

Raises:

ValueError – If any result dict is missing required keys.

Return type:

ModeComparisonResult

Structure Analysis¶

Internal structure-analysis helpers for finite-shell solvation outputs.

The routines in this module are pure NumPy/CSV helpers around the solvent state arrays already returned by q2m3.solvation.run_solvation(). They are kept out of q2m3.solvation.__init__ because they support optional post-processing workflows rather than the primary solvation API.

class q2m3.solvation.structure_analysis.RadialProfile(bin_edges, counts, shell_volumes, density_per_a3, mean_count_per_frame, cumulative_coordination, n_frames)[source]¶

Bases: object

Finite-shell oxygen radial-density profile around a solute center.

Parameters:

bin_edges (ndarray)
counts (ndarray)
shell_volumes (ndarray)
density_per_a3 (ndarray)
mean_count_per_frame (ndarray)
cumulative_coordination (ndarray)
n_frames (int)

property bin_centers: ndarray¶: Bin centers in Angstrom.

q2m3.solvation.structure_analysis.validate_state_trajectory(states)[source]¶

Return a float trajectory array with shape (n_steps, n_waters, 6).

Parameters:: states (ndarray)
Return type:: ndarray

q2m3.solvation.structure_analysis.write_state_trajectory_csv(path, states)[source]¶

Write solvent states as a long-form CSV for deterministic round-trips.

Parameters:

path (Path)
states (ndarray)

Return type:

None

q2m3.solvation.structure_analysis.load_state_trajectory_csv(path)[source]¶

Load a solvent-state trajectory CSV.

The preferred schema is long-form step,water_index,x,y,z,roll,pitch,yaw. A compact wide schema with columns like water_0_x or w0_x is also accepted for compatibility with ad hoc notebook exports.

Parameters:: path (Path)
Return type:: ndarray

q2m3.solvation.structure_analysis.solute_center(solute_coords)[source]¶

Return the arithmetic center of solute coordinates.

Parameters:: solute_coords (ndarray)
Return type:: ndarray

q2m3.solvation.structure_analysis.water_oxygen_distances(states, solute_coords)[source]¶

Return oxygen-center distances to the solute center for every frame/water.

Parameters:

states (ndarray)
solute_coords (ndarray)

Return type:

ndarray

q2m3.solvation.structure_analysis.radial_density_profile(states, solute_coords, bin_edges, *, start=0)[source]¶

Compute finite-shell oxygen radial density and cumulative coordination.

Parameters:

states (ndarray)
solute_coords (ndarray)
bin_edges (ndarray)
start (int)

Return type:

RadialProfile

q2m3.solvation.structure_analysis.shell_counts_per_frame(states, solute_coords, cutoff, *, start=0)[source]¶

Return per-frame water counts within cutoff Angstrom.

Parameters:

states (ndarray)
solute_coords (ndarray)
cutoff (float)
start (int)

Return type:

ndarray

q2m3.solvation.structure_analysis.coordination_by_cutoff(states, solute_coords, cutoffs, *, start=0)[source]¶

Return mean coordination count for each cutoff.

Parameters:

states (ndarray)
solute_coords (ndarray)
cutoffs (ndarray)
start (int)

Return type:

ndarray

q2m3.solvation.structure_analysis.select_representative_snapshot_indices(states, solute_coords, *, n_snapshots=3, cutoff=3.5, start=0)[source]¶

Select deterministic post-burn snapshots spanning shell occupancy.

Parameters:

states (ndarray)
solute_coords (ndarray)
n_snapshots (int)
cutoff (float)
start (int)

Return type:

list[int]

q2m3.solvation.structure_analysis.select_random_snapshot_indices(states, *, n_snapshots=3, start=0, stop=None, random_seed=42)[source]¶

Select reproducible random snapshot indices from a trajectory window.

Parameters:

states (ndarray)
n_snapshots (int)
start (int)
stop (int | None)
random_seed (int)

Return type:

list[int]

q2m3.solvation.structure_analysis.state_to_symbols_coords(solute_symbols, solute_coords, solvent_state)[source]¶

Convert one solvent-state frame to combined solute + TIP3P atom arrays.

Parameters:

solute_symbols (list[str])
solute_coords (ndarray)
solvent_state (ndarray)

Return type:

tuple[list[str], ndarray]

q2m3.solvation.structure_analysis.write_xyz_snapshot(path, solute_symbols, solute_coords, solvent_state, *, comment='')[source]¶

Write one solute + solvent-shell frame as an XYZ file.

Parameters:

path (Path)
solute_symbols (list[str])
solute_coords (ndarray)
solvent_state (ndarray)
comment (str)

Return type:

None

q2m3.solvation.structure_analysis.xyz_atom_count(path)[source]¶

Return the atom count declared by an XYZ file.

Parameters:: path (Path)
Return type:: int

q2m3.solvation.structure_analysis.radial_profile_to_rows(profile)[source]¶

Convert a radial profile to CSV-ready dictionaries.

Parameters:: profile (RadialProfile)
Return type:: list[dict[str, Any]]

Statistics And Plotting¶

Time statistics formatting for MC solvation simulations.

Uses Rich library for console output with tables and panels. Supports three Hamiltonian modes: hf_corrected / fixed / dynamic.

class q2m3.solvation.statistics.TimingData(quantum_compile_time=0.0, mc_loop_time=0.0, hf_times=<factory>, quantum_times=<factory>, n_mc_steps=0, n_quantum_evals=0, hamiltonian_mode='')[source]¶

Bases: object

Container for simulation timing data.

Parameters:

quantum_compile_time (float)
mc_loop_time (float)
hf_times (ndarray)
quantum_times (ndarray)
n_mc_steps (int)
n_quantum_evals (int)
hamiltonian_mode (str)

q2m3.solvation.statistics.create_timing_table(timing)[source]¶

Create Rich table for execution phase timing.

Parameters:: timing (TimingData)
Return type:: Table

q2m3.solvation.statistics.print_time_statistics(timing, console=None)[source]¶

Print formatted timing statistics to console.

Parameters:

timing (TimingData) – TimingData instance
console (Console | None) – Rich Console (creates new if None)

Return type:

None

q2m3.solvation.statistics.create_timing_data_from_result(result, compile_time, loop_time, hamiltonian_mode='')[source]¶

Create TimingData from MC loop result dictionary.

Parameters:

result (dict)
compile_time (float)
loop_time (float)
hamiltonian_mode (str)

Return type:

TimingData

Plotting utilities for MC solvation simulations.

Provides energy trajectory visualization and acceptance rate analysis plots.

q2m3.solvation.plotting.plot_energy_trajectory(mc_steps, hf_energies, quantum_steps=None, quantum_energies=None, reference_energy=None, title='Energy Trajectory', output_path=None, show=True)[source]¶

Plot energy trajectory over MC steps.

Parameters:

mc_steps (Sequence[int]) – MC step indices for HF energies
hf_energies (Sequence[float]) – HF (or total QM/MM) energies at each step
quantum_steps (Sequence[int] | None) – MC step indices where quantum evaluation was performed
quantum_energies (Sequence[float] | None) – Quantum algorithm energy estimates
reference_energy (float | None) – Reference energy for relative plotting (e.g., vacuum energy)
title (str) – Plot title
output_path (str | Path | None) – Path to save figure (None = don’t save)
show (bool) – Whether to display the plot

Returns:

Matplotlib Figure object

Return type:

Figure

q2m3.solvation.plotting.plot_acceptance_rate(mc_steps, cumulative_accepted, window_size=50, title='MC Acceptance Rate', output_path=None, show=True)[source]¶

Plot cumulative and windowed acceptance rate.

Parameters:

mc_steps (Sequence[int]) – MC step indices
cumulative_accepted (Sequence[int]) – Cumulative number of accepted moves
window_size (int) – Window size for rolling acceptance rate
title (str) – Plot title
output_path (str | Path | None) – Path to save figure
show (bool) – Whether to display

Returns:

Matplotlib Figure object

Return type:

Figure