def __init__(self,
concentration_G,
concentration_X,
concentration_Y,
dx,
dt=None,
params=None,
source_functions=None):
if dt is None:
dt = 0.1 * dx
if dt > 0.5 * dx:
warnings.warn(
"Time increment {} too large for simulation stability with grid constant {}"
.format(dt, dx))
super().__init__(dx, dt, concentration_G.shape)
if concentration_X.shape != concentration_Y.shape or concentration_X.shape != concentration_G.shape:
raise ValueError("Concentration shapes must match")
self.params = params or DEFAULT_PARAMS
self.source_functions = source_functions or {}
self.G = tf.constant(concentration_G, dtype="float64")
self.X = tf.constant(concentration_X, dtype="float64")
self.Y = tf.constant(concentration_Y, dtype="float64")
if self.dims == 1:
omega2 = self.omega_x**2
elif self.dims == 2:
omega_x, omega_y = self.omega_x, self.omega_y
omega2 = omega_x**2 + omega_y**2
elif self.dims == 3:
omega_x, omega_y, omega_z = self.omega_x, self.omega_y, self.omega_z
omega2 = omega_x**2 + omega_y**2 + omega_z**2
else:
raise ValueError('Only up to 3D supported')
delta = tf.cast(-omega2 * self.dt, 'complex128')
decay_G = tf.exp(self.params['D_G'] * delta)
decay_X = tf.exp(self.params['D_X'] * delta)
decay_Y = tf.exp(self.params['D_Y'] * delta)
def diffusion_integrator(*cons):
result = []
for decay, concentration in zip([decay_G, decay_X, decay_Y], cons):
f = self.fft(tf.cast(concentration, 'complex128'))
f *= decay
result.append(tf.cast(tf.math.real(self.ifft(f)), 'float64'))
return result
reaction_integrator_curried = lambda con_G, con_X, con_Y: reaction_integrator(
con_G, con_X, con_Y, self.dt, self.params['A'], self.params['B'],
self.params['k2'], self.params['k-2'], self.params['k5'])
self.diffusion_integrator = tf.function(diffusion_integrator)
self.reaction_integrator = tf.function(reaction_integrator_curried)
def get_trajectories(A, B, k2, k_2, k5):
G, X, Y = randn(3, 10)
Gs = [G]
Xs = [X]
Ys = [Y]
for _ in range(1000):
G, X, Y = reaction_integrator(G, X, Y, 0.02, A, B, k2, k_2, k5)
Gs.append(G)
Xs.append(X)
Ys.append(Y)
return array([Gs, Xs, Ys])
def __init__(self,
concentration_G,
concentration_X,
concentration_Y,
u,
dx,
dt=None,
params=None,
source_functions=None):
if dt is None:
dt = 0.1 * dx
if dt > 0.5 * dx:
warnings.warn(
"Time increment {} too large for simulation stability with grid constant {}"
.format(dt, dx))
super().__init__(dx, dt, concentration_G.shape)
self.params = params or DEFAULT_PARAMS
self.source_functions = source_functions or {}
self.G = tf.constant(concentration_G, 'float64')
self.X = tf.constant(concentration_X, 'float64')
self.Y = tf.constant(concentration_Y, 'float64')
#self.G = tf.constant(concentration_G, 'complex128')
#BJD added 3 lines below 9.8.2021
#self.del_G = tf.constant(concentration_G, 'float64')
#self.del_X = tf.constant(concentration_X, 'float64')
#self.del_Y = tf.constant(concentration_Y, 'float64')
G0, X0, Y0 = steady_state(self.params['A'], self.params['B'],
self.params['k2'], self.params['k-2'],
self.params['k5'])
if len(u) != self.dims:
raise ValueError(
"{0}-dimensional flow must have {0} components".format(
self.dims))
c2 = self.params["speed-of-sound"]**2
viscosity = self.params["viscosity"]
if self.dims == 1:
raise ValueError("1D not supported")
elif self.dims == 2:
self.u = tf.constant(u[0], 'float64')
self.v = tf.constant(u[1], 'float64')
omega_x, omega_y = self.omega_x, self.omega_y
omega2 = omega_x**2 + omega_y**2
omega2_x = tf.constant(
omega2 + 1 / 3 * omega_x * (omega_x + omega_y), "complex128")
omega2_y = tf.constant(
omega2 + 1 / 3 * omega_y * (omega_x + omega_y), "complex128")
decay_x = tf.exp(-viscosity * omega2_x * self.dt)
decay_y = tf.exp(-viscosity * omega2_y * self.dt)
delta = tf.constant(-omega2 * self.dt, "complex128")
decay_G = tf.exp(self.params['D_G'] * delta)
decay_X = tf.exp(self.params['D_X'] * delta)
decay_Y = tf.exp(self.params['D_Y'] * delta)
def flow_integrator(rho, u, v):
"""
Flow is integrated with respect to the total log density (rho)
"""
# Enter Fourier Domain
f_rho = self.fft(tf.cast(rho, 'complex128'))
waves_x = self.fft(tf.cast(u, 'complex128'))
waves_y = self.fft(tf.cast(v, 'complex128'))
# Viscosity and internal shear
waves_x *= decay_x
waves_y *= decay_y
# Exit Fourier Domain
u = tf.cast(self.ifft(waves_x), 'float64')
v = tf.cast(self.ifft(waves_y), 'float64')
# Calculate gradients
rho_dx = tf.cast(self.ifft(f_rho * self.kernel_dx), 'float64')
rho_dy = tf.cast(self.ifft(f_rho * self.kernel_dy), 'float64')
u_dx = tf.cast(self.ifft(waves_x * self.kernel_dx), 'float64')
u_dy = tf.cast(self.ifft(waves_x * self.kernel_dy), 'float64')
v_dx = tf.cast(self.ifft(waves_y * self.kernel_dx), 'float64')
v_dy = tf.cast(self.ifft(waves_y * self.kernel_dy), 'float64')
divergence = u_dx + v_dy
# This would handle log density continuity but it's actually handled individually for G, X and Y
# rho -= (u*rho_dx + v*rho_dy + divergence) * self.dt
# Self-advect flow
du = -u * u_dx - v * u_dy
dv = -u * v_dx - v * v_dy
# Propagate pressure waves
du -= c2 * rho_dx
dv -= c2 * rho_dy
# Apply strain
du += viscosity * (rho_dx * (u_dx + u_dx) + rho_dy *
(u_dy + v_dx) - 2 / 3 * rho_dx * divergence)
dv += viscosity * (rho_dx * (v_dx + u_dy) + rho_dy *
(v_dy + v_dy) - 2 / 3 * rho_dy * divergence)
u += du * self.dt
v += dv * self.dt
return u, v, divergence
def diffusion_advection_integrator(G, X, Y, u, v, divergence):
f_G = self.fft(tf.cast(G, 'complex128'))
f_X = self.fft(tf.cast(X, 'complex128'))
f_Y = self.fft(tf.cast(Y, 'complex128'))
#f_test_function = self.fft(tf.cast(G, 'complex128'))
#f_test_function = solver.fft(tf.constant(test_function, 'complex128'))
#test_function_dx_fft = tf.cast(self.ifft(f_test_function * self.kernel_dx), 'float64')
#test_function_dy_fft = tf.cast(self.ifft(f_test_function * self.kernel_dy), 'float64')
#test_function_nabla2_fft = np.real(solver.ifft(f_test_function * solver.kernel_laplacian).numpy())
#np.savetxt("/home/brendan/software/tf2-model-g/arrays/quiver_array34/del_G_dx.txt", test_function_dx_fft) # BJD 9.8.2021
#np.savetxt("/home/brendan/software/tf2-model-g/arrays/quiver_array34/del_G_dy.txt", test_function_dy_fft) # BJD 9.8.2021
f_G *= decay_G
f_X *= decay_X
f_Y *= decay_Y
G = tf.cast(self.ifft(f_G), 'float64')
X = tf.cast(self.ifft(f_X), 'float64')
Y = tf.cast(self.ifft(f_Y), 'float64')
G_dx = tf.cast(self.ifft(f_G * self.kernel_dx), 'float64')
G_dy = tf.cast(self.ifft(f_G * self.kernel_dy), 'float64')
X_dx = tf.cast(self.ifft(f_X * self.kernel_dx), 'float64')
X_dy = tf.cast(self.ifft(f_X * self.kernel_dy), 'float64')
Y_dx = tf.cast(self.ifft(f_Y * self.kernel_dx), 'float64')
Y_dy = tf.cast(self.ifft(f_Y * self.kernel_dy), 'float64')
G -= (u * G_dx + v * G_dy + G * divergence) * self.dt
X -= (u * X_dx + v * X_dy + X * divergence) * self.dt
Y -= (u * Y_dx + v * Y_dy + Y * divergence) * self.dt
return G, X, Y
elif self.dims == 3:
self.u = tf.constant(u[0], 'float64')
self.v = tf.constant(u[1], 'float64')
self.w = tf.constant(u[2], 'float64')
omega_x, omega_y, omega_z = self.omega_x, self.omega_y, self.omega_z
omega2 = omega_x**2 + omega_y**2 + omega_z**2
omega2_x = tf.constant(
omega2 + 1 / 3 * omega_x * (omega_x + omega_y + omega_z),
"complex128")
omega2_y = tf.constant(
omega2 + 1 / 3 * omega_y * (omega_x + omega_y + omega_z),
"complex128")
omega2_z = tf.constant(
omega2 + 1 / 3 * omega_z * (omega_x + omega_y + omega_z),
"complex128")
decay_x = tf.exp(-viscosity * omega2_x * self.dt)
decay_y = tf.exp(-viscosity * omega2_y * self.dt)
decay_z = tf.exp(-viscosity * omega2_z * self.dt)
delta = tf.constant(-omega2 * self.dt, "complex128")
decay_G = tf.exp(self.params['D_G'] * delta)
decay_X = tf.exp(self.params['D_X'] * delta)
decay_Y = tf.exp(self.params['D_Y'] * delta)
def flow_integrator(rho, u, v, w):
# Enter Fourier Domain
f_rho = self.fft(tf.cast(rho, 'complex128'))
waves_x = self.fft(tf.cast(u, 'complex128'))
waves_y = self.fft(tf.cast(v, 'complex128'))
waves_z = self.fft(tf.cast(w, 'complex128'))
# Viscosity and internal shear
waves_x *= decay_x
waves_y *= decay_y
waves_z *= decay_z
# Exit Fourier Domain
u = tf.cast(self.ifft(waves_x), 'float64')
v = tf.cast(self.ifft(waves_y), 'float64')
w = tf.cast(self.ifft(waves_z), 'float64')
# Calculate gradients
rho_dx = tf.cast(self.ifft(f_rho * self.kernel_dx), 'float64')
rho_dy = tf.cast(self.ifft(f_rho * self.kernel_dy), 'float64')
rho_dz = tf.cast(self.ifft(f_rho * self.kernel_dz), 'float64')
u_dx = tf.cast(self.ifft(waves_x * self.kernel_dx), 'float64')
u_dy = tf.cast(self.ifft(waves_x * self.kernel_dy), 'float64')
u_dz = tf.cast(self.ifft(waves_x * self.kernel_dz), 'float64')
v_dx = tf.cast(self.ifft(waves_y * self.kernel_dx), 'float64')
v_dy = tf.cast(self.ifft(waves_y * self.kernel_dy), 'float64')
v_dz = tf.cast(self.ifft(waves_y * self.kernel_dz), 'float64')
w_dx = tf.cast(self.ifft(waves_z * self.kernel_dx), 'float64')
w_dy = tf.cast(self.ifft(waves_z * self.kernel_dy), 'float64')
w_dz = tf.cast(self.ifft(waves_z * self.kernel_dz), 'float64')
divergence = u_dx + v_dy + w_dz
# This would handle log density continuity, but we do G, X and Y individually
# rho -= (u*rho_dx + v*rho_dy + w*rho_dz + divergence) * self.dt
# Self-advect flow
du = -u * u_dx - v * u_dy - w * u_dz
dv = -u * v_dx - v * v_dy - w * v_dz
dw = -u * w_dx - v * w_dy - w * w_dz
# Propagate pressure waves
du -= c2 * rho_dx
dv -= c2 * rho_dy
dw -= c2 * rho_dz
# Apply strain
du += viscosity * (rho_dx * (u_dx + u_dx) + rho_dy *
(u_dy + v_dx) + rho_dz *
(u_dz + w_dx) - 2 / 3 * rho_dx * divergence)
dv += viscosity * (rho_dx * (v_dx + u_dy) + rho_dy *
(v_dy + v_dy) + rho_dz *
(v_dz + w_dy) - 2 / 3 * rho_dy * divergence)
dw += viscosity * (rho_dx * (w_dx + u_dz) + rho_dy *
(w_dy + v_dz) + rho_dz *
(w_dz + w_dz) - 2 / 3 * rho_dz * divergence)
u += du * self.dt
v += dv * self.dt
w += dw * self.dt
return u, v, w, divergence
def diffusion_advection_integrator(G, X, Y, u, v, w, divergence):
f_G = self.fft(tf.cast(G, 'complex128'))
f_X = self.fft(tf.cast(X, 'complex128'))
f_Y = self.fft(tf.cast(Y, 'complex128'))
f_G *= decay_G
f_X *= decay_X
f_Y *= decay_Y
G = tf.cast(self.ifft(f_G), 'float64')
X = tf.cast(self.ifft(f_X), 'float64')
Y = tf.cast(self.ifft(f_Y), 'float64')
G_dx = tf.cast(self.ifft(f_G * self.kernel_dx), 'float64')
G_dy = tf.cast(self.ifft(f_G * self.kernel_dy), 'float64')
G_dz = tf.cast(self.ifft(f_G * self.kernel_dz), 'float64')
X_dx = tf.cast(self.ifft(f_X * self.kernel_dx), 'float64')
X_dy = tf.cast(self.ifft(f_X * self.kernel_dy), 'float64')
X_dz = tf.cast(self.ifft(f_X * self.kernel_dz), 'float64')
Y_dx = tf.cast(self.ifft(f_Y * self.kernel_dx), 'float64')
Y_dy = tf.cast(self.ifft(f_Y * self.kernel_dy), 'float64')
Y_dz = tf.cast(self.ifft(f_Y * self.kernel_dz), 'float64')
G -= (u * G_dx + v * G_dy + w * G_dz +
(G + G0) * divergence) * self.dt
X -= (u * X_dx + v * X_dy + w * X_dz +
(X + X0) * divergence) * self.dt
Y -= (u * Y_dx + v * Y_dy + w * Y_dz +
(Y + Y0) * divergence) * self.dt
return G, X, Y
else:
raise ValueError('Only up to 3D supported')
reaction_integrator_curried = lambda con_G, con_X, con_Y: reaction_integrator(
con_G, con_X, con_Y, self.dt, self.params['A'], self.params['B'],
self.params['k2'], self.params['k-2'], self.params['k5'])
self.reaction_integrator = tf.function(reaction_integrator_curried)
self.flow_integrator = tf.function(flow_integrator)
self.diffusion_advection_integrator = tf.function(
diffusion_advection_integrator)