QCQMC Part 7: Add Experiment classes (#354)

* Add experiment class.
* Fix fixture.
* Add check for serialization.
* Add experiment.

conftest.py

@@ -16,6 +16,7 @@
 import numpy as np
 import pytest
 
+from recirq.qcqmc import blueprint, qubit_maps
 from recirq.qcqmc.hamiltonian import (
     HamiltonianData,
     HamiltonianFileParams,
@@ -104,6 +105,30 @@ def fixture_8_qubit_ham_and_trial_wf(
     return fixture_8_qubit_ham, trial_wf
 
 
+@pytest.fixture(scope="package")
+def fixture_4_qubit_ham_trial_wf_and_blueprint(
+    fixture_4_qubit_ham_and_trial_wf,
+) -> Tuple[HamiltonianData, TrialWavefunctionData, blueprint.BlueprintData]:
+    ham_data, trial_wf_data = fixture_4_qubit_ham_and_trial_wf
+    trial_wf_params = trial_wf_data.params
+
+    blueprint_params = blueprint.BlueprintParamsTrialWf(
+        name="blueprint_test",
+        trial_wf_params=trial_wf_params,
+        n_cliffords=17,
+        qubit_partition=(
+            tuple(qubit_maps.get_qubits_a_b_reversed(n_orb=trial_wf_params.n_orb)),
+        ),
+        seed=1,
+    )
+
+    bp = blueprint.BlueprintData.build_blueprint_from_dependencies(
+        blueprint_params, dependencies={trial_wf_params: trial_wf_data}
+    )
+
+    return ham_data, trial_wf_data, bp
+
+
 def pytest_addoption(parser):
     parser.addoption("--skipslow", action="store_true", help="skips slow tests")
 

recirq/qcqmc/__init__.py

@@ -16,14 +16,16 @@
 
 from cirq.protocols.json_serialization import DEFAULT_RESOLVERS, ObjectFactory
 
-from .blueprint import (BlueprintData, BlueprintParamsRobustShadow,
-                        BlueprintParamsTrialWf)
+from .blueprint import (
+    BlueprintData,
+    BlueprintParamsRobustShadow,
+    BlueprintParamsTrialWf,
+)
+from .experiment import ExperimentData, SimulatedExperimentParams
 from .fermion_mode import FermionicMode
-from .hamiltonian import (HamiltonianData, HamiltonianFileParams,
-                          PyscfHamiltonianParams)
+from .hamiltonian import HamiltonianData, HamiltonianFileParams, PyscfHamiltonianParams
 from .layer_spec import LayerSpec
-from .trial_wf import (PerfectPairingPlusTrialWavefunctionParams,
-                       TrialWavefunctionData)
+from .trial_wf import PerfectPairingPlusTrialWavefunctionParams, TrialWavefunctionData
 
 
 @lru_cache()
@@ -40,16 +42,18 @@ def _resolve_json(cirq_type: str) -> Optional[ObjectFactory]:
     return {
         k.__name__: k
         for k in [
+            BlueprintParamsTrialWf,
+            BlueprintParamsRobustShadow,
+            BlueprintData,
+            ExperimentData,
+            FermionicMode,
             HamiltonianFileParams,
             HamiltonianData,
-            PyscfHamiltonianParams,
-            FermionicMode,
             LayerSpec,
             PerfectPairingPlusTrialWavefunctionParams,
+            PyscfHamiltonianParams,
+            SimulatedExperimentParams,
             TrialWavefunctionData,
-            BlueprintParamsTrialWf,
-            BlueprintParamsRobustShadow,
-            BlueprintData,
         ]
     }.get(cirq_type, None)
 

recirq/qcqmc/experiment.py

@@ -0,0 +1,182 @@
+# Copyright 2024 Google
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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 datetime import datetime
+from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple, Union
+
+import attrs
+import cirq
+import numpy as np
+import qsimcirq
+from pytz import timezone
+
+from recirq.qcqmc import blueprint, config, data, for_refactor
+
+
+@attrs.frozen(repr=False)
+class SimulatedExperimentParams(data.Params):
+    """Class for storing the parameters that specify an ExperimentData object.
+
+    This stage of the experiment concerns itself with executing circuits and doing
+    classical post-processing of shadow tomography data.
+
+    Args:
+        name: A `Params` name for this experiment.
+        blueprint_params: Backreference to the `BlueprintParams` preceding this stage.
+        n_samples_per_clifford: Number of circuit repetitions to take for each clifford.
+        noise_model_name: For simulation; see `utilities.get_noise_model`.
+        noise_model_params: For simulation; see `utilities.get_noise_model`.
+        seed: The random seed to use for simulation.
+        path_prefix: An optional path string prefix for the output files.
+    """
+
+    name: str
+    blueprint_params: blueprint.BlueprintParams
+    n_samples_per_clifford: int
+    noise_model_name: str
+    noise_model_params: Optional[Tuple[float, ...]] = attrs.field(
+        converter=lambda x: tuple(x) if x is not None else None, default=None
+    )
+    seed: int = 0
+    path_prefix: str = ""
+
+    @property
+    def path_string(self) -> str:
+        return (
+            self.path_prefix + config.OUTDIRS.DEFAULT_EXPERIMENT_DIRECTORY + self.name
+        )
+
+    def _json_dict_(self):
+        simple_dict = attrs.asdict(self)
+        simple_dict["blueprint_params"] = self.blueprint_params
+        return simple_dict
+
+
+@attrs.frozen(repr=False, eq=False)
+class ExperimentData(data.Data):
+    """The data defining the experimental result.
+
+    Args:
+        params: The experimental parameters.
+        raw_samples: An array of shape [n_cliffords, n_samples_per_clifford, n_qubits]
+            containing the bitstrings sampled from each of the different
+            circuits.
+        metadata: Any metadata associated with the run.
+    """
+
+    params: SimulatedExperimentParams
+    raw_samples: np.ndarray = attrs.field(converter=np.asarray)
+    metadata: Dict[str, Any] = attrs.field(factory=dict)
+
+    def _json_dict_(self):
+        simple_dict = attrs.asdict(self)
+        simple_dict["params"] = self.params
+        return simple_dict
+
+
+    @classmethod
+    def build_experiment_from_dependencies(
+        cls, params: SimulatedExperimentParams, *, dependencies: Dict[data.Params, data.Data]
+    ) -> 'ExperimentData':
+        """Builds an ExperimentData from ExperimentParams and any dependencies.
+
+        Args:
+            params: The experimental parameters.
+            dependencies: The dependencies leading up to this point (in particular the blueprint.)
+        """
+        bp = dependencies[params.blueprint_params]
+        assert isinstance(bp, blueprint.BlueprintData)
+        assert params.blueprint_params == bp.params
+
+        noise_model = for_refactor.get_noise_model(
+            params.noise_model_name, params.noise_model_params
+        )
+
+        raw_samples = get_samples_from_simulation(
+            bp.compiled_circuit,
+            bp.resolvers,
+            noise_model,
+            params.n_samples_per_clifford,
+            params.seed,
+        )
+
+        metadata = get_experimental_metadata()
+
+        return ExperimentData(params=params, raw_samples=raw_samples, metadata=metadata)
+
+
+def get_samples_from_simulation(
+    circuit: cirq.Circuit,
+    resolvers: List[cirq.ParamResolverOrSimilarType],
+    noise_model: Union[None, cirq.NoiseModel],
+    n_samples_per_clifford: int,
+    seed: Optional[int] = None,
+    simulate_single_precision: bool = config.SINGLE_PRECISION_DEFAULT,
+) -> np.ndarray:
+    """Samples the circuits and returns an array of sampled bits.
+
+    Args:
+        circuits: The shadow tomography circuits.
+        resolvers: A list of cirq parameter resolvers.
+        noise_model: An optional cirq.NoiseModel for the simulated experiment.
+        n_samples_per_clifford: The number of samples to take per clifford sample.
+        seed: An optional random seed
+        simulate_single_precision: Use single precision instead of double
+            precision for the circuit simulation.
+
+    Returns:
+        raw_samples: An array of shape [n_cliffords, n_samples_per_clifford,
+            n_qubits] containing the bitstrings sampled from each of the different
+            circuits.
+
+    """
+    simulator = qsimcirq.QSimSimulator()
+
+    sampled_bitstrings = []
+
+    if noise_model is not None:
+        circuit = cirq.Circuit(
+            noise_model.noisy_moments(circuit, sorted(circuit.all_qubits()))
+        )
+
+    for _, resolver in enumerate(resolvers):
+        results = simulator.run(
+            circuit, repetitions=n_samples_per_clifford, param_resolver=resolver
+        )
+
+        outcomes = results.measurements["all"]
+        sampled_bitstrings.append(outcomes)
+
+    raw_samples = np.stack(sampled_bitstrings)
+    return raw_samples
+
+
+def get_experimental_metadata() -> Dict[str, object]:
+    """Gets some metadata to store along with the results for an experiment.
+
+    Returns:
+        A dictionary of useful metadata including the date and time of the experiment.
+    """
+
+    date_time = datetime.now()
+    pacific_tz = timezone("US/Pacific")
+    pacific_date_time = date_time.astimezone(pacific_tz)
+
+    formatted_date_time = pacific_date_time.strftime("%m/%d/%Y, %H:%M:%S")
+
+    metadata: Dict[str, Any] = {}
+    metadata["PST_formatted_date_time"] = formatted_date_time
+    metadata["iso_formatted_date_time"] = date_time.isoformat()
+
+    return metadata

recirq/qcqmc/experiment_test.py

@@ -0,0 +1,56 @@
+# Copyright 2024 Google
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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 typing import Tuple
+
+import cirq
+
+from recirq.qcqmc.blueprint import BlueprintData
+from recirq.qcqmc.experiment import SimulatedExperimentParams, get_experimental_metadata, ExperimentData
+from recirq.qcqmc.hamiltonian import HamiltonianData
+from recirq.qcqmc.trial_wf import TrialWavefunctionData
+
+
+def test_small_experiment_raw_samples_shape(
+    fixture_4_qubit_ham_trial_wf_and_blueprint: Tuple[
+        HamiltonianData, TrialWavefunctionData, BlueprintData
+    ]
+):
+    _, _, blueprint_data = fixture_4_qubit_ham_trial_wf_and_blueprint
+
+    simulated_experiment_params = SimulatedExperimentParams(
+        name="test_1",
+        blueprint_params=blueprint_data.params,
+        noise_model_name="None",
+        noise_model_params=(0,),
+        n_samples_per_clifford=31,
+        seed=1,
+    )
+
+    experiment = ExperimentData.build_experiment_from_dependencies(
+        params=simulated_experiment_params,
+        dependencies={blueprint_data.params: blueprint_data},
+    )
+
+    raw_samples = experiment.raw_samples
+
+    assert raw_samples.shape == (17, 31, 4)
+
+    exp2 = cirq.read_json(json_text=cirq.to_json(experiment))
+    assert exp2 == experiment
+
+def test_get_experimental_metadata():
+    md = get_experimental_metadata()
+    assert md.get('PST_formatted_date_time') is not None
+    assert md.get('iso_formatted_date_time') is not None