Add single-layer UCJ energy

[hkbelagali]

No description provided.

[CLAassistant]

All committers have signed the CLA.

[hkbelagali]

@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]

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

[hkbelagali]

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.

[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.

[hkbelagali]

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.

I was trying to match symbols from the paper's equations/lemmas in the code, but I'll change the names to be more descriptive.

[kevinsung]

the lint CI failure should be fixed after merging main