From 4fd5a30ce2d7b5fe151dc6724dd8ba530ba7dc3b Mon Sep 17 00:00:00 2001 From: "Kevin J. Sung" Date: Mon, 23 Oct 2023 14:04:06 -0500 Subject: [PATCH] add indices_to_strings function --- python/ffsim/__init__.py | 1 + python/ffsim/states.py | 33 +++++++++++++++++++++++++++++++++ tests/states_test.py | 20 ++++++++++++++++++++ 3 files changed, 54 insertions(+) diff --git a/python/ffsim/__init__.py b/python/ffsim/__init__.py index 136f0a6ba..ef344891c 100644 --- a/python/ffsim/__init__.py +++ b/python/ffsim/__init__.py @@ -51,6 +51,7 @@ from ffsim.states import ( dim, dims, + indices_to_strings, one_hot, slater_determinant, slater_determinant_one_rdm, diff --git a/python/ffsim/states.py b/python/ffsim/states.py index bbec27708..8a1e894fa 100644 --- a/python/ffsim/states.py +++ b/python/ffsim/states.py @@ -128,3 +128,36 @@ def slater_determinant_one_rdm( one_rdm[(alpha_orbitals, alpha_orbitals)] = 1 one_rdm[(beta_orbitals, beta_orbitals)] = 1 return one_rdm + + +def indices_to_strings( + indices: Sequence[int], norb: int, nelec: tuple[int, int] +) -> list[str]: + """Convert statevector indices to bitstrings. + + Example: + >>> norb = 3 + >>> nelec = (2, 1) + >>> dim = ffsim.dim(norb, nelec) + >>> strings = ffsim.indices_to_strings(range(dim), norb, nelec) + >>> assert strings == [ + "011001", + "011010", + "011100", + "101001", + "101010", + "101100", + "110001", + "110010", + "110100", + ] + """ + n_alpha, n_beta = nelec + dim_b = comb(norb, n_beta, exact=True) + indices_a, indices_b = np.divmod(indices, dim_b) + strings_a = cistring.addrs2str(norb=norb, nelec=n_alpha, addrs=indices_a) + strings_b = cistring.addrs2str(norb=norb, nelec=n_beta, addrs=indices_b) + return [ + f"{string_a:0{norb}b}{string_b:0{norb}b}" + for string_a, string_b in zip(strings_a, strings_b) + ] diff --git a/tests/states_test.py b/tests/states_test.py index 3d371fcbf..0098e9be5 100644 --- a/tests/states_test.py +++ b/tests/states_test.py @@ -31,3 +31,23 @@ def test_slater_determinant(): hamiltonian = ffsim.contract.one_body_linop(one_body_tensor, norb=norb, nelec=nelec) np.testing.assert_allclose(hamiltonian @ state, eig * state) + + +def test_indices_to_strings(): + """Test converting statevector indices to strings.""" + norb = 3 + nelec = (2, 1) + + dim = ffsim.dim(norb, nelec) + strings = ffsim.indices_to_strings(range(dim), norb, nelec) + assert strings == [ + "011001", + "011010", + "011100", + "101001", + "101010", + "101100", + "110001", + "110010", + "110100", + ]