Add input validation to optimize_orbitals per Gemini Code Assist review

Addresses both high-priority comments from the automated review on PR #1442:
- validate one_body_integrals/two_body_integrals/one_rdm/two_rdm shapes
  against each other and against n_electrons parity (restricted/closed-shell
  requires an even electron count) before they reach energy(), where a
  mismatch would otherwise surface as an opaque broadcast/index error deep
  inside general_basis_change or rhf_params_to_matrix.
- validate initial_guess size against the expected nocc * nvirt parameter
  count for the same reason.

4 new regression tests, all pre-existing tests still pass.

Author: SemyonAndreyev

src/openfermion/hamiltonians/orbital_optimization.py

@@ -128,7 +128,36 @@ def optimize_orbitals(
         scipy.optimize.OptimizeResult. `result.x` is the optimal kappa
         parameter vector; `result.fun` is the optimized energy.
     """
+    if one_body_integrals.ndim != 2 or one_body_integrals.shape[0] != one_body_integrals.shape[1]:
+        raise ValueError(
+            f"one_body_integrals must be a square 2D array, got shape "
+            f"{one_body_integrals.shape}"
+        )
     n_orbitals = one_body_integrals.shape[0]
+    if two_body_integrals.shape != (n_orbitals,) * 4:
+        raise ValueError(
+            f"two_body_integrals must have shape {(n_orbitals,) * 4} to match "
+            f"one_body_integrals (n_orbitals={n_orbitals}), got "
+            f"{two_body_integrals.shape}"
+        )
+    n_spin_orbitals = 2 * n_orbitals
+    if one_rdm.shape != (n_spin_orbitals,) * 2:
+        raise ValueError(
+            f"one_rdm must have shape {(n_spin_orbitals,) * 2} (spin-orbital "
+            f"basis, 2 * n_orbitals with n_orbitals={n_orbitals}), got "
+            f"{one_rdm.shape}"
+        )
+    if two_rdm.shape != (n_spin_orbitals,) * 4:
+        raise ValueError(
+            f"two_rdm must have shape {(n_spin_orbitals,) * 4} (spin-orbital "
+            f"basis, 2 * n_orbitals with n_orbitals={n_orbitals}), got "
+            f"{two_rdm.shape}"
+        )
+    if n_electrons % 2 != 0:
+        raise ValueError(
+            f"optimize_orbitals is restricted (closed-shell) -- n_electrons "
+            f"must be even, got {n_electrons}"
+        )
     nocc = n_electrons // 2
     nvirt = n_orbitals - nocc
     if nocc <= 0 or nvirt <= 0:
@@ -159,6 +188,11 @@ def energy(params: np.ndarray) -> float:
         init_params = np.zeros(nocc * nvirt)
     else:
         init_params = np.asarray(initial_guess).flatten()
+        if init_params.size != nocc * nvirt:
+            raise ValueError(
+                f"initial_guess has {init_params.size} parameters, expected "
+                f"nocc * nvirt = {nocc} * {nvirt} = {nocc * nvirt}"
+            )
 
     sp_optimizer_options = {'disp': verbose}
     if sp_options is not None:

src/openfermion/hamiltonians/orbital_optimization_test.py

@@ -187,3 +187,57 @@ def test_optimize_orbitals_rejects_degenerate_orbital_split():
         optimize_orbitals(obi, tbi, one_rdm, two_rdm, n_electrons=4)  # all occupied, no virtuals
     with pytest.raises(ValueError):
         optimize_orbitals(obi, tbi, one_rdm, two_rdm, n_electrons=0)  # all virtual, no occupied
+
+
+def test_optimize_orbitals_rejects_mismatched_shapes():
+    """Gemini Code Assist review on PR #1442: no shape validation meant a
+    mismatched one_body_integrals/two_body_integrals/one_rdm/two_rdm/
+    n_electrons combination would fail deep inside energy() with a
+    confusing broadcast/index error instead of a clear message at the
+    call boundary."""
+    n_orbitals = 3
+    obi = np.zeros((n_orbitals, n_orbitals))
+    tbi = np.zeros((n_orbitals,) * 4)
+    one_rdm = np.zeros((2 * n_orbitals, 2 * n_orbitals))
+    two_rdm = np.zeros((2 * n_orbitals,) * 4)
+
+    with pytest.raises(ValueError, match="one_body_integrals"):
+        optimize_orbitals(np.zeros((n_orbitals, n_orbitals + 1)), tbi, one_rdm, two_rdm, 4)
+    with pytest.raises(ValueError, match="two_body_integrals"):
+        optimize_orbitals(obi, np.zeros((n_orbitals + 1,) * 4), one_rdm, two_rdm, 4)
+    with pytest.raises(ValueError, match="one_rdm"):
+        optimize_orbitals(obi, tbi, np.zeros((2 * n_orbitals + 1,) * 2), two_rdm, 4)
+    with pytest.raises(ValueError, match="two_rdm"):
+        optimize_orbitals(obi, tbi, one_rdm, np.zeros((2 * n_orbitals + 1,) * 4), 4)
+
+
+def test_optimize_orbitals_rejects_odd_electron_count():
+    """optimize_orbitals is restricted (closed-shell): n_electrons // 2
+    silently rounds an odd count down, which would optimize the wrong
+    number of occupied orbitals without any warning."""
+    n_orbitals = 3
+    obi = np.zeros((n_orbitals, n_orbitals))
+    tbi = np.zeros((n_orbitals,) * 4)
+    one_rdm = np.zeros((2 * n_orbitals, 2 * n_orbitals))
+    two_rdm = np.zeros((2 * n_orbitals,) * 4)
+    with pytest.raises(ValueError, match="even"):
+        optimize_orbitals(obi, tbi, one_rdm, two_rdm, n_electrons=3)
+
+
+def test_optimize_orbitals_rejects_mismatched_initial_guess():
+    """A caller-supplied initial_guess of the wrong length would otherwise
+    hit an IndexError deep inside rhf_params_to_matrix instead of a clear
+    message naming the expected parameter count."""
+    m = _load_h2()
+    hamiltonian = m.get_molecular_hamiltonian()
+    _, one_rdm, two_rdm = _fci_ground_state_rdms(hamiltonian, m.n_qubits)
+    with pytest.raises(ValueError, match="initial_guess"):
+        optimize_orbitals(
+            m.one_body_integrals,
+            m.two_body_integrals,
+            one_rdm,
+            two_rdm,
+            m.n_electrons,
+            initial_guess=np.zeros(5),  # wrong size for this molecule (H2/sto-3g needs 1)
+            verbose=False,
+        )