Source code for mbuild.pattern

"""mBuild pattern module."""
from itertools import product

import numpy as np

from mbuild import clone
from mbuild.coordinate_transform import force_overlap
from mbuild.utils.validation import assert_port_exists

__all__ = [
    "Pattern",
    "DiskPattern",
    "SpherePattern",
    "Random2DPattern",
    "Random3DPattern",
    "Grid2DPattern",
    "Grid3DPattern",
]


[docs]class Pattern(object): """A superclass for molecules spatial Patterns. Patterns refer to how molecules are arranged either in a box (volume) or 2-D surface. This class could serve as a superclass for different molecules patterns Attributes ---------- points : array or np.array Positions of molecules in surface or space orientations : dict, optional, default=None Orientations of ports scale : float, optional, default=None Scaling factor for the original pattern """ def __init__(self, points, orientations=None, scale=None, **kwargs): self.points = points if orientations is None: orientations = dict() self.orientations = orientations if scale is not None: self.scale(scale) def __len__(self): """Get the length of the pattern as the number of points.""" return len(self.points) def __getitem__(self, item): """Get an item from the pattern from its points.""" return self.points[item]
[docs] def scale(self, by): """Scale the points in the Pattern. Parameters ---------- by : float or np.ndarray, shape=(3,) The factor to scale by. If a scalar, scale all directions isotropically. If np.ndarray, scale each direction independently """ self.points *= np.asarray([by]) self._adjust_ports()
def _adjust_ports(self): """Adjust ports according to the provided orientations.""" for orientation, ports in self.orientations.items(): for port, point in zip(ports, self.points): port.translate(point)
[docs] def apply(self, compound, orientation="", compound_port=""): """Arrange copies of a Compound as specified by the Pattern. Parameters ---------- compound : mb.Compound mb.Compound to be applied new pattern orientation : dict, optional, default='' New orientations for ports in compound compound_port : list, optional, default=None Ports to be applied new orientations Returns ------- compound : mb.Compound mb.Compound with applied pattern """ compounds = list() if self.orientations.get(orientation): for port in self.orientations[orientation]: new_compound = clone(compound) new_port = new_compound.labels[compound_port] (new_compound, new_port["up"], port["up"]) compounds.append(new_compound) else: for point in self.points: new_compound = clone(compound) new_compound.translate(point) compounds.append(new_compound) return compounds
[docs] def apply_to_compound( self, guest, guest_port_name="down", host=None, backfill=None, backfill_port_name="up", scale=True, ): """Attach copies of a guest Compound to Ports on a host Compound. Parameters ---------- guest : mb.Compound The Compound prototype to be applied to the host Compound guest_port_name : str, optional, default='down' The name of the port located on `guest` to attach to the host host : mb.Compound, optional, default=None A Compound with available ports to add copies of `guest` to backfill : mb.Compound, optional, default=None A Compound to add to the remaining available ports on `host` after clones of `guest` have been added for each point in the pattern backfill_port_name : str, optional, default='up' The name of the port located on `backfill` to attach to the host scale : bool, optional, default=True Scale the points in the pattern to the lengths of the `host`'s `boundingbox` and shift them by the hosts mins Returns ------- guests : list of mb.Compound List of inserted guest compounds on host compound backfills : list of mb.Compound List of inserted backfill compounds on host compound """ n_ports = len(host.available_ports()) assert n_ports >= self.points.shape[0], "Not enough ports for pattern." assert_port_exists(guest_port_name, guest) box = host.get_boundingbox() if scale: self.scale(box.lengths) self.points += host.mins pattern = self.points port_positions = np.empty(shape=(n_ports, 3)) port_list = list() for port_idx, port in enumerate(host.available_ports()): port_positions[port_idx, :] = port["up"]["middle"].pos port_list.append(port) used_ports = set() # Keep track of used ports for backfilling. guests = [] for point in pattern: closest_point_idx = np.argmin( host.min_periodic_distance(point, port_positions) ) closest_port = port_list[closest_point_idx] used_ports.add(closest_port) # Attach the guest to the closest port. new_guest = clone(guest) force_overlap( new_guest, new_guest.labels[guest_port_name], closest_port ) guests.append(new_guest) # Move the port as far away as possible (simpler than removing it). # There may well be a more elegant/efficient way of doing this. port_positions[closest_point_idx, :] = np.array( [np.inf, np.inf, np.inf] ) backfills = [] if backfill: assert_port_exists(backfill_port_name, backfill) # Attach the backfilling Compound to unused ports. for port in port_list: if port not in used_ports: new_backfill = clone(backfill) # Might make sense to have a backfill_port_name option... force_overlap( new_backfill, new_backfill.labels[backfill_port_name], port, ) backfills.append(new_backfill) return guests, backfills
[docs]class Random2DPattern(Pattern): """Generate n random points on a 2D grid along z = 0. Attributes ---------- n : int Number of points to generate seed : int Seed for random number generation """ def __init__(self, n, seed=None, **kwargs): if seed: np.random.seed(seed) points = np.random.random((n, 3)) points[:, 2] = 0 super(Random2DPattern, self).__init__(points=points, **kwargs)
[docs]class Random3DPattern(Pattern): """Generate n random points on a 3D grid. Attributes ---------- n : int Number of points to generate seed : int Seed for random number generation """ def __init__(self, n, seed=None, **kwargs): if seed: np.random.seed(seed) points = np.random.random((n, 3)) super(Random3DPattern, self).__init__(points=points, **kwargs)
[docs]class Grid2DPattern(Pattern): """Generate a 2D grid (n x m) of points along z = 0. Notes ----- Points span [0,1) along x and y axes Attributes ---------- n : int Number of grid rows m : int Number of grid columns """ def __init__(self, n, m, **kwargs): points = np.zeros(shape=(n * m, 3), dtype=float) for i, j in product(range(n), range(m)): points[i * m + j, 0] = i / n points[i * m + j, 1] = j / m super(Grid2DPattern, self).__init__(points=points, **kwargs)
[docs]class Grid3DPattern(Pattern): """Generate a 3D grid (n x m x l) of points. Notes ----- Points span [0,1) along x, y, and z axes Attributes ---------- n : int Number of grid rows m : int Number of grid columns l : int Number of grid aisles """ def __init__(self, n, m, l, **kwargs): points = np.zeros(shape=(n * m * l, 3), dtype=float) for i, j, k in product(range(n), range(m), range(l)): points[i * m * l + j * l + k, 0] = i / n points[i * m * l + j * l + k, 1] = j / m points[i * m * l + j * l + k, 2] = k / l super(Grid3DPattern, self).__init__(points=points, **kwargs)
[docs]class SpherePattern(Pattern): """Generate N evenly distributed points on the unit sphere. Sphere is centered at the origin. Alrgorithm based on the 'Golden Spiral'. Code by Chris Colbert from the numpy-discussion list: http://mail.scipy.org/pipermail/numpy-discussion/2009-July/043811.html """ def __init__(self, n, **kwargs): phi = (1 + np.sqrt(5)) / 2 # the golden ratio long_incr = 2 * np.pi / phi # how much to increment the longitude dz = 2.0 / float(n) # a unit sphere has diameter 2 bands = np.arange(int(n)) # each band will have one point placed on it z = bands * dz - 1.0 + (dz / 2.0) # the height z of each band/point r = np.sqrt(1.0 - z * z) # project onto xy-plane az = bands * long_incr # azimuthal angle of point modulo 2 pi x = r * np.cos(az) y = r * np.sin(az) points = np.column_stack((x, y, z)) from mbuild.port import Port if kwargs.get("orientations") is None: ports = list() for point in points: port = Port() ports.append(port) # Make the top of the port point toward the positive x axis. port.spin(-np.pi / 2, [0, 0, 1]) # Raise up (or down) the top of the port in the z direction. port.spin(-np.arcsin(point[2]), [0, 1, 0]) # Rotate the Port along the z axis. port.spin(np.arctan2(point[1], point[0]), [0, 0, 1]) # Move the Port a bit away from the surface of the Sphere. # translate(port, point + 0.07) kwargs["orientations"] = {"normal": ports} else: raise NotImplementedError( "Custom orientation support is not yet implemented." ) super(SpherePattern, self).__init__(points=points, **kwargs)
[docs]class DiskPattern(Pattern): """Generate N evenly distributed points on the unit circle along z = 0. Disk is centered at the origin. Algorithm based on Vogel's method. Code by Alexandre Devert: http://blog.marmakoide.org/?p=1 """ def __init__(self, n, **kwargs): radius = np.sqrt(np.arange(n) / float(n)) golden_angle = np.pi * (3 - np.sqrt(5)) theta = golden_angle * np.arange(n) points = np.zeros((n, 3)) points[:, 0] = np.cos(theta) points[:, 1] = np.sin(theta) points *= radius.reshape((n, 1)) super(DiskPattern, self).__init__(points=points, **kwargs)