From a2164507d521716d25d2d523934878fec08dbae5 Mon Sep 17 00:00:00 2001
From: sebastiankloth <kloth@mailbox.org>
Date: Wed, 21 Jun 2023 13:03:05 +0200
Subject: [PATCH] Added possibility to use more images with pbc_points.

---
 mdevaluate/pbc.py | 36 +++++++++++++++++++++++-------------
 requirements.txt  |  2 +-
 2 files changed, 24 insertions(+), 14 deletions(-)

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