Merge pull request #3 from konpat/fsapt

Added the functionality to display PDB files from Psi4/F-SAPT

examples/psi4/Total.pdb

@@ -0,0 +1,26 @@
+HETATM    1 O    001 A   1      -1.389   1.930  -0.443  1.00 -4.36           O  
+HETATM    2 H    001 A   1      -0.524   1.965  -0.006  1.00 -4.36           H  
+HETATM    3 C    001 A   1      -2.007   0.764  -0.108  1.00 -2.77           C  
+HETATM    4 C    001 A   1      -1.463  -0.152   0.795  1.00 -2.77           C  
+HETATM    5 C    001 A   1      -2.148  -1.330   1.088  1.00 -2.77           C  
+HETATM    6 C    001 A   1      -3.374  -1.603   0.490  1.00 -2.77           C  
+HETATM    7 C    001 A   1      -3.914  -0.684  -0.409  1.00 -2.77           C  
+HETATM    8 C    001 A   1      -3.237   0.493  -0.710  1.00 -2.77           C  
+HETATM    9 H    001 A   1      -0.511   0.057   1.264  1.00 -2.77           H  
+HETATM   10 H    001 A   1      -1.715  -2.032   1.788  1.00 -2.77           H  
+HETATM   11 H    001 A   1      -3.902  -2.517   0.720  1.00 -2.77           H  
+HETATM   12 H    001 A   1      -4.867  -0.882  -0.881  1.00 -2.77           H  
+HETATM   13 H    001 A   1      -3.643   1.213  -1.406  1.00 -2.77           H  
+HETATM   14 O    001 A   1       1.353   1.938   0.472  1.00 -8.59           O  
+HETATM   15 H    001 A   1       1.784   2.349   1.230  1.00 -8.59           H  
+HETATM   16 C    001 A   1       2.037   0.787   0.150  1.00  1.46           C  
+HETATM   17 C    001 A   1       1.590   0.070  -0.957  1.00  1.46           C  
+HETATM   18 C    001 A   1       2.242  -1.107  -1.313  1.00  1.46           C  
+HETATM   19 C    001 A   1       3.332  -1.567  -0.575  1.00  1.46           C  
+HETATM   20 C    001 A   1       3.770  -0.840   0.529  1.00  1.46           C  
+HETATM   21 C    001 A   1       3.122   0.338   0.896  1.00  1.46           C  
+HETATM   22 H    001 A   1       0.745   0.437  -1.522  1.00  1.46           H  
+HETATM   23 H    001 A   1       1.892  -1.665  -2.170  1.00  1.46           H  
+HETATM   24 H    001 A   1       3.833  -2.481  -0.857  1.00  1.46           H  
+HETATM   25 H    001 A   1       4.614  -1.185   1.109  1.00  1.46           H  
+HETATM   26 H    001 A   1       3.460   0.903   1.757  1.00  1.46           H  

examples/psi4/example.ipynb

@@ -149,6 +149,15 @@
    "source": [
     "fortecubeview.vib('output.nh3.molden_normal_modes')"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "fortecubeview.showpdb(pdb = 'Total.pdb', erange=10.0)"
+   ]
   }
  ],
  "metadata": {
@@ -167,7 +176,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.9"
+   "version": "3.9.15"
   }
  },
  "nbformat": 4,

fortecubeview/__init__.py

@@ -1,2 +1,2 @@
-from .fortecubeview import plot, vib, colorschemes, geom
+from .fortecubeview import plot, vib, colorschemes, geom, showpdb
 

fortecubeview/cube_viewer.py

@@ -260,12 +260,14 @@ def parse_cube_files(self):
             elif m2:
                 if m2.groups()[0] == 'a':
                     label = 'Density (alpha)'
-                if m2.groups()[0] == 'b':
+                elif m2.groups()[0] == 'b':
                     label = 'Density (beta)'
-                if m2.groups()[0] == 's':
+                elif m2.groups()[0] == 's':
                     label = 'Density (spin)'
-                if m2.groups()[0] == 't':
+                elif m2.groups()[0] == 't':
                     label = 'Density (total)'
+                else:
+                    label = key
             else:
                 label = key
             labels_to_filename[label] = key

fortecubeview/fortecubeview.py

@@ -28,6 +28,32 @@ def geom(xyz = None,
         height=height,
     )
 
+def showpdb(pdb = None,
+         erange = 10.0,
+         width=400,
+         height=400):
+    """
+    A widget for rendering a psi4/fsapt PDB object.
+
+    Parameters
+    ----------
+    pdb : filename
+        A psi4 PDB file to render
+    erange : float
+        Range of displayed energy values. Values >= |erange| are given
+        maximum color saturation.
+    width : int
+        the width of the plot in pixels (default = 400)
+    height : int
+        the height of the plot in pixels (default = 400)
+    """
+    return PDBViewer(
+        pdb=pdb,
+        erange=erange,
+        width=width,
+        height=height,
+    )
+
 def plot(path='.',
          cubes=None,
          width=400,

fortecubeview/mol_viewer.py

@@ -81,3 +81,64 @@ def __init__(self,
 
         # disable scroll and display the selection widget
 #        display(widgets.HTML(widget_style))
+
+
+class PDBViewer():
+    """
+    This version reads xyz coordinates and F-SAPT energy contributions from the PDB file
+    generated by fsapt.py.
+    """
+    def __init__(self,
+                 pdb = None,
+                 erange = 10.0,
+                 path = '.',
+                 width=400,
+                 height=400,
+                 font_size=16,
+                 font_family='Helvetica'):
+
+
+        if pdb == None:
+            print(
+                f'PDBViewer: no molecule provided. The widget will not be displayed'
+            )
+            return
+
+        geom = []
+        levels = []
+
+        pdblines = [x.rstrip() for x in open(pdb,"r")]
+        for line in pdblines:
+        # fsapt.py output sometimes has no space before a minus sign
+            line = line.replace("-"," -")
+            sline = line.split()
+            if len(sline) > 0:
+                geom.append((sline[-1],float(sline[-6]),float(sline[-5]),float(sline[-4])))
+                levels.append(float(sline[-2]))
+
+        # start a Py3JSRenderer
+        renderer = Py3JSRenderer(width=width, height=height)
+
+        # add molecule to the renderer
+        renderer.set_levels(levels)
+        renderer.set_erange(erange)
+        renderer.add_molecule(geom,bohr=False, shift_to_com=False)
+
+        # output.layout.height = f'{height + 100}px'
+        widget_style = """
+        <style>
+           .jupyter-widgets-output-area .output_scroll {
+                height: unset !important;
+                border-radius: unset !important;
+                -webkit-box-shadow: unset !important;
+              box-shadow: unset !important;
+            }
+            .jupyter-widgets-output-area  {
+            height: auto !important;
+         }
+        </style>
+        """
+
+        # display the rendered and the labels
+        display(renderer.renderer)
+

fortecubeview/py3js_renderer.py

@@ -85,6 +85,8 @@ class Py3JSRenderer():
         Add a plane
     def add_box(position,width,height,depth,color,opacity=1.0,normal=(0, 0, 1))
         Add a box
+    def set_levels(atoms_list,levels_list)
+        Add energy levels to print (using a color scale) for each atom
     """
     def __init__(self, width=400, height=400):
         """
@@ -107,6 +109,7 @@ def __init__(self, width=400, height=400):
         self.atom_size = 0.5  # scaling factor for atom geometry
         self.atom_geometries = {}
         self.atom_materials = {}
+        self.atom_levels = []
         self.bond_materials = {}
         self.bond_geometry = None
         # cubefile meshes
@@ -200,6 +203,28 @@ def show_molecule(self, wfn, shift_to_com=True):
                       for i in range(natom)]
         self.add_molecule(atoms_list, bohr=True)
 
+    def set_levels(self, levels_list):
+        """
+        Set color levels for each atom
+
+        Parameters
+        ----------
+        levels_list : list(float)
+            A list of energy contribution values for each atom
+        """
+        self.atom_levels = levels_list[:]
+
+    def set_erange(self, erange):
+        """
+        Set energy range for color saturation PDB plot
+
+        Parameters
+        ----------
+        erange : float
+            Maximum |energy| to plot (corresponds to full color saturation)
+        """
+        self.erange = erange
+
     def add_molecule(self, atoms_list, bohr=False, shift_to_com=True):
         """
         Add a molecular geometry to the scene. The geometry is given as a list of atoms
@@ -237,20 +262,34 @@ def add_molecule(self, atoms_list, bohr=False, shift_to_com=True):
             Xcm, Ycm, Zcm = (0.0, 0.0, 0.0)
 
         atom_positions = collections.defaultdict(list)
-        for symbol, x, y, z in atoms_list:
-            atom_positions[symbol].append([x - Xcm, y - Ycm, z - Zcm])
+        if self.atom_levels == []:
+            for symbol, x, y, z in atoms_list:
+                atom_positions[symbol].append([x - Xcm, y - Ycm, z - Zcm])
+
         # Then add the unique atoms at all the positions
-        for atom_type in atom_positions:
-            atom_mesh = self.__get_atom_mesh((atom_type, 0.0, 0.0, 0.0))
-            clone_geom = CloneArray(original=atom_mesh,
-                                    positions=atom_positions[atom_type])
-            self.scene.add(clone_geom)
+            for atom_type in atom_positions:
+                atom_mesh = self.__get_atom_mesh((atom_type, 0.0, 0.0, 0.0))
+                clone_geom = CloneArray(original=atom_mesh,
+                                        positions=atom_positions[atom_type])
+                self.scene.add(clone_geom)
+        else:
+            for ((symbol, x, y, z), energy) in zip(atoms_list,self.atom_levels):
+                atom_positions[symbol].append([x - Xcm, y - Ycm, z - Zcm, energy])
+
+        # Then add the unique atoms at all the positions
+            for atom_type in atom_positions:
+                for atom in atom_positions[atom_type]:
+                    energy = atom[3]
+                    atom_mesh = self.__get_atom_mesh_color_by_energy((atom_type, 0.0, 0.0, 0.0), energy)
+                    clone_geom = CloneArray(original=atom_mesh,
+                                            positions=[atom[:3]])
+                    self.scene.add(clone_geom)
 
         # Add the bonds
         for i in range(len(atoms_list)):
-            atom1 = atoms_list[i]
+            atom1 = (atoms_list[i][0], atoms_list[i][1] - Xcm, atoms_list[i][2] - Ycm, atoms_list[i][3] - Zcm)
             for j in range(i + 1, len(atoms_list)):
-                atom2 = atoms_list[j]
+                atom2 = (atoms_list[j][0], atoms_list[j][1] - Xcm, atoms_list[j][2] - Ycm, atoms_list[j][3] - Zcm)
                 bond = self.__get_bond_mesh(atom1, atom2)
                 if bond:
                     self.scene.add(bond)
@@ -743,6 +782,24 @@ def __get_atom_mesh(self, atom_info):
         mesh = Mesh(geometry=geometry, material=material, position=[x, y, z])
         return mesh
 
+    def __get_atom_mesh_color_by_energy(self, atom_info, energy):
+        """
+        This function returns a Mesh object (Geometry + Material) that represents an atom
+
+        Parameters
+        ----------
+        atom_info : tuple(str, float, float, float)
+            A tuple containing the atomic symbol and coordinates of the atom using the format
+            (atomic symbol , x, y, z)
+        energy : float
+            The energy contribution for this atom (determines atom color)
+        """
+        symbol, x, y, z = atom_info
+        geometry = self.__get_atom_geometry(symbol)
+        material = self.__set_atom_material(energy)
+        mesh = Mesh(geometry=geometry, material=material, position=[x, y, z])
+        return mesh
+
     def __get_atom_geometry(self, symbol, shininess=75):
         """
         This function returns a sphere geometry object with radius proportional to the covalent atomic radius
@@ -786,6 +843,34 @@ def __get_atom_material(self, symbol, shininess=75):
         self.atom_materials[symbol] = material
         return material
 
+    def __set_atom_material(self, energy):
+        """
+        This function returns a Material object based on energy level
+
+        Parameters
+        ----------
+        energy : float
+            The energy to plot for this atom (by color)
+        """
+        erange = self.erange
+        colormin = [255, 0, 0]
+        colormax = [0, 0, 255]
+        white = [255, 255, 255]
+        thiscolor = 3*[0]
+        saturation = abs(energy / erange)
+        if saturation > 1.0: saturation = 1.0
+        if energy < 0:
+            for i in range(3):
+                thiscolor[i] = int(saturation * colormin[i] + (1.0 - saturation) * white[i])
+        else:
+            for i in range(3):
+                thiscolor[i] = int(saturation * colormax[i] + (1.0 - saturation) * white[i])
+        color = 'rgb({0[0]},{0[1]},{0[2]})'.format(thiscolor)
+        material = MeshStandardMaterial(color=color,
+                                        roughness=0.25,
+                                        metalness=0.1)
+        return material
+
     def __get_bond_mesh(self, atom1_info, atom2_info, radius=None):
         """
         This function adds a bond between two atoms