def collect_tomography_measurements(self, tomography, mini_buffer, N_msmt): """ Given the buffer of 'N_alpha' phase space points, do 'N_msmt' tomography measurements per point. Args: tomography (str): type of tomography, either 'Wigner' or 'CF' mini_buffer (list(c64)): a list of phase space points N_msmt (int): number of measurements to do per phase space point Returns: msmt (Tensor([2, batch_size, N_alpha, N_msmt], float32)): tensor of measurement outomes; msmt[0] is the first measurement in the tomography reward circuit, which disentangles the qubit and oscillator, msmt[1] is the actual tomography measurement. """ N_alpha = len(mini_buffer) M1_buffer, M2_buffer = [], [] for i in range(N_alpha): # take 1 point from the buffer and replicate it for the batch points = tf.broadcast_to(mini_buffer[i], [self.batch_size]) M1, M2 = [], [] for j in range(N_msmt): # first measure the qubit to disentangle from oscillator psi = self.info['psi_cached'] if self.tensorstate: psi, m1 = measurement(psi, self.P, sample=True) else: m1 = tf.ones([self.batch_size]) M1.append(tf.squeeze(m1)) # do 1 tomography measurement in this phase space point if tomography == 'wigner': translations = self.translate(-points) psi = tf.linalg.matvec(translations, psi) _, m2 = self.phase_estimation(psi, self.parity, angle=tf.zeros( self.batch_size), sample=True) if tomography == 'CF': translations = self.translate(points) _, m2 = self.phase_estimation(psi, translations, angle=tf.zeros( self.batch_size), sample=True) M2.append(tf.squeeze(m2)) M1_buffer.append(M1) M2_buffer.append(M2) msmt = tf.cast([M1_buffer, M2_buffer], tf.float32) if self.batch_size == 1: msmt = tf.expand_dims(msmt, axis=-1) # at this point msmt.shape = [2, N_alpha, N_msmt, B] msmt = tf.transpose(msmt, perm=[0, 3, 1, 2]) return msmt
def reward_stabilizers(self, stabilizer_translations, *args): """ Reward only on last time step with the result of measurement of one of the stabilizers using cached wavefunction (after feedback translation). """ # Calculate reward if self._elapsed_steps < self.episode_length: z = tf.zeros(self.batch_size, dtype=tf.float32) else: # TODO: this works only for 2 stabilizers, generalize mask = tf.random.categorical( tf.math.log([[0.5, 0.5] for _ in range(self.batch_size)]), 1) stabilizers = [stabilizer_translations[int(m)] for m in mask] stabilizers = tf.convert_to_tensor(stabilizers, dtype=c64) phi = tf.zeros(self.batch_size) stabilizers = self.translate(stabilizers) if self.tensorstate: psi, m = measurement(self._state, self.P, sample=True) mask = tf.squeeze(tf.where(m == 1, 1.0, 0.0)) else: psi = self.info['psi_cached'] mask = tf.ones([self.batch_size]) _, z = self.phase_estimation(psi, stabilizers, angle=phi, sample=True) z = tf.cast(z, dtype=tf.float32) z = tf.reshape(z, shape=(self.batch_size, )) * mask return z
def reward_measurement(self, sample, *args): """ Reward is simply the outcome of sigma_z measurement. """ psi, z = measurement(self.info['psi_cached'], self.P, sample) return tf.squeeze(z)
def _quantum_circuit(self, psi, action): """ Args: psi (Tensor([batch_size,N], c64)): batch of states action (dict, 'alpha' : Tensor([batch_size,2], tf.float32), 'beta' : Tensor([batch_size,2], tf.float32), 'phi' : Tensor([batch_size,1], tf.float32)) Returns: see parent class docs """ # extract parameters alpha = hf.vec_to_complex(action['alpha']) beta = hf.vec_to_complex(action['beta']) phi = action['phi'] Kraus = {} T = {'a': self.translate(alpha), 'b': self.translate(beta)} I = tf.stack([self.I] * self.batch_size) Kraus[0] = 1 / 2 * (I + self.phase(phi) * T['b']) Kraus[1] = 1 / 2 * (I - self.phase(phi) * T['b']) psi = self.simulate(psi, self.t_feedback) psi_cached = batch_dot(T['a'], psi) psi = self.simulate(psi_cached, self.t_round + self.t_idle) psi_final, msmt = measurement(psi, Kraus) return psi_final, psi_cached, msmt
def _control_circuit(self, psi, action): """ Args: psi (Tensor([batch_size,N], c64)): batch of states action (dict, 'alpha' : Tensor([batch_size,2], tf.float32), 'beta' : Tensor([batch_size,2], tf.float32), 'phi' : Tensor([batch_size,1], tf.float32), 'theta' : Tensor([batch_size,1], tf.float32)) Returns: see parent class docs """ # extract parameters alpha = hf.vec_to_complex(action['alpha']) beta = hf.vec_to_complex(action['beta']) phi = action['phi'] Rotation = self.rotate(action['theta']) Kraus = {} T = {'a': self.translate(alpha), 'b': self.translate(beta / 2.0)} Kraus[0] = 1 / 2 * (tf.linalg.adjoint(T['b']) + self.phase(phi) * T['b']) Kraus[1] = 1 / 2 * (tf.linalg.adjoint(T['b']) - self.phase(phi) * T['b']) psi = self.simulate(psi, self.t_feedback) psi = batch_dot(T['a'], psi) psi_cached = batch_dot(Rotation, psi) psi = self.simulate(psi_cached, self.t_round + self.t_idle) psi_final, msmt = measurement(psi, Kraus) return psi_final, psi_cached, msmt
def phase_estimation(self, psi, U, angle, sample=False): """ One round of phase estimation. Input: psi -- batch of state vectors; shape=[batch_size,2N] U -- unitary on which to do phase estimation. shape=(batch_size,N,N) angle -- angle along which to measure qubit. shape=(batch_size,) sample -- bool flag to sample or return expectation value Output: psi -- batch of collapsed states if sample==True, otherwise same as input psi; shape=[batch_size,2N] z -- batch of measurement outcomes if sample==True, otherwise batch of expectation values of qubit sigma_z. """ CT = self.ctrl(self.I, U) Phase = self.rotate_qb_z(tf.squeeze(angle)) psi = tf.linalg.matvec(self.hadamard, psi) psi = tf.linalg.matvec(CT, psi) psi = tf.linalg.matvec(Phase, psi) psi = tf.linalg.matvec(self.hadamard, psi) return measurement(psi, self.P, sample)
def _quantum_circuit(self, psi, action): """ Args: psi (Tensor([batch_size,N], c64)): batch of states action (dict, 'beta' : Tensor([batch_size,2], tf.float32), 'phi' : Tensor([batch_size,2], tf.float32)) Returns: see parent class docs """ # Extract parameters beta = hf.vec_to_complex(action['beta']) phase, angle = action['phi'][:,0], action['phi'][:,1] # Construct gates D = self.displace(beta/2.0) CD = self.ctrl(D, tf.linalg.adjoint(D)) R = self.rotate_qb_xy(angle, phase) flip = self.rotate_qb_xy(tf.constant(pi), tf.constant(0)) # Apply gates psi = tf.linalg.matvec(R, psi) psi = tf.linalg.matvec(CD, psi) if self._elapsed_steps < self.T-1: psi = tf.linalg.matvec(flip, psi) m = tf.ones((self.batch_size,1)) if self._elapsed_steps == self.T-1: psi, m = measurement(psi, self.P, sample=True) return psi, psi, m
def reward_fock(self, target_projector, N_msmt, error_prob): """ Reward only on last time step using the measurement of a given Fock state of the oscillator. """ if self._elapsed_steps < self.episode_length: z = tf.zeros(self.batch_size, dtype=tf.float32) else: Z = 0 for i in range(N_msmt): if self.tensorstate: # measure qubit to disentangle from oscillator psi, m = measurement(self._state, self.P, sample=True) mask = tf.squeeze(tf.where(m == 1, 1.0, 0.0)) else: psi = self._state mask = tf.ones([self.batch_size]) # sample observations p_n = expectation(psi, target_projector, reduce_batch=False) p_n = tf.reshape(tf.math.real(p_n), shape=[self.batch_size]) obs = tfp.distributions.Bernoulli(probs=p_n).sample() # sample errors and apply them to measurement outcomes e = error_prob * tf.ones_like(p_n) err = tfp.distributions.Bernoulli(probs=e).sample() obs = tf.where(err == 0, obs, 0) obs = 2 * tf.cast(obs, dtype=tf.float32) - 1 # convert to {-1,1} penalty_coeff = 1.0 Z += obs * mask - penalty_coeff * (1 - mask) z = Z / N_msmt return z
def phase_estimation(psi, U, angle, sample=False): I = ops.identity(N) angle = tf.cast(tf.reshape(angle, angle.shape + [1, 1]), c64) phase = tf.math.exp(1j * angle) Kraus = {} Kraus[0] = 1 / 2 * (I + phase * U) Kraus[1] = 1 / 2 * (I - phase * U) return measurement(psi, Kraus, sample)
def _quantum_circuit(self, psi, action): """ Args: psi (Tensor([batch_size,N], c64)): batch of states action (dict, 'alpha' : Tensor([batch_size,2], tf.float32), 'beta' : Tensor([batch_size,2], tf.float32), 'epsilon' : Tensor([batch_size,2], tf.float32)) Returns: see parent class docs """ # extract parameters alpha = hf.vec_to_complex(action['alpha']) beta = hf.vec_to_complex(action['beta']) epsilon = hf.vec_to_complex(action['epsilon']) # Construct gates Hadamard = tf.stack([self.hadamard] * self.batch_size) sx = tf.stack([self.sx] * self.batch_size) I = tf.stack([self.I] * self.batch_size) Rxp = tf.stack([self.rxp] * self.batch_size) Rxm = tf.stack([self.rxm] * self.batch_size) T, CT = {}, {} T['a'] = self.translate(alpha) T['b'] = self.translate(beta / 4.0) T['e'] = self.translate(epsilon / 2.0) CT['b'] = self.ctrl(tf.linalg.adjoint(T['b']), T['b']) CT['e'] = self.ctrl(tf.linalg.adjoint(T['e']), T['e']) # Feedback translation psi_cached = batch_dot(T['a'], psi) # Between-round wait time psi = self.simulate(psi_cached, self.t_idle) # Qubit gates psi = batch_dot(Hadamard, psi) # Conditional translation psi = batch_dot(CT['b'], psi) psi = self.simulate(psi, self.t_gate) psi = batch_dot(CT['b'], psi) # Qubit rotation psi = batch_dot(Rxp, psi) # Conditional translation psi = batch_dot(CT['e'], psi) # Qubit rotation psi = batch_dot(Rxm, psi) # Readout of finite duration psi = self.simulate(psi, self.t_read) psi, msmt = measurement(psi, self.P) psi = self.simulate(psi, self.t_read) # Feedback delay psi = self.simulate(psi, self.t_feedback) # Flip qubit conditioned on the measurement psi_final = psi * tf.cast((msmt == 1), c64) psi_final += batch_dot(sx, psi) * tf.cast((msmt == -1), c64) return psi_final, psi_cached, msmt
def _quantum_circuit(self, psi, action): """ Args: psi (Tensor([batch_size,N], c64)): batch of states action (dict, 'beta' : Tensor([batch_size,2], tf.float32), 'epsilon' : Tensor([batch_size,2], tf.float32, 'phi' : Tensor([batch_size,1]))) Returns: see parent class docs """ # extract parameters beta = hf.vec_to_complex(action['beta']) epsilon = hf.vec_to_complex(action['epsilon']) phi = tf.squeeze(action['phi']) # Construct gates Phase = self.rotate_qb_z(phi) T, CT = {}, {} T['b'] = self.translate(beta / 4.0) # 4 because it will be troterized T['e'] = self.translate(epsilon / 2.0) CT['b'] = self.ctrl(tf.linalg.adjoint(T['b']), T['b']) CT['e'] = self.ctrl(tf.linalg.adjoint(T['e']), T['e']) # Between-round wait time psi = self.simulate(psi, self.t_idle) # Rotate qubit to |+> state psi = tf.linalg.matvec(self.hadamard, psi) # Troterized conditional translation psi = tf.linalg.matvec(CT['b'], psi) psi = self.simulate(psi, self.t_gate) psi = tf.linalg.matvec(CT['b'], psi) # Qubit rotation psi = tf.linalg.matvec(self.rxp, psi) # Conditional translation psi = tf.linalg.matvec(CT['e'], psi) # Qubit rotation psi = tf.linalg.matvec(self.rxm, psi) # Troterized conditional translation psi = tf.linalg.matvec(CT['b'], psi) psi = self.simulate(psi, self.t_gate) psi = tf.linalg.matvec(CT['b'], psi) # Qubit gates psi = tf.linalg.matvec(Phase, psi) psi = tf.linalg.matvec(self.hadamard, psi) # Readout of finite duration psi = self.simulate(psi, self.t_read) psi, msmt = measurement(psi, self.P) psi = self.simulate(psi, self.t_read) # Feedback delay psi = self.simulate(psi, self.t_feedback) # Flip qubit conditioned on the measurement psi_final = tf.where(msmt == 1, psi, tf.linalg.matvec(self.sx, psi)) return psi_final, psi_final, msmt
def _control_circuit(self, psi, action): """ Args: psi (Tensor([batch_size,N], c64)): batch of states action (dict, 'alpha' : Tensor([batch_size,2], tf.float32), 'beta' : Tensor([batch_size,2], tf.float32), 'phi' : Tensor([batch_size,1], tf.float32)) Returns: see parent class docs """ # Extract parameters alpha = hf.vec_to_complex(action['alpha']) beta = hf.vec_to_complex(action['beta']) phi = action['phi'] # Construct gates Hadamard = tf.stack([self.hadamard] * self.batch_size) sx = tf.stack([self.sx] * self.batch_size) I = tf.stack([self.I] * self.batch_size) Phase = self.rotate_qb_z(tf.squeeze(phi)) T, CT = {}, {} T['a'] = self.translate(alpha) T['b'] = self.translate(beta / 2.0) CT['b'] = self.ctrl(I, T['b']) # Feedback translation psi_cached = batch_dot(T['a'], psi) # Between-round wait time psi = self.simulate(psi_cached, self.t_idle) # Qubit gates psi = batch_dot(Hadamard, psi) # Conditional translation psi = batch_dot(CT['b'], psi) psi = self.simulate(psi, self.t_gate) psi = batch_dot(CT['b'], psi) # Qubit gates psi = batch_dot(Phase, psi) psi = batch_dot(Hadamard, psi) # Readout of finite duration psi = self.simulate(psi, self.t_read) psi, msmt = measurement(psi, self.P) psi = self.simulate(psi, self.t_read) # Feedback delay psi = self.simulate(psi, self.t_feedback) # Flip qubit conditioned on the measurement psi_final = psi * tf.cast((msmt == 1), c64) psi_final += batch_dot(sx, psi) * tf.cast((msmt == -1), c64) return psi_final, psi_cached, msmt
def _quantum_circuit(self, psi, action): """ Args: psi (Tensor([batch_size,N], c64)): batch of states action (dict, 'beta' : Tensor([batch_size,2], tf.float32), 'epsilon' : Tensor([batch_size,2], tf.float32, 'phi' : Tensor([batch_size,1]))) Returns: see parent class docs """ # extract parameters beta = hf.vec_to_complex(action['beta']) epsilon = hf.vec_to_complex(action['epsilon']) phi = action['phi'] Kraus = {} T = {} T['+b'] = self.translate(beta / 2.0) T['-b'] = tf.linalg.adjoint(T['+b']) T['+e'] = self.translate(epsilon / 2.0) T['-e'] = tf.linalg.adjoint(T['+e']) chunk1 = 1j*batch_dot(T['-b'], batch_dot(T['+e'], T['+b'])) \ - 1j*batch_dot(T['-b'], batch_dot(T['-e'], T['+b'])) \ + batch_dot(T['-b'], batch_dot(T['-e'], T['-b'])) \ + batch_dot(T['-b'], batch_dot(T['+e'], T['-b'])) chunk2 = 1j*batch_dot(T['+b'], batch_dot(T['-e'], T['-b'])) \ - 1j*batch_dot(T['+b'], batch_dot(T['+e'], T['-b'])) \ + batch_dot(T['+b'], batch_dot(T['-e'], T['+b'])) \ + batch_dot(T['+b'], batch_dot(T['+e'], T['+b'])) Kraus[0] = 1 / 4 * (chunk1 + self.phase(phi) * chunk2) Kraus[1] = 1 / 4 * (chunk1 - self.phase(phi) * chunk2) psi = self.simulate(psi, self.t_round + self.t_idle) psi_final, msmt = measurement(normalize(psi), Kraus) return psi_final, psi_final, msmt
def reward_stabilizers_v2(self, beta, Delta, sample): """ Use finite-energy stabilizers instead of ideal stabilizers. These have a parameter 'Delta' controlling the squeezing level and 'beta' controlling the stabilizer magnitude (the direction is randomly sampled to be e^{0j} or e^{1j*pi/2}). See this paper by Baptiste Royer for more details: https://journals.aps.org/prl/abstract/10.1103/PhysRevLett.125.260509 """ # return 0 on all intermediate steps of the episode if self._elapsed_steps < self.episode_length: return tf.zeros(self.batch_size, dtype=tf.float32) # sample random stabilizer directions (+/-x, +/-p) mask = tf.random.uniform([self.batch_size]) theta_x = tf.where(mask > 3 / 4, 0., pi) * tf.where( mask > 1 / 2, 1., 0.) theta_p = tf.where(mask > 1 / 4, pi / 2, -pi / 2) * tf.where( mask < 1 / 2, 1., 0.) theta = theta_x + theta_p theta = tf.cast(theta, c64) if self.tensorstate: raise ValueError('Only <Oscillator> Hilbert space is supported.') eps = beta * Delta**2 # Create Kraus operators corresponding to finite-energy stabilizers D_e = self.displace((eps + 0j) * tf.math.exp(1j * theta) / 2) D_b = self.displace(-1j * beta * tf.math.exp(1j * theta) / 2) chunk1 = tf.linalg.matmul(D_b, D_e - 1j * tf.linalg.adjoint(D_e)) chunk2 = tf.linalg.matmul(tf.linalg.adjoint(D_b), \ -1j*D_e + tf.linalg.adjoint(D_e)) Kraus = {} Kraus[0] = 1 / (2 * sqrt(2)) * (chunk2 + chunk1) Kraus[1] = 1 / (2 * sqrt(2)) * (chunk2 - chunk1) psi, z = measurement(self._state, Kraus, sample=sample) return tf.squeeze(z)
def reward_Fock(self, target_projector, *args): """ Reward only on last time step using the measurement of a given Fock state of the oscillator. """ if self._elapsed_steps < self.episode_length: z = tf.zeros(self.batch_size, dtype=tf.float32) else: if self.tensorstate: # measure qubit to disentangle from oscillator psi, _ = measurement(self._state, self.P, sample=True) else: psi = self._state overlap = expectation(psi, target_projector, reduce_batch=False) z = tf.reshape(tf.math.real(overlap), shape=[self.batch_size]) obs = tfp.distributions.Bernoulli(probs=z).sample() obs = 2 * obs - 1 # convert to {-1,1} z = tf.cast(obs, dtype=tf.float32) return z
def phase_estimation(self, psi, U, angle, sample=False): """ One round of phase estimation. Input: psi -- batch of state vectors; shape=[batch_size,N] U -- unitary on which to do phase estimation. shape=(batch_size,N,N) angle -- angle along which to measure qubit. shape=(batch_size,) sample -- bool flag to sample or return expectation value Output: psi -- batch of collapsed states if sample==True, otherwise same as input psi; shape=[batch_size,N] z -- batch of measurement outcomes if sample==True, otherwise batch of expectation values of qubit sigma_z. """ Kraus = {} Kraus[0] = 1 / 2 * (self.I + self.phase(angle) * U) Kraus[1] = 1 / 2 * (self.I - self.phase(angle) * U) return measurement(psi, Kraus, sample)
def _quantum_circuit(self, psi, action): """ Args: psi (Tensor([batch_size,N], c64)): batch of states action (dict, 'alpha' : Tensor([batch_size,2], tf.float32), 'beta' : Tensor([batch_size,2], tf.float32), 'epsilon' : Tensor([batch_size,2], tf.float32)) Returns: see parent class docs """ # extract parameters alpha = hf.vec_to_complex(action['alpha']) beta = hf.vec_to_complex(action['beta']) epsilon = hf.vec_to_complex(action['epsilon']) Kraus = {} T = {} T['a'] = self.translate(alpha) T['+b'] = self.translate(beta/2.0) T['-b'] = tf.linalg.adjoint(T['+b']) T['+e'] = self.translate(epsilon/2.0) T['-e'] = tf.linalg.adjoint(T['+e']) chunk1 = batch_dot(T['-e'], T['-b']) - 1j*batch_dot(T['-e'], T['+b']) chunk2 = batch_dot(T['+e'], T['-b']) + 1j*batch_dot(T['+e'], T['+b']) Kraus[0] = 1/2/sqrt(2)*(chunk1 + chunk2) Kraus[1] = 1/2/sqrt(2)*(chunk1 - chunk2) psi = self.simulate(psi, self.t_feedback) psi_cached = batch_dot(T['a'], psi) psi = self.simulate(psi_cached, self.t_round + self.t_idle) psi_final, msmt = measurement(psi, Kraus) return psi_final, psi_cached, msmt
def _control_circuit(self, psi, action): """ Args: psi (Tensor([batch_size,N], c64)): batch of states action (dict, 'alpha' : Tensor([batch_size,2], tf.float32), 'theta' : Tensor([batch_size,10], tf.float32)) Returns: see parent class docs """ # Extract parameters alpha = hf.vec_to_complex(action['alpha']) theta = action['theta'] # Build gates displace = self.displace(alpha) snap = self.SNAP_miscalibrated(theta) # Apply gates psi = tf.linalg.matvec(displace, psi) psi = tf.linalg.matvec(snap, psi) psi = tf.linalg.matvec(tf.linalg.adjoint(displace), psi) # Either implement a measurement with a random qubit projection, or # project on the specified 'self.bit_string' sequence of qubit states. if self.bit_string is not None: s = int(self.bit_string[self._elapsed_steps]) psi = tf.linalg.matvec(self.P[s], psi) if s == 1: psi = tf.linalg.matvec(self.sx, psi) psi, norm = normalize(psi) self.norms.append(norm) msmt = tf.ones((self.batch_size, 1)) * (1 if s == 0 else -1) else: # Readout with feedback to flip the qubit psi, msmt = measurement(psi, self.P) psi = tf.where(msmt == 1, psi, tf.linalg.matvec(self.sx, psi)) return psi, psi, msmt
def reward_overlap(self, target_projector, postselect_0): """ Reward only on last time step using the overlap of the cached state with the target state. The cached state is measured prior to computing the overlap, to make sure that the oscillator and qubit are disentangled. The agent can learn to arrange things so that this measurement doesn't matter (i.e. if they are already disentangled before the measurement). """ if self._elapsed_steps < self.episode_length: z = tf.zeros(self.batch_size, dtype=tf.float32) else: if self.tensorstate: if postselect_0: # project qubit to |0> psi = tf.linalg.matvec(self.P[0], self._state) psi, _ = normalize(psi) else: # randomly project qubit with a measurement psi, _ = measurement(self._state, self.P, sample=True) else: psi = self._state overlap = expectation(psi, target_projector, reduce_batch=False) z = tf.reshape(tf.math.real(overlap), shape=[self.batch_size]) self.info['psi_cached'] = psi return z
def reward_tomography(self, target_state, window_size, tomography, sample_from_buffer): """ Reward only on last time step using the empirical approximation of the overlap of the prepared state with the target state. This overlap is computed in characteristic-function-tomography-like fashing. The state is measured prior to acquiring a tomography point on the oscillator to make sure that the oscillator and qubit are disentangled. The agent can learn to arrange things so that this measurement doesn't matter (i.e. if they are already disentangled before the measurement). """ # return 0 on all intermediate steps of the episode if self._elapsed_steps < self.episode_length: return tf.zeros(self.batch_size, dtype=tf.float32) def alpha_sample_schedule(n): return 1 def msmt_sample_schedule(n): return 1 alpha_samples = alpha_sample_schedule(self._episodes_completed) msmt_samples = msmt_sample_schedule(self._episodes_completed) # populate a new small mini-buffer for each epoch if not sample_from_buffer: mini_buffer, target_vals = self.fill_buffer(target_state, window_size, tomography, samples=alpha_samples) z = 0 for i in range(alpha_samples): if sample_from_buffer: # uniformly sample points from the large buffer; each batch # member will get its own phase space point samples, buffer_size = self.batch_size, len(self.buffer) index = tf.math.round( tf.random.uniform([samples]) * buffer_size) targets = tf.gather(self.target_vals, tf.cast(index, tf.int32)) points = tf.gather(self.buffer, tf.cast(index, tf.int32)) else: # take 1 point from the buffer and replicate it for the batch targets = tf.broadcast_to(target_vals[i], [self.batch_size]) points = tf.broadcast_to(mini_buffer[i], [self.batch_size]) M = 0 + 1e-10 Z = 0 for j in range(msmt_samples): # first measure the qubit to disentangle from oscillator psi = self.info['psi_cached'] if self.tensorstate: psi, m = measurement(psi, self.P, sample=True) mask = tf.squeeze(tf.where(m == 1, 1.0, 0.0)) else: mask = tf.ones([self.batch_size]) # do tomography in one phase space point if tomography == 'wigner': translations = self.translate(-points) psi = tf.linalg.matvec(translations, psi) _, msmt = self.phase_estimation(psi, self.parity, angle=tf.zeros( self.batch_size), sample=True) if tomography == 'characteristic_fn': translations = self.translate(points) _, msmt = self.phase_estimation(psi, translations, angle=tf.zeros( self.batch_size), sample=True) # Make a noisy Monte Carlo estimate of the overlap integral. # If using characteristic_fn, this would work only for symmetric # states (GKP, Fock etc) # Mask out trajectories where qubit was measured in |e> Z += tf.squeeze(msmt) * tf.math.sign(targets) * mask M += mask z += Z / msmt_samples z /= alpha_samples return z