Parse context with ProjectorSums

tensorflow_quantum/core/ops/batch_util.py

@@ -23,7 +23,7 @@
 #   be used end to end with engine. This current issue is that
 #   cirq.PauliSumCollector does not produce serializable gates for basis
 #   conversion.
-class TFQPauliSumCollector(cirq.work.collector.Collector):
+class TFQPauliProjectorSumCollector(cirq.work.collector.Collector):
     """Copy of cirq.PauliSumCollector with some fixes to work with engine."""
 
     def __init__(self,
@@ -261,7 +261,7 @@ def batch_calculate_expectation(circuits, param_resolvers, ops, simulator):
         if not isinstance(sub_list, (list, tuple, np.ndarray)):
             raise TypeError('elements of ops must be type list.')
         for x in sub_list:
-            if not isinstance(x, cirq.PauliSum):
+            if not isinstance(x, (cirq.PauliSum, cirq.ProjectorSum)):
                 raise TypeError('ops must contain only cirq.PauliSum objects.'
                                 ' Given: {}'.format(type(x)))
 
@@ -279,8 +279,13 @@ def batch_calculate_expectation(circuits, param_resolvers, ops, simulator):
             sim_result = simulator.simulate(c, p)
             for j, op in enumerate(op_row):
                 dm = sim_result.final_density_matrix
-                all_exp_vals[i][j] = op.expectation_from_density_matrix(
-                    dm, qubit_order, check_preconditions=False)
+                if isinstance(op, cirq.PauliSum):
+                    all_exp_vals[i][j] = op.expectation_from_density_matrix(
+                        dm, qubit_order, check_preconditions=False)
+                else:
+                    assert isinstance(op, cirq.ProjectorSum)
+                    all_exp_vals[i][j] = op.expectation_from_density_matrix(
+                        dm, qubit_order)
         else:
             # Valid observables always have real expectation values.
             all_exp_vals[i] = np.real(
@@ -359,7 +364,7 @@ def batch_calculate_sampled_expectation(circuits, param_resolvers, ops,
             continue
         circuit = cirq.resolve_parameters(c, params)
         for op_index, op in enumerate(ops[c_index]):
-            collector = TFQPauliSumCollector(
+            collector = TFQPauliProjectorSumCollector(
                 circuit, op, samples_per_term=n_samples[c_index][op_index])
             collector.collect(sampler)
             result = collector.estimated_energy().real

tensorflow_quantum/core/ops/math_ops/tfq_inner_product.cc

@@ -35,7 +35,6 @@ limitations under the License.
 namespace tfq {
 
 using ::tensorflow::Status;
-using ::tfq::proto::PauliSum;
 using ::tfq::proto::Program;
 
 typedef qsim::Cirq::GateCirq<float> QsimGate;

tensorflow_quantum/core/ops/math_ops/tfq_inner_product_grad.cc

@@ -36,7 +36,6 @@ limitations under the License.
 namespace tfq {
 
 using ::tensorflow::Status;
-using ::tfq::proto::PauliSum;
 using ::tfq::proto::Program;
 
 typedef qsim::Cirq::GateCirq<float> QsimGate;

tensorflow_quantum/core/ops/math_ops/tfq_simulate_1d_expectation.cc

@@ -85,7 +85,7 @@ class TfqSimulateMPS1DExpectationOp : public tensorflow::OpKernel {
     // is resolved.
     OP_REQUIRES_OK(context,
                    GetProgramsAndNumQubits(context, &programs, &num_qubits,
-                                           &pauli_sums, true));
+                                           &pauli_sums, nullptr, true));
 
     std::vector<SymbolMap> maps;
     OP_REQUIRES_OK(context, GetSymbolMaps(context, &maps));

tensorflow_quantum/core/ops/math_ops/tfq_simulate_1d_sampled_expectation.cc

@@ -81,7 +81,7 @@ class TfqSimulateMPS1DSampledExpectationOp : public tensorflow::OpKernel {
     std::vector<std::vector<PauliSum>> pauli_sums;
     OP_REQUIRES_OK(context,
                    GetProgramsAndNumQubits(context, &programs, &num_qubits,
-                                           &pauli_sums, true));
+                                           &pauli_sums, nullptr, true));
 
     std::vector<SymbolMap> maps;
     OP_REQUIRES_OK(context, GetSymbolMaps(context, &maps));

tensorflow_quantum/core/ops/math_ops/tfq_simulate_1d_samples.cc

@@ -66,7 +66,7 @@ class TfqSimulateMPS1DSamplesOp : public tensorflow::OpKernel {
     std::vector<int> num_qubits;
     OP_REQUIRES_OK(context,
                    GetProgramsAndNumQubits(context, &programs, &num_qubits,
-                                           nullptr, true));
+                                           nullptr, nullptr, true));
 
     // Parse symbol maps for parameter resolution in the circuits.
     std::vector<SymbolMap> maps;

tensorflow_quantum/core/ops/noise/BUILD

@@ -83,6 +83,7 @@ py_library(
 
 py_test(
     name = "noisy_expectation_op_test",
+    timeout = "long",
     srcs = ["noisy_expectation_op_test.py"],
     python_version = "PY3",
     deps = [
@@ -103,6 +104,7 @@ py_library(
 
 py_test(
     name = "noisy_sampled_expectation_op_test",
+    timeout = "long",
     srcs = ["noisy_sampled_expectation_op_test.py"],
     python_version = "PY3",
     deps = [

tensorflow_quantum/core/ops/noise/noisy_expectation_op.py

@@ -20,7 +20,12 @@
 NOISY_OP_MODULE = load_module(os.path.join("noise", "_tfq_noise_ops.so"))
 
 
-def expectation(programs, symbol_names, symbol_values, pauli_sums, num_samples):
+def expectation(programs,
+                symbol_names,
+                symbol_values,
+                pauli_sums,
+                num_samples,
+                projector_sums=None):
     """Calculate the analytic expectation values using monte-carlo trajectories.
 
     Simulate the final state of `programs` given `symbol_values` are placed
@@ -88,11 +93,21 @@ def expectation(programs, symbol_names, symbol_values, pauli_sums, num_samples):
             threads to TensorFlow. For best performance ensure that the
             quantities in `num_samples` are a multiple of the number of
             available threads.
+        projector_sums: `tf.Tensor` of strings with shape [batch_size, n_ops]
+            containing the string representation of the operators that will
+            be used on all of the circuits in the expectation calculations.
     Returns:
         `tf.Tensor` with shape [batch_size, n_ops] that holds the
             expectation value for each circuit with each op applied to it
             (after resolving the corresponding parameters in).
     """
+    # TODO(tonybruguier): Always supply a projector sum to this function and
+    # remove the special-casing below.
+    if projector_sums is None:
+        if len(pauli_sums.shape) == 1:
+            projector_sums = pauli_sums[0:0]
+        else:
+            projector_sums = pauli_sums[:, 0:0]
     return NOISY_OP_MODULE.tfq_noisy_expectation(
         programs, symbol_names, tf.cast(symbol_values, tf.float32), pauli_sums,
-        tf.cast(num_samples, dtype=tf.int32))
+        projector_sums, tf.cast(num_samples, dtype=tf.int32))