""" 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()