"""
Pythonic wrappings for mmsurf surfaces
"""
import enum
import decorator
import numpy
from schrodinger.analysis.visanalysis import volumedata
from schrodinger.analysis.visanalysis import volumedataio
from schrodinger.infra import mm
from schrodinger.infra import mmbitset
from schrodinger.infra import mmsurf
from schrodinger.structutils import analyze
from schrodinger.structutils import color
from schrodinger.utils import fileutils
# This module is first imported while Maestro is loading, so
# schrodinger.get_maestro() will return a _DummyMaestroModule if we run it now.
# Instead, we import Maestro when the first Surface object is instantiated.
maestro = DELAYED_MAESTRO_LOAD = object()
[docs]class Style(enum.IntEnum):
"""
Surface representation styles.
"""
solid = mmsurf.MMSURF_STYLE_SOLID
mesh = mmsurf.MMSURF_STYLE_MESH
dot = mmsurf.MMSURF_STYLE_DOT
[docs]class ColorFrom(enum.IntEnum):
"""
Values for surface color sources.
"""
surface = mmsurf.MMSURF_COLOR_FROM_SURFACE
vertex = mmsurf.MMSURF_COLOR_FROM_VERTEX
nearest_asl_atom = mmsurf.MMSURF_COLOR_FROM_NEAREST_ASL_ATOM
volume = mmsurf.MMSURF_COLOR_FROM_VOLUME
entry = mmsurf.MMSURF_COLOR_FROM_ENTRY
[docs]class StrEnum(str, enum.Enum):
"""
An enum class where all values are strings. All enum objects stringify to
their value.
"""
def __str__(self):
return self.value
[docs]class ColorBy(StrEnum):
"""
Values for surface color schemes.
"""
source_color = mmsurf.MMSURF_COLOR_BY_SOURCE_COLOR
partial_charge = mmsurf.MMSURF_COLOR_BY_PARTIAL_CHARGE
atom_type = mmsurf.MMSURF_COLOR_BY_ATOM_TYPE
atom_typeMM = mmsurf.MMSURF_COLOR_BY_ATOM_TYPE_MM
chain_name = mmsurf.MMSURF_COLOR_BY_CHAIN_NAME
element = mmsurf.MMSURF_COLOR_BY_ELEMENT
mol_number = mmsurf.MMSURF_COLOR_BY_MOL_NUMBER
mol_number_carbon = mmsurf.MMSURF_COLOR_BY_MOL_NUMBER_CARBON
residue_charge = mmsurf.MMSURF_COLOR_BY_RESIDUE_CHARGE
residue_hydrophobicity = mmsurf.MMSURF_COLOR_BY_RESIDUE_HYDROPHOBICITY
residue_position = mmsurf.MMSURF_COLOR_BY_RESIDUE_POSITION
residue_type = mmsurf.MMSURF_COLOR_BY_RESIDUE_TYPE
grid_property = mmsurf.MMSURF_COLOR_BY_GRID_PROPERTY
atom_color = mmsurf.MMSURF_COLOR_BY_ATOM_COLOR
cavity_depth = mmsurf.MMSURF_COLOR_BY_CAVITY_DEPTH
[docs]class MolSurfType(enum.IntEnum):
"""
Types of molecular surfaces.
"""
vdw = mmsurf.MOLSURF_VDW
extended = mmsurf.MOLSURF_EXTENDED
molecular = mmsurf.MOLSURF_MOLECULAR
@decorator.decorator
def _requires_update(func, self, *args, **kwargs):
"""
A decorator for `Surface` methods that update the visual representation of
the surface. This decorator, tells Maestro to update the workspace surface
representation and ensures that we only force one update even if a decorated
method calls another decorated method.
"""
do_update = not self._force_update
self._force_update = True
retval = func(self, *args, **kwargs)
if do_update:
self._updateMaestro()
self._force_update = False
return retval
[docs]class Surface(object):
"""
A Pythonic wrapping for mmsurf surfaces that are not associated with a
project entry. (For surfaces that are associated with a project entry, see
`ProjectSurface` below.) Surface objects can be created from an existing
mmsurf handle via `__init__`, can be read from disk via `read`, or new
surfaces can be created via `newMolecularSurface`.
"""
Style = Style
ColorBy = ColorBy
ColorFrom = ColorFrom
Color = color.Color
SURFACE_TYPE_NAME = {
MolSurfType.vdw: "van der Waals",
MolSurfType.extended: "extended radius",
MolSurfType.molecular: "molecular surface"
}
[docs] def __init__(self, handle, manage=True):
"""
:param handle: An mmsurf handle to an existing surface
:type handle: int
:param manage: If True, the mmsurf handle will be deleted when this
object is garbage collected.
:type manage: bool
"""
self._initializeMmlibs()
self._handle = handle
self._manage = manage
# Required for compatibility with _requires_update, which is required by
# ProjectSurface
self._force_update = True
@staticmethod
def _initializeMmlibs():
"""
Initialize all mmlib libraries used by this class.
"""
mmsurf.mmsurf_initialize(mm.MMERR_DEFAULT_HANDLER)
mmsurf.mmvol_initialize(mm.MMERR_DEFAULT_HANDLER)
mmsurf.mmvisio_initialize(mm.MMERR_DEFAULT_HANDLER)
@staticmethod
def _terminateMmlibs():
"""
Terminate all mmlib libraries used by this class.
"""
mmsurf.mmsurf_terminate()
mmsurf.mmvol_terminate()
mmsurf.mmvisio_terminate()
def __del__(self):
"""
When this object is garbage collected, terminate the mmlib libraries and
delete the mmsurf handle if it's managed by this object.
"""
if self._manage:
self.delete()
self._terminateMmlibs()
[docs] def delete(self):
"""
Immediately delete the mmsurf handle. After this method has been
called, any further attempts to interact with this object will result in
an MmException.
"""
mmsurf.mmsurf_delete(self._handle)
self._handle = -1
[docs] @classmethod
def newMolecularSurface(cls,
struc,
name,
asl=None,
atoms=None,
resolution=0.5,
probe_radius=None,
vdw_scaling=1.0,
mol_surf_type=MolSurfType.molecular):
"""
Create a new molecular surface for the specified surface
:param struc: The structure to create the surface for
:type proj: `schrodinger.structure.Structure`
:param name: The name of the surface.
:type name: str
:param asl: If given, the surface will only be created for atoms in the
structure that match the provided ASL. Note that only one of `asl` and
`atoms` may be given. If neither are given, then the surface will be
created for all atoms in the structure.
:type asl: str or NoneType
:param atoms: An optional list of atom numbers. If given, the surface
will only be created for the specified atoms. Note that only one of
`asl` and `atoms` may be given. If neither are given, then the
surface will be created for all atoms in the structure.
:type atoms: list or NoneType
:param resolution: The resolution of the surface, generally between
0 and 1. Smaller numbers lead to a more highly detailed surface.
:type resolution: float
:param probe_radius: The radius of the rolling sphere used to calculate
the surface. Defaults to 1.4 if `mol_surf_type` is
`MolSurfType.Molecular` or `MolSurfType.Extended`. May not be given
if `mol_surf_type` is `MolSurfType.vdw`.
:type probe_radius: float
:param vdw_scaling: If given, all atomic radii will be scaled by the
provided value before the surface is calculated.
:type vdw_scaling: float
:param mol_surf_type: The type of surface to create.
:type mol_surf_type: `MolSurfType`
:return: The new surface
:rtype: `Surface`
"""
mmsurf.mmsurf_initialize(mm.MMERR_DEFAULT_HANDLER)
handle = cls._createMolecularSurface(struc, name, asl, atoms,
resolution, probe_radius,
vdw_scaling, mol_surf_type)
surf = cls(handle)
mmsurf.mmsurf_terminate()
return surf
@classmethod
def _createMolecularSurface(cls, struc, name, asl, atoms, resolution,
probe_radius, vdw_scaling, mol_surf_type):
"""
Create a new molecular surface for the specified structure. Arguments
are the same as `newMolecularSurface` above.
:return: An mmsurf handle for the new surface
:rtype: int
"""
if probe_radius is None:
if mol_surf_type in (MolSurfType.extended, MolSurfType.molecular):
probe_radius = 1.4
elif mol_surf_type is MolSurfType.vdw:
probe_radius = 0
elif mol_surf_type is MolSurfType.vdw:
raise ValueError("May not give probe_radius for vdw surfaces.")
bs = cls._generateBitset(struc, asl, atoms)
# See mmshare/mmlibs/mmsurf/mmsurf.h for documentation on the additional
# mmsurf_molsurf() arguments
handle = mmsurf.mmsurf_molsurf(struc, bs, resolution, probe_radius,
vdw_scaling, mol_surf_type, 0)
mmsurf.mmsurf_set_name(handle, name)
surf_type = cls.SURFACE_TYPE_NAME[mol_surf_type]
mmsurf.mmsurf_set_surface_type(handle, surf_type)
return handle
@staticmethod
def _generateBitset(struc, asl, atoms):
"""
Generate a bitself that indicates which atoms to create the surface for.
Note that either `asl` or `atoms` (or neither) may be provided, but
not both. If neither are provided, the bitset will cover all atoms in
`structure`.
:param struc: The structure to generate the bitset for
:type struc: `schrodinger.structure.Structure`
:param asl: If not None, the bitset will contain only atoms that match
this ASL.
:type asl: str or NoneType
:param atoms: If not None, a list of atom numbers. The bitset will
contain only the listed atoms.
:type atoms: list or NoneType
:return: The newly generated bitset
:rtype: `mmbitset.Bitset`
"""
if asl is not None and atoms is not None:
raise ValueError("May not specify both asl and atoms.")
bs = mmbitset.Bitset(size=struc.atom_total)
if asl is not None:
atoms = analyze.evaluate_asl(struc, asl)
if atoms is None:
bs.fill()
elif not atoms:
raise ValueError("No atoms specified.")
else:
list(map(bs.set, atoms))
return bs
[docs] @classmethod
def read(cls, filename):
"""
Read surface data from a file.
:param filename: The file to read from.
:type filename: str
:return: The read surface.
:rtype: `Surface`
"""
mmsurf.mmvisio_initialize(mm.MMERR_DEFAULT_HANDLER)
handle = mmsurf.mmvisio_read_surface_from_file(filename)
surf = cls(handle)
mmsurf.mmvisio_terminate()
return surf
[docs] def write(self, filename):
"""
Write this surface to a file. Note that existing files will be
overwritten.
:param filename: The file to write to.
:type filename: str
"""
# We get strange HDF5-DIAG errors if we attempt to overwrite an existing
# file, so make sure we delete any existing file first
fileutils.force_remove(filename)
mmsurf.mmvisio_write_surface_to_file(filename, self._handle)
@property
def name(self):
"""
The surface name.
:type: str
"""
return mmsurf.mmsurf_get_name(self._handle)
@name.setter
def name(self, value):
mmsurf.mmsurf_set_name(self._handle, value)
[docs] def rename(self, value):
"""
Set the surface name. This method is provided for compatibility with
`ProjectSurface`.
"""
# For compatibility with ProjectSurface
self.name = value
@property
def volume_name(self):
"""
The volume name associated with the given surface
:type: str
"""
return mmsurf.mmsurf_get_volume_name(self._handle)
@volume_name.setter
def volume_name(self, value):
mmsurf.mmsurf_set_volume_name(self._handle, value)
@property
def isovalue(self):
"""
The isovalue for the given surface
:type: float
"""
return mmsurf.mmsurf_get_isovalue(self._handle)
@isovalue.setter
def isovalue(self, value):
mmsurf.mmsurf_set_isovalue(self._handle, value)
@property
def surface_type(self):
"""
A textual description of the type of surface.
:type: str
"""
return mmsurf.mmsurf_get_surface_type(self._handle)
@surface_type.setter
def surface_type(self, val):
mmsurf.mmsurf_set_surface_type(self._handle, val)
@property
def visible(self):
"""
Whether the surface is currently visible. This setting will be
remembered, but it will not have any effect until the surface is added
to a project and loaded into Maestro.
:type: bool
"""
vis = mmsurf.mmsurf_get_visibility(self._handle)
return bool(vis)
@visible.setter
@_requires_update
def visible(self, val):
val = int(bool(val))
mmsurf.mmsurf_set_visibility(self._handle, val)
[docs] def show(self):
"""
Sets the surface to be visible.
"""
self.visible = True
[docs] def hide(self):
"""
Hides the surface.
"""
self.visible = False
@property
def front_transparency(self):
"""
The transparency of the front of the surface (relative to the workspace
camera position). Measured on a scale from 0 (fully opaque) to 100
(fully transparent).
:type: int
"""
return mmsurf.mmsurf_get_transparency(self._handle)
@front_transparency.setter
@_requires_update
def front_transparency(self, val):
mmsurf.mmsurf_set_transparency(self._handle, val)
@property
def back_transparency(self):
"""
The transparency of the back of the surface (relative to the workspace
camera position). Measured on a scale from 0 (fully opaque) to 100
(fully transparent).
:type: int
"""
return mmsurf.mmsurf_get_transparency_back(self._handle)
@back_transparency.setter
@_requires_update
def back_transparency(self, val):
mmsurf.mmsurf_set_transparency_back(self._handle, val)
[docs] @_requires_update
def setTransparency(self, val):
"""
Set both the front and the back transparency.
:param val: The value to set the transparency to
:type val: int
"""
self.front_transparency = val
self.back_transparency = val
@property
def style(self):
"""
The visual style of the surface representation (solid, mesh, or dot).
:type: `Style`
"""
style = mmsurf.mmsurf_get_style(self._handle)
return Style(style)
@style.setter
@_requires_update
def style(self, val):
val = int(val)
mmsurf.mmsurf_set_style(self._handle, val)
@property
def darken_colors_by_cavity_depth(self):
"""
Whether the colors on the surface should be darkened based on the cavity
depth.
:type: bool
"""
val = mmsurf.mmsurf_get_darken_colors_by_cavity_depth(self._handle)
return bool(val)
@darken_colors_by_cavity_depth.setter
@_requires_update
def darken_colors_by_cavity_depth(self, val):
val = bool(val)
mmsurf.mmsurf_set_darken_colors_by_cavity_depth(self._handle, val)
@property
def color_source(self):
"""
The source of the surface colors. Note that coloring()/setColoring()
are recommended over directly manipulating `color_source`, as this will
ensure that `color_source` is set correctly.
:type: `ColorFrom`
"""
val = mmsurf.mmsurf_get_color_source(self._handle)
return ColorFrom(val)
@color_source.setter
@_requires_update
def color_source(self, val):
val = int(val)
mmsurf.mmsurf_set_color_source(self._handle, val)
@property
def color_scheme(self):
"""
The color scheme used to determine surface colors. This value may be
ignored unless `color_source` is set to `ColorFrom.NearestAslAtom`.
Note that coloring()/setColoring() are recommended over directly
manipulating `color_scheme`, as this will ensure that `color_source`
is set correctly.
:type: `ColorBy`
"""
val = mmsurf.mmsurf_get_color_scheme(self._handle)
return ColorBy(val)
@color_scheme.setter
@_requires_update
def color_scheme(self, val):
# color by
val = str(val)
mmsurf.mmsurf_set_color_scheme(self._handle, val)
@property
def color(self):
"""
The constant surface color. This value may be ignored unless
`color_source` is set to `ColorFrom.Surface` and `color_scheme` is
set to `ColorBy.SourceColor`. Note that coloring()/setColoring() are
recommended over directly manipulating `color`, as this will ensure
that `color_source` and `color_scheme` are set correctly.
:type: `Color`
"""
val = mmsurf.mmsurf_get_rgb_color(self._handle)
return color.Color(val)
@color.setter
@_requires_update
def color(self, val):
if isinstance(val, color.Color):
mmsurf.mmsurf_set_rgb_color(self._handle, val.rgb)
else:
mmsurf.mmsurf_set_color(self._handle, val)
[docs] @_requires_update
def setColoring(self, coloring):
"""
Set the surface coloring. Must be one of:
- A `ColorBy` value other than `ColorBy.SourceColor` to color based
on the nearest atom
- A `Color` value for constant coloring
- A list or numpy array containing a color for each vertex
"""
if (isinstance(coloring, ColorBy) and
coloring is not ColorBy.source_color):
# Set the color source so that the color scheme is obeyed
self.color_source = ColorFrom.nearest_asl_atom
self.color_scheme = coloring
elif isinstance(coloring, color.Color):
self.color_source = ColorFrom.surface
self.color_scheme = ColorBy.source_color
self.color = coloring
elif isinstance(coloring, (list, numpy.ndarray)):
self.color_source = ColorFrom.vertex
self.color_scheme = ColorBy.source_color
self.vertex_colors = coloring
else:
raise ValueError("Unrecognized coloring.")
[docs] def coloring(self):
"""
Return the current surface coloring. Is only guaranteed to return a
non-None value if the surface coloring was set via `setColoring`. If
the surface coloring cannot be determined, will return None.
:return: The current surface coloring
:rtype: `ColorBy`, `Color`, `numpy.ndarray`, or NoneType
"""
if (self.color_source == ColorFrom.surface and
self.color_scheme == ColorBy.source_color):
return self.color
elif (self.color_source == ColorFrom.nearest_asl_atom and
self.color_scheme != ColorBy.source_color):
return self.color_scheme
elif (self.color_source == ColorFrom.vertex and
self.color_scheme != ColorBy.source_color and
self.has_vertex_colors):
return self.vertex_colors
@property
def surface_area(self):
"""
The reported surface area of the surface
:type: float
"""
return mmsurf.mmsurf_get_surface_area(self._handle)
@property
def surface_volume(self):
"""
Compute and return surface volume of the surface. Note that unlike
surface_area, this is not cached, use appropriately.
:type: float
"""
return mmsurf.mmsurf_calculate_surface_volume(self._handle)
@property
def vertex_coords(self):
"""
A list of all vertex coordinates
:type: list
"""
return mmsurf.mmsurf_get_all_vertex_coords(self._handle)
@property
def vertex_count(self):
return mmsurf.mmsurf_get_num_vertices(self._handle)
@property
def vertex_normals(self):
"""
The normal for each vertex
:type: numpy.ndarray
"""
return mmsurf.mmsurf_get_all_vertex_normals(self._handle)
@property
def patch_count(self):
"""
The number of surface patches (i.e. triangles connecting three adjacent
vertices).
:type: int
"""
return mmsurf.mmsurf_get_num_patches(self._handle)
@property
def patch_vertices(self):
"""
A `patch_count` x 3 array containing vertex indices for each surface
patch.
:type: numpy.array
"""
return mmsurf.mmsurf_get_all_patches(self._handle)
@property
def nearest_atom_indices(self):
"""
A list of the atom indices closest to each vertex coordinate. Atom
indices are listed in a corresponding order to vertex_coords.
:type: list
"""
return mmsurf.mmsurf_get_nearest_atoms(self._handle)
def _updateMaestro(self):
"""
This method is used in `ProjectSurface` and is present here for
compatibility with the `_requires_update` decorator.
"""
# This method intentionally left blank
[docs] def smoothColors(self, colors, iterations):
"""
Given a list of vertex colors, return a list of smoothed colors. Does
not modify the surface in any way.
:param colors: A list or numpy array of the colors to smooth, where
colors are represented as either RGB or RGBA values. Note that if this
value is a numpy array, the input array will be modified in place.
:type colors: list or numpy.array
:param iterations: The number of smoothing iterations to carry out.
:type iterations: int
:return: The smoothed colors as a numpy array. If `colors` was a numpy
array, the return value will be a reference to the (modified) input
array.
:rtype: numpy.array
"""
color_len = len(colors[0])
if not isinstance(colors, numpy.ndarray):
if color_len == 3:
colors = numpy.array(colors, dtype=numpy.float32)
else:
colors = numpy.array(colors, dtype=numpy.double)
if color_len == 3:
func = mmsurf.mmsurf_smooth_colors3
elif color_len == 4:
func = mmsurf.mmsurf_smooth_colors
else:
err = ("Colors must be defined using three (r, g, b) or four "
"(r, g, b, a) terms.")
raise ValueError(err)
func(self._handle, colors, iterations)
return colors
@property
def vertex_colors(self):
"""
An array of manually specified per-vertex colors.
:type: `numpy.ndarray`
"""
return mmsurf.mmsurf_get_all_vertex_colors(self._handle)
@vertex_colors.setter
@_requires_update
def vertex_colors(self, val):
if not isinstance(val, numpy.ndarray):
val = numpy.array(val, dtype=numpy.float32)
mmsurf.mmsurf_set_all_vertex_colors(self._handle, val)
@property
def has_vertex_colors(self):
"""
Does this surface contain manually specified per-vertex colors?
:type: bool
"""
return mmsurf.mmsurf_has_vertex_colors(self._handle)
[docs] def curvatures(self, curvature_type):
"""
Return curvature values for all vertices.
:param curvature_type: mmsurf.CURVATURE_GAUSS, mmsurf.CURVATURE_MIN, mmsurf.CURVATURE_MAX, mmsurf.CURVATURE_MEAN
:type: curvature_type: mmsurf.CurvatureType enum
:rtype: numpy.array
"""
return mmsurf.mmsurf_get_curvatures(self._handle, curvature_type)
[docs] def copy(self):
"""
Create a copy of this surface. Note that this method will always return
a `Surface` object, even when a `ProjectSurface` object is copied.
:return: The copied surface
:rtype: `Surface`
"""
copy_handle = mmsurf.mmsurf_copy(self._handle)
return Surface(copy_handle)
[docs]def create_isosurface_from_grid_data(dimensions,
resolution,
origin,
isovalue,
nonzero=None,
grid=None,
surface_color=None,
name=None,
comment=None,
surface_type=None):
"""
Create a new isosurface from 3D grid data
To create a surface attached to a project entry, use::
from schrodinger import surface
from schrodinger.project.surface import ProjectSurface
pysurf = surface.create_isosurface_from_grid_data(*args, **kwargs)
projsurf = ProjectSurface.addSurfaceToProject(pysurf,
ptable,
row)
Where ptable is a Project instance and row is the desired ProjectRow
instance
:type dimensions: list
:param dimensions: The number of grid points in the X, Y and Z directions
:type resolution: list
:param resolution: The gridpoint spacing in the X, Y and Z directions
:type origin: list
:param origin: The X, Y, Z coordinates of the grid origin
:type isovalue: float
:param isovalue: The value of the isosurface
:type nonzero: iterable
:param nonzero: Each item of nonzero is an (x, y, z, value) tuple indicating
the value of the grid at the given x, y, z coordinate. All other gridpoints
are set to zero. Either nonzero or grid must be supplied but not both.
:type grid: numpy.array
:param grid: A 3-dimensional numpy.array the same size as dimensions, the
value at each point is the grid at that point. Either nonzero or grid must
be supplied but not both.
:type surface_color: str or `schrodinger.structutils.color.Color`
:param surface_color: The color of the surface. If a string, must be a color
name recognized by the Color class.
:type name: str
:param name: The name of the surface - shows in Maestro's Manage Surfaces
dialog under Surface Name
:type comment: str
:param comment: The comment for the surface - shows in Maestro's Manage
surfaces dialog under Comments
:type surface_type: str
:param surface_type: The type of the surface - shows in Maestro's Manage
Surfaces dialog under Surface Type. Note - this has nothing to do with
the molecular surface type (VDW, EXTENDED, MOLECULAR) property and is a
free text field.
:rtype: `Surface`
:return: The created isosurface
"""
Surface._initializeMmlibs()
if nonzero is None and grid is None:
raise RuntimeError('Either nonzero or grid must be given')
elif nonzero is not None and grid is not None:
raise RuntimeError('Only one of nonzero and grid can be given')
# Set the grid data for the volume the surface will be computed from
voldata = volumedata.VolumeData(N=dimensions,
resolution=resolution,
origin=origin)
if nonzero:
datapoints = voldata.getData()
for xind, yind, zind, value in nonzero:
datapoints[xind][yind][zind] = value
else:
voldata.setData(grid)
# Create infrastructure volume object
with fileutils.tempfilename(suffix='vis') as visname:
volumedataio.SaveVisFile(voldata, visname)
infravis = mmsurf.mmvisio_read_volume_from_file(visname)
# Create infrastructure surface and set properties
infrasurf = mmsurf.mmsurf_new_isosurface(infravis, isovalue, False, 0.0,
0.0, 0.0, 0.0, 0)
if comment:
mmsurf.mmsurf_set_comment(infrasurf, comment)
# Create pythonic surface and set properties
pysurf = Surface(infrasurf)
if name:
pysurf.name = name
if surface_color:
if isinstance(surface_color, str):
surface_color = color.Color(surface_color)
pysurf.setColoring(surface_color)
if surface_type:
pysurf.surface_type = surface_type
Surface._terminateMmlibs()
return pysurf