What is the best way to save a psi4 molecule object in Python?

Question

what is the best way to save and load a psi4 molecule? I tried pickle/dill but they both failed. I see a feature to save as xyz, but I'm not seeing a function to load as xyz and I'm not sure if save xyz also saves the energy and charges etc.

import psi4

# Define the molecular geometry using a string
h2_geometry = """
H 0.0 0.0 0.0
H 0.0 0.0 0.74
"""

# Create the molecule object
h2_molecule = psi4.geometry(h2_geometry)

pickle_file_path = "psi4_mol.pkl"
with open(pickle_file_path, 'wb') as file:
    pickle.dump(h2_molecule, file)

Traceback (most recent call last):
  File "<stdin>", line 2, in <module>
TypeError: cannot pickle 'psi4.core.Molecule' object

I also tried dill:

with open(pickle_file_path, 'wb') as file:
    dill.dump(h2_molecule, file)

Traceback (most recent call last):
  File "<stdin>", line 2, in <module>
  File "lib/python3.11/site-packages/dill/_dill.py", line 250, in dump
    Pickler(file, protocol, **_kwds).dump(obj)
  File "lib/python3.11/site-packages/dill/_dill.py", line 418, in dump
    StockPickler.dump(self, obj)
  File "lib/python3.11/pickle.py", line 487, in dump
    self.save(obj)
  File "lib/python3.11/site-packages/dill/_dill.py", line 412, in save
    StockPickler.save(self, obj, save_persistent_id)
  File "lib/python3.11/pickle.py", line 560, in save
    f(self, obj)  # Call unbound method with explicit self
    ^^^^^^^^^^^^
  File "lib/python3.11/pickle.py", line 887, in save_tuple
    save(element)
  File "lib/python3.11/site-packages/dill/_dill.py", line 412, in save
    StockPickler.save(self, obj, save_persistent_id)
  File "lib/python3.11/pickle.py", line 578, in save
    rv = reduce(self.proto)
         ^^^^^^^^^^^^^^^^^^
TypeError: cannot pickle 'psi4.core.Molecule' object
```

Answer 1

If you only need the geometry, molecular charge, and other related quantities, then h2_geometry.to_dict() is what you want. You can then save the returned dict with e.g. pickle.

If you also want to compute energies and such then save them to file, you want wfn.to_file and the partner function for loading, wfn.from_file. These functions also generate (and read) dicts, which are serialized forms of the entire wavefunction. You may also pass an optional filename, which will save the wavefunction object as a numpy file.

You can get the wavefunction as a return from a calculation or building it yourself, e.g.

import psi4

mol = psi4.geometry("""
H 0.0 0.0 0.0
H 0.0 0.0 0.74""")

# build wfn yourself
bas = psi4.core.BasisSet.build(mol,"ORBITAL","sto-3g",quiet=True)
wfn = psi4.core.Wavefunction(mol,bas)

# get wfn as a return
psi4.set_options({"basis":"sto-3g"})
e,wfn = psi4.energy('scf',return_wfn=True)

# either way, save to file
wfn.to_file('wfn.npy')

# load it back up as a new wfn and print the variables inside
new_wfn = psi4.core.Wavefunction.from_file('wfn.npy')
print(new_wfn.variables())