QCQMC Part 5: Add Trial Wavefunction

recirq/qcqmc/__init__.py

@@ -11,3 +11,40 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from functools import lru_cache
+from typing import Optional
+
+from cirq.protocols.json_serialization import DEFAULT_RESOLVERS, ObjectFactory
+
+from .fermion_mode import FermionicMode
+from .hamiltonian import HamiltonianData, HamiltonianFileParams, PyscfHamiltonianParams
+from .layer_spec import LayerSpec
+from .trial_wf import PerfectPairingPlusTrialWavefunctionParams, TrialWavefunctionData
+
+
+@lru_cache()
+def _resolve_json(cirq_type: str) -> Optional[ObjectFactory]:
+    """Resolve the types of `recirq.qcqmc.` json objects.
+
+    This is a Cirq JSON resolver suitable for appending to
+    `cirq.protocols.json_serialization.DEFAULT_RESOLVERS`.
+    """
+    if not cirq_type.startswith("recirq.qcqmc."):
+        return None
+
+    cirq_type = cirq_type[len("recirq.qcqmc.") :]
+    return {
+        k.__name__: k
+        for k in [
+            HamiltonianFileParams,
+            HamiltonianData,
+            PyscfHamiltonianParams,
+            FermionicMode,
+            LayerSpec,
+            PerfectPairingPlusTrialWavefunctionParams,
+            TrialWavefunctionData,
+        ]
+    }.get(cirq_type, None)
+
+
+DEFAULT_RESOLVERS.append(_resolve_json)