From 11e10274eb459b4521daf99272d3deae7cd6236e Mon Sep 17 00:00:00 2001
From: "Kevin J. Sung" <kevjsung@umich.edu>
Date: Tue, 13 Feb 2024 09:17:47 -0500
Subject: [PATCH] LUCJ: order parameters by layer (#100)

---
 python/ffsim/variational/ucj.py      | 92 ++++++++++++++++------------
 tests/optimize/linear_method_test.py |  6 +-
 2 files changed, 57 insertions(+), 41 deletions(-)

diff --git a/python/ffsim/variational/ucj.py b/python/ffsim/variational/ucj.py
index 0f3745187..9f658da32 100644
--- a/python/ffsim/variational/ucj.py
+++ b/python/ffsim/variational/ucj.py
@@ -479,30 +479,37 @@ def _ucj_from_parameters(
     diag_coulomb_mats_alpha_beta = np.zeros((n_reps, norb, norb))
     orbital_rotations = np.zeros((n_reps, norb, norb), dtype=complex)
     index = 0
-    # diag coulomb matrices, alpha-alpha
-    indices = alpha_alpha_indices
-    if indices:
-        n_params = len(indices)
-        rows, cols = zip(*indices)
-        for mat in diag_coulomb_mats_alpha_alpha:
+    for (
+        orbital_rotation,
+        diag_coulomb_mat_alpha_alpha,
+        diag_coulomb_mat_alpha_beta,
+    ) in zip(
+        orbital_rotations,
+        diag_coulomb_mats_alpha_alpha,
+        diag_coulomb_mats_alpha_beta,
+    ):
+        # orbital rotation
+        n_params = norb**2
+        orbital_rotation[:] = orbital_rotation_from_parameters(
+            params[index : index + n_params], norb
+        )
+        index += n_params
+        # diag coulomb matrix, alpha-alpha
+        if alpha_alpha_indices:
+            n_params = len(alpha_alpha_indices)
+            rows, cols = zip(*alpha_alpha_indices)
             vals = params[index : index + n_params]
-            mat[rows, cols] = vals
-            mat[cols, rows] = vals
+            diag_coulomb_mat_alpha_alpha[rows, cols] = vals
+            diag_coulomb_mat_alpha_alpha[cols, rows] = vals
             index += n_params
-    # diag coulomb matrices, alpha-beta
-    indices = alpha_beta_indices
-    if indices:
-        n_params = len(indices)
-        rows, cols = zip(*indices)
-        for mat in diag_coulomb_mats_alpha_beta:
+        # diag coulomb matrix, alpha-beta
+        if alpha_beta_indices:
+            n_params = len(alpha_beta_indices)
+            rows, cols = zip(*alpha_beta_indices)
             vals = params[index : index + n_params]
-            mat[rows, cols] = vals
-            mat[cols, rows] = vals
+            diag_coulomb_mat_alpha_beta[rows, cols] = vals
+            diag_coulomb_mat_alpha_beta[cols, rows] = vals
             index += n_params
-    # orbital rotations
-    for mat in orbital_rotations:
-        mat[:] = orbital_rotation_from_parameters(params[index : index + norb**2], norb)
-        index += norb**2
     # final orbital rotation
     final_orbital_rotation = None
     if with_final_orbital_rotation:
@@ -538,26 +545,35 @@ def _ucj_to_parameters(
         ntheta += norb**2
     theta = np.zeros(ntheta)
     index = 0
-    # diag coulomb matrices, alpha-alpha
-    indices = alpha_alpha_indices
-    if indices:
-        n_params = len(indices)
-        for mat in diag_coulomb_mats_alpha_alpha:
-            theta[index : index + n_params] = mat[tuple(zip(*indices))]
-            index += n_params
-    # diag coulomb matrices, alpha-beta
-    indices = alpha_beta_indices
-    if indices:
-        n_params = len(indices)
-        for mat in diag_coulomb_mats_alpha_beta:
-            theta[index : index + n_params] = mat[tuple(zip(*indices))]
-            index += n_params
-    # orbital rotations
-    for orbital_rotation in orbital_rotations:
-        theta[index : index + norb**2] = orbital_rotation_to_parameters(
+    for (
+        orbital_rotation,
+        diag_coulomb_mat_alpha_alpha,
+        diag_coulomb_mat_alpha_beta,
+    ) in zip(
+        orbital_rotations,
+        diag_coulomb_mats_alpha_alpha,
+        diag_coulomb_mats_alpha_beta,
+    ):
+        # orbital rotation
+        n_params = norb**2
+        theta[index : index + n_params] = orbital_rotation_to_parameters(
             orbital_rotation
         )
-        index += norb**2
+        index += n_params
+        # diag coulomb matrix, alpha-alpha
+        if alpha_alpha_indices:
+            n_params = len(alpha_alpha_indices)
+            theta[index : index + n_params] = diag_coulomb_mat_alpha_alpha[
+                tuple(zip(*alpha_alpha_indices))
+            ]
+            index += n_params
+        # diag coulomb matrix, alpha-beta
+        if alpha_beta_indices:
+            n_params = len(alpha_beta_indices)
+            theta[index : index + n_params] = diag_coulomb_mat_alpha_beta[
+                tuple(zip(*alpha_beta_indices))
+            ]
+            index += n_params
     # final orbital rotation
     if final_orbital_rotation is not None:
         theta[index : index + norb**2] = orbital_rotation_to_parameters(
diff --git a/tests/optimize/linear_method_test.py b/tests/optimize/linear_method_test.py
index 6aea5cb5e..763309379 100644
--- a/tests/optimize/linear_method_test.py
+++ b/tests/optimize/linear_method_test.py
@@ -31,10 +31,10 @@ def test_minimize_linear_method():
     hartree_fock = pyscf.scf.RHF(mol)
     hartree_fock.kernel()
 
-    # Construct UCJ operator
+    # Initialize parameters
     n_reps = 2
     n_params = ffsim.UCJOperator.n_params(hartree_fock.mol.nao_nr(), n_reps)
-    rng = np.random.default_rng(1234)
+    rng = np.random.default_rng(1804)
     x0 = rng.uniform(-10, 10, size=n_params)
 
     # Get molecular data and molecular Hamiltonian (one- and two-body tensors)
@@ -74,7 +74,7 @@ def callback(intermediate_result):
     )
     np.testing.assert_allclose(energy(result.x), result.fun)
     np.testing.assert_allclose(result.fun, -0.970773)
-    np.testing.assert_allclose(info["fun"][0], -0.834889, atol=1e-5)
+    np.testing.assert_allclose(info["fun"][0], -0.833558, atol=1e-5)
     np.testing.assert_allclose(info["fun"][-1], -0.970773, atol=1e-5)
     for params, fun in zip(info["x"], info["fun"]):
         np.testing.assert_allclose(energy(params), fun)