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