Should program_unitary work for custom gates?

[ArfatSalman]

Pre-Report Checklist

• I am running the latest versions of pyQuil and the Forest SDK
• I checked to make sure that this bug has not already been reported

Issue Description

Is program_unitary supposed to work with Custom gates? I couldn't find any docs and it feels like it should work with any circuit.

How to Reproduce

If useful, provide a numbered list of the steps that result in the error.

Otherwise, just fill out the "Code Snippet" and "Error Output" sections below.

Code Snippet

import numpy as np

from pyquil import Program
from pyquil.quil import DefGate
from pyquil.simulation.tools import program_unitary

# First we define the new gate from a matrix
sqrt_x = np.array([[ 0.5+0.5j,  0.5-0.5j],
                   [ 0.5-0.5j,  0.5+0.5j]])

# Get the Quil definition for the new gate
sqrt_x_definition = DefGate("SQRT-X", sqrt_x)
# Get the gate constructor
SQRT_X = sqrt_x_definition.get_constructor()

# Then we can use the new gate
p = Program()
p += sqrt_x_definition
p += SQRT_X(0)


program_unitary(p, 1)

Error Output

KeyError                                  Traceback (most recent call last)
Cell In[3], line 1
----> 1 program_unitary(p, 2)

File ~/Desktop/bloqs/env/lib/python3.10/site-packages/pyquil/simulation/tools.py:335, in program_unitary(program, n_qubits)
    333 for instruction in program:
    334     if isinstance(instruction, Gate):
--> 335         unitary = lifted_gate(gate=instruction, n_qubits=n_qubits)
    336         umat = unitary.dot(umat)
    337     elif isinstance(instruction, Halt):

File ~/Desktop/bloqs/env/lib/python3.10/site-packages/pyquil/simulation/tools.py:320, in lifted_gate(gate, n_qubits)
    317         else:
    318             raise TypeError("Unsupported gate modifier {}".format(mod))
--> 320 matrix = _gate_matrix(gate)
    322 return lifted_gate_matrix(matrix=matrix, qubit_inds=[q.index for q in gate.qubits], n_qubits=n_qubits)

File ~/Desktop/bloqs/env/lib/python3.10/site-packages/pyquil/simulation/tools.py:296, in lifted_gate.<locals>._gate_matrix(gate)
    294         return QUANTUM_GATES[gate.name](*gate.params)  # type: ignore
    295     else:
--> 296         return QUANTUM_GATES[gate.name]  # type: ignore
    297 else:
    298     mod = gate.modifiers[0]

KeyError: 'SQRT-X'

Environment Context

Operating System: Mac OS Ventura 13.1 (Apple M1 Pro)

Python Version (e): Python 3.10.6

Quilc Version (quilc --version): 1.23.0 [e6c0939]

QVM Version (qvm --version): 1.17.1 [cf3f91f]

Python Environment Details (pip freeze or conda list):

anyio==3.6.2
appnope==0.1.3
argon2-cffi==21.3.0
argon2-cffi-bindings==21.2.0
arrow==1.2.3
asttokens==2.2.1
attrs==20.3.0
autopep8==2.0.1
Babel==2.11.0
backcall==0.2.0
beautifulsoup4==4.11.1
black==22.12.0
bleach==6.0.0
cachetools==5.3.0
certifi==2022.12.7
cffi==1.15.1
charset-normalizer==3.0.1
cirq==1.1.0
cirq-aqt==1.1.0
cirq-core==1.1.0
cirq-google==1.1.0
cirq-ionq==1.1.0
cirq-pasqal==1.1.0
cirq-rigetti==1.1.0
cirq-web==1.1.0
click==8.1.3
comm==0.1.2
contourpy==1.0.7
cycler==0.11.0
debugpy==1.6.6
decorator==5.1.1
defusedxml==0.7.1
duet==0.2.7
entrypoints==0.4
executing==1.2.0
fastjsonschema==2.16.2
fonttools==4.38.0
fqdn==1.5.1
google-api-core==1.34.0
google-auth==2.16.0
googleapis-common-protos==1.58.0
grpcio==1.51.1
grpcio-status==1.48.2
h11==0.14.0
httpcore==0.16.3
httpx==0.23.3
idna==3.4
ipykernel==6.20.2
ipython==8.8.0
ipython-genutils==0.2.0
iso8601==1.1.0
isoduration==20.11.0
jedi==0.18.2
Jinja2==3.1.2
json5==0.9.11
jsonpointer==2.3
jsonschema==4.17.3
jupyter-events==0.6.3
jupyter_client==7.4.9
jupyter_core==5.1.5
jupyter_server==2.1.0
jupyter_server_terminals==0.4.4
jupyterlab==3.5.3
jupyterlab-pygments==0.2.2
jupyterlab_server==2.19.0
kiwisolver==1.4.4
lark==0.11.3
MarkupSafe==2.1.2
matplotlib==3.6.3
matplotlib-inline==0.1.6
mistune==2.0.4
mpmath==1.2.1
msgpack==0.6.2
mypy-extensions==0.4.3
nbclassic==0.4.8
nbclient==0.7.2
nbconvert==7.2.9
nbformat==5.7.3
nest-asyncio==1.5.6
networkx==2.8.8
notebook==6.5.2
notebook_shim==0.2.2
numpy==1.23.5
packaging==23.0
pandas==1.5.3
pandocfilters==1.5.0
parso==0.8.3
pathspec==0.11.0
pexpect==4.8.0
pickleshare==0.7.5
Pillow==9.4.0
platformdirs==2.6.2
prometheus-client==0.16.0
prompt-toolkit==3.0.36
proto-plus==1.22.2
protobuf==3.20.3
psutil==5.9.4
ptyprocess==0.7.0
pure-eval==0.2.2
py==1.11.0
pyasn1==0.4.8
pyasn1-modules==0.2.8
pycodestyle==2.10.0
pycparser==2.21
pydantic==1.10.4
Pygments==2.14.0
PyJWT==2.6.0
pyparsing==3.0.9
pyquil==3.3.2
pyrsistent==0.19.3
python-dateutil==2.8.2
python-json-logger==2.0.4
python-rapidjson==1.9
pytz==2022.7.1
PyYAML==6.0
pyzmq==25.0.0
qcs-api-client==0.21.2
requests==2.28.2
retry==0.9.2
retrying==1.3.4
rfc3339==6.2
rfc3339-validator==0.1.4
rfc3986==1.5.0
rfc3986-validator==0.1.1
rpcq==3.10.0
rsa==4.9
ruamel.yaml==0.17.21
ruamel.yaml.clib==0.2.7
scipy==1.10.0
Send2Trash==1.8.0
six==1.16.0
sniffio==1.3.0
sortedcontainers==2.4.0
soupsieve==2.3.2.post1
stack-data==0.6.2
sympy==1.11.1
terminado==0.17.1
tinycss2==1.2.1
toml==0.10.2
tomli==2.0.1
tornado==6.2
tqdm==4.64.1
traitlets==5.8.1
types-python-dateutil==2.8.19.6
types-retry==0.9.9.1
typing_extensions==4.4.0
uri-template==1.2.0
urllib3==1.26.14
wcwidth==0.2.6
webcolors==1.12
webencodings==0.5.1
websocket-client==1.4.2