Merge pull request #263 from padix-key/molvis

Add ball-and-stick model function and an corresponding example

doc/examples/scripts/structure/molecular_visualization.py

@@ -44,8 +44,6 @@
 colors[caffeine.element == "N"] = (0.0, 0.0, 0.8) # blue
 colors[caffeine.element == "O"] = (0.8, 0.0, 0.0) # red
 
-# Create a quadratic figure to ensure a correct aspect ratio
-# in the visualized molecule
 fig = plt.figure(figsize=(8.0, 8.0))
 ax = fig.add_subplot(111, projection="3d")
 graphics.plot_atoms(

doc/examples/scripts/structure/peoe_visualization.py

@@ -0,0 +1,128 @@
+"""
+Partial charge distribution
+===========================
+
+This examples shows how partial charges are distributed in a
+small molecule.
+The charges are calculated using the PEOE method [1]_.
+
+.. [1] J Gasteiger and M Marsili,
+    "Iterative partial equalization of orbital electronegativity - a
+    rapid access to atomic charges"
+    Tetrahedron, 36, 3219 - 3288 (1980).
+"""
+
+# Code source: Patrick Kunzmann
+# License: BSD 3 clause
+
+import numpy as np
+from sklearn.decomposition import PCA
+import matplotlib.pyplot as plt
+from matplotlib.colors import Normalize
+from matplotlib.cm import ScalarMappable
+import biotite.structure as struc
+import biotite.structure.info as info
+import biotite.structure.graphics as graphics
+
+
+# Acetylsalicylic acid
+MOLECULE_NAME = "AIN"
+
+# The number of iterations for the PEOE algorithm
+ITERATION_NUMBER = 6
+# The size of the element lables
+ELEMENT_FONT_SIZE = 10
+# The scaling factor of the atom 'balls'
+BALL_SCALE = 20
+# The higher this number, the more detailed are the rays
+N_RAY_STEPS = 20
+# The scaling factor of the 'ray' of charged molecules
+RAY_SCALE = 100
+# The transparency value for each 'ray ring'
+RAY_ALPHA = 0.03
+# The color map to use to depict the charge
+CMAP_NAME = "bwr_r"
+
+
+
+# Get an atom array for the selected molecule
+molecule = info.residue(MOLECULE_NAME)
+
+# Align molecule with principal component analysis:
+# The component with the least variance, i.e. the axis with the lowest
+# number of atoms lying over each other, is aligned to the z-axis,
+# which points into the plane of the figure
+pca = PCA(n_components=3)
+pca.fit(molecule.coord)
+molecule = struc.align_vectors(molecule, pca.components_[-1], [0, 0, 1])
+
+# Balls should be colored by partial charge
+charges = struc.partial_charges(molecule, ITERATION_NUMBER)
+# Later this variable stores values between 0 and 1 for use in color map
+normalized_charges = charges.copy()
+# Show no partial charge for atoms
+# that are not parametrized for the PEOE algorithm
+normalized_charges[np.isnan(normalized_charges)] = 0
+# Norm charge values to highest absolute value
+max_charge = np.max(np.abs(normalized_charges))
+normalized_charges /= max_charge
+# Transform range (-1, 1) to range (0, 1)
+normalized_charges = (normalized_charges + 1) / 2
+# Calculate colors
+color_map = plt.get_cmap(CMAP_NAME)
+colors = color_map(normalized_charges)
+
+# Ball size should be proportional to VdW radius of the respective atom
+ball_sizes = np.array(
+    [info.vdw_radius_single(e) for e in molecule.element]
+) * BALL_SCALE
+
+# Gradient of ray strength
+# The ray size is proportional to the absolute charge value
+ray_full_sizes = ball_sizes + np.abs(charges) * RAY_SCALE   
+ray_sizes = np.array([
+    np.linspace(ray_full_sizes[i], ball_sizes[i], N_RAY_STEPS, endpoint=False)
+    for i in range(molecule.array_length())
+]).T
+
+
+# The plotting begins here
+fig = plt.figure(figsize=(8.0, 6.0))
+ax = fig.add_subplot(111, projection="3d")
+
+# Plot the atoms
+# As 'axes.scatter()' uses sizes in points**2,
+# the VdW-radii as also squared
+graphics.plot_ball_and_stick_model(
+    ax, molecule, colors, ball_size=ball_sizes**2, line_width=3,
+    line_color=color_map(0.5), background_color=(.05, .05, .05), zoom=1.5
+)
+
+# Plot the element labels
+for atom in molecule:
+    ax.text(
+        *atom.coord, atom.element,
+        fontsize=ELEMENT_FONT_SIZE, color="black",
+        ha="center", va="center", zorder=100
+    )
+
+# Plots the rays
+for i in range(N_RAY_STEPS):
+    ax.scatter(
+        *molecule.coord.T, s=ray_sizes[i]**2, c=colors,
+        linewidth=0, alpha=RAY_ALPHA
+    )
+
+# Plot the colorbar
+color_bar = fig.colorbar(ScalarMappable(
+    norm=Normalize(vmin=-max_charge, vmax=max_charge),
+    cmap=color_map
+))
+color_bar.set_label("Partial charge (e)", color="white")
+color_bar.ax.yaxis.set_tick_params(color="white")
+color_bar.outline.set_edgecolor("white")
+for label in color_bar.ax.get_yticklabels():
+    label.set_color("white")
+
+fig.tight_layout()
+plt.show()

environment.yml

@@ -17,14 +17,14 @@ dependencies:
   - requests >=2.12
   - numpy >=1.13
   - scipy >=0.14
+  - scikit-learn >= 0.18.0
   - networkx >=2.0
   - pydot >=1.4
   - msgpack-python >=0.5.6
   - cython >=0.29
   - mdtraj >=1.9.3
   - matplotlib >=3.3
   - pytest >=5.2
-  - pytest-cov >=2.5
   - sphinx >=1.7
   - sphinx-gallery =0.7.0
   - numpydoc >=0.8

pyproject.toml

@@ -5,4 +5,4 @@ requires = [
     "wheel >= 0.30",
     "numpy >= 1.13",
     "cython >= 0.29"
-]
+]

src/biotite/structure/graphics/atoms.py

@@ -4,7 +4,7 @@
 
 __name__ = "biotite.structure.graphics"
 __author__ = "Patrick Kunzmann"
-__all__ = ["plot_atoms"]
+__all__ = ["plot_atoms", "plot_ball_and_stick_model"]
 
 import numpy as np
 import matplotlib.pyplot as plt
@@ -23,22 +23,18 @@ def plot_atoms(axes, atoms, colors, line_width=1.0, background_color=None,
     ----------
     axes : Axes3D
         The *Matplotlib* 3D-axes to plot the structure on.
-        For a correct aspect ratio of the visualized structure, the
-        :class:`Axes3D` must have quadratic extents on the display.
-        This can be approximately achieved by giving the :class:`Figure`
-        equal values for the width and height in its `figsize`.
-    atoms : AtomArray, length=n
+    atoms : AtomArray, shape=(n,)
         The structure to be plotted.
         The atom array must have an associated :class:`BondList`, i.e.
-        its `bonds` attribute must be defined.
+        its ``bonds`` attribute must be defined.
         One line for each bond is drawn.
     colors : ndarray, shape=(n,3) or shape=(n,4), dtype=float
         An array of RGB or RGBA colors for each atom in `atoms`.
         The values for each color channel are in the range 0 to 1.
     line_width : float, optional
         The width of the lines to be drawn.
-    background_color : string or iterable object
-        A matplotlib compatible color (color name or RGB values).
+    background_color : object
+        A *Matplotlib* compatible color (color name or RGB values).
         If set, the background is colored with the given value.
     center : tuple of (float, float, float), optional
         The coordinates of the structure that are in the center of the
@@ -108,6 +104,98 @@ def plot_atoms(axes, atoms, colors, line_width=1.0, background_color=None,
     _set_box(axes, atoms.coord, center, size, zoom)
 
 
+def plot_ball_and_stick_model(axes, atoms, colors, ball_size=200,
+                              line_color="black", line_width=1.0,
+                              background_color=None, center=None,
+                              size=None, zoom=1.0):
+    """
+    Plot an :class:`AtomArray` as *ball-and-stick* model.
+
+    The z-axis points into the screen plane.
+
+    UNSTABLE: This function is probably subject to future changes.
+    
+    Parameters
+    ----------
+    axes : Axes3D
+        The *Matplotlib* 3D-axes to plot the structure on.
+    atoms : AtomArray, length=n
+        The structure to be plotted.
+        The atom array must have an associated :class:`BondList`, i.e.
+        its ``bonds`` attribute must be defined.
+        One line for each bond is drawn.
+    colors : ndarray, shape=(n,3) or shape=(n,4), dtype=float
+        An array of RGB or RGBA colors for each atom in `atoms`, that
+        is used to color the *balls* of the model.
+        The values for each color channel are in the range 0 to 1.
+    ball_size : int or iterable of int, shape=(n,)
+        The size of the *balls* in the model in *pt*.
+        Either a single value for all atoms or an iterable object of
+        values for each atom.
+    line_color : object
+        A *Matplotlib* compatible color value for the *sticks* of the
+        model.
+    line_width : float, optional
+        The width of the *sticks* in the model in *pt*.
+    background_color : object
+        A *Matplotlib* compatible color (color name or RGB values).
+        If set, the background is colored with the given value.
+    center : tuple of (float, float, float), optional
+        The coordinates of the structure that are in the center of the
+        plot.
+        By default the complete molecule is centered.
+    size : float, optional
+        The size of each dimension in the plot.
+        The limits of the :class:`Axes3D` are set to
+        ``(center - size/(2*zoom)), (center + size/(2*zoom))``.
+    zoom : float, optional
+        Zoom in or zoom out.
+
+        - ``> 1.0``: Zoom in.
+        - ``< 1.0``: Zoom out.
+    
+    Notes
+    -----
+    This is a very simple visualization tools for quick visual analysis
+    of a structure.
+    For publication-ready molecular images the usage of a dedicated
+    molecular visualization tool is recommended.
+    """
+    if not isinstance(axes, Axes3D):
+        raise ValueError("The given axes mut be an 'Axes3D'")
+    if atoms.bonds is None:
+        raise ValueError("The atom array must have an associated bond list")
+
+    # Calculating connections between atoms
+    line_coord = [
+        (atoms.coord[index1], atoms.coord[index2])
+        for index1, index2 in atoms.bonds.as_array()[:,:2]
+    ]
+
+    # Plot sticks
+    # Use 'Line3DCollection' for higher efficiency
+    sticks = Line3DCollection(
+        line_coord, color=line_color, linewidths=line_width
+    )
+    axes.add_collection(sticks)
+
+    # Plot balls
+    axes.scatter(
+        *atoms.coord.T, s=ball_size, c=colors, linewidth=0, alpha=1
+    )
+
+    # Set viewing angle
+    axes.azim = -90
+    axes.elev = 90
+    # Remove frame
+    axes.axis("off")
+    # Set background color
+    if background_color is not None:
+        axes.set_facecolor(background_color)
+        axes.get_figure().set_facecolor(background_color)
+    _set_box(axes, atoms.coord, center, size, zoom)
+
+
 def _set_box(axes, coord, center, size, zoom):
     """
     This ensures an approximately equal aspect ratio in a 3D plot under
@@ -130,7 +218,7 @@ def _set_box(axes, coord, center, size, zoom):
     axes.set_xlim(center[0] - size/(2*zoom), center[0] + size/(2*zoom))
     axes.set_ylim(center[1] - size/(2*zoom), center[1] + size/(2*zoom))
     axes.set_zlim(center[2] - size/(2*zoom), center[2] + size/(2*zoom))
+
     # Make the axis lengths of the 'plot box' equal
     # The 'plot box' is not visible due to 'axes.axis("off")'
-
     axes.set_box_aspect([1,1,1])

src/biotite/structure/transform.py

@@ -243,7 +243,7 @@ def align_vectors(atoms, origin_direction, target_direction,
     At first the transformation (translation and rotation), that is
     necessary to align the origin vector to the target vector is
     calculated.
-    This means means, that the application of the transformation on the
+    This means, that the application of the transformation on the
     origin vector would give the target vector.
     Then the same transformation is applied to the given
     atoms/coordinates. 
@@ -264,8 +264,8 @@ def align_vectors(atoms, origin_direction, target_direction,
     Returns
     -------
     transformed : Atom or AtomArray or AtomArrayStack or ndarray, shape=(3,) or shape=(n,3) or shape=(m,n,3)
-        A copy of the input atoms or coordinates, rotated about the
-        given axis.
+        A copy of the input atoms or coordinates with the applied
+        transformation.
         
     See Also
     --------