From 9c3034e7daa1c8afc6f298292b7814eaea20d877 Mon Sep 17 00:00:00 2001
From: Eric Berquist <eric.berquist@gmail.com>
Date: Sun, 11 Jul 2021 15:16:12 -0400
Subject: [PATCH] QCSchemaWriter: write MO coefficients and energies

---
 cclib/io/qcschemawriter.py | 44 +++++++++++++++++++++++++++++++++++++-
 1 file changed, 43 insertions(+), 1 deletion(-)

diff --git a/cclib/io/qcschemawriter.py b/cclib/io/qcschemawriter.py
index 3172f33dc..83979742e 100644
--- a/cclib/io/qcschemawriter.py
+++ b/cclib/io/qcschemawriter.py
@@ -17,6 +17,8 @@
 )
 from cclib.parser.utils import convertor, find_package
 
+import numpy as np
+
 _found_qcschema = find_package("qcschema")
 if _found_qcschema:
     import qcschema
@@ -81,9 +83,11 @@ def as_dict(self, validate=True):
             # FIXME
             method = "HF"
 
+        basis_set_name = metadata["basis_set"].lower()
+
         qcschema_dict["model"] = {
             "method": method.lower(),
-            "basis": metadata["basis_set"].lower(),
+            "basis": basis_set_name,
         }
 
         scf_total_energy = convertor(self.ccdata.scfenergies[-1], "eV", "hartree")
@@ -183,6 +187,44 @@ def as_dict(self, validate=True):
                 }
             )
 
+        qcschema_dict["wavefunction"] = {
+            "basis": {"name": basis_set_name, "center_data": {}, "atom_map": []},
+            # TODO in latest schema version
+            # "restricted": bool,
+        }
+
+        has_beta = len(self.ccdata.homos) == 2
+
+        if hasattr(self.ccdata, "gbasis"):
+            pass
+
+        # FIXME Again, since we currently do not know if the (natural) orbital
+        # coefficients and eigenvalues come from diagonalizing a (correlated)
+        # density matrix, we assume that they are from SCF.
+        if hasattr(self.ccdata, "moenergies"):
+            qcschema_dict["wavefunction"]["scf_eigenvalues_a"] = self.ccdata.moenergies[
+                0
+            ].tolist()
+            if has_beta:
+                qcschema_dict["wavefunction"][
+                    "scf_eigenvalues_b"
+                ] = self.ccdata.moenergies[1].tolist()
+
+        if hasattr(self.ccdata, "mocoeffs"):
+            mocoeffs_a = self.ccdata.mocoeffs[0]
+            # Sometimes we don't parse all MOs and fill missing ones with NaN,
+            # which isn't allowed by QCSchema.
+            if not np.isnan(mocoeffs_a).any():
+                qcschema_dict["wavefunction"]["scf_orbitals_a"] = (
+                    mocoeffs_a.transpose().tolist()
+                )
+            if has_beta:
+                mocoeffs_b = self.ccdata.mocoeffs[1]
+                if not np.isnan(mocoeffs_b).any():
+                    qcschema_dict["wavefunction"]["scf_orbitals_b"] = (
+                        mocoeffs_b.transpose().tolist()
+                    )
+
         if driver == "energy":
             return_result = return_energy
         elif driver == "gradient":