diff --git a/python/ffsim/states/states.py b/python/ffsim/states/states.py index a1e151a77..b78d12cb2 100644 --- a/python/ffsim/states/states.py +++ b/python/ffsim/states/states.py @@ -43,6 +43,18 @@ class StateVector: norb: int nelec: int | tuple[int, int] + def __array__(self, dtype=None, copy=None): + # TODO in Numpy 2.0 this can be simplified to + # return np.array(self.vec, dtype=dtype, copy=copy) + if copy: + if dtype is None: + return self.vec.copy() + else: + return self.vec.astype(dtype, copy=True) + if dtype is None: + return self.vec + return self.vec.astype(dtype, copy=False) + def dims(norb: int, nelec: tuple[int, int]) -> tuple[int, int]: """Get the dimensions of the FCI space. diff --git a/tests/python/states/states_test.py b/tests/python/states/states_test.py index 0cab96937..96a04fe81 100644 --- a/tests/python/states/states_test.py +++ b/tests/python/states/states_test.py @@ -511,3 +511,12 @@ def test_slater_determinant_one_rdm_diff_rotation( expected = ffsim.rdm(vec, norb, nelec, spin_summed=spin_summed) np.testing.assert_allclose(rdm, expected, atol=1e-12) + + +def test_state_vector_array(): + """Test StateVector's __array__ method.""" + norb = 5 + nelec = (3, 2) + vec = ffsim.random.random_state_vector(ffsim.dim(norb, nelec), seed=3556) + state_vec = ffsim.StateVector(vec, norb, nelec) + assert np.array_equal(np.abs(state_vec), np.abs(vec))