Source code for mdevaluate.pbc


from collections import OrderedDict

import numpy as np
import numba

from scipy.spatial import cKDTree
from itertools import product

from .logging import logger

[docs]def pbc_diff_old(v1, v2, box): """ Calculate the difference of two vestors, considering optional boundary conditions. """ if box is None: v = v1 - v2 else: v = v1 % box - v2 % box v -= (v > box / 2) * box v += (v < -box / 2) * box return v
[docs]def pbc_diff(v1, v2=None, box=None): if box is None: out = v1 - v2 elif len(getattr(box, 'shape', [])) == 1: out = pbc_diff_rect(v1, v2, box) elif len(getattr(box, 'shape', [])) == 2: out = pbc_diff_tric(v1, v2, box) else: raise NotImplementedError("cannot handle box") return out
[docs]def pbc_diff_rect(v1, v2, box): """ Calculate the difference of two vectors, considering periodic boundary conditions. """ if v2 is None: v = v1 else: v = v1 -v2 s = v / box v = box * (s - s.round()) return v
[docs]def pbc_diff_tric(v1, v2=None, box=None): """ difference vector for arbitrary pbc Args: box_matrix: CoordinateFrame.box """ if len(box.shape) == 1: box = np.diag(box) if v1.shape == (3,): v1 = v1.reshape((1,3)) #quick 'n dirty if v2.shape == (3,): v2 = v2.reshape((1,3)) if box is not None: r3 = np.subtract(v1, v2) r2 = np.subtract(r3, (np.rint(np.divide(r3[:,2],box[2][2])))[:,np.newaxis] * box[2][np.newaxis,:]) r1 = np.subtract(r2, (np.rint(np.divide(r2[:,1],box[1][1])))[:,np.newaxis] * box[1][np.newaxis,:]) v = np.subtract(r1, (np.rint(np.divide(r1[:,0],box[0][0])))[:,np.newaxis] * box[0][np.newaxis,:]) else: v = v1 - v2 return v
[docs]def pbc_dist(a1,a2,box = None): return ((pbc_diff(a1,a2,box)**2).sum(axis=1))**0.5
[docs]def pbc_extend(c, box): """ in: c is frame, box is frame.box out: all atoms in frame and their perio. image (shape => array(len(c)*27,3)) """ c=np.asarray(c) if c.shape == (3,): c = c.reshape((1,3)) #quick 'n dirty comb = np.array([np.asarray(i) for i in product([0,-1,1],[0,-1,1],[0,-1,1])]) b_matrices = comb[:,:,np.newaxis]*box[np.newaxis,:,:] b_vectors = b_matrices.sum(axis=1)[np.newaxis,:,:] return (c[:,np.newaxis,:]+b_vectors)
[docs]def pbc_kdtree(v1,box, leafsize = 32, compact_nodes = False, balanced_tree = False): """ kd_tree with periodic images box - whole matrix rest: optional optimization """ r0 = cKDTree(pbc_extend(v1,box).reshape((-1,3)),leafsize ,compact_nodes ,balanced_tree) return r0
[docs]def pbc_kdtree_query(v1,v2,box,n = 1): """ kd_tree query with periodic images """ r0, r1 = pbc_kdtree(v1,box).query(v2,n) r1 = r1 // 27 return r0, r1
[docs]def pbc_backfold_rect(act_frame,box_matrix): """ mimics "trjconv ... -pbc atom -ur rect" folds coords of act_frame in cuboid """ af=np.asarray(act_frame) if af.shape == (3,): act_frame = act_frame.reshape((1,3)) #quick 'n dirty b = box_matrix c = np.diag(b)/2 af = pbc_diff(np.zeros((1,3)),af-c,b) return af + c
[docs]def pbc_backfold_compact(act_frame,box_matrix): """ mimics "trjconv ... -pbc atom -ur compact" folds coords of act_frame in wigner-seitz-cell (e.g. dodecahedron) """ c = act_frame box = box_matrix ctr = box.sum(0)/2 c=np.asarray(c) shape = c.shape if shape == (3,): c = c.reshape((1,3)) shape = (1,3) #quick 'n dirty comb = np.array([np.asarray(i) for i in product([0,-1,1],[0,-1,1],[0,-1,1])]) b_matrices = comb[:,:,np.newaxis]*box[np.newaxis,:,:] b_vectors = b_matrices.sum(axis=1)[np.newaxis,:,:] sc = c[:,np.newaxis,:]+b_vectors w = np.argsort((((sc)-ctr)**2).sum(2),1)[:,0] return sc[range(shape[0]),w]
[docs]@numba.jit(nopython=True) def pbc_diff_numba(ri, rj, box): v = ri % box - rj % box v -= (v > box / 2) * box v += (v < -box / 2) * box return v
[docs]def whole(frame): """ Apply ``-pbc whole`` to a CoordinateFrame. """ residue_ids = frame.coordinates.atom_subset.residue_ids box = frame.box.diagonal() coms = np.array([ np.bincount(residue_ids, weights=c * frame.masses)[1:] / np.bincount(residue_ids, weights=frame.masses)[1:] for c in frame.T ]).T[residue_ids - 1] cor = np.zeros_like(frame) cd = frame - coms n, d = np.where(cd > box / 2 * 0.9) cor[n, d] = -box[d] n, d = np.where(cd < -box / 2 * 0.9) cor[n, d] = box[d] duomask = np.bincount(residue_ids)[1:][residue_ids - 1] == 2 if np.any(duomask): duomask[::2] = False cor[duomask] = 0 return frame + cor
NOJUMP_CACHESIZE = 128
[docs]def nojump(frame, usecache=True): """ Return the nojump coordinates of a frame, based on a jump matrix. """ selection = frame.selection reader = frame.coordinates.frames if usecache: if not hasattr(reader, '_nojump_cache'): reader._nojump_cache = OrderedDict() # make sure to use absolute (non negative) index abstep = frame.step % len(frame.coordinates) i0s = [x for x in reader._nojump_cache if x <= abstep] if len(i0s) > 0: i0 = max(i0s) delta = reader._nojump_cache[i0] i0 += 1 else: i0 = 0 delta = 0 delta = delta + np.array(np.vstack( [m[i0:abstep + 1].sum(axis=0) for m in reader.nojump_matrixes] ).T) * frame.box.diagonal() reader._nojump_cache[abstep] = delta while len(reader._nojump_cache) > NOJUMP_CACHESIZE: reader._nojump_cache.popitem(last=False) delta = delta[selection, :] else: delta = np.array(np.vstack( [m[:frame.step + 1, selection].sum(axis=0) for m in reader.nojump_matrixes] ).T) * frame.box.diagonal() return frame - delta
[docs]def pbc_points(coordinates, box, thickness=0, index=False, inclusive=True, center=None): """ Returns the points their first periodic images. Does not fold them back into the box. Thickness 0 means all 27 boxes. Positive means the box+thickness. Negative values mean that less than the box is returned. index=True also returns the indices with indices of images being their originals values. inclusive=False returns only images, does not work with thickness <= 0 """ if center is None: center = box/2 allcoordinates = np.copy(coordinates) indices = np.tile(np.arange(len(coordinates)),(27)) for x in range(-1, 2, 1): for y in range(-1, 2, 1): for z in range(-1, 2, 1): vv = np.array([x, y, z], dtype=float) if not (vv == 0).all() : allcoordinates = np.concatenate((allcoordinates, coordinates + vv*box), axis=0) if thickness != 0: mask = np.all(allcoordinates < center+box/2+thickness, axis=1) allcoordinates = allcoordinates[mask] indices = indices[mask] mask = np.all(allcoordinates > center-box/2-thickness, axis=1) allcoordinates = allcoordinates[mask] indices = indices[mask] if not inclusive and thickness > 0: allcoordinates = allcoordinates[len(coordinates):] indices = indices[len(coordinates):] if index: return (allcoordinates, indices) return allcoordinates