def compute_p2(P_stencil,n,grid_size):
batch_size = P_stencil.shape[0]
K = map_2_to_1(grid_size=grid_size)
pi = np.pi
theta_x = np.array(([i * 2 * pi / n for i in range(-n // (grid_size*2) + 1, n // (grid_size*2) + 1)]))
theta_y = np.array([i * 2 * pi / n for i in range(-n // (grid_size*2) + 1, n // (grid_size*2) + 1)])
num_modes = theta_x.shape[0]
X, Y = np.meshgrid(np.arange(-1, 2), np.arange(-1, 2))
with tf.device("gpu:0"):
P = tf.zeros((len(theta_y), len(theta_x),batch_size,grid_size**2,(grid_size//2)**2),dtype=tf.complex128)
modes = np.array([[np.exp(-1j * (tx * X + ty * Y)) for tx in theta_x] for ty in theta_y])
fourier_component = tf.to_complex128(np.tile(modes, (batch_size,1,1,1,1)))
for ic in range(grid_size//2):
i = 2*ic #ic is the index on the coarse grid, and i is the index on the fine grid
for jc in range(grid_size//2):
j = 2*jc #jc is the index on the coarse grid, and j is the index on the fine grid
J = int(grid_size//2*jc+ic)
for k in range(3):
for m in range(3):
I = int(K[i,j,k,m])
a = fourier_component[:,:,:,k, m] * tf.reshape(P_stencil[:, ic, jc, k, m],(-1,1,1))
a = tf.transpose(a,(1,2,0))
P = P + tf.to_complex128(
tf.scatter_nd(tf.constant(idx_array((I,J,int(batch_size),num_modes))),
tf.ones(batch_size*(num_modes**2)),
tf.constant([num_modes,num_modes,batch_size, grid_size**2, (grid_size//2)**2])))*tf.reshape(a, (theta_x.shape[0],theta_y.shape[0], batch_size, 1, 1))
return P
def partial_trace(tensor, dims_subsystems):
'''
Take the partial trace, leaving H_A, where H = H_A \otimes H_B \otimes H_N
# Einsum does not allow you to do traces so we need this hacky method.
# reorder the axes and take the trace of the final two.
# These loops does the following (for 2 subsystems):
# first loop
# permute axes 0,a,b,c,a,b,c -> 0,a,a,c,b,b,c
# permute axes 0,a,a,c,b,b,c -> 0,a,a,b,b,c,c
# perform permutation
# second loop
# trace 0,a,a,b,b,c,c -> 0,a,a,b,b
# trace 0,a,a,b,b -> 0,a,a
:param tensor: NxN density matrix
:param dims_subsystems: Dimensions of subsystems
:return: Reduced density matrix
'''
dims = np.tile(dims_subsystems, 2)
# insert batch dimension
dims = np.insert(dims, [0], [-1])
tensor = tf.reshape(tensor, dims)
idx = list(range(len(dims)))
n = len(dims_subsystems)
for i in range(n - 1):
idx[n + i + 1], idx[2 + i] = idx[2 + i], idx[n + i + 1]
tensor = tf.transpose(tensor, perm=idx)
for _ in range(n - 1):
tensor = tf.trace(tensor)
return tf.to_complex128(tensor)
def to_complex(self, x):
if self.dtype(x) in (np.complex64, np.complex128):
return x
if self.dtype(x) == np.float64:
return tf.to_complex128(x)
else:
return tf.to_complex64(x)
def build_graph(self,
hilbert_space_composition,
epsilon=0.1,
complex_weights=True,
device='CPU'):
# Assert first that the model is compiled properly
assert self.dim != None, 'Data dimension unknown, specify as model argument or run QubitLearning.get_statistics'
assert self.n_samples != None, 'Number of samples unknown, specify as model argument or run QubitLearning.get_statistics'
assert isinstance(
hilbert_space_composition,
(list, tuple, np.ndarray
)), 'Pass list, tuple or ndarray of Hilbert space dimensions'
assert len(
hilbert_space_composition
) > 0, 'Hilbert space must consist of at least one dimension'
assert all(
isinstance(d, int) for d in hilbert_space_composition
), 'Expected list of integers for the dimensions of the composite Hilbert spaces'
assert all(d >= 2 for d in hilbert_space_composition
), 'Composite Hilbert spaces must have dim(H_i) >= 2'
assert (
hilbert_space_composition[0] == self.c
), 'dim(H_A) must be equal to the number of classes = {}, but dim(H_A) = {}'.format(
self.c, hilbert_space_composition[0])
# Build the model
with tf.name_scope('model'):
# TODO: X now has shape 1 but needs shape -1, rest of functions needs to be adapted
dims = np.array([d for d in hilbert_space_composition])
# self.x = tf.placeholder(dtype=tf.float32, shape=(None, self.dim))
# The data density matrix is C x C where C is the number of classes
self.eta = tf.placeholder(dtype=tf.complex128,
shape=(None, self.c, self.c))
if len(hilbert_space_composition) == 1:
print('Building non-entangled model')
self.MODEL_TYPE = 'single'
generators = self.SU_generators(dims[0])
self.x = tf.placeholder(dtype=tf.float32,
shape=(None, self.dim))
# For single system we just take the SU(C) matrices with C^2 - 1 free parameters
weights = tf.get_variable("weights",
(self.dim, dims[0]**2 - 1),
dtype=tf.float32)
# Multiply the weights with the inputs to get a field phi for each vector parameter.
phi = tf.matmul(self.x, weights)
phi_norm = tf.sqrt(
tf.reshape(tf.einsum('nj,nj->n', phi, tf.conj(phi)),
(-1, 1)))
phi = tf.tanh(phi_norm) * phi / phi_norm
# Construct the model density matrix from the SU(C) generators.
rho_red = 0.5 * (
tf.eye(dims[0], dims[0], dtype=tf.complex128) + tf.einsum(
'ni,jki->njk', tf.to_complex128(phi), generators))
else:
print('Building entangled model')
if len(hilbert_space_composition) == 2:
self.MODEL_TYPE = 'bipartite'
elif len(hilbert_space_composition) > 2:
self.MODEL_TYPE = 'multipartite'
# In the bi- or multipartite case the number of parameters is equal to the size the combined Hilbert space
if self.learn_hamiltonian is None:
# weights = tf.get_variable("weights", (self.dim, (np.product(dims, keepdims=True))), dtype=tf.float32)
if complex_weights:
self.x = tf.placeholder(dtype=tf.complex128,
shape=(None, self.dim))
weights_r = tf.get_variable(
"weights_r",
(self.dim, (np.product(dims, keepdims=True))),
dtype=tf.float32)
weights_c = tf.get_variable(
"weights_c",
(self.dim, (np.product(dims, keepdims=True))),
dtype=tf.float32)
weights = tf.to_complex128(
tf.complex(weights_r, weights_c))
else:
self.x = tf.placeholder(dtype=tf.float32,
shape=(None, self.dim))
weights = tf.get_variable(
"weights",
(self.dim, (np.product(dims, keepdims=True))),
dtype=tf.float32)
phi = tf.matmul(self.x, weights)
phi_norm = tf.sqrt(
tf.reshape(tf.einsum('nj,nj->n', phi, tf.conj(phi)),
(-1, 1)))
phi /= phi_norm
rho = tf.einsum('ni,nj->nij', phi, tf.conj(phi))
# Trace out the entangled systems to get a C x C density matrix
rho_red = self.partial_trace(rho, dims)
else:
self.x = tf.placeholder(dtype=tf.complex128,
shape=(None, self.dim))
I = tf.ones((tf.shape(self.x)[0], 1), dtype=tf.complex128)
H = tf.zeros(shape=(np.product(dims), np.product(dims)),
dtype=tf.complex128)
num_weights = 0
for flag, name, f in self.hamiltonian.field_hamiltonian_iterator(
):
if flag:
weight = tf.get_variable("weight_{}".format(name),
(self.dim, 1),
dtype=tf.float32)
h = tf.matmul(self.x,
tf.to_complex128(weight),
name="field_{}".format(name))
H += tf.einsum('ni,jk->njk', h, tf.constant(f))
num_weights += 1
else:
H += tf.einsum('ni,jk->njk', I, tf.constant(f))
for flag, name, f in self.hamiltonian.coupling_hamiltonian_iterator(
):
if flag:
weight = tf.get_variable("weight_{}".format(name),
(self.dim, 1),
dtype=tf.float32)
h = tf.matmul(self.x,
tf.to_complex128(weight),
name="coupling_{}".format(name))
H += tf.einsum('ni,jk->njk', h, tf.constant(f))
num_weights += 1
else:
H += tf.einsum('ni,jk->njk', I, tf.constant(f))
energies, phi = tf.linalg.eigh(H)
energies = tf.identity(energies, name='energies')
if self.rank_one_approx:
rho = tf.einsum('ni,nj->nij', phi[:, 0],
tf.conj(phi[:, 0]))
else:
rho = self.matrix_exp(H)
rho /= tf.reshape(tf.trace(rho), (-1, 1, 1))
rho_red = self.partial_trace(rho, dims)
with tf.name_scope('quantities'):
tf.identity(energies[:, 0], name='gs_energy')
tf.identity(energies[:, 0] - energies[:, 1],
name='width')
lams, _ = tf.linalg.eigh(rho_red)
lams = tf.identity(lams, name='lams')
tf.reduce_sum(-lams * tf.log(lams),
axis=1,
name='vn_entropy')
with tf.name_scope('negative_quantum_log_likelihood'):
# if lambda_op:
# gamma_0 = tf.constant(np.array([[1, 0], [0, 0]]), dtype=tf.complex128)
# gamma_1 = tf.constant(np.array([[0, 0], [0, 1]]), dtype=tf.complex128)
#
# self.negqll =- tf.reduce_sum(self.q_y[0] * tf.real(tf.log(tf.trace(tf.einsum('nij,jk->nik', rho_red, gamma_0)))))
# self.negqll -= tf.reduce_sum(self.q_y[1] * tf.real(tf.log(tf.trace(tf.einsum('nij,jk->nik', rho_red, gamma_1)))))
if self.c == 2:
self.negqll = -tf.reduce_sum(
tf.real(
tf.trace(
tf.matmul(self.eta,
self.matrix_log_2x2(rho_red)))))
else:
self.negqll = -tf.reduce_sum(
tf.real(
tf.trace(tf.matmul(self.eta,
self.matrix_log(rho_red)))))
tf.summary.scalar('negative_quantum_log_likelihood', self.negqll)
with tf.name_scope('optimizer'):
# Get the optimizer
opt = tf.train.GradientDescentOptimizer(learning_rate=epsilon)
# opt = tf.train.AdagradOptimizer(learning_rate=epsilon)
# self.train_step = opt.minimize(self.negqll, var_list=[weights])
if self.MODEL_TYPE == 'bipartite' or self.MODEL_TYPE == 'multipartite':
if self.learn_hamiltonian:
self.train_step = opt.minimize(
self.negqll,
var_list=[var for var in tf.trainable_variables()])
else:
if complex_weights:
self.train_step = opt.minimize(
self.negqll, var_list=[weights_r, weights_c])
else:
self.train_step = opt.minimize(self.negqll,
var_list=[weights])
else:
self.train_step = opt.minimize(self.negqll, var_list=[weights])
self._initialize_session(device)
with tf.name_scope('predictor'):
# if lambda_op:
# self.probabilities = tf.trace(tf.einsum('nij,jk->nik', rho_red, gamma_0))
# self.probabilities = tf.stack([self.probabilities, 1 - self.probabilities], axis=1)
if self.c == 2:
self.probabilities = tf.stack(
[rho_red[:, 0, 0], rho_red[:, 1, 1]], axis=1)
else:
_, v = tf.linalg.eigh(rho_red)
self.probabilities = tf.abs(v[:, :, 0])**2
self.FLAG_GRAPH = True
writer = SummaryWriter(log_dir='runs/' + run_name)
periodic: bool = args.bc == 'periodic'
input_transforms = model.periodic_input_transform if periodic else model.dirichlet_input_transform
output_transforms = model.periodic_output_transform if periodic else model.dirichlet_output_tranform
with tf.device(DEVICE):
m = PNetwork() if not args.simple else PNetworkSimple()
root = tf.train.Checkpoint(optimizer=optimizer, model=m, optimizer_step=tf.train.get_or_create_global_step())
blackbox_model = blackbox.black_box_periodic_output_transform if periodic else blackbox.black_box_dirichlet_output_transform
with tf.device(DEVICE):
pi = tf.constant(np.pi)
ci = tf.to_complex128(1j)
A_stencils_test, A_matrices_test, S_matrices_test, num_of_modes = utils.get_A_S_matrices(num_test_samples, np.pi,
grid_size, n_test)
with tf.device(DEVICE):
A_stencils_test = tf.convert_to_tensor(A_stencils_test, dtype=tf.double)
A_matrices_test = tf.convert_to_tensor(A_matrices_test, dtype=tf.complex128)
S_matrices_test = tf.reshape(tf.convert_to_tensor(S_matrices_test, dtype=tf.complex128),
(-1, num_of_modes, num_of_modes, grid_size ** 2, grid_size ** 2))
A_stencils_train = np.array(utils.two_d_stencil(num_training_samples))
n_train_list = [16, 16, 32]
initial_epsi = 1e-0
def call(self,
inputs,
black_box=False,
index=None,
pos=-1.,
phase='Training'):
with tf.device(self.device):
batch_size = inputs.shape[0]
right_contributions_input = tf.gather(
params=inputs,
indices=[i for i in range(1, self.grid_size, 2)],
axis=1)
right_contributions_input = tf.gather(
params=right_contributions_input,
indices=[i for i in range(0, self.grid_size, 2)],
axis=2)
idx = [(i - 1) % self.grid_size
for i in range(0, self.grid_size, 2)]
left_contributions_input = tf.gather(params=inputs,
indices=idx,
axis=1)
left_contributions_input = tf.gather(
params=left_contributions_input,
indices=[i for i in range(0, self.grid_size, 2)],
axis=2)
left_contributions_input = tf.reshape(
tensor=left_contributions_input,
shape=(-1, self.grid_size // 2, self.grid_size // 2, 3, 3))
up_contributions_input = tf.gather(
params=inputs,
indices=[i for i in range(0, self.grid_size, 2)],
axis=1)
up_contributions_input = tf.gather(
params=up_contributions_input,
indices=[i for i in range(1, self.grid_size, 2)],
axis=2)
up_contributions_input = tf.reshape(tensor=up_contributions_input,
shape=(-1, self.grid_size // 2,
self.grid_size // 2, 3,
3))
down_contributions_input = tf.gather(
params=inputs,
indices=[i for i in range(0, self.grid_size, 2)],
axis=1)
down_contributions_input = tf.gather(
params=down_contributions_input, indices=idx, axis=2)
down_contributions_input = tf.reshape(
tensor=down_contributions_input,
shape=(-1, self.grid_size // 2, self.grid_size // 2, 3, 3))
#
center_contributions_input = tf.gather(
params=inputs,
indices=[i for i in range(0, self.grid_size, 2)],
axis=1)
center_contributions_input = tf.gather(
params=center_contributions_input,
indices=[i for i in range(0, self.grid_size, 2)],
axis=2)
center_contributions_input = tf.reshape(
tensor=center_contributions_input,
shape=(-1, self.grid_size // 2, self.grid_size // 2, 3, 3))
inputs_combined = tf.concat([
right_contributions_input, left_contributions_input,
up_contributions_input, down_contributions_input,
center_contributions_input
], 0)
flattended = tf.reshape(inputs_combined, (-1, 9))
temp = (self.grid_size // 2)**2
# bug then augmented with doubled grid size
flattended = tf.concat([
flattended[:batch_size * temp],
flattended[temp * batch_size:temp * 2 * batch_size],
flattended[temp * 2 * batch_size:temp * 3 * batch_size],
flattended[temp * 3 * batch_size:temp * 4 * batch_size],
flattended[temp * 4 * batch_size:]
], -1)
# x = self.linear0(flattended)
# x = tf.nn.relu(x)
# for i in range(1, self.num_layers, 2):
# x1 = getattr(self, "bias_1%i" % i) + x
# x1 = getattr(self, "linear%i" % i)(x1)
# x1 = x1 + getattr(self, "bias_2%i" % i) + x1
# x1 = tf.nn.relu(x1)
# x1 = x1 + getattr(self, "bias_3%i" % i) + x1
# x1 = getattr(self, "linear%i" % (i + 1))(x1)
# x1 = tf.multiply(x1, getattr(self, "multiplier_%i" % i))
# x = x + x1
# x = x + getattr(self, "bias_4%i" % i)
# x = tf.nn.relu(x)
x = self.dense0(flattended)
for i in range(1, self.num_layers):
x = getattr(self, "dense%i" % 1)(x)
x = self.output_layer(x)
if index is not None:
indices = tf.constant([[index]])
updates = [tf.to_double(pos)]
shape = tf.constant([2 * 2 * 2 * 8])
scatter = tf.scatter_nd(indices, updates, shape)
x = self.new_output + tf.reshape(scatter, (-1, 2, 2, 8))
ld_contribution = x[:, :, :, 0]
left_contributions_output = x[:, :, :, 1]
lu_contribution = x[:, :, :, 2]
down_contributions_output = x[:, :, :, 3]
up_contributions_output = x[:, :, :, 4]
ones = tf.ones_like(up_contributions_output)
right_contributions_output = x[:, :, :, 6]
rd_contribution = x[:, :, :, 5]
ru_contribution = x[:, :, :, 7]
first_row = tf.concat([
tf.expand_dims(ld_contribution, -1),
tf.expand_dims(left_contributions_output, -1),
tf.expand_dims(lu_contribution, -1)
], -1)
second_row = tf.concat([
tf.expand_dims(down_contributions_output, -1),
tf.expand_dims(ones, -1),
tf.expand_dims(up_contributions_output, -1)
], -1)
third_row = tf.concat([
tf.expand_dims(rd_contribution, -1),
tf.expand_dims(right_contributions_output, -1),
tf.expand_dims(ru_contribution, -1)
], -1)
output = tf.stack([first_row, second_row, third_row], 0)
output = tf.transpose(output, (1, 2, 3, 0, 4))
if not black_box:
return tf.to_complex128(output)
else:
x = tf.reshape(
x, (-1, self.grid_size // 2, self.grid_size // 2, 4))
if black_box:
up_contributions_output = tf.gather(
inputs, [i for i in range(0, self.grid_size, 2)], axis=1)
up_contributions_output = tf.gather(
up_contributions_output,
[i for i in range(1, self.grid_size, 2)],
axis=2)
up_contributions_output = -tf.reduce_sum(
up_contributions_output[:, :, :, :, 0],
axis=-1) / tf.reduce_sum(
up_contributions_output[:, :, :, :, 1], axis=-1)
left_contributions_output = tf.gather(inputs, idx, axis=1)
left_contributions_output = tf.gather(
left_contributions_output,
[i for i in range(0, self.grid_size, 2)],
axis=2)
left_contributions_output = -tf.reduce_sum(
left_contributions_output[:, :, :, 2, :],
axis=-1) / tf.reduce_sum(
left_contributions_output[:, :, :, 1, :], axis=-1)
right_contributions_output = tf.gather(
inputs, [i for i in range(1, self.grid_size, 2)], axis=1)
right_contributions_output = tf.gather(
right_contributions_output,
[i for i in range(0, self.grid_size, 2)],
axis=2)
right_contributions_output = -tf.reduce_sum(
right_contributions_output[:, :, :, 0, :],
axis=-1) / tf.reduce_sum(
right_contributions_output[:, :, :, 1, :], axis=-1)
down_contributions_output = tf.gather(
inputs, [i for i in range(0, self.grid_size, 2)], axis=1)
down_contributions_output = tf.gather(
down_contributions_output, idx, axis=2)
down_contributions_output = -tf.reduce_sum(
down_contributions_output[:, :, :, :, 2],
axis=-1) / tf.reduce_sum(
down_contributions_output[:, :, :, :, 1], axis=-1)
else:
jm1 = [(i - 1) % (self.grid_size // 2)
for i in range(self.grid_size // 2)]
jp1 = [(i + 1) % (self.grid_size // 2)
for i in range(self.grid_size // 2)]
right_contributions_output = x[:, :, :, 0] / (
tf.gather(x[:, :, :, 1], jp1, axis=1) + x[:, :, :, 0])
left_contributions_output = x[:, :, :, 1] / (
x[:, :, :, 1] + tf.gather(x[:, :, :, 0], jm1, axis=1))
up_contributions_output = x[:, :, :, 2] / (
x[:, :, :, 2] + tf.gather(x[:, :, :, 3], jp1, axis=2))
down_contributions_output = x[:, :, :, 3] / (
tf.gather(x[:, :, :, 2], jm1, axis=2) + x[:, :, :, 3])
ones = tf.ones_like(down_contributions_output)
# based on rule 2 given rule 1:
up_right_contribution = tf.gather(
inputs, [i for i in range(1, self.grid_size, 2)], axis=1)
up_right_contribution = tf.gather(
up_right_contribution,
[i for i in range(1, self.grid_size, 2)],
axis=2)
up_right_contribution = up_right_contribution[:, :, :, 0, 1]
right_up_contirbution = tf.gather(
inputs, [i for i in range(1, self.grid_size, 2)], axis=1)
right_up_contirbution = tf.gather(
right_up_contirbution,
[i for i in range(1, self.grid_size, 2)],
axis=2)
right_up_contirbution_additional_term = right_up_contirbution[:, :, :,
0, 0]
right_up_contirbution = right_up_contirbution[:, :, :, 1, 0]
ru_center_ = tf.gather(inputs,
[i for i in range(1, self.grid_size, 2)],
axis=1)
ru_center_ = tf.gather(ru_center_,
[i for i in range(1, self.grid_size, 2)],
axis=2)
ru_center_ = ru_center_[:, :, :, 1, 1]
ru_contribution = -tf.expand_dims((right_up_contirbution_additional_term +
tf.multiply(right_up_contirbution, right_contributions_output) + \
tf.multiply(up_right_contribution,
up_contributions_output)) / ru_center_, -1)
up_left_contribution = tf.gather(inputs, idx, axis=1)
up_left_contribution = tf.gather(
up_left_contribution, [i for i in range(1, self.grid_size, 2)],
axis=2)
up_left_contribution = up_left_contribution[:, :, :, 2, 1]
left_up_contirbution = tf.gather(inputs, idx, axis=1)
left_up_contirbution = tf.gather(
left_up_contirbution, [i for i in range(1, self.grid_size, 2)],
axis=2)
left_up_contirbution_addtional_term = left_up_contirbution[:, :, :,
2, 0]
left_up_contirbution = left_up_contirbution[:, :, :, 1, 0]
lu_center_ = tf.gather(inputs, idx, axis=1)
lu_center_ = tf.gather(lu_center_,
[i for i in range(1, self.grid_size, 2)],
axis=2)
lu_center_ = lu_center_[:, :, :, 1, 1]
lu_contribution = -tf.expand_dims((left_up_contirbution_addtional_term +
tf.multiply(up_left_contribution, up_contributions_output) + \
tf.multiply(left_up_contirbution,
left_contributions_output)) / lu_center_, -1)
down_left_contribution = tf.gather(inputs, idx, axis=1)
down_left_contribution = tf.gather(down_left_contribution,
idx,
axis=2)
down_left_contribution = down_left_contribution[:, :, :, 2, 1]
left_down_contirbution = tf.gather(inputs, idx, axis=1)
left_down_contirbution = tf.gather(left_down_contirbution,
idx,
axis=2)
left_down_contirbution_additional_term = left_down_contirbution[:, :, :,
2,
2]
left_down_contirbution = left_down_contirbution[:, :, :, 1, 2]
ld_center_ = tf.gather(inputs, idx, axis=1)
ld_center_ = tf.gather(ld_center_, idx, axis=2)
ld_center_ = ld_center_[:, :, :, 1, 1]
ld_contribution = -tf.expand_dims((left_down_contirbution_additional_term +
tf.multiply(down_left_contribution, down_contributions_output) + \
tf.multiply(left_down_contirbution,
left_contributions_output)) / ld_center_, -1)
down_right_contribution = tf.gather(
inputs, [i for i in range(1, self.grid_size, 2)], axis=1)
down_right_contribution = tf.gather(down_right_contribution,
idx,
axis=2)
down_right_contribution = down_right_contribution[:, :, :, 0, 1]
right_down_contirbution = tf.gather(
inputs, [i for i in range(1, self.grid_size, 2)], axis=1)
right_down_contirbution = tf.gather(right_down_contirbution,
idx,
axis=2)
right_down_contirbution_addtional_term = right_down_contirbution[:, :, :,
0,
2]
right_down_contirbution = right_down_contirbution[:, :, :, 1, 2]
rd_center_ = tf.gather(inputs,
[i for i in range(1, self.grid_size, 2)],
axis=1)
rd_center_ = tf.gather(rd_center_, idx, axis=2)
rd_center_ = rd_center_[:, :, :, 1, 1]
rd_contribution = -tf.expand_dims((right_down_contirbution_addtional_term + tf.multiply(
down_right_contribution, down_contributions_output) + \
tf.multiply(right_down_contirbution,
right_contributions_output)) / rd_center_, -1)
first_row = tf.concat([
ld_contribution,
tf.expand_dims(left_contributions_output, -1), lu_contribution
], -1)
second_row = tf.concat([
tf.expand_dims(down_contributions_output, -1),
tf.expand_dims(ones, -1),
tf.expand_dims(up_contributions_output, -1)
], -1)
third_row = tf.concat([
rd_contribution,
tf.expand_dims(right_contributions_output, -1), ru_contribution
], -1)
output = tf.stack([first_row, second_row, third_row], 0)
output = tf.transpose(output, (1, 2, 3, 0, 4))
return tf.to_complex128(output)
def body_inside(rho, next_term, k, iter, my_L):
help = tf.add(k, one)
help2 = (delta_t / tf.to_complex128(tf.complex(help * iter, zero)))
temp = tf.matmul(help2 * my_L, next_term)
return [rho + temp, temp, help, iter, my_L]
def call(self,
inputs,
black_box=False,
index=None,
pos=-1.,
phase='Training'):
with tf.device("gpu:0"):
batch_size = inputs.shape[0]
right_contributions_input = tf.gather(
inputs, [i for i in range(2, self.grid_size, 2)], axis=1)
right_contributions_input = tf.gather(
right_contributions_input,
[i for i in range(1, self.grid_size, 2)],
axis=2)
idx = [i for i in range(0, self.grid_size - 1, 2)]
left_contributions_input = tf.gather(inputs, idx, axis=1)
left_contributions_input = tf.gather(
left_contributions_input,
[i for i in range(1, self.grid_size, 2)],
axis=2)
left_contributions_input = tf.reshape(
left_contributions_input,
(-1, self.grid_size // 2, self.grid_size // 2, 3, 3))
up_contributions_input = tf.gather(
inputs, [i for i in range(1, self.grid_size, 2)], axis=1)
up_contributions_input = tf.gather(
up_contributions_input,
[i for i in range(2, self.grid_size, 2)],
axis=2)
up_contributions_input = tf.reshape(
up_contributions_input,
(-1, self.grid_size // 2, self.grid_size // 2, 3, 3))
down_contributions_input = tf.gather(
inputs, [i for i in range(1, self.grid_size, 2)], axis=1)
down_contributions_input = tf.gather(down_contributions_input,
idx,
axis=2)
down_contributions_input = tf.reshape(
down_contributions_input,
(-1, self.grid_size // 2, self.grid_size // 2, 3, 3))
#
center_contributions_input = tf.gather(
inputs, [i for i in range(1, self.grid_size, 2)], axis=1)
center_contributions_input = tf.gather(
center_contributions_input,
[i for i in range(1, self.grid_size, 2)],
axis=2)
center_contributions_input = tf.reshape(
center_contributions_input,
(-1, self.grid_size // 2, self.grid_size // 2, 3, 3))
inputs_combined = tf.concat([
right_contributions_input, left_contributions_input,
up_contributions_input, down_contributions_input,
center_contributions_input
], 0)
inputs_combined_temp = tf.concat([
tf.reduce_sum(inputs_combined, axis=3),
tf.reduce_sum(inputs_combined, axis=4)
],
axis=-1)
inputs_combined_temp = tf.reshape(inputs_combined_temp, (-1, 6))
inputs_combined_temp_fixed = tf.where(
tf.equal(inputs_combined_temp,
0), 1e20 * tf.ones_like(inputs_combined_temp),
inputs_combined_temp)
#
flattended = tf.reshape(inputs_combined, (-1, 9))
# flattended = tf.concat([flattended, inputs_combined_temp,tf.reduce_sum(flattended,1,keepdims=True)],-1)
# flattended = tf.check_numerics(tf.concat([flattended,tf.reciprocal(flattended)], axis=-1),"sdf")
# flattended = tf.concat([flattended, tf.reciprocal(flattended)], -1)
temp = (self.grid_size // 2)**2
flattended = tf.concat([
flattended[:batch_size * temp],
flattended[temp * batch_size:temp * 2 * batch_size],
flattended[temp * 2 * batch_size:temp * 3 * batch_size],
flattended[temp * 3 * batch_size:temp * 4 * batch_size],
flattended[temp * 4 * batch_size:]
], -1)
x = self.linear0(flattended)
x = tf.nn.relu(x)
for i in range(1, self.num_layers, 2):
x1 = getattr(self, "bias_1%i" % i) + x
x1 = getattr(self, "linear%i" % i)(x1)
x1 = x1 + getattr(self, "bias_2%i" % i) + x1
x1 = tf.nn.relu(x1)
x1 = x1 + getattr(self, "bias_3%i" % i) + x1
x1 = getattr(self, "linear%i" % (i + 1))(x1)
x1 = tf.multiply(x1, getattr(self, "multiplier_%i" % i))
x = x + x1
x = x + getattr(self, "bias_4%i" % i)
x = tf.nn.relu(x)
x = self.output_layer(x)
if index is not None:
indices = tf.constant([[index]])
updates = [tf.to_double(pos)]
shape = tf.constant([2 * 2 * 2 * 8])
scatter = tf.scatter_nd(indices, updates, shape)
x = self.new_output + tf.reshape(scatter, (-1, 2, 2, 8))
ld_contribution = x[:, :, :, 0]
left_contributions_output = x[:, :, :, 1]
lu_contribution = x[:, :, :, 2]
down_contributions_output = x[:, :, :, 3]
up_contributions_output = x[:, :, :, 4]
ones = tf.ones_like(up_contributions_output)
right_contributions_output = x[:, :, :, 6]
rd_contribution = x[:, :, :, 5]
ru_contribution = x[:, :, :, 7]
first_row = tf.concat([
tf.expand_dims(ld_contribution, -1),
tf.expand_dims(left_contributions_output, -1),
tf.expand_dims(lu_contribution, -1)
], -1)
second_row = tf.concat([
tf.expand_dims(down_contributions_output, -1),
tf.expand_dims(ones, -1),
tf.expand_dims(up_contributions_output, -1)
], -1)
third_row = tf.concat([
tf.expand_dims(rd_contribution, -1),
tf.expand_dims(right_contributions_output, -1),
tf.expand_dims(ru_contribution, -1)
], -1)
output = tf.stack([first_row, second_row, third_row], 0)
output = tf.transpose(output, (1, 2, 3, 0, 4))
if not black_box:
return tf.to_complex128(output)
else:
x = tf.reshape(
x, (-1, self.grid_size // 2, self.grid_size // 2, 4))
if black_box:
up_contributions_output = tf.gather(
inputs, [i for i in range(1, self.grid_size, 2)], axis=1)
up_contributions_output = tf.gather(
up_contributions_output,
[i for i in range(2, self.grid_size, 2)],
axis=2)
up_contributions_output = -tf.reduce_sum(
up_contributions_output[:, :, :, :, 0],
axis=-1) / tf.reduce_sum(
up_contributions_output[:, :, :, :, 1], axis=-1)
left_contributions_output = tf.gather(inputs, idx, axis=1)
left_contributions_output = tf.gather(
left_contributions_output,
[i for i in range(1, self.grid_size, 2)],
axis=2)
left_contributions_output = -tf.reduce_sum(
left_contributions_output[:, :, :, 2, :],
axis=-1) / tf.reduce_sum(
left_contributions_output[:, :, :, 1, :], axis=-1)
right_contributions_output = tf.gather(
inputs, [i for i in range(2, self.grid_size, 2)], axis=1)
right_contributions_output = tf.gather(
right_contributions_output,
[i for i in range(1, self.grid_size, 2)],
axis=2)
right_contributions_output = -tf.reduce_sum(
right_contributions_output[:, :, :, 0, :],
axis=-1) / tf.reduce_sum(
right_contributions_output[:, :, :, 1, :], axis=-1)
down_contributions_output = tf.gather(
inputs, [i for i in range(1, self.grid_size, 2)], axis=1)
down_contributions_output = tf.gather(
down_contributions_output, idx, axis=2)
down_contributions_output = -tf.reduce_sum(
down_contributions_output[:, :, :, :, 2],
axis=-1) / tf.reduce_sum(
down_contributions_output[:, :, :, :, 1], axis=-1)
else:
jm1 = [(i - 0) % (self.grid_size // 2)
for i in range(self.grid_size // 2 - 1)]
jp1 = [(i + 1) % (self.grid_size // 2)
for i in range(self.grid_size // 2 - 1)]
right_contributions_output = x[:, :-1, :, 0] / (
tf.gather(x[:, :, :, 1], jp1, axis=1) + x[:, :-1, :, 0])
left_contributions_output = x[:, 1:, :, 1] / (
x[:, 1:, :, 1] + tf.gather(x[:, :, :, 0], jm1, axis=1))
up_contributions_output = x[:, :, :-1, 2] / (
x[:, :, :-1, 2] + tf.gather(x[:, :, :, 3], jp1, axis=2))
down_contributions_output = x[:, :, 1:, 3] / (
tf.gather(x[:, :, :, 2], jm1, axis=2) + x[:, :, 1:, 3])
#complete right with black box:
right_contributions_output_bb = tf.gather(
inputs, [i for i in range(2, self.grid_size, 2)], axis=1)
right_contributions_output_bb = tf.gather(
right_contributions_output_bb,
[i for i in range(1, self.grid_size, 2)],
axis=2)
right_contributions_output_bb = -tf.reduce_sum(
right_contributions_output_bb[:, :, :, 0, :],
axis=-1) / tf.reduce_sum(
right_contributions_output_bb[:, :, :, 1, :], axis=-1)
right_contributions_output_bb = tf.reshape(
right_contributions_output_bb[:, -1, :], (1, 1, -1))
right_contributions_output = tf.concat([
right_contributions_output, right_contributions_output_bb
],
axis=1)
left_contributions_output_bb = tf.gather(inputs, idx, axis=1)
left_contributions_output_bb = tf.gather(
left_contributions_output_bb,
[i for i in range(1, self.grid_size, 2)],
axis=2)
left_contributions_output_bb = -tf.reduce_sum(
left_contributions_output_bb[:, :, :, 2, :],
axis=-1) / tf.reduce_sum(
left_contributions_output_bb[:, :, :, 1, :], axis=-1)
left_contributions_output_bb = tf.reshape(
left_contributions_output_bb[:, 0, :], (1, 1, -1))
left_contributions_output = tf.concat(
[left_contributions_output_bb, left_contributions_output],
axis=1)
up_contributions_output_bb = tf.gather(
inputs, [i for i in range(1, self.grid_size, 2)], axis=1)
up_contributions_output_bb = tf.gather(
up_contributions_output_bb,
[i for i in range(2, self.grid_size, 2)],
axis=2)
up_contributions_output_bb = -tf.reduce_sum(
up_contributions_output_bb[:, :, :, :, 0],
axis=-1) / tf.reduce_sum(
up_contributions_output_bb[:, :, :, :, 1], axis=-1)
up_contributions_output_bb = tf.reshape(
up_contributions_output_bb[:, :, -1], (1, -1, 1))
up_contributions_output = tf.concat(
[up_contributions_output, up_contributions_output_bb],
axis=-1)
down_contributions_output_bb = tf.gather(
inputs, [i for i in range(1, self.grid_size, 2)], axis=1)
down_contributions_output_bb = tf.gather(
down_contributions_output_bb, idx, axis=2)
down_contributions_output_bb = -tf.reduce_sum(
down_contributions_output_bb[:, :, :, :, 2],
axis=-1) / tf.reduce_sum(
down_contributions_output_bb[:, :, :, :, 1], axis=-1)
down_contributions_output_bb = tf.reshape(
down_contributions_output_bb[:, :, 0], (1, -1, 1))
down_contributions_output = tf.concat(
[down_contributions_output_bb, down_contributions_output],
axis=-1)
ones = tf.ones_like(down_contributions_output)
idx = [i for i in range(0, self.grid_size - 1, 2)]
#based on rule 2 given rule 1:
# x,y = np.ix_([3, 1], [1, 3])
up_right_contribution = tf.gather(
inputs, [i for i in range(2, self.grid_size, 2)], axis=1)
up_right_contribution = tf.gather(
up_right_contribution,
[i for i in range(2, self.grid_size, 2)],
axis=2)
up_right_contribution = up_right_contribution[:, :, :, 0, 1]
right_up_contirbution = tf.gather(
inputs, [i for i in range(2, self.grid_size, 2)], axis=1)
right_up_contirbution = tf.gather(
right_up_contirbution,
[i for i in range(2, self.grid_size, 2)],
axis=2)
right_up_contirbution_additional_term = right_up_contirbution[:, :, :,
0, 0]
right_up_contirbution = right_up_contirbution[:, :, :, 1, 0]
ru_center_ = tf.gather(inputs,
[i for i in range(2, self.grid_size, 2)],
axis=1)
ru_center_ = tf.gather(ru_center_,
[i for i in range(2, self.grid_size, 2)],
axis=2)
ru_center_ = ru_center_[:, :, :, 1, 1]
ru_contribution = -tf.expand_dims((right_up_contirbution_additional_term+
tf.multiply(right_up_contirbution,right_contributions_output) +\
tf.multiply(up_right_contribution,up_contributions_output))/ru_center_, -1)
# x,y = np.ix_([3, 1], [3, 1])
up_left_contribution = tf.gather(inputs, idx, axis=1)
up_left_contribution = tf.gather(
up_left_contribution, [i for i in range(2, self.grid_size, 2)],
axis=2)
up_left_contribution = up_left_contribution[:, :, :, 2, 1]
left_up_contirbution = tf.gather(inputs, idx, axis=1)
left_up_contirbution = tf.gather(
left_up_contirbution, [i for i in range(2, self.grid_size, 2)],
axis=2)
left_up_contirbution_addtional_term = left_up_contirbution[:, :, :,
2, 0]
left_up_contirbution = left_up_contirbution[:, :, :, 1, 0]
lu_center_ = tf.gather(inputs, idx, axis=1)
lu_center_ = tf.gather(lu_center_,
[i for i in range(2, self.grid_size, 2)],
axis=2)
lu_center_ = lu_center_[:, :, :, 1, 1]
lu_contribution = -tf.expand_dims((left_up_contirbution_addtional_term+
tf.multiply(up_left_contribution , up_contributions_output) + \
tf.multiply(left_up_contirbution , left_contributions_output)) / lu_center_, -1)
# x,y = np.ix_([1, 3], [3, 1])
down_left_contribution = tf.gather(inputs, idx, axis=1)
down_left_contribution = tf.gather(down_left_contribution,
idx,
axis=2)
down_left_contribution = down_left_contribution[:, :, :, 2, 1]
left_down_contirbution = tf.gather(inputs, idx, axis=1)
left_down_contirbution = tf.gather(left_down_contirbution,
idx,
axis=2)
left_down_contirbution_additional_term = left_down_contirbution[:, :, :,
2,
2]
left_down_contirbution = left_down_contirbution[:, :, :, 1, 2]
ld_center_ = tf.gather(inputs, idx, axis=1)
ld_center_ = tf.gather(ld_center_, idx, axis=2)
ld_center_ = ld_center_[:, :, :, 1, 1]
ld_contribution = -tf.expand_dims((left_down_contirbution_additional_term+
tf.multiply(down_left_contribution , down_contributions_output) + \
tf.multiply(left_down_contirbution , left_contributions_output)) / ld_center_,-1)
# x,y = np.ix_([1, 3], [1, 3])
down_right_contribution = tf.gather(
inputs, [i for i in range(2, self.grid_size, 2)], axis=1)
down_right_contribution = tf.gather(down_right_contribution,
idx,
axis=2)
down_right_contribution = down_right_contribution[:, :, :, 0, 1]
right_down_contirbution = tf.gather(
inputs, [i for i in range(2, self.grid_size, 2)], axis=1)
right_down_contirbution = tf.gather(right_down_contirbution,
idx,
axis=2)
right_down_contirbution_addtional_term = right_down_contirbution[:, :, :,
0,
2]
right_down_contirbution = right_down_contirbution[:, :, :, 1, 2]
rd_center_ = tf.gather(inputs,
[i for i in range(2, self.grid_size, 2)],
axis=1)
rd_center_ = tf.gather(rd_center_, idx, axis=2)
rd_center_ = rd_center_[:, :, :, 1, 1]
rd_contribution = -tf.expand_dims((right_down_contirbution_addtional_term+tf.multiply(down_right_contribution , down_contributions_output) + \
tf.multiply(right_down_contirbution, right_contributions_output)) / rd_center_,-1)
first_row = tf.concat([
ld_contribution,
tf.expand_dims(left_contributions_output, -1), lu_contribution
], -1)
second_row = tf.concat([
tf.expand_dims(down_contributions_output, -1),
tf.expand_dims(ones, -1),
tf.expand_dims(up_contributions_output, -1)
], -1)
third_row = tf.concat([
rd_contribution,
tf.expand_dims(right_contributions_output, -1), ru_contribution
], -1)
output = tf.stack([first_row, second_row, third_row], 0)
output = tf.transpose(output, (1, 2, 3, 0, 4))
return tf.to_complex128(output)
