Orbital optimization: Reshape params outside of JAX

[kevinsung]

This improves performance but something is wrong. The optimization with complex numbers gives a worse result than before.

>       np.testing.assert_allclose(energy, -108.58613393502857)
E       AssertionError:
E       Not equal to tolerance rtol=1e-07, atol=0
E
E       Mismatched elements: 1 / 1 (100%)
E       Max absolute difference among violations: 5.03683129e-05
E       Max relative difference among violations: 4.63855845e-07
E        ACTUAL: array(-108.586084)
E        DESIRED: array(-108.586134)

tests/python/variational/orbital_optimization_test.py:86: AssertionError

[q-inho]

I'm not confident whether it is appropriate to comment on this PR draft, but I'd like to share a quick thought.
From my scheme, I think current _generator_to_parameters misses that each variational parameter controls two elements of the anti‑Hermitian generator.
Do you think multiply the gradient contributions by 2 and to account for the sign of the imaginary part when converting back to the parameter vector would solve the issue?

[kevinsung]

I'm not confident whether it is appropriate to comment on this PR draft, but I'd like to share a quick thought. From my scheme, I think current _generator_to_parameters misses that each variational parameter controls two elements of the anti‑Hermitian generator. Do you think multiply the gradient contributions by 2 and to account for the sign of the imaginary part when converting back to the parameter vector would solve the issue?

Sorry, I don't quite understand what you mean, but feel free to open a pull request if you have a solution 🙂.

[kevinsung]

It turns out the issue was jax-ml/jax#4891.