Add single-layer UCJ energy by hkbelagali · Pull Request #684 · qiskit-community/ffsim

hkbelagali · 2026-08-03T18:42:08Z

No description provided.

…ations yet

…nsatze

CLAassistant · 2026-08-03T18:42:17Z

All committers have signed the CLA.

…ing norb**4

…out in CI

hkbelagali · 2026-08-14T00:54:27Z

@kevinsung I was looking into why the test cases are getting stuck in CI, I believe it can be narrowed down to a deadlock in jax 0.10.2. This error does not occur on my laptop when I run the test cases, but if I restrict to 4 cores like the GitHub Actions runners, then I am able to reproduce the deadlock in the UCJ algorithm implementation. This code also produces the same deadlock on jax 0.10.2, but works fine with 0.9.2.

import time
import numpy as np
import jax
import jax.numpy as jnp

jax.config.update("jax_enable_x64", True)

n_calls = 16
batch = 1024
n  = 5 

rng = np.random.default_rng(0)
A = jnp.asarray(rng.normal(size=(batch, n, n)) + 1j * rng.normal(size=(batch, n, n)))

def f(t):
    return sum(jnp.real(jnp.sum(jnp.linalg.det(A * jnp.exp(1j * (t + k)))))
               for k in range(n_calls))
c = jax.jit(f).lower(0.3).compile()
t = time.time(); o = jax.block_until_ready(c(0.3))

print(f"OK exec {time.time()-t:.2f}s val={float(o):.3f}")

JAX_PLATFORMS=cpu taskset -c 0-3 python test.py

on jax==0.9.2, this prints OK exec 0.01s val=-3054.291, but it never finishes running on jax==0.10.2. I think this is because jax parallelized LAPACK operations in 0.10.0, and XLA's source code here has a comment on the safety of this. I believe the fix for this is also live on the jax main branch right now according to this PR. The deadlock disappeared when I used a nightly build of jax. Would it be possible to temporarily pin jax<0.10 until the next release comes out?

kevinsung · 2026-08-14T15:53:33Z

Sure, you can go ahead and edit pyproject.toml to restrict to working JAX versions.

hkbelagali · 2026-08-14T16:04:21Z

Sounds good, thanks!

kevinsung

@hkbelagali Thanks for the contribution! My first request is that you use the newly introduced rotate_one_body_tensor and rotate_two_body_tensor functions from https://github.com/qiskit-community/ffsim/blob/main/python/ffsim/linalg/util.py. I think these can replace the _propagate_through_orbital_rotations and _propagate_spin_sector_tensor functions you introduced here. Note, however, that the orbital rotation convention is transposed from your convention (please check this), so you either need to pass u.T.conj() everywhere, or rework your logic to align with the ffsim convention.

kevinsung

To keep things simple for now, let's get rid of the high-level dispatcher functions like ucj_energy, ucj_energy_and_grad, and optimize_ucj_energy and just force the user to use the appropriate function for their operator and Hamiltonian type.

…raction_pairs arg for energy calculation

kevinsung

I wonder if the variable names with single letters like q, h, and g can be made more descriptive. If you can't think of better names, it's fine though.

kevinsung · 2026-08-18T01:43:32Z

Needs a copyright header.

kevinsung · 2026-08-18T01:43:52Z

Needs a copyright header.

kevinsung · 2026-08-18T01:52:08Z

    "Topic :: Software Development :: Libraries :: Python Modules",
 ]
 dependencies = [
-    "jax",


jax 0.11.1 was just released, let's try reverting this

kevinsung · 2026-08-18T02:17:22Z

+def test_ucj_energy_spin_balanced_n2():
+    """Compare fermionic backpropagation against statevector simulation."""
+    mol = pyscf.gto.Mole()
+    mol.build(
+        atom=[["N", (0, 0, 0)], ["N", (0, 0, 1.0)]],
+        basis="sto-6g",
+        symmetry="Dooh",
+    )
+    scf = pyscf.scf.RHF(mol).run()
+
+    n_frozen = 2
+    active_space = range(n_frozen, mol.nao_nr())
+    ccsd = pyscf.cc.CCSD(
+        scf, frozen=[i for i in range(mol.nao_nr()) if i not in active_space]
+    ).run()
+
+    mol_data = ffsim.MolecularData.from_scf(scf, active_space=active_space)
+    mol_hamiltonian = mol_data.hamiltonian
+    norb = mol_data.norb
+    nelec = mol_data.nelec
+    assert norb == 8
+    assert nelec == (5, 5)
+
+    pairs_aa = [(p, p) for p in range(norb-1)]
+    pairs_ab = [(p, p) for p in range(norb-1)]
+
+    ucj_op = ffsim.UCJOpSpinBalanced.from_t_amplitudes(ccsd.t2, t1=ccsd.t1, interaction_pairs=(pairs_aa, pairs_ab), n_reps=1)
+
+    backprop_energy = ffsim.ucj_energy_spin_balanced(ucj_op, mol_hamiltonian, nelec)
+    statevector_energy = _statevector_energy(ucj_op, mol_hamiltonian, norb, nelec)
+
+    np.testing.assert_allclose(backprop_energy, statevector_energy)
+
+    optimized_ucj_op, result = ffsim.optimize_ucj_energy_spin_balanced(
+        ucj_op,
+        mol_hamiltonian,
+        nelec,
+        options={"maxiter": 1},
+        return_optimize_result=True,
+    )
+    optimized_backprop_energy = ffsim.ucj_energy_spin_balanced(
+        optimized_ucj_op, mol_hamiltonian, nelec
+    )
+    optimized_statevector_energy = _statevector_energy(
+        optimized_ucj_op, mol_hamiltonian, norb, nelec
+    )
+
+    np.testing.assert_allclose(optimized_backprop_energy, result.fun)
+    np.testing.assert_allclose(optimized_backprop_energy, optimized_statevector_energy)
+
+
+def test_ucj_energy_spin_unbalanced_n2():
+    """Compare spin-unbalanced backpropagation against statevector simulation."""
+    mol = pyscf.gto.Mole()
+    mol.build(
+        atom=[["N", (0, 0, 0)], ["N", (0, 0, 1.0)]],
+        basis="sto-6g",
+        symmetry="Dooh",
+    )
+    scf = pyscf.scf.RHF(mol).run()
+
+    n_frozen = 2
+    active_space = range(n_frozen, mol.nao_nr())
+    mol_data = ffsim.MolecularData.from_scf(scf, active_space=active_space)
+    mol_hamiltonian = mol_data.hamiltonian
+    norb = mol_data.norb
+    nelec = mol_data.nelec
+
+    ucj_op = ffsim.random.random_ucj_op_spin_unbalanced(
+        norb,
+        n_reps=1,
+        with_final_orbital_rotation=True,
+        diag_coulomb_scale=0.5,
+        seed=RNG,
+    )
+
+    backprop_energy = ffsim.ucj_energy_spin_unbalanced(ucj_op, mol_hamiltonian, nelec)
+    statevector_energy = _statevector_energy(ucj_op, mol_hamiltonian, norb, nelec)
+
+    np.testing.assert_allclose(backprop_energy, statevector_energy)
+
+
+def test_ucj_energy_spinless_n2():
+    """Compare spinless backpropagation against statevector simulation."""
+    mol = pyscf.gto.Mole()
+    mol.build(
+        atom=[["N", (0, 0, 0)], ["N", (0, 0, 1.0)]],
+        basis="sto-6g",
+        symmetry="Dooh",
+    )
+    scf = pyscf.scf.RHF(mol).run()
+
+    n_frozen = 2
+    active_space = range(n_frozen, mol.nao_nr())
+    mol_data = ffsim.MolecularData.from_scf(scf, active_space=active_space)
+    mol_hamiltonian = ffsim.MolecularHamiltonianSpinless(
+        one_body_tensor=mol_data.hamiltonian.one_body_tensor,
+        two_body_tensor=mol_data.hamiltonian.two_body_tensor,
+        constant=mol_data.hamiltonian.constant,
+    )
+    norb = mol_data.norb
+    nelec = mol_data.nelec[0]
+
+    ucj_op = ffsim.random.random_ucj_op_spinless(
+        norb,
+        n_reps=1,
+        with_final_orbital_rotation=True,
+        diag_coulomb_scale=0.5,
+        seed=RNG,
+    )
+
+    backprop_energy = ffsim.ucj_energy_spinless(ucj_op, mol_hamiltonian, nelec)
+    statevector_energy = _statevector_energy(ucj_op, mol_hamiltonian, norb, nelec)
+
+    np.testing.assert_allclose(backprop_energy, statevector_energy)
+
+
+def test_optimize_ucj_energy_spin_unbalanced():
+    """Compare spin-unbalanced optimization objective against statevector simulation."""
+    norb = 3
+    nelec = (1, 1)
+    mol_hamiltonian = ffsim.random.random_molecular_hamiltonian(norb, seed=RNG)
+    ucj_op = ffsim.random.random_ucj_op_spin_unbalanced(
+        norb,
+        n_reps=1,
+        with_final_orbital_rotation=True,
+        diag_coulomb_scale=0.5,
+        seed=RNG,
+    )
+
+    optimized_ucj_op, result = ffsim.optimize_ucj_energy_spin_unbalanced(
+        ucj_op,
+        mol_hamiltonian,
+        nelec,
+        options={"maxiter": 1},
+        return_optimize_result=True,
+    )
+    optimized_backprop_energy = ffsim.ucj_energy_spin_unbalanced(
+        optimized_ucj_op, mol_hamiltonian, nelec
+    )
+    optimized_statevector_energy = _statevector_energy(
+        optimized_ucj_op, mol_hamiltonian, norb, nelec
+    )
+
+    np.testing.assert_allclose(optimized_backprop_energy, result.fun)
+    np.testing.assert_allclose(optimized_backprop_energy, optimized_statevector_energy)
+
+
+def test_optimize_ucj_energy_spinless():
+    """Compare spinless optimization objective against statevector simulation."""
+    norb = 4
+    nelec = 2
+    mol_hamiltonian = ffsim.random.random_molecular_hamiltonian_spinless(norb, seed=RNG)
+    ucj_op = ffsim.random.random_ucj_op_spinless(
+        norb,
+        n_reps=1,
+        with_final_orbital_rotation=True,
+        diag_coulomb_scale=0.5,
+        seed=RNG,
+    )
+
+    optimized_ucj_op, result = ffsim.optimize_ucj_energy_spinless(
+        ucj_op,
+        mol_hamiltonian,
+        nelec,
+        options={"maxiter": 1},
+        return_optimize_result=True,
+    )
+    optimized_backprop_energy = ffsim.ucj_energy_spinless(
+        optimized_ucj_op, mol_hamiltonian, nelec
+    )
+    optimized_statevector_energy = _statevector_energy(
+        optimized_ucj_op, mol_hamiltonian, norb, nelec
+    )
+
+    np.testing.assert_allclose(optimized_backprop_energy, result.fun)
+    np.testing.assert_allclose(optimized_backprop_energy, optimized_statevector_energy)


Let's simplify these 5 tests into 3 tests that are more consistent. One test for each spin variant. Each test should

Generate a random UCJ operator (no need to use a molecule)

Compute the energy and compare it to the state vector reference

Optimize

Compute the energy and compare it to the state vector reference

Assert that the optimized energy is lower than the starting energy

Be parameterized for a few different values of interaction pairs, such as [], None, [(1, 2), (3, 4)]

These tests should exercise the default path of not specifying return_optimize_result. Then, add separate (cheaper) tests exercising the return_optimize_result=True path.

kevinsung · 2026-08-18T02:25:06Z

+    np.testing.assert_allclose(optimized_backprop_energy, optimized_statevector_energy)
+
+
+def test_ucj_energy_chunk_size_nondivisor():


Split this into three tests, one for each spin variant.

kevinsung · 2026-08-18T02:27:05Z

+    np.testing.assert_allclose(chunked, unchunked)
+
+
+def test_ucj_energy_and_grad():


Split this into three consistent tests, one for each spin variant. Extract a helper function to compute the finite difference gradient, and compute the whole gradient instead of just index 0. If this is too expensive, perhaps do it just for a few randomly chosen indices.

kevinsung · 2026-08-18T17:49:48Z

+    return jax.jit(jax.value_and_grad(energy, argnums=0))
+
+
+def _spin_balanced_jastrow_phase(same, diff, norb):


Please add type annotations to all function signatures, including the helper functions.

kevinsung · 2026-08-18T18:23:30Z

+def _interaction_pairs_key(
+    interaction_pairs: Sequence[tuple[int, int]],
+) -> tuple[tuple[int, int], ...]:
+    """Functools requires arguments to be hashable."""


Suggested change

"""Functools requires arguments to be hashable."""

"""functools.cache requires arguments to be hashable."""

and similarly elsewhere

kevinsung · 2026-08-18T19:57:51Z

+    )
+
+
+def ucj_energy_and_grad_spin_balanced(


I noticed that this function isn't used in the guide, and it's also not used in optimize_ucj_energy_spin_balanced. Let's replace both ucj_energy_and_grad_spin_balanced and optimize_ucj_energy_spin_balanced with a single function that returns the callable that can be passed to scipy. Something like

def ucj_energy_and_grad_func_spin_balanced( hamiltonian, nelec, *, norb, interaction_pairs=None, with_final_orbital_rotation=False, occupied_orbitals=None, chunk_size=None, ) -> Callable[[np.ndarray], tuple[float, np.ndarray]]:

Then, the guide should use this function to construct the callable, and pass it scipy.optimize.minimize. This gives the user more control over the optimization and lets us remove some more code from this file.

Same for the other spin variants.

hkbelagali added 11 commits July 27, 2026 22:42

Initial skeleton of the code, not sure if I can combine the implement…

85a23c7

…ations yet

Typing mistake in interaction_pairs

c513f80

Port fermionic backpropagation code

702ff04

Rewrite a docstring

7fe162b

Add the functions here

4bd55f1

Add statevector test case

2fb5f77

Add an initial guide here

d21e64b

Add fermionic backprop support for spinless and spin unbalanced UCJ a…

5601c3b

…nsatze

Inline some helpers to condense the file

104de0b

Compress helper methods

cf8b691

Add spin-unbalanced example to the guide

8cfa5d3

hkbelagali and others added 8 commits August 7, 2026 10:15

Merge branch 'main' into main

7a52ec5

Add term canonicalization, and remove restriction on chunk size divid…

2efbe17

…ing norb**4

Expose gradient helpers to be public

4cf3a44

Add random operator test cases for gradient helpers

ea0a5be

Change some stuff to pass mypy and ruff

8e2a697

Add the notebook here

ab1f899

Fix ruff in the notebook

edb4527

Change the header size in the notebook

43ef56d

hkbelagali marked this pull request as ready for review August 8, 2026 02:38

hkbelagali and others added 9 commits August 8, 2026 17:02

Merge branch 'main' into main

ba070ee

Merge branch 'main' into main

2e2e796

Add chunking to optimization loops and reduce maxiter to prevent time…

7264efe

…out in CI

Merge branch 'main' into main

639a038

Merge branch 'main' into main

0c8c317

Merge branch 'main' into main

95dc4f0

Forgot a comma to pass ruff

c097173

Merge branch 'main' into main

853d6f2

Merge branch 'main' into main

bd2346f

Merge branch 'main' into main

5129de0

Temporarily pin jax<0.10 to avoid parallelized LAPACK deadlocks

82f9a62

kevinsung reviewed Aug 14, 2026

View reviewed changes

Remove dispatcher functions

3746e20

kevinsung reviewed Aug 14, 2026

View reviewed changes

Comment thread python/ffsim/variational/ucj_energy.py Outdated

hkbelagali and others added 6 commits August 17, 2026 13:46

Remove references to dispatcher functions

b1dd639

Remove interaction_pairs parameter from here

9ede48f

Switch to ffsim linalg helpers for orbital rotations, and remove inte…

f91ecae

…raction_pairs arg for energy calculation

Test that interaction_pairs are handled properly

6a757a8

Merge branch 'main' into main

1022431

Fix linting issues

b4145ab

kevinsung reviewed Aug 18, 2026

View reviewed changes

		np.testing.assert_allclose(optimized_backprop_energy, optimized_statevector_energy)


		def test_ucj_energy_chunk_size_nondivisor():

		np.testing.assert_allclose(chunked, unchunked)


		def test_ucj_energy_and_grad():

		return jax.jit(jax.value_and_grad(energy, argnums=0))


		def _spin_balanced_jastrow_phase(same, diff, norb):

	"""Functools requires arguments to be hashable."""
	"""functools.cache requires arguments to be hashable."""

Uh oh!

Conversation

hkbelagali commented Aug 3, 2026

Uh oh!

CLAassistant commented Aug 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

hkbelagali commented Aug 14, 2026

Uh oh!

kevinsung commented Aug 14, 2026

Uh oh!

hkbelagali commented Aug 14, 2026

Uh oh!

kevinsung left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

kevinsung left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

kevinsung left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

CLAassistant commented Aug 3, 2026 •

edited

Loading

kevinsung left a comment •

edited

Loading