def potential_energy(pos, edges_half, neg_edges_half, edges,
ljatom_diameter_tf):
with tf.name_scope("potential_energy"):
dcut = tf.constant(2.5, dtype=pos.dtype, name="dcut")
elj = 1.0
four = 4.0
one = 1.0
two = 2.0
dcut_6 = tf.pow(dcut, 6, name="dcut_6")
dcut_12 = tf.pow(dcut, 12, name="dcut_12")
energy_shift = four * elj * (one / dcut_12 - one / dcut_6)
distances = tf.compat.v1.vectorized_map(
fn=lambda atom_pos: pos - atom_pos, elems=pos)
distances = tf.compat.v1.where_v2(distances > edges_half,
distances - edges,
distances,
name="edges_half_where")
distances = tf.compat.v1.where_v2(distances < neg_edges_half,
distances + edges,
distances,
name="neg_edges_half_where")
magnitude = common.magnitude(distances)
d_6 = tf.pow(ljatom_diameter_tf, 6, name="energy_d_6")
r_6 = tf.pow(magnitude, 6, name="energy_r_6")
ljpair = four * elj * (d_6 / r_6) * ((d_6 / r_6) - one) - energy_shift
ret = tf.reduce_sum(ljpair, name="energy_reduce_sum") / two
return ret
def lj_force(pos, edges_half, neg_edges_half, edges, forces_zeroes_tf,
ljatom_diameter_tf):
with tf.name_scope("lj_force"):
distances = tf.compat.v1.vectorized_map(
fn=lambda atom_pos: pos - atom_pos, elems=pos)
distances = tf.compat.v1.where_v2(distances > edges_half,
distances - edges,
distances,
name="where_edges_half")
distances = tf.compat.v1.where_v2(distances < neg_edges_half,
distances + edges,
distances,
name="where_neg_edges_half")
magnitude = common.magnitude(distances)
twelve = tf.math.pow(ljatom_diameter_tf, 12.0,
name="diam_12_pow") / tf.math.pow(
magnitude, 12.0, name="mag_12_pow")
six = tf.math.pow(ljatom_diameter_tf, 6.0,
name="diam_6_pow") / tf.math.pow(
magnitude, 6.0, name="mag_6_pow")
mag_squared = 1.0 / tf.math.pow(magnitude, 2, name="mag_sqr_pow")
slice_forces = distances * (48.0 * 1.0 *
(((twelve - 0.5 * six) * mag_squared)))
# handle case distances pos - atom_pos == 0, causing inf and nan to appear in that position
# can't see a way to remove that case in all above computations, easier to do it all at once at the end
filter = tf.math.logical_or(tf.math.is_nan(slice_forces),
magnitude < (ljatom_diameter_tf * 2.5),
name="or")
filtered = tf.compat.v1.where_v2(filter,
forces_zeroes_tf,
slice_forces,
name="where_or")
forces = tf.math.reduce_sum(filtered, axis=0)
return forces
def _particle_electrostatic_force(simul_box, ion_dict):
"""
force on the particles (electrostatic)
parallel calculation of forces (uniform case)
"""
with tf.name_scope("particle_electrostatic_force"):
distances = common.wrap_vectorize(
fn=lambda atom_pos: ion_dict[interface.ion_pos_str] - atom_pos,
elems=ion_dict[interface.ion_pos_str])
z_distances = distances[:, :,
-1] # get z-axis value #TODO: Remove the need for third axis/pulling out z dimension => see if faster way
abs_z_distances = tf.math.abs(z_distances)
r1 = tf.math.sqrt(0.5 + ((z_distances / simul_box.lx) *
(z_distances / simul_box.lx)))
r2 = tf.math.sqrt(0.25 + ((z_distances / simul_box.lx) *
(z_distances / simul_box.lx)))
E_z = 4 * tf.math.atan(4 * abs_z_distances * r1 / simul_box.lx)
factor = tf.compat.v1.where_v2(z_distances >= 0.0,
_tf_one,
_tf_neg_one,
name="where_factor")
hcsh = (4 / simul_box.lx) * (1 / (r1 * (0.5 + r1)) - 1 / (r2 * r2)) * z_distances + factor * E_z + \
16 * abs_z_distances * (simul_box.lx / (simul_box.lx * simul_box.lx + 16 * z_distances * z_distances * r1 * r1)) * \
(abs_z_distances * z_distances / (simul_box.lx * simul_box.lx * r1) + factor * r1) # MATHEMATICAL
# print("hcsh.shape", hcsh.shape)
#h1.z = h1.z + 2 * ion[i].q * (ion[j].q / (box.lx * box.lx)) * 0.5 * (1 / ion[i].epsilon + 1 / ion[j].epsilon) * hcsh
one_over_ep = 1 / ion_dict[interface.ion_epsilon_str]
q_over_lx_sq = ion_dict[interface.ion_charges_str] / (simul_box.lx *
simul_box.lx)
vec_one_over_ep = common.wrap_vectorize(
fn=lambda epsilon_j: one_over_ep + epsilon_j, elems=one_over_ep)
# print("vec_one_over_ep.shape", vec_one_over_ep.shape)
vec_q_over_lx_sq = common.wrap_vectorize(
fn=lambda q_j: ion_dict[interface.ion_charges_str] * q_j,
elems=q_over_lx_sq)
# print("vec_q_over_lx_sq.shape", vec_q_over_lx_sq.shape)
h1_z = 2 * vec_q_over_lx_sq * 0.5 * vec_one_over_ep * hcsh
h1_z = tf.math.reduce_sum(h1_z, axis=1, keepdims=True)
# print("h1_z.shape", h1_z.shape)
# h1 =h1+ ((temp_vec ^ ((-1.0) / r3)) ^ ((-0.5) * ion[i].q * ion[j].q * (1 / ion[i].epsilon + 1 / ion[j].epsilon)));
wrapped_distances = common.wrap_distances_on_edges(
simul_box, distances)
r = common.magnitude(
wrapped_distances, keepdims=True
) # keep third dimension to divide third dim in wrapped_distances later
r3 = tf.math.pow(r, 3)
vec_q_mul = common.wrap_vectorize(
fn=lambda q_j: ion_dict[interface.ion_charges_str] * q_j,
elems=ion_dict[interface.ion_charges_str])
a = _zero_nans(
wrapped_distances * ((-1.0) / r3)
) # r3 can have zeroes in it, so remove the nans that come from div by zero
b = ((-0.5) * vec_q_mul * vec_one_over_ep)
# print("a.shape", a.shape)
# print("b.shape", b.shape)
# print("a * b[:,:,tf.newaxis].shape", (a * b[:,:,tf.newaxis]).shape)
h1 = tf.math.reduce_sum(
a * b[:, :, tf.newaxis],
axis=1,
keepdims=False,
name="sum_a_times_b") #TODO: remove need for newaxis here
# print("h1.shape", h1.shape)
h1_x_y = h1[:, 0:2] #TODO: replace this junk with better impl
c = h1[:, 2:3] + h1_z
con = tf.concat(values=[h1_x_y, c],
axis=1,
name="x_y_and_c_concatenate")
return con * utility.scalefactor
def _electrostatic_wall_force(simul_box, ion_dict, wall_dictionary):
"""
ion interacting via electrostatic force with discrete planar wall
"""
with tf.name_scope("electrostatic_wall_force"):
wall_distances = common.wrap_vectorize(
fn=lambda atom_pos: atom_pos - wall_dictionary["posvec"],
elems=ion_dict[interface.ion_pos_str])
wall_z_dist = wall_distances[:, :, -1] # get z-axis value
factor = tf.compat.v1.where_v2(wall_z_dist >= 0.0,
_tf_one,
_tf_neg_one,
name="where_factor")
r1_rightwall = tf.math.sqrt(0.5 + (wall_z_dist / simul_box.lx) *
(wall_z_dist / simul_box.lx))
r2_rightwall = tf.math.sqrt(0.25 + (wall_z_dist / simul_box.lx) *
(wall_z_dist / simul_box.lx))
E_z_rightwall = 4 * tf.math.atan(
4 * tf.math.abs(wall_z_dist) * r1_rightwall / simul_box.lx)
hcsh_rightwall = (4 / simul_box.lx) * (1 / (r1_rightwall * (0.5 + r1_rightwall)) - 1 / (r2_rightwall * r2_rightwall)) * wall_z_dist + factor * E_z_rightwall +\
16 * tf.math.abs(wall_z_dist) * (simul_box.lx / (simul_box.lx * simul_box.lx + 16 * wall_z_dist * wall_z_dist * r1_rightwall * r1_rightwall)) *\
(tf.math.abs(wall_z_dist) * wall_z_dist / (simul_box.lx * simul_box.lx * r1_rightwall) + factor * r1_rightwall)
# h1_rightwall.z = h1_rightwall.z + 2 * ion[i].q * (wall_dummy.q / (box.lx * box.lx)) * 0.5 * (1 / ion[i].epsilon + 1 / wall_dummy.epsilon) * hcsh_rightwall;
ion_one_over_ep = 1 / ion_dict[
interface.ion_epsilon_str] # 1 / ion[i].epsilon
wall_one_over_ep = 1 / wall_dictionary[
"epsilon"] # 1 / wall_dummy.epsilon
q_over_lx_sq = wall_dictionary["q"] / (
simul_box.lx * simul_box.lx) # (wall_dummy.q / (box.lx * box.lx))
vec_one_over_ep = common.wrap_vectorize(
fn=lambda ion_eps: wall_one_over_ep + ion_eps,
elems=ion_one_over_ep
) # (1 / ion[i].epsilon + 1 / wall_dummy.epsilon)
vec_q_over_lx_sq = common.wrap_vectorize(
fn=lambda q_j: q_over_lx_sq * q_j,
elems=ion_dict[interface.ion_charges_str]
) # ion[i].q * (wall_dummy.q / (box.lx * box.lx))
h1_z = 2 * vec_q_over_lx_sq * 0.5 * (vec_one_over_ep) * hcsh_rightwall
h1_z = tf.math.reduce_sum(h1_z, axis=1, keepdims=True, name="sum_h1_z")
# h1_rightwall = h1_rightwall+ ((temp_vec_rightwall ^ ((-1.0) / r3_rightwall)) ^ ((-0.5) * ion[i].q * wall_dummy.q * (1 / ion[i].epsilon + 1 / wall_dummy.epsilon)));
wrapped_distances = common.wrap_distances_on_edges(
simul_box, wall_distances)
r = common.magnitude(
wrapped_distances, keepdims=True
) # keep third dimension to divide third dim in wrapped_distances later
r3 = tf.math.pow(r, 3.0, name="r_3")
vec_q_mul = common.wrap_vectorize(
fn=lambda q_j: wall_dictionary["q"] * q_j,
elems=ion_dict[interface.ion_charges_str])
a = _zero_nans(wrapped_distances * ((-1.0) / r3)) * (
(-0.5) * vec_q_mul * vec_one_over_ep)[:, :, tf.newaxis]
h1 = tf.math.reduce_sum(a, axis=1, keepdims=False, name="sum_a_mul_b")
z = h1[:, 2:3] + h1_z
con = tf.concat(values=[h1[:, 0:2], z],
axis=1,
name="h1x_y_and_h1_z_concatenate")
return con * utility.scalefactor
def kinetic_energy(vel, ljatom_diameter_tf):
with tf.name_scope("kinetic_energy"):
# half = tf.constant(0.5, dtype=tf.float64)
magnitude = common.magnitude(vel)
return tf.reduce_sum(0.5 * ljatom_diameter_tf * magnitude * magnitude)