Add trial wavefunction.

recirq/qcqmc/__init__.py

@@ -11,3 +11,47 @@
 # 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 .hamiltonian import (
+    HamiltonianData,
+    LoadFromFileHamiltonianParams,
+    PyscfHamiltonianParams,
+)
+from .trial_wf import (
+    FermionicMode,
+    LayerSpec,
+    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 [
+            LoadFromFileHamiltonianParams,
+            HamiltonianData,
+            PyscfHamiltonianParams,
+            FermionicMode,
+            LayerSpec,
+            PerfectPairingPlusTrialWavefunctionParams,
+            TrialWavefunctionData,
+        ]
+    }.get(cirq_type, None)
+
+
+DEFAULT_RESOLVERS.append(_resolve_json)

recirq/qcqmc/conftest.py

@@ -0,0 +1,92 @@
+from typing import Tuple
+
+import numpy as np
+import pytest
+
+from recirq.qcqmc.hamiltonian import (
+    HamiltonianData,
+    LoadFromFileHamiltonianParams,
+    build_hamiltonian_from_file,
+)
+from recirq.qcqmc.trial_wf import (
+    PerfectPairingPlusTrialWavefunctionParams,
+    TrialWavefunctionData,
+    _get_qubits_a_b_reversed,
+    build_pp_plus_trial_wavefunction,
+)
+
+
+@pytest.fixture(scope="package")
+def fixture_4_qubit_ham() -> HamiltonianData:
+    params = LoadFromFileHamiltonianParams(
+        name="test hamiltonian 4 qubits", integral_key="fh_sto3g", n_orb=2, n_elec=2
+    )
+
+    hamiltonian_data = build_hamiltonian_from_file(params)
+
+    return hamiltonian_data
+
+
+@pytest.fixture(scope="package")
+def fixture_8_qubit_ham() -> HamiltonianData:
+    params = LoadFromFileHamiltonianParams(
+        name="test hamiltonian 8 qubits", integral_key="h4_sto3g", n_orb=4, n_elec=4
+    )
+
+    hamiltonian_data = build_hamiltonian_from_file(params)
+
+    return hamiltonian_data
+
+
+@pytest.fixture(scope="package")
+def fixture_12_qubit_ham() -> HamiltonianData:
+    params = LoadFromFileHamiltonianParams(
+        name="test hamiltonian 12 qubits",
+        integral_key="diamond_dzvp/cas66",
+        n_orb=6,
+        n_elec=6,
+        do_eri_restore=True,
+    )
+
+    hamiltonian_data = build_hamiltonian_from_file(params)
+
+    return hamiltonian_data
+
+
+@pytest.fixture(scope="package")
+def fixture_4_qubit_ham_and_trial_wf(
+    fixture_4_qubit_ham: HamiltonianData,
+) -> Tuple[HamiltonianData, TrialWavefunctionData]:
+    params = PerfectPairingPlusTrialWavefunctionParams(
+        name="pp_test_wf_1",
+        hamiltonian_params=fixture_4_qubit_ham.params,
+        heuristic_layers=tuple(),
+        do_pp=True,
+        restricted=True,
+    )
+
+    trial_wf = build_pp_plus_trial_wavefunction(
+        params, dependencies={fixture_4_qubit_ham.params: fixture_4_qubit_ham}
+    )
+
+    return fixture_4_qubit_ham, trial_wf
+
+
+@pytest.fixture(scope="package")
+def fixture_8_qubit_ham_and_trial_wf(
+    fixture_8_qubit_ham: HamiltonianData,
+) -> Tuple[HamiltonianData, TrialWavefunctionData]:
+    params = PerfectPairingPlusTrialWavefunctionParams(
+        name="pp_test_wf_qchem",
+        hamiltonian_params=fixture_8_qubit_ham.params,
+        heuristic_layers=tuple(),
+        initial_orbital_rotation=None,
+        initial_two_body_qchem_amplitudes=np.asarray([0.3, 0.4]),
+        do_optimization=False,
+    )
+
+    trial_wf = build_pp_plus_trial_wavefunction(
+        params, dependencies={fixture_8_qubit_ham.params: fixture_8_qubit_ham}
+    )
+
+    return fixture_8_qubit_ham, trial_wf