fix: Handle quil_rs FunctionCallExpressions

rigetti · Sep 17, 2023 · ab551d5 · ab551d5
1 parent 79989c6
commit ab551d5
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 9 deletions.
diff --git a/pyquil/quilatom.py b/pyquil/quilatom.py
@@ -449,6 +449,19 @@ def _convert_to_py_expression(
             return BinaryExp._from_rs_infix_expression(expression.to_infix())
         if expression.is_address():
             return MemoryReference._from_rs_memory_reference(expression.to_address())
+        if expression.is_function_call():
+            fc = expression.to_function_call()
+            parameter = _convert_to_py_expression(fc.expression)
+            if fc.function == quil_rs_expr.ExpressionFunction.Cis:
+                return quil_cis(parameter)
+            if fc.function == quil_rs_expr.ExpressionFunction.Cosine:
+                return quil_cos(parameter)
+            if fc.function == quil_rs_expr.ExpressionFunction.Exponent:
+                return quil_exp(parameter)
+            if fc.function == quil_rs_expr.ExpressionFunction.Sine:
+                return quil_sin(parameter)
+            if fc.function == quil_rs_expr.ExpressionFunction.SquareRoot:
+                return quil_sqrt(parameter)
         if expression.is_prefix():
             prefix = expression.to_prefix()
             py_expression = _convert_to_py_expression(prefix.expression)

diff --git a/test/unit/__snapshots__/test_quilbase.ambr b/test/unit/__snapshots__/test_quilbase.ambr
@@ -240,10 +240,10 @@
 # name: TestDefGate.test_out[Params]
   '''
   DEFGATE ParameterizedGate(%X) AS MATRIX:
-  	%X, 0, 0, 0
-  	0, %X, 0, 0
-  	0, 0, %X, 0
-  	0, 0, 0, %X
+  	cos(%X), 0, 0, 0
+  	0, cos(%X), 0, 0
+  	0, 0, cos(%X), 0
+  	0, 0, 0, cos(%X)
 
   '''
 # ---
@@ -260,10 +260,10 @@
 # name: TestDefGate.test_str[Params]
   '''
   DEFGATE ParameterizedGate(%X) AS MATRIX:
-  	%X, 0, 0, 0
-  	0, %X, 0, 0
-  	0, 0, %X, 0
-  	0, 0, 0, %X
+  	cos(%X), 0, 0, 0
+  	0, cos(%X), 0, 0
+  	0, 0, cos(%X), 0
+  	0, 0, 0, cos(%X)
 
   '''
 # ---

diff --git a/test/unit/test_quilbase.py b/test/unit/test_quilbase.py
@@ -6,6 +6,7 @@
 import pytest
 from syrupy.assertion import SnapshotAssertion
 
+from pyquil.quilatom import quil_cos
 from pyquil.gates import X
 from pyquil.quil import Program
 from pyquil.quilbase import (
@@ -162,7 +163,7 @@ def test_compile(self, program: Program, compiler: QPUCompiler):
     ("name", "matrix", "parameters"),
     [
         ("NoParamGate", np.eye(4), []),
-        ("ParameterizedGate", np.diag([Parameter("X")] * 4), [Parameter("X")]),
+        ("ParameterizedGate", np.diag([quil_cos(Parameter("X"))] * 4), [Parameter("X")]),
     ],
     ids=("No-Params", "Params"),
 )