validate UCJ dataclasses (#256)

python/ffsim/variational/ucj_spin_balanced.py

@@ -13,7 +13,7 @@
 from __future__ import annotations
 
 import itertools
-from dataclasses import dataclass
+from dataclasses import InitVar, dataclass
 from typing import cast
 
 import numpy as np
@@ -93,6 +93,54 @@ class UCJOpSpinBalanced:
     diag_coulomb_mats: np.ndarray  # shape: (n_reps, 2, norb, norb)
     orbital_rotations: np.ndarray  # shape: (n_reps, norb, norb)
     final_orbital_rotation: np.ndarray | None = None  # shape: (norb, norb)
+    validate: InitVar[bool] = True
+    rtol: InitVar[float] = 1e-5
+    atol: InitVar[float] = 1e-8
+
+    def __post_init__(self, validate: bool, rtol: float, atol: float):
+        if validate:
+            if self.diag_coulomb_mats.ndim != 4 or self.diag_coulomb_mats.shape[1] != 2:
+                raise ValueError(
+                    "diag_coulomb_mats should have shape (n_reps, 2, norb, norb). "
+                    f"Got shape {self.diag_coulomb_mats.shape}."
+                )
+            if self.orbital_rotations.ndim != 3:
+                raise ValueError(
+                    "orbital_rotations should have shape (n_reps, norb, norb). "
+                    f"Got shape {self.orbital_rotations.shape}."
+                )
+            if (
+                self.final_orbital_rotation is not None
+                and self.final_orbital_rotation.ndim != 2
+            ):
+                raise ValueError(
+                    "final_orbital_rotation should have shape (norb, norb). "
+                    f"Got shape {self.final_orbital_rotation.shape}."
+                )
+            if self.diag_coulomb_mats.shape[0] != self.orbital_rotations.shape[0]:
+                raise ValueError(
+                    "diag_coulomb_mats and orbital_rotations should have the same "
+                    "first dimension. "
+                    f"Got {self.diag_coulomb_mats.shape[0]} and "
+                    f"{self.orbital_rotations.shape[0]}."
+                )
+            if not all(
+                linalg.is_real_symmetric(mats[0], rtol=rtol, atol=atol)
+                and linalg.is_real_symmetric(mats[1], rtol=rtol, atol=atol)
+                for mats in self.diag_coulomb_mats
+            ):
+                raise ValueError(
+                    "Diagonal Coulomb matrices were not all real symmetric."
+                )
+            if not all(
+                linalg.is_unitary(orbital_rotation, rtol=rtol, atol=atol)
+                for orbital_rotation in self.orbital_rotations
+            ):
+                raise ValueError("Orbital rotations were not all unitary.")
+            if self.final_orbital_rotation is not None and not linalg.is_unitary(
+                self.final_orbital_rotation, rtol=rtol, atol=atol
+            ):
+                raise ValueError("Final orbital rotation was not unitary.")
 
     @property
     def norb(self):

python/ffsim/variational/ucj_spin_unbalanced.py

@@ -13,7 +13,7 @@
 from __future__ import annotations
 
 import itertools
-from dataclasses import dataclass
+from dataclasses import InitVar, dataclass
 from typing import cast
 
 import numpy as np
@@ -95,6 +95,59 @@ class UCJOpSpinUnbalanced:
     diag_coulomb_mats: np.ndarray  # shape: (n_reps, 3, norb, norb)
     orbital_rotations: np.ndarray  # shape: (n_reps, 2, norb, norb)
     final_orbital_rotation: np.ndarray | None = None  # shape: (2, norb, norb)
+    validate: InitVar[bool] = True
+    rtol: InitVar[float] = 1e-5
+    atol: InitVar[float] = 1e-8
+
+    def __post_init__(self, validate: bool, rtol: float, atol: float):
+        if validate:
+            if self.diag_coulomb_mats.ndim != 4 or self.diag_coulomb_mats.shape[1] != 3:
+                raise ValueError(
+                    "diag_coulomb_mats should have shape (n_reps, 3, norb, norb). "
+                    f"Got shape {self.diag_coulomb_mats.shape}."
+                )
+            if self.orbital_rotations.ndim != 4 or self.orbital_rotations.shape[1] != 2:
+                raise ValueError(
+                    "orbital_rotations should have shape (n_reps, 2, norb, norb). "
+                    f"Got shape {self.orbital_rotations.shape}."
+                )
+            if (
+                self.final_orbital_rotation is not None
+                and self.final_orbital_rotation.ndim != 3
+            ):
+                raise ValueError(
+                    "final_orbital_rotation should have shape (2, norb, norb). "
+                    f"Got shape {self.final_orbital_rotation.shape}."
+                )
+            if self.diag_coulomb_mats.shape[0] != self.orbital_rotations.shape[0]:
+                raise ValueError(
+                    "diag_coulomb_mats and orbital_rotations should have the same "
+                    "first dimension. "
+                    f"Got {self.diag_coulomb_mats.shape[0]} and "
+                    f"{self.orbital_rotations.shape[0]}."
+                )
+            if not all(
+                linalg.is_real_symmetric(mats[0], rtol=rtol, atol=atol)
+                and linalg.is_real_symmetric(mats[2], rtol=rtol, atol=atol)
+                for mats in self.diag_coulomb_mats
+            ):
+                raise ValueError(
+                    "alpha-alpha and beta-beta diagonal Coulomb matrices were not all "
+                    "real symmetric."
+                )
+            if not all(
+                linalg.is_unitary(orbital_rotation[0], rtol=rtol, atol=atol)
+                and linalg.is_unitary(orbital_rotation[1], rtol=rtol, atol=atol)
+                for orbital_rotation in self.orbital_rotations
+            ):
+                raise ValueError("Orbital rotations were not all unitary.")
+            if self.final_orbital_rotation is not None and not (
+                linalg.is_unitary(self.final_orbital_rotation[0], rtol=rtol, atol=atol)
+                and linalg.is_unitary(
+                    self.final_orbital_rotation[1], rtol=rtol, atol=atol
+                )
+            ):
+                raise ValueError("Final orbital rotation was not unitary.")
 
     @property
     def norb(self):

python/ffsim/variational/ucj_spinless.py

@@ -13,7 +13,7 @@
 from __future__ import annotations
 
 import itertools
-from dataclasses import dataclass
+from dataclasses import InitVar, dataclass
 from typing import cast
 
 import numpy as np
@@ -82,6 +82,53 @@ class UCJOpSpinless:
     diag_coulomb_mats: np.ndarray  # shape: (n_reps, norb, norb)
     orbital_rotations: np.ndarray  # shape: (n_reps, norb, norb)
     final_orbital_rotation: np.ndarray | None = None  # shape: (norb, norb)
+    validate: InitVar[bool] = True
+    rtol: InitVar[float] = 1e-5
+    atol: InitVar[float] = 1e-8
+
+    def __post_init__(self, validate: bool, rtol: float, atol: float):
+        if validate:
+            if self.diag_coulomb_mats.ndim != 3:
+                raise ValueError(
+                    "diag_coulomb_mats should have shape (n_reps, norb, norb). "
+                    f"Got shape {self.diag_coulomb_mats.shape}."
+                )
+            if self.orbital_rotations.ndim != 3:
+                raise ValueError(
+                    "orbital_rotations should have shape (n_reps, norb, norb). "
+                    f"Got shape {self.orbital_rotations.shape}."
+                )
+            if (
+                self.final_orbital_rotation is not None
+                and self.final_orbital_rotation.ndim != 2
+            ):
+                raise ValueError(
+                    "final_orbital_rotation should have shape (norb, norb). "
+                    f"Got shape {self.final_orbital_rotation.shape}."
+                )
+            if self.diag_coulomb_mats.shape[0] != self.orbital_rotations.shape[0]:
+                raise ValueError(
+                    "diag_coulomb_mats and orbital_rotations should have the same "
+                    "first dimension. "
+                    f"Got {self.diag_coulomb_mats.shape[0]} and "
+                    f"{self.orbital_rotations.shape[0]}."
+                )
+            if not all(
+                linalg.is_real_symmetric(mat, rtol=rtol, atol=atol)
+                for mat in self.diag_coulomb_mats
+            ):
+                raise ValueError(
+                    "Diagonal Coulomb matrices were not all real symmetric."
+                )
+            if not all(
+                linalg.is_unitary(orbital_rotation, rtol=rtol, atol=atol)
+                for orbital_rotation in self.orbital_rotations
+            ):
+                raise ValueError("Orbital rotations were not all unitary.")
+            if self.final_orbital_rotation is not None and not linalg.is_unitary(
+                self.final_orbital_rotation, rtol=rtol, atol=atol
+            ):
+                raise ValueError("Final orbital rotation was not unitary.")
 
     @property
     def norb(self):

tests/python/variational/ucj_spin_balanced_test.py

@@ -171,3 +171,69 @@ def test_t_amplitudes_restrict_indices():
     )
 
     assert ffsim.approx_eq(operator, other_operator, rtol=1e-12)
+
+
+def test_validate():
+    rng = np.random.default_rng(335)
+    n_reps = 3
+    norb = 4
+    eye = np.eye(norb)
+    diag_coulomb_mats = np.stack([np.stack([eye, eye]) for _ in range(n_reps)])
+    orbital_rotations = np.stack([eye for _ in range(n_reps)])
+
+    _ = ffsim.UCJOpSpinBalanced(
+        diag_coulomb_mats=rng.standard_normal(10),
+        orbital_rotations=orbital_rotations,
+        validate=False,
+    )
+
+    _ = ffsim.UCJOpSpinBalanced(
+        diag_coulomb_mats=rng.standard_normal((n_reps, 2, norb, norb)),
+        orbital_rotations=orbital_rotations,
+        atol=10,
+    )
+
+    with pytest.raises(ValueError, match="shape"):
+        _ = ffsim.UCJOpSpinBalanced(
+            diag_coulomb_mats=rng.standard_normal(10),
+            orbital_rotations=orbital_rotations,
+        )
+    with pytest.raises(ValueError, match="shape"):
+        _ = ffsim.UCJOpSpinBalanced(
+            diag_coulomb_mats=diag_coulomb_mats,
+            orbital_rotations=rng.standard_normal(10),
+        )
+    with pytest.raises(ValueError, match="shape"):
+        _ = ffsim.UCJOpSpinBalanced(
+            diag_coulomb_mats=diag_coulomb_mats,
+            orbital_rotations=orbital_rotations,
+            final_orbital_rotation=rng.standard_normal(10),
+        )
+    with pytest.raises(ValueError, match="dimension"):
+        _ = ffsim.UCJOpSpinBalanced(
+            diag_coulomb_mats=diag_coulomb_mats,
+            orbital_rotations=np.concatenate([orbital_rotations, orbital_rotations]),
+        )
+    with pytest.raises(ValueError, match="symmetric"):
+        _ = ffsim.UCJOpSpinBalanced(
+            diag_coulomb_mats=rng.standard_normal((n_reps, 2, norb, norb)),
+            orbital_rotations=orbital_rotations,
+        )
+    with pytest.raises(ValueError, match="unitary"):
+        _ = ffsim.UCJOpSpinBalanced(
+            diag_coulomb_mats=diag_coulomb_mats,
+            orbital_rotations=rng.standard_normal((n_reps, norb, norb)),
+        )
+    with pytest.raises(ValueError, match="unitary"):
+        _ = ffsim.UCJOpSpinBalanced(
+            diag_coulomb_mats=diag_coulomb_mats,
+            orbital_rotations=orbital_rotations,
+            final_orbital_rotation=rng.standard_normal((norb, norb)),
+        )
+    with pytest.raises(ValueError, match="shape"):
+        _ = ffsim.UCJOpSpinBalanced(
+            diag_coulomb_mats=np.stack(
+                [np.stack([eye, eye, eye]) for _ in range(n_reps)]
+            ),
+            orbital_rotations=orbital_rotations,
+        )

tests/python/variational/ucj_spin_unbalanced_test.py

@@ -170,3 +170,74 @@ def test_t_amplitudes_restrict_indices():
     )
 
     assert ffsim.approx_eq(operator, other_operator, rtol=1e-12)
+
+
+def test_validate():
+    rng = np.random.default_rng(335)
+    n_reps = 3
+    norb = 4
+    eye = np.eye(norb)
+    diag_coulomb_mats = np.stack([np.stack([eye, eye, eye]) for _ in range(n_reps)])
+    orbital_rotations = np.stack([np.stack([eye, eye]) for _ in range(n_reps)])
+
+    _ = ffsim.UCJOpSpinUnbalanced(
+        diag_coulomb_mats=rng.standard_normal(10),
+        orbital_rotations=orbital_rotations,
+        validate=False,
+    )
+
+    _ = ffsim.UCJOpSpinUnbalanced(
+        diag_coulomb_mats=rng.standard_normal((n_reps, 3, norb, norb)),
+        orbital_rotations=orbital_rotations,
+        atol=10,
+    )
+
+    with pytest.raises(ValueError, match="shape"):
+        _ = ffsim.UCJOpSpinUnbalanced(
+            diag_coulomb_mats=rng.standard_normal(10),
+            orbital_rotations=orbital_rotations,
+        )
+    with pytest.raises(ValueError, match="shape"):
+        _ = ffsim.UCJOpSpinUnbalanced(
+            diag_coulomb_mats=diag_coulomb_mats,
+            orbital_rotations=rng.standard_normal(10),
+        )
+    with pytest.raises(ValueError, match="shape"):
+        _ = ffsim.UCJOpSpinUnbalanced(
+            diag_coulomb_mats=diag_coulomb_mats,
+            orbital_rotations=orbital_rotations,
+            final_orbital_rotation=rng.standard_normal(10),
+        )
+    with pytest.raises(ValueError, match="dimension"):
+        _ = ffsim.UCJOpSpinUnbalanced(
+            diag_coulomb_mats=diag_coulomb_mats,
+            orbital_rotations=np.concatenate([orbital_rotations, orbital_rotations]),
+        )
+    with pytest.raises(ValueError, match="symmetric"):
+        _ = ffsim.UCJOpSpinUnbalanced(
+            diag_coulomb_mats=rng.standard_normal((n_reps, 3, norb, norb)),
+            orbital_rotations=orbital_rotations,
+        )
+    with pytest.raises(ValueError, match="unitary"):
+        _ = ffsim.UCJOpSpinUnbalanced(
+            diag_coulomb_mats=diag_coulomb_mats,
+            orbital_rotations=rng.standard_normal((n_reps, 2, norb, norb)),
+        )
+    with pytest.raises(ValueError, match="unitary"):
+        _ = ffsim.UCJOpSpinUnbalanced(
+            diag_coulomb_mats=diag_coulomb_mats,
+            orbital_rotations=orbital_rotations,
+            final_orbital_rotation=rng.standard_normal((2, norb, norb)),
+        )
+    with pytest.raises(ValueError, match="shape"):
+        _ = ffsim.UCJOpSpinUnbalanced(
+            diag_coulomb_mats=np.stack([np.stack([eye, eye]) for _ in range(n_reps)]),
+            orbital_rotations=orbital_rotations,
+        )
+    with pytest.raises(ValueError, match="shape"):
+        _ = ffsim.UCJOpSpinUnbalanced(
+            diag_coulomb_mats=diag_coulomb_mats,
+            orbital_rotations=np.stack(
+                [np.stack([eye, eye, eye]) for _ in range(n_reps)]
+            ),
+        )