diff --git a/mdevaluate/pbc.py b/mdevaluate/pbc.py index fb6b50f..84340be 100644 --- a/mdevaluate/pbc.py +++ b/mdevaluate/pbc.py @@ -198,28 +198,38 @@ def nojump(frame, usecache=True): [m[:frame.step + 1, selection].sum(axis=0) for m in reader.nojump_matrixes] ).T) * frame.box.diagonal() return frame - delta - -def pbc_points(coordinates, box, thickness=0, index=False, shear=False): + + +def pbc_points(coordinates, box, thickness=0, index=False, shear=False, + extra_image=False): """ 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. + index=True also returns the indices with indices of images being their + originals values. """ if shear: - box[2,0] = box[2,0]%box[0,0] - grid = np.array([ [i, j, k] for k in [-1,0,1] for j in [1,0,-1] for i in [-1,0,1] ]) - coordinatesPBC = np.concatenate([coordinates+v@box for v in grid], axis=0) + box[2, 0] = box[2, 0] % box[0, 0] + if shear or extra_image: + grid = np.array([[i, j, k] for k in [-2, -1, 0, 1, 2] + for j in [2, 1, 0, -1, -2] for i in [-2, -1, 0, 1, 2]]) + indices = np.tile(np.arange(len(coordinates)), 125) + else: + grid = np.array([[i, j, k] for k in [-1, 0, 1] for j in [1, 0, -1] + for i in [-1, 0, 1]]) + indices = np.tile(np.arange(len(coordinates)), 27) + coordinates_pbc = np.concatenate([coordinates+v@box for v in grid], axis=0) size = np.diag(box) - indices = np.tile(np.arange(len(coordinates)),(27)) + if thickness != 0: - mask = np.all(coordinatesPBC > -thickness, axis=1) - coordinatesPBC = coordinatesPBC[mask] + mask = np.all(coordinates_pbc > -thickness, axis=1) + coordinates_pbc = coordinates_pbc[mask] indices = indices[mask] - mask = np.all(coordinatesPBC < size+thickness, axis=1) - coordinatesPBC = coordinatesPBC[mask] + mask = np.all(coordinates_pbc < size+thickness, axis=1) + coordinates_pbc = coordinates_pbc[mask] indices = indices[mask] if index: - return coordinatesPBC, indices - return coordinatesPBC + return coordinates_pbc, indices + return coordinates_pbc diff --git a/requirements.txt b/requirements.txt index 30294a4..e51b8a7 100644 --- a/requirements.txt +++ b/requirements.txt @@ -6,5 +6,5 @@ pandas dask sqlalchemy psycopg2-binary -ipyvlume +ipyvolume jupyter_contrib_nbextensions