From 8d20bbe32f1010ec816deeee384e633d4c0a741c Mon Sep 17 00:00:00 2001
From: "Kevin J. Sung" <kevjsung@umich.edu>
Date: Sat, 6 Jul 2024 12:03:13 -0400
Subject: [PATCH] add MolecularData.to_fcidump (#272)

---
 python/ffsim/molecular_data.py      | 62 +++++++++++++++++++----------
 tests/python/molecular_data_test.py | 22 ++--------
 2 files changed, 45 insertions(+), 39 deletions(-)

diff --git a/python/ffsim/molecular_data.py b/python/ffsim/molecular_data.py
index 8c1c1f4b8..04f7e23e3 100644
--- a/python/ffsim/molecular_data.py
+++ b/python/ffsim/molecular_data.py
@@ -124,27 +124,6 @@ def mole(self) -> pyscf.gto.Mole:
             symmetry=self.symmetry,
         )
 
-    @staticmethod
-    def from_fcidump(file: str | bytes | os.PathLike) -> MolecularData:
-        """Initialize a MolecularData from an FCIDUMP file.
-
-        Args:
-            file: The FCIDUMP file path.
-        """
-        data = pyscf.tools.fcidump.read(file, verbose=False)
-        n_electrons = data["NELEC"]
-        spin = data["MS2"]
-        n_alpha = (n_electrons + spin) // 2
-        n_beta = (n_electrons - spin) // 2
-        return MolecularData(
-            core_energy=data["ECORE"],
-            one_body_integrals=data["H1"],
-            two_body_integrals=data["H2"],
-            norb=data["NORB"],
-            nelec=(n_alpha, n_beta),
-            spin=spin,
-        )
-
     @staticmethod
     def from_scf(
         hartree_fock: pyscf.scf.hf.SCF, active_space: Iterable[int] | None = None
@@ -400,3 +379,44 @@ def as_array_tuple_or_none(val):
             dipole_integrals=as_array_or_none(data.get("dipole_integrals")),
             orbital_symmetries=data.get("orbital_symmetries"),
         )
+
+    def to_fcidump(self, file: str | bytes | os.PathLike) -> None:
+        """Save data to disk in FCIDUMP format.
+
+        .. note::
+            The FCIDUMP format does not retain all information stored in the
+            MolecularData object. To serialize a MolecularData object losslessly, use
+            the :func:`to_json` method to save to JSON format.
+
+        Args:
+            file: The file path to save to.
+        """
+        pyscf.tools.fcidump.from_integrals(
+            file,
+            h1e=self.one_body_integrals,
+            h2e=self.two_body_integrals,
+            nuc=self.core_energy,
+            nmo=self.norb,
+            nelec=self.nelec,
+        )
+
+    @staticmethod
+    def from_fcidump(file: str | bytes | os.PathLike) -> MolecularData:
+        """Initialize a MolecularData from an FCIDUMP file.
+
+        Args:
+            file: The FCIDUMP file path.
+        """
+        data = pyscf.tools.fcidump.read(file, verbose=False)
+        n_electrons = data["NELEC"]
+        spin = data["MS2"]
+        n_alpha = (n_electrons + spin) // 2
+        n_beta = (n_electrons - spin) // 2
+        return MolecularData(
+            core_energy=data["ECORE"],
+            one_body_integrals=data["H1"],
+            two_body_integrals=data["H2"],
+            norb=data["NORB"],
+            nelec=(n_alpha, n_beta),
+            spin=spin,
+        )
diff --git a/tests/python/molecular_data_test.py b/tests/python/molecular_data_test.py
index fbaeeeffd..b2bb9c177 100644
--- a/tests/python/molecular_data_test.py
+++ b/tests/python/molecular_data_test.py
@@ -138,8 +138,8 @@ def test_json_open_shell(tmp_path: pathlib.Path):
         _assert_mol_data_equal(loaded_mol_data, mol_data)
 
 
-def test_from_fcidump(tmp_path: pathlib.Path):
-    """Test loading from FCIDUMP."""
+def test_fcidump(tmp_path: pathlib.Path):
+    """Test saving to and loading from FCIDUMP."""
     mol = pyscf.gto.Mole()
     mol.build(
         atom=[["N", (0, 0, 0)], ["N", (1.0, 0, 0)]],
@@ -151,14 +151,7 @@ def test_from_fcidump(tmp_path: pathlib.Path):
     active_space = range(n_frozen, mol.nao_nr())
     scf = pyscf.scf.RHF(mol).run()
     mol_data = ffsim.MolecularData.from_scf(scf, active_space=active_space)
-    pyscf.tools.fcidump.from_integrals(
-        tmp_path / "test.fcidump",
-        h1e=mol_data.one_body_integrals,
-        h2e=mol_data.two_body_integrals,
-        nuc=mol_data.core_energy,
-        nmo=mol_data.norb,
-        nelec=mol_data.nelec,
-    )
+    mol_data.to_fcidump(tmp_path / "test.fcidump")
     loaded_mol_data = ffsim.MolecularData.from_fcidump(tmp_path / "test.fcidump")
     assert loaded_mol_data.norb == mol_data.norb
     assert loaded_mol_data.nelec == mol_data.nelec
@@ -174,14 +167,7 @@ def test_from_fcidump(tmp_path: pathlib.Path):
     )
     scf = pyscf.scf.ROHF(mol).run()
     mol_data = ffsim.MolecularData.from_scf(scf, active_space=active_space)
-    pyscf.tools.fcidump.from_integrals(
-        tmp_path / "test.fcidump",
-        h1e=mol_data.one_body_integrals,
-        h2e=mol_data.two_body_integrals,
-        nuc=mol_data.core_energy,
-        nmo=mol_data.norb,
-        nelec=mol_data.nelec,
-    )
+    mol_data.to_fcidump(tmp_path / "test.fcidump")
     loaded_mol_data = ffsim.MolecularData.from_fcidump(tmp_path / "test.fcidump")
     assert loaded_mol_data.norb == mol_data.norb
     assert loaded_mol_data.nelec == mol_data.nelec