mirror of
https://github.com/OpenFreeEnergy/openfe.git
synced 2026-06-05 14:54:37 +08:00
103 lines
3.2 KiB
Python
103 lines
3.2 KiB
Python
# This code is part of OpenFE and is licensed under the MIT license.
|
|
# For details, see https://github.com/OpenFreeEnergy/openfe
|
|
|
|
import click
|
|
from openfecli import OFECommandPlugin
|
|
from openfecli.parameters import MOL, MAPPER, OUTPUT_FILE_AND_EXT
|
|
|
|
|
|
def allow_two_molecules(ctx, param, value):
|
|
"""click callback to require that --mol is specified exactly twice"""
|
|
if len(value) != 2:
|
|
raise click.BadParameter("Must specify --mol exactly twice.")
|
|
return value
|
|
|
|
|
|
@click.command(
|
|
"atommapping",
|
|
short_help="Check the atom mapping of a given pair of ligands"
|
|
)
|
|
@MOL.parameter(multiple=True, callback=allow_two_molecules, required=True,
|
|
help=MOL.kwargs['help'] + " Must be specified twice.")
|
|
@MAPPER.parameter(required=True)
|
|
@OUTPUT_FILE_AND_EXT.parameter(
|
|
help=OUTPUT_FILE_AND_EXT.kwargs['help'] + " (PNG format)"
|
|
)
|
|
def atommapping(mol, mapper, output):
|
|
"""
|
|
This provides tools for looking at a specific atommapping.
|
|
"""
|
|
# note that the text of the docstring will be the help when you run
|
|
# `openfe atommapping --help`.
|
|
molA_str, molB_str = mol
|
|
molA = MOL.get(molA_str)
|
|
molB = MOL.get(molB_str)
|
|
mapper_cls = MAPPER.get(mapper)
|
|
mapper_obj = mapper_cls()
|
|
file, ext = OUTPUT_FILE_AND_EXT.get(output)
|
|
if file:
|
|
atommapping_visualize_main(mapper_obj, molA, molB, file, ext)
|
|
else:
|
|
atommapping_print_dict_main(mapper_obj, molA, molB)
|
|
|
|
|
|
def generate_mapping(mapper, molA, molB):
|
|
"""Utility method to extract a single mapping from a mapper.
|
|
|
|
Parameters
|
|
----------
|
|
mapper : :class:`.LigandAtomMapper`
|
|
the mapper to use to generate the mapping
|
|
molA, molB : :class:`.SmallMoleculeComponent`
|
|
molecules to map between
|
|
|
|
Returns
|
|
------
|
|
:class:`.LigandAtomMapping` :
|
|
the mapping generated by the mapper; errors if there is not exactly
|
|
one mapping generated
|
|
"""
|
|
mappings = list(mapper.suggest_mappings(molA, molB))
|
|
if len(mappings) != 1:
|
|
raise click.UsageError(
|
|
f"Found {len(mappings)} mappings; this command requires a mapper "
|
|
"to provide exactly 1 mapping"
|
|
)
|
|
return mappings[0]
|
|
|
|
|
|
def atommapping_print_dict_main(mapper, molA, molB):
|
|
"""Main function for generating and printing out the mapping"""
|
|
mapping = generate_mapping(mapper, molA, molB)
|
|
print(mapping.componentA_to_componentB)
|
|
|
|
|
|
def atommapping_visualize_main(mapper, molA, molB, file, ext):
|
|
from rdkit.Chem import Draw
|
|
from gufe.visualization import mapping_visualization as vis
|
|
|
|
mapping = generate_mapping(mapper, molA, molB)
|
|
ext_to_artist = {
|
|
"png": Draw.rdMolDraw2D.MolDraw2DCairo(600, 300, 300, 300),
|
|
}
|
|
try:
|
|
artist = ext_to_artist[ext]
|
|
except KeyError:
|
|
raise click.BadParameter(
|
|
f"Unknown file format: '{ext}'. The following formats are "
|
|
"supported: " + ", ".join([f"'{ext}'" for ext in ext_to_artist])
|
|
)
|
|
|
|
contents = vis.draw_mapping(mapping.componentA_to_componentB,
|
|
mapping.componentA.to_rdkit(),
|
|
mapping.componentB.to_rdkit(), d2d=artist)
|
|
|
|
file.write(contents)
|
|
|
|
|
|
PLUGIN = OFECommandPlugin(
|
|
command=atommapping,
|
|
section="hidden",
|
|
requires_ofe=(0, 0, 1),
|
|
)
|