Merge pull request #4 from DHRUVJ2003/selection_support

Added selection support for scale.py and flattenz.py

flattenZ.py

@@ -17,8 +17,28 @@ def getOptions():
 
 def flattenZ(opts, mol):
     coords = mol['atoms']['coords']['3d']
+    # check if the user has any atoms selected
+    any_selected = False
+    if 'selected' in mol['atoms']:
+        for item in mol['atoms']['selected']:
+            if item:
+                any_selected = True
+                break  # we have *some* atoms selected
+    indices = []
+    atomic_numbers = mol['atoms']['elements']['number']
+
+    for i in range(len(atomic_numbers)):
+        if not any_selected:
+            indices.append(i)
+            # or only do selected atoms
+        elif mol['atoms']['selected'][i]:
+            indices.append(i)
+    j = 0
+
     for i in range(0, len(coords), 3):
-        coords[i+2] = 0.0
+        if j in indices:
+            coords[i+2] = 0.0
+        j = j+1
 
     return mol
 
@@ -35,6 +55,7 @@ def runCommand():
     result['cjson'] = flattenZ(opts, opts['cjson'])
     return result
 
+
 if __name__ == "__main__":
     parser = argparse.ArgumentParser('Flatten Z axis.')
     parser.add_argument('--debug', action='store_true')

scale.py

@@ -47,12 +47,30 @@ def scale(opts, mol):
     xScale = float(opts['X Scale'])
     yScale = float(opts['Y Scale'])
     zScale = float(opts['Z Scale'])
-
+    # check if the user has any atoms selected
+    any_selected = False
+    if 'selected' in mol['atoms']:
+        for item in mol['atoms']['selected']:
+            if item:
+                any_selected = True
+                break  # we have *some* atoms selected
+    indices = []
+    atomic_numbers = mol['atoms']['elements']['number']
+
+    for i in range(len(atomic_numbers)):
+        if not any_selected:
+            indices.append(i)
+            # or only do selected atoms
+        elif mol['atoms']['selected'][i]:
+            indices.append(i)
+    j = 0
     coords = mol['atoms']['coords']['3d']
     for i in range(0, len(coords), 3):
-        coords[i] = coords[i] * xScale
-        coords[i + 1] = coords[i + 1] * yScale
-        coords[i + 2] = coords[i + 2] * zScale
+        if j in indices:
+            coords[i] = coords[i] * xScale
+            coords[i + 1] = coords[i + 1] * yScale
+            coords[i + 2] = coords[i + 2] * zScale
+        j = j+1
 
     return mol
 
@@ -69,6 +87,7 @@ def runCommand():
     result['cjson'] = scale(opts, opts['cjson'])
     return result
 
+
 if __name__ == "__main__":
     parser = argparse.ArgumentParser('Scale molecular coordinates.')
     parser.add_argument('--debug', action='store_true')