Improved the behaviour of the atom selection plugin and its connection

with the molecular viewer plugin
Redesigned the partial charges plugin

MDANSE/Framework/Plugins/AtomSelectionPlugin.py

@@ -168,7 +168,7 @@ class AtomSelectionPlugin(UserDefinitionPlugin):
         pub.subscribe(self.msg_select_atoms_from_viewer, 'msg_select_atoms_from_viewer')
 
         UserDefinitionPlugin.__init__(self,parent,size=(800,500))
-        
+                
     def build_panel(self):
                                                 
         self._mainPanel = wx.ScrolledWindow(self, wx.ID_ANY, size=self.GetSize())
@@ -274,6 +274,9 @@ class AtomSelectionPlugin(UserDefinitionPlugin):
         self.parent.mgr.Update()
         
         self.set_trajectory(self.dataproxy.data)
+
+        if self.parent.selectedAtoms:
+            self.insert_keyword_values('atom_picked',sorted(self.parent.selectedAtoms))     
                                 
     def set_trajectory(self,trajectory):
 
@@ -312,11 +315,7 @@ class AtomSelectionPlugin(UserDefinitionPlugin):
         self.selectionTextCtrl.SetValue(self._query.get_expression())
         
         self.display_selection_summary()
-        
-    def close(self):
                 
-        pub.sendMessage('msg_clear_selection',plugin=self)                                
-        
     def on_add_operator(self, event=None):
 
         operator = event.GetEventObject().GetLabel()

MDANSE/Framework/Plugins/MolecularViewerPlugin.py

@@ -30,6 +30,8 @@ Created on Apr 14, 2015
 :author: Gael Goret, Bachir Aoun, Eric C. Pellegrini
 '''
 
+import os
+
 import numpy
 
 import vtk
@@ -47,6 +49,7 @@ from MDANSE.MolecularDynamics.Trajectory import sorted_atoms
 
 from MDANSE.Framework.Plugins.DataPlugin import get_data_plugin 
 from MDANSE.Framework.Plugins.ComponentPlugin import ComponentPlugin
+from MDANSE.Framework.UserDefinitionStore import UD_STORE
 
 # The colour for a selected atom (R,G,B,Alpha).
 RGB_COLOURS = {}
@@ -403,6 +406,7 @@ class MolecularViewerPanel(ComponentPlugin):
 
         popupMenu.AppendSeparator()
         selectionBox = popupMenu.Append(wx.ID_ANY, 'Show/hide selection box',kind=wx.ITEM_CHECK)
+        saveSelection = popupMenu.Append(wx.ID_ANY, 'Save selection')
         clearSelection = popupMenu.Append(wx.ID_ANY, 'Clear selection')
 
         popupMenu.AppendSeparator()
@@ -411,6 +415,7 @@ class MolecularViewerPanel(ComponentPlugin):
         
         popupMenu.Bind(wx.EVT_MENU, self.on_toggle_selection_box, selectionBox)
         popupMenu.Bind(wx.EVT_MENU, self.on_clear_selection, clearSelection)
+        popupMenu.Bind(wx.EVT_MENU, self.on_save_selection, saveSelection)
         popupMenu.Bind(wx.EVT_MENU, self.parallel_proj_onoff, parallelProjection)
         popupMenu.Check(parallelProjection.GetId(), self.camera.GetParallelProjection())
         
@@ -480,6 +485,11 @@ class MolecularViewerPanel(ComponentPlugin):
     def picked_atoms(self):
         
         return self.__pickedAtoms
+    
+    @property
+    def selectedAtoms(self):
+        
+        return self.__pickedAtoms
 
     def change_frame_rate(self, laps):
 
@@ -640,9 +650,11 @@ class MolecularViewerPanel(ComponentPlugin):
             self.pick_atoms([idx],new=(not obj.GetShiftKey()))
 
     def on_info_pick(self, obj, evt = None):
+        
+        
         if not self._trajectoryLoaded:
             return
-        
+                
         pos = obj.GetEventPosition()
         
         self.picker.AddPickList(self.picking_domain)
@@ -689,7 +701,41 @@ class MolecularViewerPanel(ComponentPlugin):
         self._polydata.Modified()
         
         self._iren.Render()
+
+    def on_save_selection(self,event=None):
+        
+        if not self._trajectoryLoaded:
+            return
+        
+        if not self.__pickedAtoms:
+            return
+        
+        d = wx.TextEntryDialog(self,"Enter selection name","New selection")
+ 
+        # If the new element dialog is closed by clicking on OK. 
+        if d.ShowModal() == wx.ID_CANCEL:
+            return
+ 
+        # Get rid of wxpython unicode string formatting
+        name = str(d.GetValue())
+                         
+        if not name:
+            return
+        
+        target = os.path.basename(self._trajectory.filename)
+        
+        if UD_STORE.has_definition(target,"atom_selection",name):
+            LOGGER('There is already a user-definition with that name.','error',['console'])
+            return
+                  
+        UD_STORE.set_definition(target,"atom_selection",name,{'indexes':sorted(self.__pickedAtoms)})
+                 
+        UD_STORE.save()
+                 
+        pub.sendMessage("msg_set_ud")
         
+        LOGGER('User definition %r successfully set.' % name,'info',['console'])
+                                                         
     def show_selection(self, selection):
         
         if not self._trajectoryLoaded:

MDANSE/Framework/Plugins/PartialChargesPlugin.py

@@ -37,6 +37,10 @@ import wx
 import wx.aui as wxaui
 import wx.grid as wxgrid
 
+from MMTK import Atom, AtomCluster, Molecule
+from MMTK.NucleicAcids import NucleotideChain
+from MMTK.Proteins import PeptideChain, Protein
+
 from MDANSE import LOGGER
 from MDANSE.Framework.Plugins.UserDefinitionPlugin import UserDefinitionPlugin
 
@@ -78,6 +82,8 @@ class PartialChargesPlugin(UserDefinitionPlugin):
         self.parent.mgr.Update()
         
         self.set_trajectory(self.dataproxy.data)
+                
+        self._mgr.Update()
 
     def set_trajectory(self,trajectory):
 
@@ -85,41 +91,74 @@ class PartialChargesPlugin(UserDefinitionPlugin):
 
         self._target = os.path.basename(self._trajectory.filename)
 
-        self._grid.CreateGrid(self._trajectory.universe.numberOfAtoms(),3)
-
+        self._contents = {}
+        for obj in self._trajectory.universe.objectList():
+        
+            objname = obj.name
+            if isinstance(obj, (PeptideChain,Protein,NucleotideChain)):
+                for res in obj.residues():
+                    resname = res.type.symbol.strip()
+                    for at in res.atomList():
+                        descr = (objname,resname,at.name)
+                        d = self._contents.setdefault(descr,[0.0,[]])
+                        d[1].append(at.index)
+            elif isinstance(obj, (Molecule,AtomCluster)):
+                for at in obj.atomList():
+                    descr = (objname,at.name)
+                    d = self._contents.setdefault(descr,[0.0,[]])
+                    d[1].append(at.index)
+            else:
+                descr = (at.name,)
+                d = self._contents.setdefault(descr,[0.0,[]])
+                d[1].append(at.index)
+            
+        self._grid.CreateGrid(0,2)
+        
         self._grid.SetRowLabelSize(1)
         
-        self._grid.SetColFormatNumber(0)
         self._grid.SetColFormatNumber(1)
-        self._grid.SetColFormatNumber(2)
                 
         roAttr = wxgrid.GridCellAttr()
         roAttr.SetReadOnly(True)
         roAttr.SetBackgroundColour(wx.Colour(220,220,220))
         self._grid.SetColAttr(0,roAttr)
-        self._grid.SetColAttr(1,roAttr)
         
         floatAttr = wxgrid.GridCellAttr()
         floatAttr.SetRenderer(wxgrid.GridCellFloatRenderer())
         floatAttr.SetEditor(wxgrid.GridCellFloatEditor())
-        self._grid.SetColAttr(2,floatAttr)
+        self._grid.SetColAttr(1,floatAttr)
         
-        self._grid.SetColLabelValue(0,'index')
-        self._grid.SetColLabelValue(1,'name')
-        self._grid.SetColLabelValue(2,'charge')
-
-        atoms = sorted(self._trajectory.universe.atomList(), key=operator.attrgetter('index'))
-        for idx, at in enumerate(atoms):
-            self._grid.SetCellValue(idx,0,str(idx))
-            self._grid.SetCellValue(idx,1,at.name)
+        self._grid.SetColLabelValue(0,'name')
+        self._grid.SetColLabelValue(1,'charge')
+
+        for i,(k,v) in enumerate(sorted(self._contents.items())):
+            self._grid.AppendRows(1)
+            self._grid.SetCellValue(i,0,".".join(k))
+            self._grid.SetCellValue(i,1,str(v[0]))
+            
+        self._grid.Bind(wx.EVT_SIZE, self.OnSize)
+ 
+    def OnSize(self, event):
+         
+        width,_ = self._grid.GetClientSizeTuple()
+        w = 4*width/5
+        self._grid.SetColSize(0, w)                
+        self._grid.SetColSize(1, width-w-5)
         
+        self._grid.ForceRefresh()
+                                            
     def validate(self):
         
         charges = {}
-        for i in range(self._trajectory.universe.numberOfAtoms()):
-            val = self._grid.GetCellValue(i,2)
-            if val:
-                charges[i] = float(val)
+        for i in range(self._grid.GetNumberRows()):
+            name = tuple(self._grid.GetCellValue(i,0).split('.'))
+            charge = self._grid.GetCellValue(i,1)
+            
+            val = self._contents.get(name,None)
+            if val is None:
+                continue
+            for idx in val[1]:
+                charges[idx] = float(charge)
                 
         if not charges:
             LOGGER("No partial charges defined.", "error", ["dialog"])