Allow symbol for coupling_mhz in CouplerPulse gate

cirq-google/cirq_google/experimental/ops/coupler_pulse.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Optional, Tuple
+from typing import AbstractSet, Any, Optional, Tuple
 
 import numpy as np
 
@@ -51,7 +51,7 @@ class CouplerPulse(cirq.ops.Gate):
     def __init__(
         self,
         hold_time: cirq.Duration,
-        coupling_mhz: float,
+        coupling_mhz: cirq.TParamVal,
         rise_time: Optional[cirq.Duration] = cirq.Duration(nanos=8),
         padding_time: Optional[cirq.Duration] = cirq.Duration(nanos=2.5),
     ):
@@ -92,6 +92,32 @@ def __str__(self) -> str:
             + f'padding_time={self.padding_time})'
         )
 
+    def _is_parameterized_(self) -> bool:
+        return (
+            cirq.is_parameterized(self.hold_time)
+            or cirq.is_parameterized(self.rise_time)
+            or cirq.is_parameterized(self.padding_time)
+            or cirq.is_parameterized(self.coupling_mhz)
+        )
+
+    def _parameter_names_(self: Any) -> AbstractSet[str]:
+        return (
+            cirq.parameter_names(self.hold_time)
+            | cirq.parameter_names(self.coupling_mhz)
+            | cirq.parameter_names(self.rise_time)
+            | cirq.parameter_names(self.padding_time)
+        )
+
+    def _resolve_parameters_(
+        self, resolver: cirq.ParamResolverOrSimilarType, recursive: bool
+    ) -> 'CouplerPulse':
+        return CouplerPulse(
+            hold_time=cirq.resolve_parameters(self.hold_time, resolver, recursive=recursive),
+            coupling_mhz=cirq.resolve_parameters(self.coupling_mhz, resolver, recursive=recursive),
+            rise_time=cirq.resolve_parameters(self.rise_time, resolver, recursive=recursive),
+            padding_time=cirq.resolve_parameters(self.padding_time, resolver, recursive=recursive),
+        )
+
     def _value_equality_values_(self) -> Any:
         return self.hold_time, self.coupling_mhz, self.rise_time, self.padding_time
 

cirq-google/cirq_google/experimental/ops/coupler_pulse_test.py

@@ -11,6 +11,8 @@
 # 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.
+import pytest
+import sympy
 
 import cirq
 import cirq_google.experimental.ops.coupler_pulse as coupler_pulse
@@ -112,3 +114,86 @@ def test_coupler_pulse_circuit_diagram():
 1: ───/‾‾(10 ns@25.0MHz)‾‾\───
 """,
     )
+
+
+@pytest.mark.parametrize(
+    'gate, resolver, expected',
+    [
+        (
+            coupler_pulse.CouplerPulse(
+                hold_time=cirq.Duration(nanos=sympy.Symbol('t_ns')), coupling_mhz=10
+            ),
+            {'t_ns': 50},
+            coupler_pulse.CouplerPulse(hold_time=cirq.Duration(nanos=50), coupling_mhz=10),
+        ),
+        (
+            coupler_pulse.CouplerPulse(
+                hold_time=cirq.Duration(nanos=50), coupling_mhz=sympy.Symbol('g')
+            ),
+            {'g': 10},
+            coupler_pulse.CouplerPulse(hold_time=cirq.Duration(nanos=50), coupling_mhz=10),
+        ),
+    ],
+)
+def test_coupler_pulse_resolution(gate, resolver, expected):
+    assert cirq.resolve_parameters(gate, resolver) == expected
+
+
+@pytest.mark.parametrize(
+    'gate, param_names',
+    [
+        (
+                coupler_pulse.CouplerPulse(
+                    hold_time=cirq.Duration(nanos=sympy.Symbol('t_ns')), coupling_mhz=10
+                ),
+                {'t_ns'},
+        ),
+        (
+                coupler_pulse.CouplerPulse(
+                    hold_time=cirq.Duration(nanos=50), coupling_mhz=sympy.Symbol('g')
+                ),
+                {'g'},
+        ),
+        (
+                coupler_pulse.CouplerPulse(
+                    hold_time=cirq.Duration(nanos=sympy.Symbol('t_ns')), coupling_mhz=sympy.Symbol('g')
+                ),
+                {'g', 't_ns'},
+        ),
+    ],
+)
+def test_coupler_pulse_parameter_names(gate, param_names):
+    assert cirq.parameter_names(gate) == param_names
+
+
+@pytest.mark.parametrize(
+    'gate, is_parameterized',
+    [
+        (
+                coupler_pulse.CouplerPulse(
+                    hold_time=cirq.Duration(nanos=50), coupling_mhz=10
+                ),
+                False,
+        ),
+        (
+                coupler_pulse.CouplerPulse(
+                    hold_time=cirq.Duration(nanos=sympy.Symbol('t_ns')), coupling_mhz=10
+                ),
+                True,
+        ),
+        (
+                coupler_pulse.CouplerPulse(
+                    hold_time=cirq.Duration(nanos=50), coupling_mhz=sympy.Symbol('g')
+                ),
+                True,
+        ),
+        (
+                coupler_pulse.CouplerPulse(
+                    hold_time=cirq.Duration(nanos=sympy.Symbol('t_ns')), coupling_mhz=sympy.Symbol('g')
+                ),
+                True,
+        ),
+    ],
+)
+def test_coupler_pulse_is_parameterized(gate, is_parameterized):
+    assert cirq.is_parameterized(gate) == is_parameterized