def calculate_coord_numbers(traj, atoms, cutoff):
"""Compute the coordination numbers for `mask` atoms at all times in `traj`.
Args:
- traj (ndarray n_frames x n_atoms x 3)
- mask (ndarray bool n_atoms)
- atoms (ase.Atoms)
- cutoff (float, distance units)
- skin (float, distance units, default: 0)
Returns:
ndarray of int, n_frames x n_atoms
"""
n_atoms = len(atoms)
# Prealloc buffers
out = np.full(shape=(len(traj), n_atoms), fill_value=-1, dtype=np.int)
distbuf = np.empty(shape=(n_atoms, n_atoms))
neighborbuf = np.empty(shape=(n_atoms, n_atoms), dtype=np.bool)
pbcc = PBCCalculator(atoms.cell)
for f_idex, frame in enumerate(tqdm(traj)):
pbcc.pairwise_distances(frame, out=distbuf)
np.less_equal(distbuf, cutoff, out=neighborbuf)
np.sum(neighborbuf, axis=1, out=out[f_idex])
out -= 1 # Previous sum always counted atom itself in its CN, which is wrong
assert np.min(out) >= 0
return out
def func(atoms, **kwargs):
nonlocal pbcc, dmat, connmat, newtags, layer_mask
# preallocate buffers
if pbcc is None:
pbcc = PBCCalculator(atoms.cell)
dmat = np.empty(shape=(len(atoms), len(atoms)))
connmat = np.empty(shape=(len(atoms), len(atoms)), dtype=np.bool)
newtags = np.empty(shape=len(atoms), dtype=np.int)
layer_mask = np.empty(shape=len(atoms), dtype=np.bool)
tags = groupfunc(atoms, **kwargs)
layers = np.unique(tags)
layers.sort()
newtags.fill(-1)
pbcc.pairwise_distances(atoms.positions, out=dmat)
np.less_equal(dmat, cutoff, out=connmat)
agreegrp_conns = []
nexttag = 0
for layer in layers:
np.equal(tags, layer, out=layer_mask)
layer_conrows = connmat[layer_mask]
layer_conmat = layer_conrows[:, layer_mask]
n_groups_layer, group_tags = connected_components(layer_conmat,
directed=False)
group_tags += nexttag
newtags[layer_mask] = group_tags
neighbor_groups = newtags[np.logical_or.reduce(layer_conrows,
axis=0)]
agreegrp_conns.append(neighbor_groups)
nexttag += n_groups_layer
agreegrp_connmat = np.zeros(shape=(nexttag + 1, nexttag + 1),
dtype=np.bool)
for agreegrp, neighbors in enumerate(agreegrp_conns):
agreegrp_connmat[agreegrp, neighbors] = True
agreegrp_connmat = agreegrp_connmat[:-1, :-1]
agreegrp_connmat |= agreegrp_connmat.T
return newtags, np.arange(nexttag), agreegrp_connmat
def cfunc(sn):
jl = sn.jump_lag.copy()
jl -= 1.0 # Center it around 1 since that's the minimum lag, 1 frame
jl /= jump_lag_sigma
np.square(jl, out=jl)
jl *= -0.5
np.exp(jl, out=jl) # exp correctly takes the -infs to 0
jl[sn.jump_lag > jump_lag_cutoff] = 0.
# Distance term
pbccalc = PBCCalculator(sn.structure.cell)
dists = pbccalc.pairwise_distances(sn.centers)
dmat = dists.copy()
# We want to strongly boost the similarity of *very* close sites
dmat /= distance_sigma
np.square(dmat, out=dmat)
dmat *= -0.5
np.exp(dmat, out=dmat)
return (sn.p_ij + jump_lag_coeff * jl) * (distance_coeff * dmat +
(1 - distance_coeff))
def _get_sites_to_merge(self, st, coordinating_mask = None):
sn = st.site_network
# -- Compute jump statistics
if not sn.has_attribute('n_ij'):
ja = JumpAnalysis()
ja.run(st)
pos = sn.centers
if coordinating_mask is None:
coordinating_mask = sn.static_mask
else:
assert not np.any(coordinating_mask & sn.mobile_mask)
# -- Build images
mobile_idex = np.where(sn.mobile_mask)[0][0]
one_mobile_structure = sn.structure[coordinating_mask]
one_mobile_structure.extend(sn.structure[mobile_idex])
mobile_idex = -1
one_mobile_structure.set_calculator(self.calculator)
interpolation_coeffs = np.linspace(0, 1, self.n_driven_images)
energies = np.empty(shape = self.n_driven_images)
# -- Decide on pairs to check
pbcc = PBCCalculator(sn.structure.cell)
dists = pbcc.pairwise_distances(pos)
# At the start, all within distance cutoff are mergable
mergable = dists <= self.maximum_pairwise_distance
mergable &= sn.n_ij >= self.minimum_jumps_mergable
# -- Check pairs' barriers
# Symmetric, and diagonal is trivially true. Combinations avoids those cases.
jbuf = pos[0].copy()
first_calculate = True
mergable_pairs = (p for p in itertools.combinations(range(sn.n_sites), r = 2) if mergable[p] or mergable[p[1], p[0]])
n_mergable = (np.sum(mergable) - sn.n_sites) // 2
for i, j in tqdm(mergable_pairs, total = n_mergable):
jbuf[:] = pos[j]
# Get minimage
_ = pbcc.min_image(pos[i], jbuf)
# Do coordinate driving
vector = jbuf - pos[i]
for image_i in range(self.n_driven_images):
one_mobile_structure.positions[mobile_idex] = vector
one_mobile_structure.positions[mobile_idex] *= interpolation_coeffs[image_i]
one_mobile_structure.positions[mobile_idex] += pos[i]
energies[image_i] = one_mobile_structure.get_potential_energy()
first_calculate = False
# Check barrier
barrier_idex = np.argmax(energies)
forward_barrier = energies[barrier_idex] - energies[0]
backward_barrier = energies[barrier_idex] - energies[-1]
# If it's an actual maxima barrier between them, then we want to
# check its height
if barrier_idex != 0 and barrier_idex != self.n_driven_images - 1:
mergable[i, j] = forward_barrier <= self.barrier_threshold
mergable[j, i] = backward_barrier <= self.barrier_threshold
# Otherwise, if there's no maxima between them, they are in the same
# basin.
# Get mergable groups
n_merged_sites, labels = connected_components(
mergable,
directed = True,
connection = 'strong'
)
# MergeSites will check pairwise distances; we just need to make it the
# right format.
merge_groups = []
for lbl in range(n_merged_sites):
merge_groups.append(np.where(labels == lbl)[0])
return merge_groups