@Richard's answer to Going to try to move some of my scipy/numpy calculation to a new GPU, how to avoid disappointing results? is quite helpful, and as promised I've added a simple running example below of the kind of project I've started.
While there may be "high level libraries" out there that can do this, I always code/script my projects myself first (to the extent feasible) in order to get insight/appreciation into the problem.
I have a layer of atoms on a surface and am trying to do a 2D Frenkel–Kontorova-like energy minimization using at first simplex then once that works, with simulated annealing. The number of atoms could range from 100 to 10,000.
I'll be looking at deviations from coincident lattices (see also related code golf).
Each evaluation of energy is calculated from the atomic positions with respect to the substrate lattice and nearest-neighbor bond distances (and later, angles as well) within the layer, and will be done with numpy. In the future I may need to call scipy.special functions also.
The script below first builds the problem and stores it in the object bob for convenience (plotting, interrogation, rendering with Blender, etc.) but then extracts a few constants that are stored in args and three arrays:
atom_positionsan N x 3 array of positions of atoms, order arbitrary except that the first element is the centermost of the group. These are the values that a minimization routine will adjust.atoms_1andatoms_2are numpy indices for the atom_position array that are used to identify all pairs of atoms for which bond energies should be calculated; they are in this case nearest neighbors and to avoid double counting there are three per atom on the interior and fewer on the edges/corners. The order is again arbitrary.
There is an alternative way to formulate the numpy position array for minimization by making it a masked array. It might be slightly larger to add an extra "dummy atom" all the way around, and bond energies would be calculated by simple shifts, something like x_masked[1:, :] - x_masked[:-1, :] letting the predefined masks deal with atoms on the edge with fewer bonds.
note 1: Later I will be introducing defects in the lattice and including honeycomb structures (e.g. Xenes). In this case indexing has the advantage that you can introduce these quickly by editing the index arrays. Trying to do the defects in a regular array may become impossible.
note 2: The example script calls a generic Scipy minimization routine for convenience. I will ultimately be scripting a simulated annealing algorithm and experimenting to find a good annealing schedule, but that can't really happen until I get this running faster, likely by GPU augmentation of a PC.
Question: How amenable is this 2D Frenkel–Kontorova-like energy minimization problem in Python to the use of a modest PC + GPU? Is heavy reliance on indexing going to bite me later as I scale it up to 10,000 atoms?
Tiny example (runs in seconds on a laptop) with nmax=6. With nmax=20 for example you get 1261 atoms and a much more challenging problem to calculate.
bob = Problem('bob', sig=0.2, k_bond=20, a_sub=a_substrate, a=a_lattice, h=height,
R=R, origin=origin, nmax=6)
note 3: per request I have tried to do some profiling using scaline but the results seem trivial, to minimize energy most time is spent in the scipy minimizer and the line of numpy in my function mini_me that calculates the energy to minimize. Results are here https://pastebin.com/GfneKJhB and in nicely formatted html here https://pastebin.com/ZJAL8gqv
note 4: colors of the atoms and the bonds reflect their relative energies.
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection
from matplotlib import colors as mcolors
from scipy.optimize import minimize
import time
class Problem():
def __init__(self, name=None, nmax=12, a_sub=1., k_vert=1., a=1.,
h=1.0, k_bond=2., R=0., origin=np.zeros(2), sig=0., seed=42):
self.name = str(name)
self.nmax = int(nmax)
self.a_sub = float(a_sub)
twopi = 2 * np.pi
r3o2 = np.sqrt(3) / 2
g = (twopi / (self.a_sub * r3o2)) * np.array([[0, 1], [r3o2, 0.5], [-r3o2, 0.5]]).T
self.g_sub = g[::-1] # for purposes of calculating hexagonal substrate potential
self.k_vert = float(k_vert)
self.a = float(a)
self.h = float(h)
self.k_bond = float(k_bond)
self.R = float(R)
self.origin = np.array(origin, dtype=float)
self.sig = float(sig)
self.seed = int(seed)
self.initial_hex_lattice(origin=origin)
np.random.seed(self.seed)
self.random = self.sig * np.random.random(self.x0_init.shape)
self.x_initial = self.x0_init + self.random
self.x = self.x_initial.copy() # this array will be updated
# self.x_final = None
self.n_atoms = self.x_initial.shape[0]
self.define_nearest_neighbor_bonds()
self.calculate_bond_energies()
self.calculate_surface_energies()
def initial_hex_lattice(self, origin):
"""generates an N x 3 array of atom locations, origin atom is first, wight
Gaussian noise in 3D. The reduction to 1/6 slice by hex-hex symmetry is optional"""
i, j = np.mgrid[-self.nmax:self.nmax+1, -self.nmax:self.nmax+1]
keep = (np.abs(i + j) <= self.nmax) * ~((i == 0) * (j == 0)) # add origin back later
i, j = [thing[keep] for thing in (i, j)]
r3o2 = np.sqrt(3) / 2.
x = self.a * (i + 0.5 * j)
y = self.a * r3o2 * j
s, c = [f(np.radians(self.R)) for f in (np.sin, np.cos)]
x, y = c * x - s * y, c * y + s * x
z = np.zeros_like(x)
# origin = np.array([0, 0, self.h])
self.x0_init = origin + np.vstack([np.zeros(3), np.stack([x, y, z], axis=1)])
self.i = np.hstack(([0], i))
self.j = np.hstack(([0], j))
self.ij = np.stack((self.i, self.j), axis=1)
def define_nearest_neighbor_bonds(self):
"""calculate index pairs for all nearest-neighbor bonds in the hexagonal
lattice"""
a_bonds, b_bonds, c_bonds = [], [], []
n_atom = np.arange(self.n_atoms)
A, B, C, D = (self.i < self.nmax, self.i + self.j < self.nmax,
self.j < self.nmax, self.i > -self.nmax)
self.A, self.B, self.C, self.D = A, B, C, D
self.atoms_1a = n_atom[A * B]
self.atoms_1b = n_atom[B * C]
self.atoms_1c = n_atom[C * D]
self.atoms_2aij = self.ij[self.atoms_1a] + np.array([1, 0]) # to the right
self.atoms_2bij = self.ij[self.atoms_1b] + np.array([0, 1]) # to the above-rith
self.atoms_2cij = self.ij[self.atoms_1c] + np.array([-1, 1]) # to the above-left
self.atoms_2a = np.array([np.where(np.all(self.ij == atom , axis=1))[0][0]
for atom in self.atoms_2aij])
self.atoms_2b = np.array([np.where(np.all(self.ij == atom , axis=1))[0][0]
for atom in self.atoms_2bij])
self.atoms_2c = np.array([np.where(np.all(self.ij == atom , axis=1))[0][0]
for atom in self.atoms_2cij])
self.atoms_1 = np.hstack((self.atoms_1a, self.atoms_1b, self.atoms_1c))
self.atoms_2 = np.hstack((self.atoms_2a, self.atoms_2b, self.atoms_2c))
def calculate_bond_energies(self, which='x'):
"""harmonic oscilator potential for bond lengths, minimum at 'a'
with spring contsant 'k_bond'"""
array = getattr(self, which)
vectors = array[self.atoms_2] - array[self.atoms_1]
self.bond_lengths = np.sqrt((vectors**2).sum(axis=1))
self.bond_energies = self.k_bond * (self.bond_lengths - self.a)**2
def calculate_surface_energies(self, which='x'):
""" applies smooth sinusoidal-like hexagonal pattern with lattice constant 'a'
to a N x 2 grid of x-y points. Minimum at origin."""
array = getattr(self, which)
E = np.cos((array[:, :2, None] * self.g_sub).sum(axis=1))
result = (1.5 + E.sum(axis=1)) / 4.5
surface_energy_xy = 1. - (1.5 + E.sum(axis=-1)) / 4.5 # "sinusolidal" in x, y
surface_energy_z = self.k_vert * (self.x[:, 2] - self.h)**2 # harmonic osc. in z
self.surface_energies = surface_energy_xy + surface_energy_z
def get_energy(self, which='x'):
self.calculate_bond_energies(which)
self.calculate_surface_energies(which)
return self.bond_energies.sum() + self.surface_energies.sum()
def potential_yx(self, yx):
""" applies smooth sinusoidal-like hexagonal pattern with lattice constant 'a'
to a 2 x M x N y-x grid of points. Minimum at origin."""
E = np.cos((yx[:, None, ...] * self.g_sub[::-1][..., None, None]).sum(axis=0)) # sum over x, y then cos()
result = 1. - (1.5 + E.sum(axis=0)) / 4.5 # then sum over three directions
return result
def plot_it(self, hw_pixels=250, cmap_sub='gray', cmap_atoms='inferno',
cmap_bonds='jet', linewidth=2.5, vmin=0., vmax=1.2, which='x'):
array = getattr(self, which)
hw = 1.05 * np.abs(array[:, :2]).max() # largest cartesian distance of an atom from center + 10%
self.plot_scale = hw / hw_pixels # 0.05 angstroms per pixel
self.yx = self.plot_scale * np.stack(np.mgrid[-hw_pixels:hw_pixels+1,
-hw_pixels:hw_pixels+1], axis=0)
self.substrate = self.potential_yx(yx=self.yx)
self.plot_extent = [self.yx[1].min(), self.yx[1].max(),
self.yx[0].min(), self.yx[0].max()]
print('self.substrate.shape: ', self.substrate.shape)
# energy histograms
a, b = np.histogram(self.surface_energies, bins=(self.n_atoms >> 2))
c, d = np.histogram(self.bond_energies, bins=(self.n_atoms >> 2))
# color coded bond lines
my_cmap = plt.get_cmap(cmap_bonds)
colors = my_cmap(self.bond_energies / self.bond_energies.max()) # normalized for cmap
linez = np.stack((self.x[bob.atoms_1][:, :2], self.x[bob.atoms_2][:, :2]), axis=1)
ln_coll = LineCollection(linez, colors=colors, linewidths=linewidth)
fig = plt.figure(constrained_layout=False, figsize=[9, 7.5])
gs = fig.add_gridspec(nrows=4, ncols=5, left=0.08, right=0.92,
bottom=0.07, top=0.96, hspace=0.15)
ax1 = fig.add_subplot(gs[:3, :])
# fig, ax = plt.subplots(1, 1, figsize=[9, 7.5])
ax1.imshow(self.substrate, origin='lower', extent=self.plot_extent,
cmap=cmap_sub, vmin=vmin, vmax=vmax) # surface in background
ax1.add_collection(ln_coll) # then bonds
x, y = self.x.T[:2]
ax1.scatter(x, y, c=self.surface_energies, cmap=cmap_atoms, s=60,
edgecolors='k', zorder=2) # then atoms on top of the bonds
ax1.set_ylabel('Angstroms')
ax2 = fig.add_subplot(gs[3:, 1:4])
ax2.plot(b[1:], a, label='surface energies')
ax2.plot(d[1:], c, label='bond energies')
ax2.legend()
ax2.set_yscale('log')
ax2.set_xlabel('energy')
plt.show()
def ang(i, j, k, l):
root3 = np.sqrt(3)
th_ij = np.arctan2(root3 * j, 2*i+j)
th_kl = np.arctan2(root3 * l, 2*k+l)
return np.degrees(th_ij - th_kl)
#### NEXT
# print old/new mean origin
# print old/new mean rotation
# print old/new mean bond length
# is there any collective behavior?
# BEGIN
a_lattice = 3.5
a_substrate = 2.9
height = a_lattice
R = 9.
origin = [0, 0, height]
bob = Problem('bob', sig=0.2, k_bond=20, a_sub=a_substrate, a=a_lattice, h=height,
R=R, origin=origin, nmax=6)
bob.plot_it()
def mini_me(atom_positions_flattened, args):
g_sub, k_vert, a, h, k_bond, atoms_1, atoms_2 = args
atom_positions = atom_positions_flattened.reshape(-1, 3)
# bond energies
vectors = atom_positions[atoms_2] - atom_positions[atoms_1]
bond_lengths = np.sqrt((vectors**2).sum(axis=1))
bond_energy = k_bond * ((bond_lengths - a)**2).sum()
# surface energies
E = np.cos((atom_positions[:, :2, None] * g_sub).sum(axis=1))
result = (1.5 + E.sum(axis=1)) / 4.5
surface_energies_xy = 1. - (1.5 + E.sum(axis=-1)) / 4.5 # "sinusolidal-like" in x, y
surface_energies_z = k_vert * (atom_positions[:, 2] - h)**2 # harmonic osc. in z
surface_energy = (surface_energies_xy + surface_energies_z).sum()
return bond_energy + surface_energy
atom_positions = bob.x.copy()
parameters = ('g_sub', 'k_vert' , 'a', 'h', 'k_bond', 'atoms_1', 'atoms_2')
args = [getattr(bob, attr) for attr in parameters]
if True:
print('start minimization!')
print('mini_me(atom_positions, args): ', mini_me(atom_positions, args))
maxiter = 1000
options = {'maxiter': maxiter}
t_start = time.process_time()
wow = minimize(mini_me, atom_positions.flatten(), args=args,
tol=1e-2, options=options) # method='Nelder-Mead', bounds=bounds
process_time = time.process_time() - t_start
print('wow.fun: ', wow.fun)
print('wow.message: ', wow.message)
print('process time: ', process_time)
print('number of iterations: ', wow.nit)
print('iterartions per second: ', wow.nit / process_time)
bob.x_final = wow.x.reshape(-1, 3) # insert final results back into object
print('bob.get_energy(which="x_initial"): ', round(bob.get_energy('x_initial'), 3))
print('bob.get_energy(which="x_final"): ', round(bob.get_energy(which='x_final'), 3))
bob.plot_it(which='x_final')
