This chapter is a long-form FAQ for things that have actually broken during development of the package, in roughly the order they’re likely to surface. If your problem isn’t here, the test suite is the second place to look — every public function has at least one test, and the tests are short enough to read in a few minutes.
jnp.linalg.eigh complains about complex numbersYou probably have float32 JAX. surfacehop_jax enables float64 on
import via
import jax
jax.config.update("jax_enable_x64", True)
in the package’s __init__.py. If something else in your code (or an
import order quirk) disables it after we’ve enabled it, eigenvalue
problems on the electronic generator can underflow / lose precision.
Put import surfacehop_jax early, ideally before any other JAX-using
imports.
ImportError: cannot import name 'jaxlib'pip resolved jax to a newer version than jaxlib can supply on
your platform. Pinning both versions explicitly usually fixes it:
pip install "jax==0.4.30" "jaxlib==0.4.30"
Then reinstall surfacehop_jax on top.
import jax
print(jax.devices()) # [CpuDevice(id=0)] when you wanted CUDA
Most common cause: a CUDA-toolkit/driver mismatch. JAX’s install page
has a compatibility table. Less common: you installed both the CPU and
CUDA wheels and the CPU one was picked. pip uninstall jax jaxlib and
reinstall just the CUDA wheel.
A few likely causes:
dt.history.total_energy[-1] - history.total_energy[0]
over the whole trajectory, not differences across single steps near
frustrated hops.H(x) is asymmetric,
jnp.linalg.eigh returns real eigenvalues anyway (it implicitly
symmetrises), but the energy you compute via the eigenvectors
doesn’t match what the propagator integrates. Make $H$ Hermitian.The hop probability g is being computed correctly, but always coming
out near zero. The probable cause is the velocity used in vdotd:
if v is essentially zero (e.g., the trajectory started at rest at
the FC point in a high-symmetry coordinate system), the
v·d_{ij} term in the generator is zero and population never
transfers off the initial state, so no hops happen. Either start with
non-zero velocity (Wigner sample with p0 != 0, or perturb), or
verify that the Hamiltonian actually has an inter-state coupling at
the initial geometry.
The trajectory walked too close to a conical intersection and eigh
returned ill-defined eigenvectors. Mitigations:
(x, v, energies) every few steps
near the divergence and identify the step where it goes wrong.Plots of the eigenvector components or NAC vectors show sudden sign
flips between steps. The propagator’s internal phase tracking
(fix_eigenvector_phase) keeps phases consistent step-to-step
inside the trajectory, but it doesn’t fix global orientation. If you
plot raw NAC components across a long trajectory you’ll see flips
where the phase tracking was indifferent (e.g., near an avoided
crossing where eigenvector continuity isn’t unique).
The trajectory itself is fine — observables and energies are
sign-invariant. If you’re plotting NAC components for diagnostics,
plot |NAC| instead.
Either you don’t have enough trajectories ($N \lesssim 100$ for a
photochemistry problem gives visibly noisy curves), or you’re computing
the active-state fraction of a poorly-mixed problem (the fraction is
inherently noisier than $\langle|c|^2\rangle$ for small $N$). Use more
trajectories or plot population instead of active_state.
This is FSSH over-coherence, not a bug — it’s the standard symptom near a conical intersection. Switch on the Granucci–Persico EDC (Chapter 8).
If the curves disagree even with EDC on, the implementation may have a real bug. Sanity checks:
tests/
Tully-1 test verifies this).If either fails, report it; we’d want to dig in.
run_ensemble fails with shape errorThe most common message is something like
TypeError: dot_general requires contracting dimensions to have ...
or similar. The cause is almost always init_states doesn’t have a
leading n_traj axis on every field. Check:
print(init.x.shape) # should be (n_traj, ndim)
print(init.coeffs.shape) # should be (n_traj, nel) complex
print(init.state.shape) # should be (n_traj,)
The correct way to build this is jax.vmap(initialize, in_axes=...)
over the Wigner-sampled (Q0, V0) arrays; see Chapter 7.
ConcretizationTypeError when building LinearVibronicCouplingIn versions ≤ 1.0 the LVC __post_init__ symmetry check called
float(jnp.max(...)) on the coupling tensor, which fails when the
tensor is a JAX tracer (i.e., we’re inside jax.grad or jax.jit).
Fixed in v1.1: the check is wrapped in try / except
jax.errors.ConcretizationTypeError so it’s a no-op during tracing.
You’re now responsible for passing a symmetric coupling tensor — the
standard pattern is to set coupling[i, j] and coupling[j, i]
together.
If you still see this error, you’re running an old version (pip show
surfacehop_jax). Upgrade.
jax.grad returns zeroYou’re differentiating a step-function observable — almost always the
active-state fraction
(active_state == i).mean(). Active-state fractions involve discrete
hop decisions, so their derivative is zero almost everywhere. Use a
smooth observable like jnp.mean(jnp.abs(coeffs[:, i])**2) instead.
jax.grad returns NaNMost often: the trajectory hit a singular point (an exact CoIn) or a
0/0 ratio in the momentum-rescaling routine. The latter is guarded
against in the implementation (the “double-where” pattern around the
discriminant), so check the former: print the energies at every step
and see if any pair is anomalously close.
If the NaN persists, the culprit is sometimes the matrix exponential on a generator with very large $|G|\,\Delta t$. Reduce $\Delta t$.
Run the sanity check:
g = jax.grad(loss)(theta0)
eps = 1e-3
g_fd = (loss(theta0 + eps) - loss(theta0 - eps)) / (2 * eps)
If they disagree by more than ~1%, common causes are:
| Non-smooth observable: use $ | c | ^2$, not active-state fraction. |
jnp.sign(x) instead of two-branch
jnp.where; jnp.heaviside; Python if clauses depending on the
coordinate.First call to a function takes minutes. The trace got large — common causes:
n_steps or n_traj is very large at first call. Start with a
smaller call to warm up the cache, then scale up.vmap of a scan. This is
fine but the trace is deeper. Wrap with jax.jit so the trace cost
is paid once.H is a neural
network with millions of parameters, jacrev will dominate trace
time. Consider whether you really need the full Jacobian or could
use a forward-mode jacfwd instead.If the curve is flipped (high momentum stays on the lower adiabat,
low momentum hops to the upper), the sign of the b term in the hop
probability is wrong. The correct formula is
$b = +2\,\mathrm{Re}(\rho_{ij}^{*}\,(v \cdot d_{ij}))$. The
tests/test_tully1_benchmark.py test catches this.
A few possible reasons:
pyrazine_4mode() uses the MCTDH-Heidelberg tutorial
set. If you’re comparing to a paper using a different set, expect
quantitative differences.test_wigner.py fails after a refactorThe most likely cause is a regression of the mass-factor bug discussed in Chapter 7. The widths $(\sigma_q, \sigma_p)$ should each contain a factor of $\sqrt{m}$. The test sanity-checks empirical against analytic widths to within a few percent.
This package is a research tool, and the test suite and documentation are the primary support resources. If you find a real bug, the maintainer welcomes a minimum-working-example issue or PR.
Things that are in scope to fix:
ConcretizationTypeError, a NaN gradient
in a place it shouldn’t).Things that aren’t in scope:
decoherence.py are the right pattern.