A short script for generating networks using OpenFE

gen_networks.py

import yaml
import json
import argparse
from typing import List

import numpy as np
from rdkit import Chem

from openfe import SmallMoleculeComponent
from openfe.setup import LomapAtomMapper, lomap_scorers, LigandNetwork
from openfe.setup.ligand_network_planning import (generate_minimal_spanning_network,
                                                  generate_radial_network)
from openmm import unit
from collections import OrderedDict


parser = argparse.ArgumentParser(
    description="Write out a network based on a set of input molecules",
)


parser.add_argument(
    "--mols_sdf",
    type=str,
    help=("Path to the input molecules SDF file. "
          "Note: if star graph is generate, the first entry will be used as "
          "the central ligand"),
)


parser.add_argument(
    "--network_type",
    type=str,
    help=("Network to be generated, either star (star graph), "
          "or mst (minimum spanning graph)"),
)


parser.add_argument(
    "--outfile",
    type=str,
    help="Name of output yaml file",
)


parser.add_argument(
   "--outtype",
   type=str,
   default="json",
   help=("type of output file to write, "
         "either `json` (default), `yaml` (PLB-like), or `graphml`")
)


def _load_ligands(filepath: str) -> List[SmallMoleculeComponent]:
    """
    Load a list of ligands from an SDF file
    
    Parameters
    ----------
    filepath : str
      A str path to the SDF file to be read
    
    Returns
    -------
    List[SmallMoleculeComponent]
      A list of SmallMoleculeComponent objects for each molecule in the
      SDF file.
    """
    supp = Chem.SDMolSupplier(filepath, removeHs=False)
    return [SmallMoleculeComponent(m) for m in supp]
    

def _gen_network(mols: List[SmallMoleculeComponent],
                 network_type: str) -> LigandNetwork:
    """
    Generate a network (of type defined by network_type), from the input
    mols using the LomapAtomMapper.
    
    Load a list of ligands from an SDF file
    
    Parameters
    ----------
    mols : List[SmallMoleculeComponent]
      List of SmallMoleculeComponents to generate a network from. Please
      note that the first molecule will be used as the central ligand
      when building star graphs.
    network_type : str
      The network type to build. Currently only supports "star" (a star
      graph), or "mst" (a minimum spanning graph).
    
    Returns
    -------
    LigandNetwork
      A network of ligands with each edges defining an alchemical
      transformation.
    """
    mappers = [
        LomapAtomMapper(time=60, threed=True, element_change=True, max3d=1),
    ]

    if network_type == "star":
        net = generate_radial_network(
                ligands=mols[1:], central_ligand=mols[0], mappers=mappers,
                scorer=lomap_scorers.default_lomap_score)
    elif network_type == "mst":
        net = generate_minimal_spanning_network(
                mols, mappers, lomap_scorers.default_lomap_score)
    else:
        errmsg = f"{network_type} is not a currently supported network"
        raise ValueError(errmsg)

    return net
        

def _write_network(network: LigandNetwork, mols: List[SmallMoleculeComponent],
                   network_type: str, outfile: str, outtype: str) -> None:
    """
    Write out a network to yaml.
    
    Parameters
    ----------
    network : LigandNetwork
      A ligand transformation network to be written out.
    mols : List[SmallMoleculeComponent]
      List of ligands, as ordered in the input SDF file.
    network_type : str
      The type of network which was created.
    outfile : str
      Name of the serialised network file to be written.
    outtype : str
      Type of file to be written out. Can either be json (default),
      yaml, or graphml (which dumps out the network object directly).
    """

    def get_index(mols, name):
        for i, mol in enumerate(mols):
            if mol.name == name:
                return i

    output = {}
    output['mapper'] = "LOMAP (time=60, threed=True, element_change=True, max3d=1)"
    output['planner'] = f"openfe {network_type} graph"
    
    edges_dict = dict()
    
    for edge in network.edges:
        edge_name = f'edge_{edge.componentA.name}_{edge.componentB.name}'
        mapping = [list(i) for i in edge.componentA_to_componentB.items()]
        edges_dict[edge_name] = {
            'ligand_a': edge.componentA.name,
            'ligand_a_id': get_index(mols, edge.componentA.name),
            'ligand_b': edge.componentB.name,
            'ligand_b_id': get_index(mols, edge.componentB.name),
            'atom mapping': edge.componentA_to_componentB,
            'score': f"{edge.annotations['score']:.3f}",}

    output['edges'] = edges_dict
    
    if outtype == "json":
        with open(outfile, 'w') as json_file:
            json.dump(output, json_file)
    elif outtype == "yaml":
        with open(outfile, 'w') as yaml_file:
            yaml.dump(output, yaml_file,
                      default_flow_style=None, sort_keys=False)
    elif outtype == "graphml":
        with open(outfile, 'w') as writer:
            writer.write(network.to_graphml())
    else:
        raise ValueError("Unknown serialisation type")


if __name__ == "__main__":
    args = parser.parse_args()
    
    mols = _load_ligands(args.mols_sdf)
    network = _gen_network(mols, args.network_type)
    _write_network(network, mols, args.network_type, args.outfile, args.outtype)