fixed bug #10 by adding the ability to reset modeMap and Gaussian cir…

…cuit, and bug #11 by adding np.reshape for pure states (#12)

strawberryfields/backends/base.py

@@ -127,8 +127,13 @@ class ModeMap:
     Simple internal class for maintaining a map of existing modes.
     """
     def __init__(self, num_subsystems):
+        self._init = num_subsystems
         self._map = [k for k in range(num_subsystems)]
 
+    def reset(self):
+        """reset the modemap to the initial state"""
+        self._map = [k for k in range(self._init)]
+
     def _single_mode_valid(self, mode):
         if mode is None:
             return False

strawberryfields/backends/fockbackend/backend.py

@@ -79,6 +79,7 @@ def begin_circuit(self, num_subsystems, cutoff_dim=None, hbar=2, pure=True, **kw
         elif not isinstance(pure, bool):
             raise ValueError("Argument 'pure' must be either True or False")
 
+        self._init_modes = num_subsystems
         self.qreg = QReg(num_subsystems, cutoff_dim, hbar, pure)
         self._modeMap = ModeMap(num_subsystems)
 
@@ -122,7 +123,8 @@ def reset(self, pure=True, **kwargs):
         Args:
             pure (bool): whether to use a pure state representation upon reset
         """
-        self.qreg.reset(pure)
+        self._modeMap.reset()
+        self.qreg.reset(pure, num_subsystems=self._init_modes)
 
     def prepare_vacuum_state(self, mode):
         """Prepare the vacuum state on the specified mode.

strawberryfields/backends/fockbackend/circuit.py

@@ -127,15 +127,20 @@ def _apply_channel(self, kraus_ops, modes):
             self._state = sum(states)
 
 
-    def reset(self, pure=None):
+    def reset(self, pure=None, num_subsystems=None):
         """Resets the simulation state.
 
         Args:
-            pure (bool, optional): Sets the purity setting. Default is unchanged
+            pure (bool, optional): Sets the purity setting. Default is unchanged.
+            num_subsystems (int, optional): Sets the number of modes in the reset
+                circuit. Default is unchanged.
         """
         if pure is not None:
             self._pure = pure
 
+        if num_subsystems is not None:
+            self._num_modes = num_subsystems
+
         if self._pure:
             self._state = ops.vacuumState(self._num_modes, self._trunc)
         else:

strawberryfields/backends/gaussianbackend/backend.py

@@ -44,6 +44,7 @@ def begin_circuit(self, num_subsystems, cutoff_dim=None, hbar=2, pure=None, **kw
                 By default, :math:`\hbar=2`. See :ref:`conventions` for more details.
         """
         # pylint: disable=attribute-defined-outside-init
+        self._init_modes = num_subsystems
         self.circuit = GaussianModes(num_subsystems, hbar)
 
     def add_mode(self, n=1):
@@ -84,10 +85,7 @@ def reset(self, pure=True, **kwargs):
         """
         Resets the circuit state back to an all-vacuum state.
         """
-
-        num_modes = self.circuit.get_modes()
-        for mode in num_modes:
-            self.circuit.loss(0.0, mode)
+        self.circuit.reset(self._init_modes)
 
     def prepare_thermal_state(self, nbar, mode):
         """

strawberryfields/backends/gaussianbackend/gaussiancircuit.py

@@ -50,12 +50,8 @@ def __init__(self, num_subsystems, hbar):
         if not isinstance(num_subsystems, int):
             raise ValueError("Number of modes must be an integer")
 
-        self.nmat = np.zeros((num_subsystems, num_subsystems), dtype=complex)
-        self.mmat = np.zeros((num_subsystems, num_subsystems), dtype=complex)
-        self.mean = np.zeros(num_subsystems, dtype=complex)
-        self.nlen = num_subsystems
-        self.active = list(np.arange(num_subsystems, dtype=int))
         self.hbar = hbar
+        self.reset(num_subsystems)
 
     def add_mode(self, n=1):
         """add mode to the circuit"""
@@ -90,6 +86,20 @@ def del_mode(self, modes):
             self.loss(0.0, mode)
             self.active[mode] = None
 
+    def reset(self, num_subsystems=None):
+        """Resets the simulation state.
+
+        Args:
+            num_subsystems (int, optional): Sets the number of modes in the reset
+                circuit. Default is unchanged.
+        """
+        if num_subsystems is not None:
+            self.nlen = num_subsystems
+
+        self.nmat = np.zeros((self.nlen, self.nlen), dtype=complex)
+        self.mmat = np.zeros((self.nlen, self.nlen), dtype=complex)
+        self.mean = np.zeros(self.nlen, dtype=complex)
+        self.active = list(np.arange(self.nlen, dtype=int))
 
     def get_modes(self):
         """return the modes currently active"""

strawberryfields/backends/states.py

@@ -451,7 +451,7 @@ def all_fock_probs(self, **kwargs):
         # pylint: disable=unused-argument
         if self._pure:
             s = np.ravel(self.ket()) # into 1D array
-            return (s * s.conj()).real
+            return np.reshape((s * s.conj()).real, [self._cutoff]*self._modes)
 
         s = self.dm()
         num_axes = len(s.shape)

strawberryfields/backends/tfbackend/backend.py

@@ -97,6 +97,7 @@ def begin_circuit(self, num_subsystems, cutoff_dim=None, hbar=2, pure=True, **kw
                 self._modemap = ModeMap(num_subsystems)
                 circuit = QReg(self._graph, num_subsystems, cutoff_dim, hbar, pure, batch_size)
 
+        self._init_modes = num_subsystems
         self.circuit = circuit
 
     def reset(self, pure=True, **kwargs):
@@ -118,7 +119,8 @@ def reset(self, pure=True, **kwargs):
             self._graph = tf.get_default_graph()
 
         with tf.name_scope('Reset'):
-            self.circuit.reset(pure, graph=self._graph)
+            self._modemap.reset()
+            self.circuit.reset(pure, graph=self._graph, num_subsystems=self._init_modes)
 
     def get_cutoff_dim(self):
         """Returns the Hilbert space cutoff dimension used.

strawberryfields/backends/tfbackend/circuit.py

@@ -185,13 +185,15 @@ def add_mode(self, num_modes):
         self._update_state(new_state)
         self._num_modes += num_modes
 
-    def reset(self, pure=True, graph=None):
+    def reset(self, pure=True, graph=None, num_subsystems=None):
         """
         Resets the state of the circuit to have all modes in vacuum.
         Args:
             pure (bool): If True, the reset circuit will represent its state as a pure state. If False, the representation will be mixed.
             graph: If this is an instance of tf.Graph, then the underlying graph (and any associated attributes) is replaced with this supplied graph. Otherwise, the same underlying
             graph (and all its defined operations) will be kept.
+            num_subsystems (int, optional): Sets the number of modes in the reset
+                circuit. Default is unchanged.
 
         Returns:
             None
@@ -204,6 +206,10 @@ def reset(self, pure=True, graph=None):
             self._make_vac_states()
             self._state_history = []
             self._cache = {}
+
+        if num_subsystems is not None:
+            self._num_modes = num_subsystems
+
         with self._graph.as_default():
             single_mode_vac = self._single_mode_pure_vac if pure else self._single_mode_mixed_vac
             if self._num_modes == 1: