added first bits to include a scheme for selecting structures based on their representation

new/codered.py

@@ -275,7 +275,7 @@ class CodeRed:
         return
 
 
-    def red(self, red_type=None, surface='reduced', representation="positions", representation_params={}, selection="symmetry", selection_params={}, **reduction_params):
+    def red(self, red_type=None, surface='reduced', representation=None, representation_params={}, selector=None, selector_params={}, **reduction_params):
         #TODO: process reduction_params to set up and execute the right reduction technique
         #      e.g. symmetry reduction, mesh projection reduction, soap distance reduction, ...
 
@@ -288,7 +288,9 @@ class CodeRed:
             structures = self.structures.copy()
 
         # transform structures according to representation setting
-        if representation in ["positions", "pos", "coordinates"]:
+        if representation==None:
+            structures_repr = structures.copy()
+        elif representation in ["positions", "pos", "coordinates"]:
             structures_repr = representations.positions(structures, **representation_params)
         elif representation in ["box", "boxes", "boxgrid", "grid"]:
             structures_repr = representations.boxgrid(structures, **representation_params)
@@ -299,6 +301,16 @@ class CodeRed:
                 print("This is not a valid representation!")
                 raise
 
+        # get a structure selection from the selector
+        if selector==None:
+            selection = np.ones(len(self.structures), dtype='bool')
+        else:
+            try:
+                selection = selector(structures_repr, **selection_params)
+            except TypeError:
+                print("This is not a valid selector!")
+                raise
+
         ## perform reduction according to the set reduction type
         #if red_type=='symmetry':
         #    # generate informative output
@@ -322,12 +334,6 @@ class CodeRed:
         #        if red_array[i]:
         #            structures_new.append(self.structures[i])
         #
-        #elif red_type=='sidechains':
-        #    # generate informative output
-        #    print("Generating side chains and corresponding minimum distance reduction...")
-        #
-        #    # generate structures including side chains
-        #    structures_new = self.make_sidechains()
 
         # save reduced structure list
         self.structures = structures_new

new/representations.py

@@ -17,7 +17,7 @@ def positions(structures, cartesian=False):
     return representation
 
 
-def boxgrid(structures, n_boxes_a=1, n_boxes_b=1, n_boxes_c=1, cartesian=False):
+def boxgrid(structures, n_boxes_a=1, n_boxes_b=1, n_boxes_c=1, cartesian=False, return_structures=False):
     """
     This representation splits the bounding box of the structures into a grid of
     boxes and represents each position by the center of the sub-box it is in.
@@ -44,8 +44,27 @@ def boxgrid(structures, n_boxes_a=1, n_boxes_b=1, n_boxes_c=1, cartesian=False):
     grid = np.array([grid_aa, grid_bb, grid_cc]).reshape(3,len(ga)*len(gb)*len(gc)).T
 
     # project positions on the grid
-    representation = np.copy(positions).reshape(-1,3)
-    for i in range(len(representation)):
-        representation[i] = grid[np.argmin(np.linalg.norm(grid-representation[i],axis=-1))]
-    
-    return representation.reshape(positions.shape)
+    projections = np.copy(positions).reshape(-1,3)
+    for i in range(len(projections)):
+        projections[i] = grid[np.argmin(np.linalg.norm(grid-projections[i],axis=-1))]
+
+    # transform to desired output
+    if return_structures:
+        # reshape to fit (N_structures, N_atoms, 3) shape
+        new_positions = projections.reshape(positions.shape)
+        representation = []
+        for i in range(len(structures)):
+            copy = structures[i].copy()
+            # set cartesian or relative coordinates
+            if cartesian:
+                copy.set_positions(new_positions[i])
+            else:
+                copy.set_scaled_positions(new_positions[i])
+            # add to representation list
+            representation.append(copy)
+
+    else:
+        # reshape to fit (N_structures, N_atoms, 3) shape
+        representation = projections.reshape(positions.shape)
+
+    return representation