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
