Exemplo n.º 1
0
def extend(model, nodes, atoms):
    """Returns mapping indices and an :class:`.AtomMap`."""

    try:
        n_atoms = model.numAtoms()
        is3d = model.is3d()
    except AttributeError:
        raise ValueError('model must be an NMA instance')

    try:
        n_nodes = nodes.numAtoms()
        i_nodes = nodes.iterAtoms()
    except AttributeError:
        raise ValueError('nodes must be an Atomic instance')

    if n_atoms != n_nodes:
        raise ValueError('atom numbers must be the same')

    if not nodes in atoms:
        raise ValueError('nodes must be a subset of atoms')

    atom_indices = []
    indices = []
    get = HierView(atoms).getResidue

    for i, node in enumerate(i_nodes):
        res = get(node.getChid() or None, node.getResnum(),
                  node.getIcode() or None,
                  node.getSegname() or None)
        if res is None:
            raise ValueError('atoms must contain a residue for all atoms')
        atom_indices.append(res._getIndices())
        if is3d:
            indices.append(list(range(i * 3, (i + 1) * 3)) * len(res))
        else:
            indices.append([i] * len(res))
    atom_indices = np.concatenate(atom_indices)
    indices = np.concatenate(indices)

    try:
        ag = atoms.getAtomGroup()
    except AttributeError:
        ag = atoms
    atommap = AtomMap(ag,
                      atom_indices,
                      atoms.getACSIndex(),
                      title=str(atoms),
                      intarrays=True)
    return indices, atommap
Exemplo n.º 2
0
def showProtein(*atoms, **kwargs):
    """Show protein representation using :meth:`~mpl_toolkits.mplot3d.Axes3D`.
    This function is designed for generating a quick view of the contents of a
    :class:`~.AtomGroup` or :class:`~.Selection`.

    Protein atoms matching ``"calpha"`` selection are displayed using solid
    lines by picking a random and unique color per chain.  Line with can
    be adjusted using *lw* argument, e.g. ``lw=12``. Default width is 4.
    Chain colors can be overwritten using chain identifier as in ``A='green'``.

    Water molecule oxygen atoms are represented by red colored circles.  Color
    can be changed using *water* keyword argument, e.g. ``water='aqua'``.
    Water marker and size can be changed using *wmarker* and *wsize* keywords,
    defaults values are ``wmarker='.', wsize=6``.

    Hetero atoms matching ``"hetero and noh"`` selection are represented by
    circles and unique colors are picked at random on a per residue basis.
    Colors can be customized using residue name as in ``NAH='purple'``.  Note
    that this will color all distinct residues with the same name in the same
    color.  Hetero atom marker and size can be changed using *hmarker* and
    *hsize* keywords, default values are ``hmarker='o', hsize=6``.

    ProDy will set the size of axis so the representation is not distorted when
    the shape of figure window is close to a square.  Colors are picked at
    random, except for water oxygens which will always be colored red."""

    alist = atoms
    for atoms in alist:
        if not isinstance(atoms, Atomic):
            raise TypeError('atoms must be an Atomic instance')
    import matplotlib.pyplot as plt
    from mpl_toolkits.mplot3d import Axes3D
    cf = plt.gcf()
    show = None
    for child in cf.get_children():
        if isinstance(child, Axes3D):
            show = child
            break
    if show is None:
        show = Axes3D(cf)
    from matplotlib import colors
    cnames = dict(colors.cnames)
    wcolor = kwargs.get('water', 'red').lower()
    avoid = np.array(colors.hex2color(cnames.pop(wcolor, cnames.pop('red'))))
    for cn, val in cnames.items():  # PY3K: OK
        clr = np.array(colors.hex2color(val))
        if clr.sum() > 2.4:
            cnames.pop(cn)
        elif np.abs(avoid - clr).sum() <= 0.6:
            cnames.pop(cn)
    cnames = list(cnames)
    import random
    random.shuffle(cnames)
    min_ = list()
    max_ = list()
    for atoms in alist:
        if isinstance(atoms, AtomGroup):
            title = atoms.getTitle()
        else:
            title = atoms.getAtomGroup().getTitle()
        calpha = atoms.select('calpha')
        if calpha:
            for ch in HierView(calpha, chain=True):
                xyz = ch._getCoords()
                chid = ch.getChid()
                show.plot(xyz[:, 0],
                          xyz[:, 1],
                          xyz[:, 2],
                          label=title + '_' + chid,
                          color=kwargs.get(chid, cnames.pop()).lower(),
                          lw=kwargs.get('lw', 4))
        water = atoms.select('water and noh')
        if water:
            xyz = atoms.select('water')._getCoords()
            show.plot(xyz[:, 0],
                      xyz[:, 1],
                      xyz[:, 2],
                      label=title + '_water',
                      color=wcolor,
                      ls='None',
                      marker=kwargs.get('wmarker', '.'),
                      ms=kwargs.get('wsize', 6))
        hetero = atoms.select('not protein and not nucleic and not water')
        if hetero:
            for res in HierView(hetero).iterResidues():
                xyz = res._getCoords()
                resname = res.getResname()
                resnum = str(res.getResnum())
                chid = res.getChid()
                show.plot(xyz[:, 0],
                          xyz[:, 1],
                          xyz[:, 2],
                          ls='None',
                          color=kwargs.get(resname, cnames.pop()).lower(),
                          label=title + '_' + chid + '_' + resname + resnum,
                          marker=kwargs.get('hmarker', 'o'),
                          ms=kwargs.get('hsize', 6))
        xyz = atoms._getCoords()
        min_.append(xyz.min(0))
        max_.append(xyz.max(0))

    show.set_xlabel('x')
    show.set_ylabel('y')
    show.set_zlabel('z')
    min_ = np.array(min_).min(0)
    max_ = np.array(max_).max(0)
    center = (max_ + min_) / 2
    half = (max_ - min_).max() / 2
    show.set_xlim3d(center[0] - half, center[0] + half)
    show.set_ylim3d(center[1] - half, center[1] + half)
    show.set_zlim3d(center[2] - half, center[2] + half)
    if kwargs.get('legend', False):
        show.legend(prop={'size': 10})
    if SETTINGS['auto_show']:
        showFigure()
    return show
Exemplo n.º 3
0
def showProtein(*atoms, **kwargs):
    """Show protein representation using :meth:`~mpl_toolkits.mplot3d.Axes3D`.
    This function is designed for generating a quick view of the contents of a
    :class:`~.AtomGroup` or :class:`~.Selection`.    

    Protein atoms matching ``"calpha"`` selection are displayed using solid
    lines by picking a random and unique color per chain.  Line with can
    be adjusted using *lw* argument, e.g. ``lw=12``. Default width is 4.
    Chain colors can be overwritten using chain identifier as in ``A='green'``.

    Water molecule oxygen atoms are represented by red colored circles.  Color
    can be changed using *water* keyword argument, e.g. ``water='aqua'``.
    Water marker and size can be changed using *wmarker* and *wsize* keywords,
    defaults values are ``wmarker='.', wsize=6``.

    Hetero atoms matching ``"hetero and noh"`` selection are represented by
    circles and unique colors are picked at random on a per residue basis.
    Colors can be customized using residue name as in ``NAH='purple'``.  Note
    that this will color all distinct residues with the same name in the same
    color.  Hetero atom marker and size can be changed using *hmarker* and
    *hsize* keywords, default values are ``hmarker='o', hsize=6``.

    ProDy will set the size of axis so the representation is not distorted when
    the shape of figure window is close to a square.  Colors are picked at
    random, except for water oxygens which will always be colored red.
    
    *** Interactive 3D Rendering in Jupyter Notebook ***
    
    If py3Dmol has been imported then it will be used instead to display 
    an interactive viewer.  See :func:`view3D`
    
    
    """

    from prody.dynamics.mode import Mode

    method = kwargs.pop('draw', None)
    modes = kwargs.get('mode', None)
    scale = kwargs.get('scale', 100)

    # modes need to be specifically a list or a tuple (cannot be an array)
    if modes is None:
        n_modes = 0
    else:
        modes = wrapModes(modes)
        n_modes = len(modes)

    if method is None:
        import sys
        if 'py3Dmol' in sys.modules:
            method = 'py3Dmol'
        else:
            method = 'matplotlib'
    method = method.lower()

    alist = atoms
    for atoms in alist:
        if not isinstance(atoms, Atomic):
            raise TypeError('atoms must be an Atomic instance')

    if n_modes and n_modes != len(alist):
        raise RuntimeError(
            'the number of proteins ({0}) does not match that of the modes ({1}).'
            .format(len(alist), n_modes))

    if '3dmol' in method:
        mol = view3D(*alist, **kwargs)
        return mol
    else:
        kwargs.pop('mode', None)
        kwargs.pop('scale', 100)

        import matplotlib.pyplot as plt
        from mpl_toolkits.mplot3d import Axes3D
        cf = plt.gcf()
        show = None
        for child in cf.get_children():
            if isinstance(child, Axes3D):
                show = child
                break
        if show is None:
            show = Axes3D(cf)
        from matplotlib import colors
        cnames = dict(colors.cnames)
        wcolor = kwargs.get('water', 'red').lower()
        avoid = np.array(
            colors.hex2color(cnames.pop(wcolor, cnames.pop('red'))))
        for cn, val in cnames.copy().items():  # PY3K: OK
            clr = np.array(colors.hex2color(val))
            if clr.sum() > 2.4:
                cnames.pop(cn)
            elif np.abs(avoid - clr).sum() <= 0.6:
                cnames.pop(cn)
        cnames = list(cnames)
        import random
        random.shuffle(cnames)
        cnames_copy = list(cnames)
        min_ = list()
        max_ = list()
        for i, atoms in enumerate(alist):
            if isinstance(atoms, AtomGroup):
                title = atoms.getTitle()
            else:
                title = atoms.getAtomGroup().getTitle()
            calpha = atoms.select('calpha')
            if calpha:
                partition = False
                mode = modes[i] if n_modes else None
                if mode is not None:
                    is3d = False
                    try:
                        arr = mode.getArray()
                        is3d = mode.is3d()
                        n_nodes = mode.numAtoms()
                    except AttributeError:
                        arr = mode
                        is3d = len(arr) == len(calpha) * 3
                        n_nodes = len(arr) // 3 if is3d else len(arr)
                    if n_nodes != len(calpha):
                        raise RuntimeError(
                            'size mismatch between the protein ({0} residues) and the mode ({1} nodes).'
                            .format(len(calpha), n_nodes))
                    partition = not is3d

                if partition:
                    xyz = calpha._getCoords()
                    chids = calpha.getChids()
                    rbody = []
                    last_sign = np.sign(arr[0])
                    last_chid = chids[0]
                    rcolor = ['red', 'red', 'blue']
                    n = 1
                    for i, a in enumerate(arr):
                        s = np.sign(a)
                        ch = chids[i]
                        if s == 0: s = last_sign
                        if last_sign != s or i == len(
                                arr) - 1 or last_chid != ch:
                            if last_chid == ch:
                                rbody.append(i)
                            show.plot(xyz[rbody, 0],
                                      xyz[rbody, 1],
                                      xyz[rbody, 2],
                                      label=title + '_regid%d' % n,
                                      color=rcolor[int(last_sign + 1)],
                                      lw=kwargs.get('lw', 4))
                            rbody = []
                            n += 1
                            last_sign = s
                            last_chid = ch
                        rbody.append(i)
                else:
                    for ch in HierView(calpha, chain=True):
                        xyz = ch._getCoords()
                        chid = ch.getChid()
                        if len(cnames) == 0:
                            cnames = list(cnames_copy)
                        show.plot(xyz[:, 0],
                                  xyz[:, 1],
                                  xyz[:, 2],
                                  label=title + '_' + chid,
                                  color=kwargs.get(chid, cnames.pop()).lower(),
                                  lw=kwargs.get('lw', 4))

                    if mode is not None:
                        from prody.utilities.drawtools import drawArrow3D
                        XYZ = calpha._getCoords()
                        arr = arr.reshape((n_nodes, 3))
                        XYZ2 = XYZ + arr * scale
                        for i, xyz in enumerate(XYZ):
                            xyz2 = XYZ2[i]
                            mutation_scale = kwargs.pop('mutation_scale', 10)
                            drawArrow3D(xyz,
                                        xyz2,
                                        mutation_scale=mutation_scale,
                                        **kwargs)

            water = atoms.select('water and noh')
            if water:
                xyz = atoms.select('water')._getCoords()
                show.plot(xyz[:, 0],
                          xyz[:, 1],
                          xyz[:, 2],
                          label=title + '_water',
                          color=wcolor,
                          ls='None',
                          marker=kwargs.get('wmarker', '.'),
                          ms=kwargs.get('wsize', 6))
            hetero = atoms.select(
                'not protein and not nucleic and not water and not dummy')
            if hetero:
                for res in HierView(hetero).iterResidues():
                    xyz = res._getCoords()
                    resname = res.getResname()
                    resnum = str(res.getResnum())
                    chid = res.getChid()
                    if len(cnames) == 0:
                        cnames = list(cnames_copy)
                    show.plot(xyz[:, 0],
                              xyz[:, 1],
                              xyz[:, 2],
                              ls='None',
                              color=kwargs.get(resname, cnames.pop()).lower(),
                              label=title + '_' + chid + '_' + resname +
                              resnum,
                              marker=kwargs.get('hmarker', 'o'),
                              ms=kwargs.get('hsize', 6))
            xyz = atoms._getCoords()
            min_.append(xyz.min(0))
            max_.append(xyz.max(0))

        show.set_xlabel('x')
        show.set_ylabel('y')
        show.set_zlabel('z')
        min_ = np.array(min_).min(0)
        max_ = np.array(max_).max(0)
        center = (max_ + min_) / 2
        half = (max_ - min_).max() / 2
        show.set_xlim3d(center[0] - half, center[0] + half)
        show.set_ylim3d(center[1] - half, center[1] + half)
        show.set_zlim3d(center[2] - half, center[2] + half)
        if kwargs.get('legend', False):
            show.legend(prop={'size': 10})
        if SETTINGS['auto_show']:
            showFigure()
        return show