Repository: rigetti/pyquil (python)

Issue:
QVM returns only MEASURE'd memory regions
Issue Description
-----------------

In #873, we attempted to enable the QVM (and, later, the QPU) to let the user do post-execution analysis on all memory regions, rather than just on the contents of `ro`. However, in https://github.com/rigetti/pyquil/blob/master/pyquil/api/_qvm.py#L544 we filter the memory regions that the QVM returns down to only those into which it `MEASURE`s. This is unexpected behavior, and this restriction should be lifted.

How to Reproduce
----------------

### Code Snippet

```python
from pyquil import get_qc, Program

qc = get_qc('2q-qvm') 
qc.run(Program("DECLARE ro BIT\nDECLARE not_ro BIT\nMOVE ro 1\nMOVE not_ro 1\nMEASURE 0 ro[0]").wrap_in_numshots_loop(10)) 
qc.qam.read_memory(region_name="not_ro")
```

### Error Output

```
KeyError: 'not_ro'
```

Environment Context
-------------------

Operating System: 

Python Version (`python -V`): Python 3.7.4

Quilc Version (`quilc --version`): 1.15.3 [25b95cb]

QVM Version (`qvm --version`): 1.15.2 [1b3d43a]

Python Environment Details (`pip freeze` or `conda list`): 💤


Discussion:
Even within the 'ro' memory region, it only returns those indices that are written to by measurements:
```
from pyquil import Program, get_qc
from pyquil.gates import CNOT, X, MEASURE

p = Program()
ro = p.declare("ro", "BIT", 4)
p += CNOT(0, 1)
p += X(2)
p += CNOT(2, 3)
p += MEASURE(0, ro[0])
p += MEASURE(3, ro[3])
p.wrap_in_numshots_loop(10)

qc = get_qc('9q-square-qvm')
ex = qc.compile(p)
result = qc.run(ex)
print(result)
```
gives
```
[[0 1]
 [0 1]
 [0 1]
 [0 1]
 [0 1]
 [0 1]
 [0 1]
 [0 1]
 [0 1]
 [0 1]]
```
(Python 3.7.7, quilc 1.18.0, qvm 1.17.0)
Though it's possible cause of this case might be more linked to #1194

Patch:
diff --git a/pyquil/api/_base_connection.py b/pyquil/api/_base_connection.py
--- a/pyquil/api/_base_connection.py
+++ b/pyquil/api/_base_connection.py
@@ -600,7 +600,7 @@ def _qvm_run(
         measurement_noise: Optional[Tuple[float, float, float]],
         gate_noise: Optional[Tuple[float, float, float]],
         random_seed: Optional[int],
-    ) -> np.ndarray:
+    ) -> Dict[str, np.ndarray]:
         """
         Run a Forest ``run`` job on a QVM.
 
@@ -611,7 +611,7 @@ def _qvm_run(
         )
         response = post_json(self.session, self.sync_endpoint + "/qvm", payload)
 
-        ram = response.json()
+        ram: Dict[str, np.ndarray] = {key: np.array(val) for key, val in response.json().items()}
 
         for k in ram.keys():
             ram[k] = np.array(ram[k])
diff --git a/pyquil/api/_qam.py b/pyquil/api/_qam.py
--- a/pyquil/api/_qam.py
+++ b/pyquil/api/_qam.py
@@ -41,6 +41,8 @@ class QAM(ABC):
     pretend to be a QPI-compliant quantum computer.
     """
 
+    _memory_results: Dict[str, np.ndarray]
+
     @_record_call
     def __init__(self) -> None:
         self.reset()
@@ -63,7 +65,7 @@ def load(
         self._executable: Optional[
             Union[QuiltBinaryExecutableResponse, PyQuilExecutableResponse]
         ] = executable
-        self._memory_results: Optional[Dict[str, np.ndarray]] = defaultdict(lambda: None)
+        self._memory_results = defaultdict(lambda: None)
         self.status = "loaded"
         return self
 
diff --git a/pyquil/api/_qvm.py b/pyquil/api/_qvm.py
--- a/pyquil/api/_qvm.py
+++ b/pyquil/api/_qvm.py
@@ -538,7 +538,10 @@ def load(self, executable: Union[Program, PyQuilExecutableResponse]) -> "QVM":
                     "`Program`. You provided {}".format(type(executable))
                 )
 
-        return cast("QVM", super().load(executable))
+        qvm = cast("QVM", super().load(executable))
+        for region in executable.declarations.keys():
+            self._memory_results[region] = np.ndarray((executable.num_shots, 0), dtype=np.int64)
+        return qvm
 
     @_record_call
     def run(self) -> "QVM":
@@ -565,7 +568,7 @@ def run(self) -> "QVM":
 
         quil_program = self.augment_program_with_memory_values(quil_program)
 
-        self._memory_results = self.connection._qvm_run(
+        results = self.connection._qvm_run(
             quil_program=quil_program,
             classical_addresses=classical_addresses,
             trials=trials,
@@ -573,9 +576,7 @@ def run(self) -> "QVM":
             gate_noise=self.gate_noise,
             random_seed=self.random_seed,
         )
-
-        if "ro" not in self._memory_results or len(self._memory_results["ro"]) == 0:
-            self._memory_results["ro"] = np.zeros((trials, 0), dtype=np.int64)
+        self._memory_results.update(results)
 
         return self
 
diff --git a/pyquil/pyqvm.py b/pyquil/pyqvm.py
--- a/pyquil/pyqvm.py
+++ b/pyquil/pyqvm.py
@@ -211,7 +211,7 @@ def __init__(
         # private implementation details
         self._qubit_to_ram: Optional[Dict[int, int]] = None
         self._ro_size: Optional[int] = None
-        self._memory_results: Optional[Dict[str, np.ndarray]] = None
+        self._memory_results = {}
 
         self.rs = np.random.RandomState(seed=seed)
         self.wf_simulator = quantum_simulator_type(n_qubits=n_qubits, rs=self.rs)
@@ -226,7 +226,7 @@ def load(self, executable: Union[Program, PyQuilExecutableResponse]) -> "PyQVM":
         # initialize program counter
         self.program = program
         self.program_counter = 0
-        self._memory_results = None
+        self._memory_results = {}
 
         # clear RAM, although it's not strictly clear if this should happen here
         self.ram = {}
diff --git a/pyquil/quil.py b/pyquil/quil.py
--- a/pyquil/quil.py
+++ b/pyquil/quil.py
@@ -40,7 +40,7 @@
 from rpcq.messages import NativeQuilMetadata
 
 from pyquil._parser.parser import run_parser
-
+from pyquil.gates import MEASURE, RESET
 from pyquil.noise import _check_kraus_ops, _create_kraus_pragmas, pauli_kraus_map
 from pyquil.quilatom import (
     Label,
@@ -58,7 +58,6 @@
     unpack_classical_reg,
     unpack_qubit,
 )
-from pyquil.gates import MEASURE, RESET
 from pyquil.quilbase import (
     DefGate,
     Gate,
@@ -100,7 +99,6 @@
     match_calibration,
 )
 
-
 InstructionDesignator = Union[
     AbstractInstruction,
     DefGate,
@@ -140,6 +138,9 @@ def __init__(self, *instructions: InstructionDesignator):
         # method.  It is marked as None whenever new instructions are added.
         self._synthesized_instructions: Optional[List[AbstractInstruction]] = None
 
+        # "ro" is always implicitly declared
+        self._declarations: Dict[str, Declare] = {"ro": Declare("ro", "BIT")}
+
         self.inst(*instructions)
 
         # Filled in with quil_to_native_quil
@@ -166,6 +167,11 @@ def frames(self) -> Dict[Frame, DefFrame]:
         """ A mapping from Quil-T frames to their definitions. """
         return self._frames
 
+    @property
+    def declarations(self) -> Dict[str, Declare]:
+        """ A mapping from declared region names to their declarations. """
+        return self._declarations
+
     def copy_everything_except_instructions(self) -> "Program":
         """
         Copy all the members that live on a Program object.
@@ -293,6 +299,9 @@ def inst(self, *instructions: InstructionDesignator) -> "Program":
             elif isinstance(instruction, AbstractInstruction):
                 self._instructions.append(instruction)
                 self._synthesized_instructions = None
+
+                if isinstance(instruction, Declare):
+                    self._declarations[instruction.name] = instruction
             else:
                 raise TypeError("Invalid instruction: {}".format(instruction))
 


Test patch:
diff --git a/pyquil/tests/test_quil.py b/pyquil/tests/test_quil.py
--- a/pyquil/tests/test_quil.py
+++ b/pyquil/tests/test_quil.py
@@ -243,10 +243,26 @@ def test_prog_init():
     assert p.out() == ("DECLARE ro BIT[1]\nX 0\nMEASURE 0 ro[0]\n")
 
 
-def test_classical_regs():
+def test_classical_regs_implicit_ro():
     p = Program()
-    p.inst(Declare("ro", "BIT", 2), X(0)).measure(0, MemoryReference("ro", 1))
-    assert p.out() == ("DECLARE ro BIT[2]\nX 0\nMEASURE 0 ro[1]\n")
+    p.inst(Declare("reg", "BIT", 2), X(0)).measure(0, MemoryReference("reg", 1))
+    assert p.out() == "DECLARE reg BIT[2]\nX 0\nMEASURE 0 reg[1]\n"
+    assert p.declarations == {
+        "ro": Declare("ro", "BIT", 1),
+        "reg": Declare("reg", "BIT", 2),
+    }
+
+
+def test_classical_regs_explicit_ro():
+    p = Program()
+    p.inst(Declare("ro", "BIT", 2), Declare("reg", "BIT", 2), X(0)).measure(
+        0, MemoryReference("reg", 1)
+    )
+    assert p.out() == "DECLARE ro BIT[2]\nDECLARE reg BIT[2]\nX 0\nMEASURE 0 reg[1]\n"
+    assert p.declarations == {
+        "ro": Declare("ro", "BIT", 2),
+        "reg": Declare("reg", "BIT", 2),
+    }
 
 
 def test_simple_instructions():
diff --git a/pyquil/tests/test_qvm.py b/pyquil/tests/test_qvm.py
--- a/pyquil/tests/test_qvm.py
+++ b/pyquil/tests/test_qvm.py
@@ -6,6 +6,7 @@
 from pyquil import Program
 from pyquil.api import ForestConnection, QVM
 from pyquil.api._compiler import _extract_program_from_pyquil_executable_response
+from pyquil.api._errors import QVMError
 from pyquil.gates import MEASURE, X, CNOT, H
 from pyquil.quilbase import Declare, MemoryReference
 
@@ -55,7 +56,44 @@ def test_qvm_run_only_pqer(forest: ForestConnection):
     assert np.mean(bitstrings) > 0.8
 
 
-def test_qvm_run_no_measure(forest: ForestConnection):
+def test_qvm_run_region_declared_and_measured(forest: ForestConnection):
+    qvm = QVM(connection=forest)
+    p = Program(Declare("reg", "BIT"), X(0), MEASURE(0, MemoryReference("reg")))
+    nq = PyQuilExecutableResponse(program=p.out(), attributes={"num_shots": 100})
+    qvm.load(nq).run().wait()
+    bitstrings = qvm.read_memory(region_name="reg")
+    assert bitstrings.shape == (100, 1)
+
+
+def test_qvm_run_region_declared_not_measured(forest: ForestConnection):
+    qvm = QVM(connection=forest)
+    p = Program(Declare("reg", "BIT"), X(0))
+    nq = PyQuilExecutableResponse(program=p.out(), attributes={"num_shots": 100})
+    qvm.load(nq).run().wait()
+    bitstrings = qvm.read_memory(region_name="reg")
+    assert bitstrings.shape == (100, 0)
+
+
+# For backwards compatibility, we support omitting the declaration for "ro" specifically
+def test_qvm_run_region_not_declared_is_measured_ro(forest: ForestConnection):
+    qvm = QVM(connection=forest)
+    p = Program(X(0), MEASURE(0, MemoryReference("ro")))
+    nq = PyQuilExecutableResponse(program=p.out(), attributes={"num_shots": 100})
+    qvm.load(nq).run().wait()
+    bitstrings = qvm.read_memory(region_name="ro")
+    assert bitstrings.shape == (100, 1)
+
+
+def test_qvm_run_region_not_declared_is_measured_non_ro(forest: ForestConnection):
+    qvm = QVM(connection=forest)
+    p = Program(X(0), MEASURE(0, MemoryReference("reg")))
+    nq = PyQuilExecutableResponse(program=p.out(), attributes={"num_shots": 100})
+
+    with pytest.raises(QVMError, match='Bad memory region name "reg" in MEASURE'):
+        qvm.load(nq).run().wait()
+
+
+def test_qvm_run_region_not_declared_not_measured_ro(forest: ForestConnection):
     qvm = QVM(connection=forest)
     p = Program(X(0))
     nq = PyQuilExecutableResponse(program=p.out(), attributes={"num_shots": 100})
@@ -64,6 +102,14 @@ def test_qvm_run_no_measure(forest: ForestConnection):
     assert bitstrings.shape == (100, 0)
 
 
+def test_qvm_run_region_not_declared_not_measured_non_ro(forest: ForestConnection):
+    qvm = QVM(connection=forest)
+    p = Program(X(0))
+    nq = PyQuilExecutableResponse(program=p.out(), attributes={"num_shots": 100})
+    qvm.load(nq).run().wait()
+    assert qvm.read_memory(region_name="reg") is None
+
+
 def test_roundtrip_pyquilexecutableresponse(compiler):
     p = Program(H(10), CNOT(10, 11))
     pqer = compiler.native_quil_to_executable(p)