mol_generator.py

"""
Using the schrodinger Python modules, read a template .mae file to
generate the coordgen template C++ files.
"""

import sys
import os
import textwrap
import argparse
import subprocess

from schrodinger import structure

NAME = 'CoordgenTemplates'

# String that creates a molecule and adds it to a container
# of sketcherMinimizerMolecule* at index i.
#
# C++ programmers: `{` and `}` are escaped by duplication, so
# `{{` means `{`.
_MOLECULE = """
{{
    auto molecule = new sketcherMinimizerMolecule();
    std::array<std::tuple<int, float, float>, {atom_total}> atoms = {{{{
{atoms}
    }}}};
    std::array<std::array<int, 3>, {bond_total}> bonds = {{{{
{bonds}
    }}}};

    add_atoms(molecule, atoms);
    add_bonds(molecule, bonds);

    molecules[i] = molecule;
    ++i;
}}
"""

_DOC = """
///
// Find the templates for coordinate generation
//
// Autogenerated, do not edit.
//
//   $SCHRODINGER/run {script_name} {template_file}
//
// generated using {template_file} version {git_hash}.
//
"""

_HEADER = """
#pragma once

{doc}

#include <vector>

class sketcherMinimizerMolecule;
namespace schrodinger {{

///
// Create a new vector of sketcherMinimizerMolecule*. Caller
// owns the sketcherMinimizerMolecule objects.
std::vector<sketcherMinimizerMolecule*> coordgen_templates();

}}
"""

_IMPLEMENTATION = """
{doc}

#include "{name}.h"

#include <array>

#include "sketcherMinimizerMolecule.h"
#include "sketcherMinimizerAtom.h"
#include "sketcherMinimizerBond.h"

using std::array;
using std::tuple;

template <typename T>
void add_atoms(sketcherMinimizerMolecule* molecule, const T& atoms)
{{
    for (const auto& a: atoms) {{
        auto atom = molecule->addNewAtom();
        atom->setAtomicNumber(std::get<0>(a));
        atom->setCoordinates(sketcherMinimizerPointF(std::get<1>(a), std::get<2>(a)));
    }}
}}

template <typename T>
void add_bonds(sketcherMinimizerMolecule* molecule, const T& bonds)
{{
    for (const auto& b: bonds) {{
        auto* from_atom = molecule->getAtoms().at(b[0]);
        auto* to_atom = molecule->getAtoms().at(b[1]);
        auto bond = molecule->addNewBond(from_atom, to_atom);
        bond->setBondOrder(b[2]);
    }}
}}

namespace schrodinger {{

std::vector<sketcherMinimizerMolecule*> coordgen_templates()
{{
    std::vector<sketcherMinimizerMolecule*> molecules({total});
    size_t i = 0;

    {body}

    return molecules;
}}

}}

"""


def get_mol_def(st):
    """
    Use _MOLECULE to define `st` as a sketcherMinimizerMolecule
    """
    atoms = (
        f'        tuple<int, float, float>({a.atomic_number}, {a.x}f, {a.y}f)'
        for a in st.atom)
    atoms = ',\n'.join(atoms)

    bonds = (
        f'        {{ {b.atom1.index - 1}, {b.atom2.index - 1}, {b.order} }}'
        for b in st.bond)
    bonds = ',\n'.join(bonds)

    t = _MOLECULE.format(
        bonds=bonds,
        atoms=atoms,
        atom_total=st.atom_total,
        bond_total=len(st.bond))
    return t


def main(args=None):
    parser = argparse.ArgumentParser(args)
    parser.add_argument('template_file')

    opts = parser.parse_args()

    git_hash = subprocess.check_output(
        ['git', 'log', '-n', '1', '--pretty=format:%H', opts.template_file])
    git_hash = git_hash.decode().strip()[:20]
    template_dir, template_base_name = os.path.split(opts.template_file)

    doc = _DOC.format(
        script_name=os.path.basename(__file__),
        template_file=template_base_name,
        git_hash=git_hash)

    header = _HEADER.format(doc=doc)

    total = structure.count_structures(opts.template_file)
    body = ''

    with structure.StructureReader(opts.template_file) as r:
        for st in r:
            mol_def = get_mol_def(st)
            body += mol_def
    body = textwrap.indent(body, '    ')

    implementation = _IMPLEMENTATION.format(
        doc=doc, name=NAME, body=body, total=total)

    header_path = os.path.join(template_dir, f'{NAME}.h')
    with open(header_path, 'w') as fh:
        fh.write(header)

    implementation_path = os.path.join(template_dir, f'{NAME}.cpp')
    with open(implementation_path, 'w') as fh:
        fh.write(implementation)


if __name__ == '__main__':
    main()