diff --git a/python/ffsim/variational/ucj_spin_balanced.py b/python/ffsim/variational/ucj_spin_balanced.py index 235bfc11d..194314f3c 100644 --- a/python/ffsim/variational/ucj_spin_balanced.py +++ b/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): diff --git a/python/ffsim/variational/ucj_spin_unbalanced.py b/python/ffsim/variational/ucj_spin_unbalanced.py index c1cba42a6..e7e289c74 100644 --- a/python/ffsim/variational/ucj_spin_unbalanced.py +++ b/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): diff --git a/python/ffsim/variational/ucj_spinless.py b/python/ffsim/variational/ucj_spinless.py index 9c9d45c6a..40699fc03 100644 --- a/python/ffsim/variational/ucj_spinless.py +++ b/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): diff --git a/tests/python/variational/ucj_spin_balanced_test.py b/tests/python/variational/ucj_spin_balanced_test.py index b476ef5ae..a4a1615a8 100644 --- a/tests/python/variational/ucj_spin_balanced_test.py +++ b/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, + ) diff --git a/tests/python/variational/ucj_spin_unbalanced_test.py b/tests/python/variational/ucj_spin_unbalanced_test.py index e46d1192e..c6319093b 100644 --- a/tests/python/variational/ucj_spin_unbalanced_test.py +++ b/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)] + ), + ) diff --git a/tests/python/variational/ucj_spinless_test.py b/tests/python/variational/ucj_spinless_test.py index 1bb04e27f..488499046 100644 --- a/tests/python/variational/ucj_spinless_test.py +++ b/tests/python/variational/ucj_spinless_test.py @@ -176,3 +176,62 @@ 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([eye for _ in range(n_reps)]) + orbital_rotations = np.stack([eye for _ in range(n_reps)]) + + _ = ffsim.UCJOpSpinless( + diag_coulomb_mats=rng.standard_normal(10), + orbital_rotations=orbital_rotations, + validate=False, + ) + + _ = ffsim.UCJOpSpinless( + diag_coulomb_mats=rng.standard_normal((n_reps, norb, norb)), + orbital_rotations=orbital_rotations, + atol=10, + ) + + with pytest.raises(ValueError, match="shape"): + _ = ffsim.UCJOpSpinless( + diag_coulomb_mats=rng.standard_normal(10), + orbital_rotations=orbital_rotations, + ) + with pytest.raises(ValueError, match="shape"): + _ = ffsim.UCJOpSpinless( + diag_coulomb_mats=diag_coulomb_mats, + orbital_rotations=rng.standard_normal(10), + ) + with pytest.raises(ValueError, match="shape"): + _ = ffsim.UCJOpSpinless( + diag_coulomb_mats=diag_coulomb_mats, + orbital_rotations=orbital_rotations, + final_orbital_rotation=rng.standard_normal(10), + ) + with pytest.raises(ValueError, match="dimension"): + _ = ffsim.UCJOpSpinless( + diag_coulomb_mats=diag_coulomb_mats, + orbital_rotations=np.concatenate([orbital_rotations, orbital_rotations]), + ) + with pytest.raises(ValueError, match="symmetric"): + _ = ffsim.UCJOpSpinless( + diag_coulomb_mats=rng.standard_normal((n_reps, norb, norb)), + orbital_rotations=orbital_rotations, + ) + with pytest.raises(ValueError, match="unitary"): + _ = ffsim.UCJOpSpinless( + diag_coulomb_mats=diag_coulomb_mats, + orbital_rotations=rng.standard_normal((n_reps, norb, norb)), + ) + with pytest.raises(ValueError, match="unitary"): + _ = ffsim.UCJOpSpinless( + diag_coulomb_mats=diag_coulomb_mats, + orbital_rotations=orbital_rotations, + final_orbital_rotation=rng.standard_normal((norb, norb)), + )