Source code for schrodinger.application.matsci.polymer

"""
Module containing functionalities related to polymers

Copyright Schrodinger, LLC. All rights reserved.
"""
from schrodinger.application.matsci import amorphous
from schrodinger.application.matsci import clusterstruct
from schrodinger.application.matsci import msprops


def create_polymer_name(mol):
    """
    Create a polymer name for the passed molecule by finding its monomers.
    Returns None if the molecule is not a polymer or has unexpected composition,
    e.g. more than one initiator.
    The format is "a-{initiator}-w-{terminator}-poly({list of other monomers})"
    Example: a-chloro-w-fluoro-poly(caprolactam-chloroprene)

    :param `structure._Molecule` mol: The molecule to get the name for

    :rtype: str or None
    :return: The polymer name, or None if the molecule is not a polymer
    """

    initiator = ''
    terminator = ''
    monomers = set()
    for atom in mol.atom:
        monomer_name = atom.property.get(msprops.MONOMER_NAME_PROP)
        if monomer_name is None:
            continue

        if atom.pdbres.rstrip() == amorphous.INI:
            if initiator:
                if initiator != monomer_name:
                    return None  # More than one initiator monomer
            else:
                initiator = monomer_name
        elif atom.pdbres.rstrip() == amorphous.TRM:
            if terminator:
                if terminator != monomer_name:
                    return None
            else:
                terminator = monomer_name
        else:
            monomers.add(monomer_name)

    if not all([initiator, terminator, monomers]):
        # Not a polymer (not one we made at least)
        return None

    parts = []
    if initiator != 'aH':
        # [1:] removes 'a' or 'w' from the beginning of the name
        parts.extend(['a', initiator[1:]])
    if terminator != 'wH':
        parts.extend(['w', terminator[1:]])
    if parts:
        parts.append('')  # Adds an extra dash before "poly"

    monomers_str = '-'.join(sorted(monomers))
    polymer_name = '-'.join(parts) + f'poly({monomers_str})'
    return polymer_name


def find_polymer_species(structs):
    """
    Convenience function to find polymer species in the passed structures.
    Molecules that are not polymers are grouped by SMILES.
    See `clusterstruct.find_species` for more information

    :param iterable structs: The structures to get species for

    :rtype: dict
    :return: Keys are unique SMILES strings or group names for species, values
        are SpeciesData objects for that species
    """
    return clusterstruct.find_species(structs,
                                      naming_function=create_polymer_name)