define __array__ for StateVector (#280)

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.

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))