diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 98ad6d6e..faf3af05 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -26,7 +26,7 @@ jobs: cache: "pip" - run: "pip install -e '.[dev]'" - run: mypy src - black: + ruff-format: runs-on: ubuntu-22.04 steps: - uses: actions/checkout@v3 @@ -35,7 +35,7 @@ jobs: python-version: "3.11" cache: "pip" - run: "pip install -e '.[dev]'" - - run: black --check . + - run: ruff format --check . pytest: # https://ericmjl.github.io/blog/2021/12/30/better-conda-environments-on-github-actions/ runs-on: ubuntu-22.04 diff --git a/.gitignore b/.gitignore index df7ea3fa..7812d1db 100644 --- a/.gitignore +++ b/.gitignore @@ -131,4 +131,5 @@ dmypy.json # Version. src/cgexplore/_version.py testff_iterations.json -.vscode \ No newline at end of file +.vscode +docs/source/_autosummary \ No newline at end of file diff --git a/docs/Makefile b/docs/Makefile new file mode 100644 index 00000000..258497cf --- /dev/null +++ b/docs/Makefile @@ -0,0 +1,20 @@ +# Minimal makefile for Sphinx documentation +# + +# You can set these variables from the command line, and also +# from the environment for the first two. +SPHINXOPTS ?= -W --keep-going +SPHINXBUILD ?= sphinx-build +SOURCEDIR = source +BUILDDIR = build + +# Put it first so that "make" without argument is like "make help". +help: + @$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) + +.PHONY: help Makefile + +# Catch-all target: route all unknown targets to Sphinx using the new +# "make mode" option. $(O) is meant as a shortcut for $(SPHINXOPTS). +%: Makefile + @$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) \ No newline at end of file diff --git a/docs/source/_templates/class.rst b/docs/source/_templates/class.rst new file mode 100644 index 00000000..67273242 --- /dev/null +++ b/docs/source/_templates/class.rst @@ -0,0 +1,35 @@ +{{ fullname | escape | underline}} + +.. currentmodule:: {{ module }} + +.. autoclass:: {{ objname }} + :members: + :inherited-members: + :undoc-members: + :show-inheritance: + + + {% block methods %} + {% if methods %} + .. rubric:: {{ _('Methods') }} + + .. autosummary:: + :nosignatures: + {% for item in methods %} + {%- if not item.startswith('_') and item not in inherited_members %} + ~{{ name }}.{{ item }} + {%- endif -%} + {%- endfor %} + {% endif %} + {% endblock %} + + {% block attributes %} + {% if attributes %} + .. rubric:: {{ _('Attributes') }} + + .. autosummary:: + {% for item in attributes %} + ~{{ name }}.{{ item }} + {%- endfor %} + {% endif %} + {% endblock %} diff --git a/docs/source/_templates/module.rst b/docs/source/_templates/module.rst new file mode 100644 index 00000000..7acbe67b --- /dev/null +++ b/docs/source/_templates/module.rst @@ -0,0 +1,68 @@ +{{ fullname | escape | underline}} + +.. automodule:: {{ fullname }} + + {% block attributes %} + {% if attributes %} + .. rubric:: Module attributes + + .. autosummary:: + :toctree: + {% for item in attributes %} + {{ item }} + {%- endfor %} + {% endif %} + {% endblock %} + + {% block functions %} + {% if functions %} + .. rubric:: {{ _('Functions') }} + + .. autosummary:: + :toctree: + :nosignatures: + {% for item in functions %} + {{ item }} + {%- endfor %} + {% endif %} + {% endblock %} + + {% block classes %} + {% if classes %} + .. rubric:: {{ _('Classes') }} + + .. autosummary:: + :toctree: + :template: class.rst + :nosignatures: + {% for item in classes %} + {{ item }} + {%- endfor %} + {% endif %} + {% endblock %} + + {% block exceptions %} + {% if exceptions %} + .. rubric:: {{ _('Exceptions') }} + + .. autosummary:: + :toctree: + {% for item in exceptions %} + {{ item }} + {%- endfor %} + {% endif %} + {% endblock %} + +{% block modules %} +{% if modules %} +.. rubric:: Modules + +.. autosummary:: + :toctree: + :template: module.rst + :recursive: +{% for item in modules %} + {{ item }} +{%- endfor %} +{% endif %} +{% endblock %} diff --git a/docs/source/cgexplore.rst b/docs/source/cgexplore.rst new file mode 100644 index 00000000..6d3a6b7f --- /dev/null +++ b/docs/source/cgexplore.rst @@ -0,0 +1,17 @@ +cgexplore +========= + + + + +Beads +----- + +.. toctree:: + :maxdepth: 1 + + CgBeads <_autosummary/cgexplore.CgBeads> + + + +This docs are a huge WIP. \ No newline at end of file diff --git a/docs/source/conf.py b/docs/source/conf.py new file mode 100644 index 00000000..1e3df380 --- /dev/null +++ b/docs/source/conf.py @@ -0,0 +1,45 @@ +# Configuration file for the Sphinx documentation builder. +# +# For the full list of built-in configuration values, see the documentation: +# https://www.sphinx-doc.org/en/master/usage/configuration.html + +# -- Project information ----------------------------------------------------- +# https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information +from __future__ import annotations + +project = "CGExplore" +project_copyright = "2023, Andrew Tarzia" +author = "Andrew Tarzia" + +# -- General configuration --------------------------------------------------- +# https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration + +extensions = [ + "sphinx.ext.doctest", + "sphinx.ext.napoleon", + "sphinx.ext.autosummary", + "sphinx.ext.intersphinx", + "sphinx.ext.viewcode", + "sphinx_copybutton", +] + +autosummary_imported_members = True + +autodoc_typehints = "description" +autodoc_member_order = "groupwise" +autoclass_content = "class" + +intersphinx_mapping = { + "python": ("https://docs.python.org/3", None), +} + + +templates_path = ["_templates"] +exclude_patterns: list[str] = [] + + +# -- Options for HTML output ------------------------------------------------- +# https://www.sphinx-doc.org/en/master/usage/configuration.html#options-for-html-output + +html_theme = "furo" +html_static_path = ["_static"] diff --git a/docs/source/index.rst b/docs/source/index.rst new file mode 100644 index 00000000..3536192a --- /dev/null +++ b/docs/source/index.rst @@ -0,0 +1,49 @@ +.. toctree:: + :hidden: + :caption: CGExplore + :maxdepth: 2 + + cgexplore + +.. toctree:: + :hidden: + :maxdepth: 2 + :caption: Modules: + + Modules + +============ +Introduction +============ + +| GitHub: https://www.github.com/andrewtarzia/CGExplore + + +:mod:`.CGExplore` is a Python library for doing something... + + + + +Installation +------------ + +TODO + + +Examples +-------- + +TODO + + +Acknowledgements +---------------- + +I developed much of this code when working in the Pavan group (https://www.gmpavanlab.com/). + +Indices and tables +------------------ + +* :ref:`genindex` +* :ref:`modindex` +* :ref:`search` diff --git a/docs/source/modules.rst b/docs/source/modules.rst new file mode 100644 index 00000000..27a8b68a --- /dev/null +++ b/docs/source/modules.rst @@ -0,0 +1,9 @@ +Modules +======= + +.. autosummary:: + :toctree: _autosummary + :template: module.rst + :recursive: + + cgexplore diff --git a/first_paper_example/analysis.py b/first_paper_example/analysis.py index 0576034c..ea10d7ec 100644 --- a/first_paper_example/analysis.py +++ b/first_paper_example/analysis.py @@ -41,7 +41,7 @@ def analyse_cage( conformer, name, output_dir, - force_field, + forcefield, node_element, ligand_element, database, @@ -157,7 +157,7 @@ def analyse_cage( # This is matched to the existing analysis code. I recommend # generalising in the future. - ff_targets = force_field.get_targets() + ff_targets = forcefield.get_targets() if "6P8" in name: torsions = "toff" else: @@ -206,8 +206,8 @@ def analyse_cage( elif ("b1", "m1", "b1") in (cp, tuple(reversed(cp))): clangle = at.angle.value_in_unit(openmm.unit.degrees) - force_field_dict = { - "ff_id": force_field.get_identifier(), + forcefield_dict = { + "ff_id": forcefield.get_identifier(), "torsions": torsions, "vdws": "von", "clbb_bead1": "", @@ -235,7 +235,7 @@ def analyse_cage( "min_b2b_distance": min_b2b_distance, "radius_gyration": radius_gyration, "max_diameter": max_diameter, - "force_field_dict": force_field_dict, + "forcefield_dict": forcefield_dict, } with open(output_file, "w") as f: json.dump(res_dict, f, indent=4) @@ -624,8 +624,8 @@ def data_to_array(json_files, output_dir): row["c2bb_name"] = c2bb_name row["topology"] = t_str row["ff_name"] = ff_name - row["torsions"] = res_dict["force_field_dict"]["torsions"] - row["vdws"] = res_dict["force_field_dict"]["vdws"] + row["torsions"] = res_dict["forcefield_dict"]["torsions"] + row["vdws"] = res_dict["forcefield_dict"]["vdws"] row["run_number"] = 0 row["cltopo"] = int(clbb_name[0]) @@ -633,18 +633,18 @@ def data_to_array(json_files, output_dir): "2p3" ) or t_str in cage_topology_options("2p4"): cltitle = "3C1" if row["cltopo"] == 3 else "4C1" - row["c2r0"] = res_dict["force_field_dict"]["c2r0"] - row["c2angle"] = res_dict["force_field_dict"]["c2angle"] + row["c2r0"] = res_dict["forcefield_dict"]["c2r0"] + row["c2angle"] = res_dict["forcefield_dict"]["c2angle"] row["target_bite_angle"] = (row["c2angle"] - 90) * 2 elif t_str in cage_topology_options("3p4"): cltitle = "4C1" - row["c3r0"] = res_dict["force_field_dict"]["c3r0"] - row["c3angle"] = res_dict["force_field_dict"]["c3angle"] + row["c3r0"] = res_dict["forcefield_dict"]["c3r0"] + row["c3angle"] = res_dict["forcefield_dict"]["c3angle"] row["cltitle"] = cltitle - row["clr0"] = res_dict["force_field_dict"]["clr0"] - row["clangle"] = res_dict["force_field_dict"]["clangle"] + row["clr0"] = res_dict["forcefield_dict"]["clr0"] + row["clangle"] = res_dict["forcefield_dict"]["clangle"] row["bbpair"] = clbb_name + c2bb_name + ff_name row["optimised"] = res_dict["optimised"] diff --git a/first_paper_example/generation.py b/first_paper_example/generation.py index 25bfdeea..a776b416 100644 --- a/first_paper_example/generation.py +++ b/first_paper_example/generation.py @@ -7,7 +7,7 @@ """ -import itertools +import itertools as it import logging import os @@ -36,7 +36,7 @@ def optimise_cage( molecule, name, output_dir, - force_field, + forcefield, platform, database, ): @@ -75,7 +75,7 @@ def optimise_cage( ) return ensemble.get_lowest_e_conformer() - assigned_system = force_field.assign_terms(molecule, name, output_dir) + assigned_system = forcefield.assign_terms(molecule, name, output_dir) ensemble = Ensemble( base_molecule=molecule, @@ -99,7 +99,7 @@ def optimise_cage( conformer = run_optimisation( assigned_system=AssignedSystem( molecule=temp_molecule, - force_field_terms=assigned_system.force_field_terms, + forcefield_terms=assigned_system.forcefield_terms, system_xml=assigned_system.system_xml, topology_xml=assigned_system.topology_xml, bead_set=assigned_system.bead_set, @@ -119,13 +119,13 @@ def optimise_cage( # building blocks. logging.info(f"optimisation of shifted structures of {name}") for test_molecule in yield_shifted_models( - temp_molecule, force_field, kicks=(1, 2, 3, 4) + temp_molecule, forcefield, kicks=(1, 2, 3, 4) ): try: conformer = run_optimisation( assigned_system=AssignedSystem( molecule=test_molecule, - force_field_terms=assigned_system.force_field_terms, + forcefield_terms=assigned_system.forcefield_terms, system_xml=assigned_system.system_xml, topology_xml=assigned_system.topology_xml, bead_set=assigned_system.bead_set, @@ -144,8 +144,8 @@ def optimise_cage( # Collect and optimise structures nearby in phase space. logging.info(f"optimisation of nearby structures of {name}") neighbour_library = get_neighbour_library( - ffnum=int(force_field.get_identifier()), - fftype=force_field.get_prefix(), + ffnum=int(forcefield.get_identifier()), + fftype=forcefield.get_prefix(), ) for test_molecule in yield_near_models( molecule=molecule, @@ -156,7 +156,7 @@ def optimise_cage( conformer = run_optimisation( assigned_system=AssignedSystem( molecule=test_molecule, - force_field_terms=assigned_system.force_field_terms, + forcefield_terms=assigned_system.forcefield_terms, system_xml=assigned_system.system_xml, topology_xml=assigned_system.topology_xml, bead_set=assigned_system.bead_set, @@ -176,7 +176,7 @@ def optimise_cage( name=name, assigned_system=AssignedSystem( molecule=ensemble.get_lowest_e_conformer().molecule, - force_field_terms=assigned_system.force_field_terms, + forcefield_terms=assigned_system.forcefield_terms, system_xml=assigned_system.system_xml, topology_xml=assigned_system.topology_xml, bead_set=assigned_system.bead_set, @@ -212,7 +212,7 @@ def optimise_cage( conformer = run_optimisation( assigned_system=AssignedSystem( molecule=md_conformer.molecule, - force_field_terms=assigned_system.force_field_terms, + forcefield_terms=assigned_system.forcefield_terms, system_xml=assigned_system.system_xml, topology_xml=assigned_system.topology_xml, bead_set=assigned_system.bead_set, @@ -264,44 +264,44 @@ def build_populations( logging.info(f"running population {population}") popn_dict = populations[population] topologies = popn_dict["topologies"] - force_fields = tuple(popn_dict["fflibrary"].yield_forcefields()) + forcefields = tuple(popn_dict["fflibrary"].yield_forcefields()) logging.info(f"there are {len(topologies)} topologies") - logging.info(f"there are {len(force_fields)} ffs") - logging.info(f"building {len(force_fields) * len(topologies)} cages") - popn_iterator = itertools.product(topologies, force_fields) - for cage_topo_str, force_field in popn_iterator: + logging.info(f"there are {len(forcefields)} ffs") + logging.info(f"building {len(forcefields) * len(topologies)} cages") + popn_iterator = it.product(topologies, forcefields) + for cage_topo_str, forcefield in popn_iterator: c2_precursor = popn_dict["c2"] cl_precursor = popn_dict["cl"] name = ( f"{cage_topo_str}_{cl_precursor.get_name()}_" f"{c2_precursor.get_name()}_" - f"f{force_field.get_identifier()}" + f"f{forcefield.get_identifier()}" ) # Write out force field. - force_field.write_human_readable(calculation_output) + forcefield.write_human_readable(calculation_output) # Optimise building blocks. c2_name = ( - f"{c2_precursor.get_name()}_f{force_field.get_identifier()}" + f"{c2_precursor.get_name()}_f{forcefield.get_identifier()}" ) c2_building_block = optimise_ligand( molecule=c2_precursor.get_building_block(), name=c2_name, output_dir=calculation_output, - force_field=force_field, + forcefield=forcefield, platform=None, ) c2_building_block.write(str(ligand_output / f"{c2_name}_optl.mol")) cl_name = ( - f"{cl_precursor.get_name()}_f{force_field.get_identifier()}" + f"{cl_precursor.get_name()}_f{forcefield.get_identifier()}" ) cl_building_block = optimise_ligand( molecule=cl_precursor.get_building_block(), name=cl_name, output_dir=calculation_output, - force_field=force_field, + forcefield=forcefield, platform=None, ) cl_building_block.write(str(ligand_output / f"{cl_name}_optl.mol")) @@ -317,7 +317,7 @@ def build_populations( molecule=cage, name=name, output_dir=calculation_output, - force_field=force_field, + forcefield=forcefield, platform=platform, database=database, ) @@ -331,7 +331,7 @@ def build_populations( conformer=conformer, name=name, output_dir=calculation_output, - force_field=force_field, + forcefield=forcefield, node_element=node_element, ligand_element=ligand_element, database=database, diff --git a/first_paper_example/plot_cages.py b/first_paper_example/plot_cages.py index a4de2cc5..445ba43c 100644 --- a/first_paper_example/plot_cages.py +++ b/first_paper_example/plot_cages.py @@ -1055,13 +1055,11 @@ def generate_movies(figure_output): def naming_convention_map(old_name, tors="toff"): - """ - This only applies because of the change of naming convention. + """This only applies because of the change of naming convention. In future, users should stick to the `new` naming convention. """ - if "_f" in old_name: # You do not need to convert this. return old_name @@ -1170,7 +1168,8 @@ def main(): struct_output, struct_figure_output, ) - raise SystemExit("not implemented from here.") + msg = "not implemented from here." + raise SystemExit(msg) webapp_csv( all_data, figure_output, diff --git a/first_paper_example/plot_maps.py b/first_paper_example/plot_maps.py index 24408dee..a4c1a7fc 100644 --- a/first_paper_example/plot_maps.py +++ b/first_paper_example/plot_maps.py @@ -7,7 +7,7 @@ """ -import itertools +import itertools as it import logging import math import os @@ -15,7 +15,6 @@ import matplotlib as mpl import matplotlib.pyplot as plt -import numpy as np from analysis import ( Xc_map, angle_str, @@ -28,6 +27,7 @@ stoich_map, topology_labels, ) +from cgexplore.utilities import draw_pie from env_set import calculations, figures, outputdata logging.basicConfig( @@ -319,48 +319,6 @@ def selectivity_map(all_data, figure_output): plt.close() -def draw_pie(colours, xpos, ypos, size, ax): - """From: - https://stackoverflow.com/questions/56337732/how-to-plot-scatter- - pie-chart-using-matplotlib. - - """ - num_points = len(colours) - if num_points == 1: - ax.scatter( - xpos, - ypos, - c=colours[0], - edgecolors="k", - s=size, - ) - else: - ratios = [1 / num_points for i in range(num_points)] - assert sum(ratios) <= 1, "sum of ratios needs to be < 1" - - markers = [] - previous = 0 - # calculate the points of the pie pieces - for color, ratio in zip(colours, ratios, strict=True): - this = 2 * np.pi * ratio + previous - x = [0, *np.cos(np.linspace(previous, this, 100)).tolist(), 0] - y = [0, *np.sin(np.linspace(previous, this, 100)).tolist(), 0] - xy = np.column_stack([x, y]) - previous = this - markers.append( - { - "marker": xy, - "s": np.abs(xy).max() ** 2 * np.array(size), - "facecolor": color, - "edgecolors": "k", - } - ) - - # scatter each of the pie pieces to create pies - for marker in markers: - ax.scatter(xpos, ypos, **marker) - - def selfsort_legend(all_data, figure_output): logging.info("running selfsort_legend") @@ -403,7 +361,7 @@ def selfsort_map(all_data, figure_output): io1 = sorted(set(all_data[cols_to_iter[0]])) io2 = sorted(set(all_data[cols_to_iter[1]])) io3 = sorted(set(all_data[cols_to_iter[2]])) - for tor, vdw, cltitle in itertools.product(io1, io2, io3): + for tor, vdw, cltitle in it.product(io1, io2, io3): data = all_data[all_data[cols_to_iter[0]] == tor] data = data[data[cols_to_iter[1]] == vdw] data = data[data[cols_to_iter[2]] == cltitle] @@ -412,9 +370,7 @@ def selfsort_map(all_data, figure_output): uo1 = sorted(set(data[cols_to_map[0]])) uo2 = sorted(set(data[cols_to_map[1]])) - for (i, cla), (_j, ba) in itertools.product( - enumerate(uo1), enumerate(uo2) - ): + for (i, cla), (_j, ba) in it.product(enumerate(uo1), enumerate(uo2)): plot_data = data[data[cols_to_map[0]] == cla] plot_data = plot_data[plot_data[cols_to_map[1]] == ba] @@ -481,7 +437,7 @@ def kinetic_selfsort_map(all_data, figure_output): io1 = sorted(set(all_data[cols_to_iter[0]])) io2 = sorted(set(all_data[cols_to_iter[1]])) io3 = sorted(set(all_data[cols_to_iter[2]])) - for tor, vdw, cltitle in itertools.product(io1, io2, io3): + for tor, vdw, cltitle in it.product(io1, io2, io3): data = all_data[all_data[cols_to_iter[0]] == tor] data = data[data[cols_to_iter[1]] == vdw] data = data[data[cols_to_iter[2]] == cltitle] @@ -490,9 +446,7 @@ def kinetic_selfsort_map(all_data, figure_output): uo1 = sorted(set(data[cols_to_map[0]])) uo2 = sorted(set(data[cols_to_map[1]])) - for (i, cla), (_j, ba) in itertools.product( - enumerate(uo1), enumerate(uo2) - ): + for (i, cla), (_j, ba) in it.product(enumerate(uo1), enumerate(uo2)): plot_data = data[data[cols_to_map[0]] == cla] plot_data = plot_data[plot_data[cols_to_map[1]] == ba] diff --git a/justfile b/justfile index b3cc499a..4945e874 100644 --- a/justfile +++ b/justfile @@ -17,7 +17,7 @@ check: (set -x; ruff . ) echo - ( set -x; black --check . ) + ( set -x; ruff format --check . ) echo ( set -x; mypy src ) @@ -27,7 +27,15 @@ check: test $error = 0 + # Auto-fix code issues. fix: - black . + ruff format . ruff --fix . + + +# Build docs. +docs: + rm -rf docs/source/_autosummary + make -C docs html + echo Docs are in $PWD/docs/build/html/index.html diff --git a/pyproject.toml b/pyproject.toml index 577b1d83..dd15bf08 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -24,15 +24,16 @@ readme = "README.md" [project.optional-dependencies] dev = [ "ruff", - "black", "mypy", "pip-tools", - "pytest", - "pytest-benchmark", + "pytest==7.4.2", "pytest-datadir", "pytest-lazy-fixture", "pytest-cov", + "sphinx", + "sphinx-copybutton", "twine", + "furo", ] [project.urls] @@ -47,13 +48,15 @@ where = [ "src", ] -[tool.black] -line-length = 79 - [tool.ruff] line-length = 79 -extend-select = ["I"] -ignore = ["ANN101", "COM812", "ISC001"] +exclude = [ + "first_paper_example", +] + +[tool.ruff.lint] +select = ["ALL"] +ignore = ["ANN101", "ANN401", "COM812", "ISC001", "G004", "PTH123"] [tool.ruff.lint.pydocstyle] convention = "google" @@ -71,6 +74,7 @@ convention = "google" "S101", "INP001", "T201", + "PLR0913", ] "docs/source/conf.py" = ["D100", "INP001"] @@ -81,12 +85,9 @@ testpaths = [ python_files = [ "test_*.py", "*_test.py", - "benchmark_*.py", - "*_benchmark.py", ] python_functions = [ "test_*", - "benchmark_*", ] [tool.mypy] @@ -94,6 +95,12 @@ show_error_codes = true implicit_optional = false warn_no_return = true strict_optional = true +disallow_untyped_defs = true +disallow_incomplete_defs = true +check_untyped_defs = true +disallow_untyped_decorators = true +warn_unreachable = true +disallow_any_generics = false [[tool.mypy.overrides]] module = [ diff --git a/src/cgexplore/__init__.py b/src/cgexplore/__init__.py index 0b2524b8..e797e3b8 100644 --- a/src/cgexplore/__init__.py +++ b/src/cgexplore/__init__.py @@ -1,18 +1,18 @@ -from .angles import * # noqa -from .assigned_system import * # noqa -from .beads import * # noqa -from .bonds import * # noqa -from .databases import * # noqa -from .ensembles import * # noqa -from .errors import * # noqa -from .forcefield import * # noqa -from .generation_utilities import * # noqa -from .geom import * # noqa -from .martini import * # noqa -from .molecule_construction import * # noqa -from .nonbonded import * # noqa -from .openmm_optimizer import * # noqa -from .shape import * # noqa -from .torsions import * # noqa -from .utilities import * # noqa -from .visualisation import * # noqa +from .angles import * # noqa: F403, D104 +from .assigned_system import * # noqa: F403 +from .beads import * # noqa: F403 +from .bonds import * # noqa: F403 +from .databases import * # noqa: F403 +from .ensembles import * # noqa: F403 +from .errors import * # noqa: F403 +from .forcefield import * # noqa: F403 +from .generation_utilities import * # noqa: F403 +from .geom import * # noqa: F403 +from .martini import * # noqa: F403 +from .molecule_construction import * # noqa: F403 +from .nonbonded import * # noqa: F403 +from .openmm_optimizer import * # noqa: F403 +from .shape import * # noqa: F403 +from .torsions import * # noqa: F403 +from .utilities import * # noqa: F403 +from .visualisation import * # noqa: F403 diff --git a/src/cgexplore/angles.py b/src/cgexplore/angles.py index f5b3fd2d..f1880c22 100644 --- a/src/cgexplore/angles.py +++ b/src/cgexplore/angles.py @@ -2,14 +2,14 @@ """Module for handling angles.""" -import itertools +import itertools as it import logging from collections import abc from dataclasses import dataclass import stk from openmm import openmm -from rdkit.Chem import AllChem as rdkit +from rdkit.Chem import AllChem from .errors import ForceFieldUnitError from .utilities import convert_pyramid_angle @@ -20,12 +20,13 @@ ) -def angle_k_unit(): - return openmm.unit.kilojoules_per_mole / openmm.unit.radian**2 +_angle_k_unit = openmm.unit.kilojoules_per_mole / openmm.unit.radian**2 @dataclass class Angle: + """Class containing angle defintion.""" + atom_names: tuple[str, ...] atom_ids: tuple[int, ...] angle: openmm.unit.Quantity @@ -37,6 +38,8 @@ class Angle: @dataclass class CosineAngle: + """Class containing cosine-angle defintion.""" + atom_names: tuple[str, ...] atom_ids: tuple[int, ...] n: int @@ -48,6 +51,8 @@ class CosineAngle: @dataclass class TargetAngle: + """Defines a target angle to search for in a molecule.""" + type1: str type2: str type3: str @@ -57,19 +62,33 @@ class TargetAngle: angle: openmm.unit.Quantity angle_k: openmm.unit.Quantity + def vector_key(self) -> str: + """Return key for vector defining this target term.""" + return f"{self.type1}{self.type2}{self.type3}" + + def vector(self) -> tuple[float, float]: + """Return vector defining this target term.""" + return ( + self.angle.value_in_unit(openmm.unit.degrees), + self.angle_k.value_in_unit(_angle_k_unit), + ) + def human_readable(self) -> str: + """Return human-readable definition of this target term.""" return ( f"{self.__class__.__name__}(" f"{self.type1}{self.type2}{self.type3}, " f"{self.element1}{self.element2}{self.element3}, " f"{self.angle.in_units_of(openmm.unit.degrees)}, " - f"{self.angle_k.in_units_of(angle_k_unit())}, " + f"{self.angle_k.in_units_of(_angle_k_unit)}, " ")" ) @dataclass class TargetAngleRange: + """Defines a target angle and ranges in parameters to search for.""" + type1: str type2: str type3: str @@ -79,8 +98,9 @@ class TargetAngleRange: angles: tuple[openmm.unit.Quantity] angle_ks: tuple[openmm.unit.Quantity] - def yield_angles(self): - for angle, k in itertools.product(self.angles, self.angle_ks): + def yield_angles(self) -> abc.Iterable[TargetAngle]: + """Find angles matching target.""" + for angle, k in it.product(self.angles, self.angle_ks): yield TargetAngle( type1=self.type1, type2=self.type2, @@ -93,24 +113,10 @@ def yield_angles(self): ) -@dataclass -class TargetPyramidAngle(TargetAngle): - opposite_angle: openmm.unit.Quantity - - def human_readable(self) -> str: - return ( - f"{self.__class__.__name__}(" - f"{self.type1}{self.type2}{self.type3}, " - f"{self.element1}{self.element2}{self.element3}, " - f"{self.angle.in_units_of(openmm.unit.degrees)}, " - f"{self.opposite_angle.in_units_of(openmm.unit.degrees)}, " - f"{self.angle_k.in_units_of(angle_k_unit())}, " - ")" - ) - - @dataclass class TargetCosineAngle: + """Defines a target angle to search for in a molecule.""" + type1: str type2: str type3: str @@ -121,19 +127,34 @@ class TargetCosineAngle: b: int angle_k: openmm.unit.Quantity + def vector_key(self) -> str: + """Return key for vector defining this target term.""" + return f"{self.type1}{self.type2}{self.type3}" + + def vector(self) -> tuple[float, float, float]: + """Return vector defining this target term.""" + return ( + self.n, + self.b, + self.angle_k.value_in_unit(openmm.unit.kilojoules_per_mole), + ) + def human_readable(self) -> str: + """Return human-readable definition of this target term.""" return ( f"{self.__class__.__name__}(" f"{self.type1}{self.type2}{self.type3}, " f"{self.element1}{self.element2}{self.element3}, " f"{self.n}, {self.b}, " - f"{self.angle_k.in_units_of(angle_k_unit())}, " + f"{self.angle_k.in_units_of(openmm.unit.kilojoules_per_mole)}, " ")" ) @dataclass class TargetCosineAngleRange: + """Defines a target angle and ranges in parameters to search for.""" + type1: str type2: str type3: str @@ -144,8 +165,9 @@ class TargetCosineAngleRange: bs: tuple[int] angle_ks: tuple[openmm.unit.Quantity] - def yield_angles(self): - for n, b, k in itertools.product(self.ns, self.bs, self.angle_ks): + def yield_angles(self) -> abc.Iterable[TargetCosineAngle]: + """Find angles matching target.""" + for n, b, k in it.product(self.ns, self.bs, self.angle_ks): yield TargetCosineAngle( type1=self.type1, type2=self.type2, @@ -159,8 +181,29 @@ def yield_angles(self): ) +@dataclass +class TargetPyramidAngle(TargetAngle): + """Defines a target angle to search for in a molecule.""" + + opposite_angle: openmm.unit.Quantity + + def human_readable(self) -> str: + """Return human-readable definition of this target term.""" + return ( + f"{self.__class__.__name__}(" + f"{self.type1}{self.type2}{self.type3}, " + f"{self.element1}{self.element2}{self.element3}, " + f"{self.angle.in_units_of(openmm.unit.degrees)}, " + f"{self.opposite_angle.in_units_of(openmm.unit.degrees)}, " + f"{self.angle_k.in_units_of(_angle_k_unit)}, " + ")" + ) + + @dataclass class PyramidAngleRange: + """Defines a target angle and ranges in parameters to search for.""" + type1: str type2: str type3: str @@ -170,8 +213,9 @@ class PyramidAngleRange: angles: tuple[openmm.unit.Quantity] angle_ks: tuple[openmm.unit.Quantity] - def yield_angles(self): - for angle, k in itertools.product(self.angles, self.angle_ks): + def yield_angles(self) -> abc.Iterable[TargetPyramidAngle]: + """Find angles matching target.""" + for angle, k in it.product(self.angles, self.angle_ks): try: opposite_angle = openmm.unit.Quantity( value=convert_pyramid_angle( @@ -181,7 +225,7 @@ def yield_angles(self): ) except AttributeError: msg = f"{self} in angles does not have units for parameters" - raise ForceFieldUnitError(msg) + raise ForceFieldUnitError(msg) # noqa: B904, TRY200 yield TargetPyramidAngle( type1=self.type1, @@ -198,12 +242,15 @@ def yield_angles(self): @dataclass class FoundAngle: + """Define a found forcefield term.""" + atoms: tuple[stk.Atom, ...] atom_ids: tuple[int, ...] def find_angles(molecule: stk.Molecule) -> abc.Iterator[FoundAngle]: - paths = rdkit.FindAllPathsOfLengthN( + """Find angles based on bonds in molecule.""" + paths = AllChem.FindAllPathsOfLengthN( mol=molecule.to_rdkit_mol(), length=3, useBonds=False, @@ -219,6 +266,8 @@ def find_angles(molecule: stk.Molecule) -> abc.Iterator[FoundAngle]: @dataclass class TargetMartiniAngle: + """Defines a target angle to search for in a molecule.""" + type1: str type2: str type3: str @@ -230,19 +279,22 @@ class TargetMartiniAngle: angle_k: openmm.unit.Quantity def human_readable(self) -> str: + """Return human-readable definition of this target term.""" return ( f"{self.__class__.__name__}(" f"{self.type1}{self.type2}{self.type3}, " f"{self.element1}{self.element2}{self.element3}, " f"{self.funct}, " f"{self.angle.in_units_of(openmm.unit.degrees)}, " - f"{self.angle_k.in_units_of(angle_k_unit())}, " + f"{self.angle_k.in_units_of(_angle_k_unit)}, " ")" ) @dataclass class MartiniAngleRange: + """Defines a target angle and ranges in parameters to search for.""" + type1: str type2: str type3: str @@ -253,8 +305,9 @@ class MartiniAngleRange: angles: tuple[openmm.unit.Quantity] angle_ks: tuple[openmm.unit.Quantity] - def yield_angles(self): - for angle, k in itertools.product(self.angles, self.angle_ks): + def yield_angles(self) -> abc.Iterable[TargetMartiniAngle]: + """Find angles matching target.""" + for angle, k in it.product(self.angles, self.angle_ks): yield TargetMartiniAngle( type1=self.type1, type2=self.type2, diff --git a/src/cgexplore/assigned_system.py b/src/cgexplore/assigned_system.py index c164a9cf..7ac33a69 100644 --- a/src/cgexplore/assigned_system.py +++ b/src/cgexplore/assigned_system.py @@ -1,4 +1,3 @@ -#!/usr/bin/env python # Distributed under the terms of the MIT License. """Module for system classes with assigned terms. @@ -23,8 +22,10 @@ class ForcedSystem: + """A system with forces assigned.""" + molecule: stk.Molecule - force_field_terms: dict[str, tuple] + forcefield_terms: dict[str, tuple] vdw_bond_cutoff: int def _available_forces(self, force_type: str) -> openmm.Force: @@ -41,7 +42,7 @@ def _available_forces(self, force_type: str) -> openmm.Force: return available[force_type] def _add_bonds(self, system: openmm.System) -> openmm.System: - forces = self.force_field_terms["bond"] + forces = self.forcefield_terms["bond"] force_types = {i.force for i in forces} for force_type in force_types: if "Martini" in force_type: @@ -66,12 +67,12 @@ def _add_bonds(self, system: openmm.System) -> openmm.System: ) except AttributeError: msg = f"{assigned_force} in bonds does not have units" - raise ForceFieldUnitError(msg) + raise ForceFieldUnitError(msg) # noqa: TRY200, B904 return system def _add_angles(self, system: openmm.System) -> openmm.System: - forces = self.force_field_terms["angle"] + forces = self.forcefield_terms["angle"] force_types = {i.force for i in forces} for force_type in force_types: if "Martini" in force_type: @@ -113,12 +114,12 @@ def _add_angles(self, system: openmm.System) -> openmm.System: ) except AttributeError: msg = f"{assigned_force} in angles does not have units" - raise ForceFieldUnitError(msg) + raise ForceFieldUnitError(msg) # noqa: TRY200, B904 return system def _add_torsions(self, system: openmm.System) -> openmm.System: - forces = self.force_field_terms["torsion"] + forces = self.forcefield_terms["torsion"] force_types = {i.force for i in forces} for force_type in force_types: if "Martini" in force_type: @@ -144,7 +145,7 @@ def _add_torsions(self, system: openmm.System) -> openmm.System: ) except AttributeError: msg = f"{assigned_force} in torsions does not have units" - raise ForceFieldUnitError(msg) + raise ForceFieldUnitError(msg) # noqa: TRY200, B904 return system @@ -153,7 +154,7 @@ def _add_nonbondeds(self, system: openmm.System) -> openmm.System: (i.get_atom1().get_id(), i.get_atom2().get_id()) for i in self.molecule.get_bonds() ] - forces = self.force_field_terms["nonbonded"] + forces = self.forcefield_terms["nonbonded"] force_types = {i.force for i in forces} for force_type in force_types: force_function = self._available_forces(force_type) @@ -175,7 +176,7 @@ def _add_nonbondeds(self, system: openmm.System) -> openmm.System: except AttributeError: msg = f"{assigned_force} in nonbondeds does not have units" - raise ForceFieldUnitError(msg) + raise ForceFieldUnitError(msg) # noqa: TRY200, B904 try: # This method MUST be after terms are assigned. @@ -185,10 +186,13 @@ def _add_nonbondeds(self, system: openmm.System) -> openmm.System: ) except OpenMMException: msg = f"{force_type} is missing a definition for a particle." - raise ForceFieldUnitError(msg) + raise ForceFieldUnitError(msg) # noqa: TRY200, B904 return system + def _add_atoms(self, system: openmm.System) -> openmm.System: + raise NotImplementedError + def _add_forces(self, system: openmm.System) -> openmm.System: system = self._add_bonds(system) system = self._add_angles(system) @@ -196,22 +200,31 @@ def _add_forces(self, system: openmm.System) -> openmm.System: return self._add_nonbondeds(system) def get_openmm_topology(self) -> app.topology.Topology: + """Return OpenMM.Topology object.""" raise NotImplementedError def get_openmm_system(self) -> openmm.System: + """Return OpenMM.System object.""" raise NotImplementedError @dataclass(frozen=True) class AssignedSystem(ForcedSystem): + """A system with forces assigned.""" + molecule: stk.Molecule - force_field_terms: dict[str, tuple] + forcefield_terms: dict[str, tuple] system_xml: pathlib.Path topology_xml: pathlib.Path bead_set: dict[str, CgBead] vdw_bond_cutoff: int mass: float = 10 + def _add_atoms(self, system: openmm.System) -> openmm.System: + for _atom in self.molecule.get_atoms(): + system.addParticle(self.mass) + return system + def _get_topology_xml_string(self, molecule: stk.Molecule) -> str: ff_str = "\n\n" @@ -251,7 +264,6 @@ def _get_topology_xml_string(self, molecule: stk.Molecule) -> str: re_str += " \n\n" ff_str += at_str - ff_str += re_str ff_str += "\n" return ff_str @@ -263,6 +275,7 @@ def _write_topology_xml(self, molecule: stk.Molecule) -> None: f.write(ff_str) def get_openmm_topology(self) -> app.topology.Topology: + """Return OpenMM.Topology object.""" topology = app.topology.Topology() chain = topology.addChain() residue = topology.addResidue(name="ALL", chain=chain) @@ -299,10 +312,9 @@ def get_openmm_topology(self) -> app.topology.Topology: return topology def get_openmm_system(self) -> openmm.System: - self._write_topology_xml(self.molecule) - forcefield = app.ForceField(self.topology_xml) - topology = self.get_openmm_topology() - system = forcefield.createSystem(topology) + """Return OpenMM.System object.""" + system = openmm.System() + system = self._add_atoms(system) system = self._add_forces(system) with open(self.system_xml, "w") as f: @@ -313,13 +325,10 @@ def get_openmm_system(self) -> openmm.System: @dataclass(frozen=True) class MartiniSystem(ForcedSystem): - """ - Assign a system using martini_openmm. - - """ + """Assign a system using martini_openmm.""" molecule: stk.Molecule - force_field_terms: dict[str, tuple] + forcefield_terms: dict[str, tuple] system_xml: pathlib.Path topology_itp: pathlib.Path vdw_bond_cutoff: int @@ -357,9 +366,9 @@ def _get_atoms_string( atoms_string += "\n" return atoms_string - def _get_bonds_string(self, molecule: stk.Molecule) -> str: + def _get_bonds_string(self) -> str: string = "[bonds]\n; i j funct length\n" - forces = self.force_field_terms["bond"] + forces = self.forcefield_terms["bond"] for assigned_force in forces: force_type = assigned_force.force if force_type != "MartiniDefinedBond": @@ -380,9 +389,9 @@ def _get_bonds_string(self, molecule: stk.Molecule) -> str: string += "\n" return string - def _get_angles_string(self, molecule: stk.Molecule) -> str: + def _get_angles_string(self) -> str: string = "[angles]\n; i j k funct ref.angle force_k\n" - forces = self.force_field_terms["angle"] + forces = self.forcefield_terms["angle"] for assigned_force in forces: force_type = assigned_force.force @@ -407,22 +416,22 @@ def _get_angles_string(self, molecule: stk.Molecule) -> str: string += "\n" return string - def _get_torsions_string(self, molecule: stk.Molecule) -> str: + def _get_torsions_string(self) -> str: string = "[dihedrals]\n; i j k l funct ref.angle force_k\n" - forces = self.force_field_terms["torsion"] + forces = self.forcefield_terms["torsion"] force_types = {i.force for i in forces} for force_type in force_types: if force_type != "MartiniDefinedTorsion": continue - print(force_type) - raise SystemExit() + print(force_type) # noqa: T201 + raise SystemExit string += "\n" return string - def _get_contraints_string(self, molecule: stk.Molecule) -> str: + def _get_contraints_string(self) -> str: string = "[constraints]\n; i j funct length\n" - for constraint in self.force_field_terms["constraints"]: + for constraint in self.forcefield_terms["constraints"]: string += ( f" {constraint[0]} {constraint[1]} {constraint[2]} " f"{constraint[3]} {constraint[4]}" @@ -430,9 +439,9 @@ def _get_contraints_string(self, molecule: stk.Molecule) -> str: string += "\n" return string - def _get_exclusions_string(self, molecule: stk.Molecule) -> str: + def _get_exclusions_string(self) -> str: string = "[exclusions]\n; i j funct length\n" - for constraint in self.force_field_terms["constraints"]: + for constraint in self.forcefield_terms["constraints"]: string += ( f" {constraint[0]} {constraint[1]} {constraint[2]} " f"{constraint[3]} {constraint[4]}" @@ -450,10 +459,10 @@ def _write_topology_itp(self, molecule: stk.Molecule) -> None: ) atoms_string = self._get_atoms_string(molecule, molecule_name) - bonds_string = self._get_bonds_string(molecule) - angles_string = self._get_angles_string(molecule) - torsions_string = self._get_torsions_string(molecule) - constraints_string = self._get_contraints_string(molecule) + bonds_string = self._get_bonds_string() + angles_string = self._get_angles_string() + torsions_string = self._get_torsions_string() + constraints_string = self._get_contraints_string() string += atoms_string string += bonds_string @@ -465,10 +474,12 @@ def _write_topology_itp(self, molecule: stk.Molecule) -> None: f.write(string) def get_openmm_topology(self) -> app.topology.Topology: + """Return OpenMM.Topology object.""" self._write_topology_itp(self.molecule) return MartiniTopology(self.topology_itp).get_openmm_topology() def get_openmm_system(self) -> openmm.System: + """Return OpenMM.System object.""" self._write_topology_itp(self.molecule) topology = MartiniTopology(self.topology_itp) system = topology.get_openmm_system() diff --git a/src/cgexplore/beads.py b/src/cgexplore/beads.py index 77cb2234..750b2489 100644 --- a/src/cgexplore/beads.py +++ b/src/cgexplore/beads.py @@ -14,6 +14,8 @@ @dataclass class CgBead: + """Define a coarse-grained bead.""" + element_string: str bead_type: str bead_class: str @@ -24,6 +26,7 @@ def get_cgbead_from_type( bead_type: str, bead_set: dict[str, CgBead], ) -> CgBead: + """Get CgBead class from matching to bead type.""" return bead_set[bead_type] @@ -31,6 +34,7 @@ def get_cgbead_from_element( estring: str, bead_set: dict[str, CgBead], ) -> CgBead: + """Get CgBead class from matching to element string.""" for i in bead_set: bead = bead_set[i] if bead.element_string == estring: @@ -40,6 +44,7 @@ def get_cgbead_from_element( def periodic_table() -> dict[str, int]: + """Periodic table of elements.""" return { "H": 1, "He": 2, @@ -142,11 +147,13 @@ def periodic_table() -> dict[str, int]: def string_to_atom_number(string: str) -> int: + """Convert atom string to atom number.""" return periodic_table()[string] def bead_library_check(bead_library: tuple[CgBead]) -> None: - print(bead_library) + """Check bead library for bad definitions.""" + logging.info(bead_library) logging.info(f"there are {len(bead_library)} beads") used_names = tuple(i.bead_class for i in bead_library) counts = Counter(used_names) diff --git a/src/cgexplore/bonds.py b/src/cgexplore/bonds.py index 55a7bf85..d9c2019f 100644 --- a/src/cgexplore/bonds.py +++ b/src/cgexplore/bonds.py @@ -1,4 +1,3 @@ -#!/usr/bin/env python # Distributed under the terms of the MIT License. """Module for handling bonds. @@ -7,8 +6,9 @@ """ -import itertools +import itertools as it import logging +from collections import abc from dataclasses import dataclass import stk @@ -20,12 +20,13 @@ ) -def bond_k_unit(): - return openmm.unit.kilojoules_per_mole / openmm.unit.nanometer**2 +_bond_k_unit = openmm.unit.kilojoules_per_mole / openmm.unit.nanometer**2 @dataclass class Bond: + """Class containing bond defintion.""" + atom_names: tuple[str, ...] atom_ids: tuple[int, ...] bond_r: openmm.unit.Quantity @@ -37,6 +38,8 @@ class Bond: @dataclass class TargetBond: + """Defines a target term to search for in a molecule.""" + type1: str type2: str element1: str @@ -45,19 +48,33 @@ class TargetBond: bond_k: openmm.unit.Quantity funct: int = 0 + def vector_key(self) -> str: + """Return key for vector defining this target term.""" + return f"{self.type1}{self.type2}" + + def vector(self) -> tuple[float, float]: + """Return vector defining this target term.""" + return ( + self.bond_r.value_in_unit(openmm.unit.angstrom), + self.bond_k.value_in_unit(_bond_k_unit), + ) + def human_readable(self) -> str: + """Return human-readable definition of this target term.""" return ( f"{self.__class__.__name__}(" f"{self.type1}{self.type2}, " f"{self.element1}{self.element2}, " f"{self.bond_r.in_units_of(openmm.unit.angstrom)}, " - f"{self.bond_k.in_units_of(bond_k_unit())}, " + f"{self.bond_k.in_units_of(_bond_k_unit)}, " ")" ) @dataclass class TargetBondRange: + """Defines a target term and ranges in parameters to search for.""" + type1: str type2: str element1: str @@ -65,8 +82,9 @@ class TargetBondRange: bond_rs: tuple[openmm.unit.Quantity] bond_ks: tuple[openmm.unit.Quantity] - def yield_bonds(self): - for r, k in itertools.product(self.bond_rs, self.bond_ks): + def yield_bonds(self) -> abc.Iterable[TargetBond]: + """Find interactions matching target.""" + for r, k in it.product(self.bond_rs, self.bond_ks): yield TargetBond( type1=self.type1, type2=self.type2, @@ -77,8 +95,43 @@ def yield_bonds(self): ) +@dataclass +class TargetPairedBondRange: + """Defines a target term and ranges in parameters to search for.""" + + type1s: tuple[str] + type2s: tuple[str] + element1s: tuple[str] + element2s: tuple[str] + bond_rs: tuple[openmm.unit.Quantity] + bond_ks: tuple[openmm.unit.Quantity] + + def yield_bonds(self) -> abc.Iterable[TargetBond]: + """Find interactions matching target.""" + raise NotImplementedError + for r, k in it.product(self.bond_rs, self.bond_ks): # type: ignore[unreachable] + for type1, type2, element1, element2 in zip( + self.type1s, + self.type2s, + self.element1s, + self.element2s, + strict=True, + ): + print(type1, type2, element1, element2) # noqa: T201 + yield TargetBond( + type1=type1, + type2=type2, + element1=element1, + element2=element2, + bond_k=k, + bond_r=r, + ) + + @dataclass class TargetMartiniBond: + """Defines a target angle to search for in a molecule.""" + type1: str type2: str element1: str @@ -88,19 +141,22 @@ class TargetMartiniBond: bond_k: openmm.unit.Quantity def human_readable(self) -> str: + """Return human-readable definition of this target term.""" return ( f"{self.__class__.__name__}(" f"{self.type1}{self.type2}, " f"{self.element1}{self.element2}, " f"{self.funct}," f"{self.bond_r.in_units_of(openmm.unit.angstrom)}, " - f"{self.bond_k.in_units_of(bond_k_unit())}, " + f"{self.bond_k.in_units_of(_bond_k_unit)}, " ")" ) @dataclass class MartiniBondRange: + """Defines a target bond and ranges in parameters to search for.""" + type1: str type2: str element1: str @@ -109,8 +165,9 @@ class MartiniBondRange: bond_rs: tuple[openmm.unit.Quantity] bond_ks: tuple[openmm.unit.Quantity] - def yield_bonds(self): - for r, k in itertools.product(self.bond_rs, self.bond_ks): + def yield_bonds(self) -> abc.Iterable[TargetMartiniBond]: + """Find bonds matching target.""" + for r, k in it.product(self.bond_rs, self.bond_ks): yield TargetMartiniBond( type1=self.type1, type2=self.type2, diff --git a/src/cgexplore/databases.py b/src/cgexplore/databases.py index 0ac05480..491ece44 100644 --- a/src/cgexplore/databases.py +++ b/src/cgexplore/databases.py @@ -1,3 +1,8 @@ +"""Module for databasing, with AtomLite. + +Author: Andrew Tarzia +""" + import logging import pathlib import sqlite3 @@ -16,6 +21,7 @@ class AtomliteDatabase: """Holds an atomlite database with some useful methods.""" def __init__(self, db_file: pathlib.Path) -> None: + """Initialize database.""" self._db_file = db_file self._db = atomlite.Database(db_file) @@ -24,6 +30,7 @@ def get_num_entries(self) -> int: return self._db.num_entries() def add_molecule(self, molecule: stk.Molecule, key: str) -> None: + """Add molecule to database as entry.""" entry = atomlite.Entry.from_rdkit( key=key, molecule=molecule.to_rdkit_mol(), @@ -34,15 +41,18 @@ def add_molecule(self, molecule: stk.Molecule, key: str) -> None: self._db.update_entries(entry) def get_entries(self) -> abc.Iterator[atomlite.Entry]: + """Get all entries.""" return self._db.get_entries() def get_entry(self, key: str) -> atomlite.Entry: + """Get specific entry.""" if not self._db.has_entry(key): msg = f"{key} not in database" raise RuntimeError(msg) return self._db.get_entry(key) # type: ignore[return-value] def get_molecule(self, key: str) -> stk.Molecule: + """Get a molecule.""" rdkit_molecule = atomlite.json_to_rdkit(self.get_entry(key).molecule) return stk.BuildingBlock.init_from_rdkit_mol(rdkit_molecule) @@ -51,6 +61,7 @@ def add_properties( key: str, property_dict: dict[str, atomlite.Json], ) -> None: + """Add properties to an entry by key.""" self._db.update_properties( atomlite.PropertyEntry(key=key, properties=property_dict) ) @@ -61,6 +72,7 @@ def get_property( property_key: str, property_type: type, ) -> atomlite.Json: + """Get the properties of an entry.""" if property_type is bool: value = self._db.get_bool_property( # type: ignore[assignment] key=key, @@ -94,9 +106,11 @@ def get_property( return value # type: ignore[return-value] def has_molecule(self, key: str) -> bool: + """Check if database has a molecule by key.""" return bool(self._db.has_entry(key)) def remove_entry(self, key: str) -> None: + """Remove an entry by key.""" self._db.remove_entries(keys=key) def keep_if( @@ -104,6 +118,7 @@ def keep_if( column: str, value: str | float, ) -> abc.Iterator[atomlite.Entry]: + """Filter database entries by properties.""" for entry in self.get_entries(): if entry.properties[column] == value: yield entry diff --git a/src/cgexplore/ensembles.py b/src/cgexplore/ensembles.py index ed687765..6db232d2 100644 --- a/src/cgexplore/ensembles.py +++ b/src/cgexplore/ensembles.py @@ -1,4 +1,3 @@ -#!/usr/bin/env python # Distributed under the terms of the MIT License. """Module for ensemble and trajectory classes. @@ -8,7 +7,7 @@ """ import json -import os +import pathlib from collections import abc from dataclasses import dataclass @@ -18,6 +17,8 @@ @dataclass class Conformer: + """Define conformer information.""" + molecule: stk.Molecule energy_decomposition: dict conformer_id: int | None = None @@ -26,30 +27,35 @@ class Conformer: @dataclass class Timestep: + """Define timestep information.""" + molecule: stk.Molecule timestep: float class Ensemble: - def __init__( + """Class to contain ensemble information.""" + + def __init__( # noqa: PLR0913 self, base_molecule: stk.Molecule, base_mol_path: str, conformer_xyz: str, data_json: str, - overwrite: bool, + overwrite: bool, # noqa: FBT001 ) -> None: + """Initialize Ensemble class.""" self._base_molecule = base_molecule self._molecule_num_atoms = base_molecule.get_num_atoms() - self._base_mol_path = base_mol_path - self._conformer_xyz = conformer_xyz - self._data_json = data_json + self._base_mol_path = pathlib.Path(base_mol_path) + self._conformer_xyz = pathlib.Path(conformer_xyz) + self._data_json = pathlib.Path(data_json) if overwrite: self._base_molecule.write(self._base_mol_path) - if os.path.exists(self._conformer_xyz): - os.remove(self._conformer_xyz) - if os.path.exists(self._data_json): - os.remove(self._data_json) + if self._conformer_xyz.exists(): + self._conformer_xyz.unlink() + if self._data_json.exists(): + self._data_json.unlink() self._data = {} self._trajectory = {} else: @@ -57,9 +63,11 @@ def __init__( self._trajectory = self.load_trajectory() def get_num_conformers(self) -> int: + """Get number of conformers in ensemble.""" return len(self._data) def write_conformers_to_file(self) -> None: + """Write conformers to xyz file.""" with open(self._conformer_xyz, "w") as f: for conf in self._trajectory: xyz_string = self._trajectory[conf] @@ -68,12 +76,15 @@ def write_conformers_to_file(self) -> None: json.dump(self._data, f, indent=4) def add_conformer(self, conformer: Conformer, source: str) -> None: + """Add a conformer to ensemble.""" if conformer.conformer_id is None: conf_id = self.get_num_conformers() else: conf_id = conformer.conformer_id - assert conf_id not in self._trajectory + if conf_id in self._trajectory: + msg = f"{conf_id} is already in trajectory" + raise RuntimeError(msg) conformer_label = f"conf {conf_id}" # Add structure to XYZ file. @@ -91,6 +102,7 @@ def add_conformer(self, conformer: Conformer, source: str) -> None: self._data[conf_id] = conf_data def load_trajectory(self) -> dict[int, list[str]]: + """Load trajectory.""" num_atoms = self._molecule_num_atoms trajectory = {} with open(self._conformer_xyz) as f: @@ -139,10 +151,12 @@ def _yield_from_xyz(self) -> abc.Iterator[Conformer]: ) def yield_conformers(self) -> abc.Iterator[Conformer]: + """Yield conformers.""" for conf_id in self._trajectory: yield self.get_conformer(conf_id) def get_lowest_e_conformer(self) -> Conformer: + """Get lowest energy conformer.""" temp_energy_conformerid: int | str try: temp_energy_conformerid = 0 @@ -164,6 +178,7 @@ def get_lowest_e_conformer(self) -> Conformer: return self.get_conformer(min_energy_conformerid) def get_conformer(self, conf_id: int | str) -> Conformer: + """Get a specific conformer.""" if conf_id not in self._data: if str(conf_id) not in self._data: msg = ( @@ -172,8 +187,8 @@ def get_conformer(self, conf_id: int | str) -> Conformer: f"coming. Current types: {[type(i) for i in self._data]}" ) raise ValueError(msg) - else: - conf_id = str(conf_id) + + conf_id = str(conf_id) conf_lines = self._trajectory[int(conf_id)] extracted_id = int(conf_lines[1].strip().split()[1]) @@ -205,12 +220,15 @@ def get_conformer(self, conf_id: int | str) -> Conformer: ) def get_base_molecule(self) -> stk.Molecule: + """Get the base molecule defining the ensemble.""" return self._base_molecule def get_molecule_num_atoms(self) -> int: + """Get the number of atoms in the molecule in the ensemble.""" return self._molecule_num_atoms def load_data(self) -> dict[int, dict]: + """Load ensemble data.""" try: with open(self._data_json) as f: return json.load(f) @@ -218,10 +236,12 @@ def load_data(self) -> dict[int, dict]: return {} def __str__(self) -> str: + """Return a string representation of the Ensemble.""" return ( f"{self.__class__.__name__}(" f"num_confs={self.get_num_conformers()})" ) def __repr__(self) -> str: + """Return a string representation of the Ensemble.""" return str(self) diff --git a/src/cgexplore/errors.py b/src/cgexplore/errors.py index 61571510..29ce84ad 100644 --- a/src/cgexplore/errors.py +++ b/src/cgexplore/errors.py @@ -1,4 +1,3 @@ -#!/usr/bin/env python # Distributed under the terms of the MIT License. """Module for containing exceptions. @@ -17,8 +16,8 @@ class ForceFieldUnitError(Exception): - pass + """Error found in units of forcefield term.""" class ForceFieldUnavailableError(Exception): - pass + """Error found assigning forcefield term.""" diff --git a/src/cgexplore/forcefield.py b/src/cgexplore/forcefield.py index 9b5a9be1..e82596c8 100644 --- a/src/cgexplore/forcefield.py +++ b/src/cgexplore/forcefield.py @@ -1,4 +1,3 @@ -#!/usr/bin/env python # Distributed under the terms of the MIT License. """Module for containing forcefields. @@ -7,9 +6,10 @@ """ -import itertools +import itertools as it import logging import pathlib +from collections import abc from heapq import nsmallest import numpy as np @@ -43,99 +43,10 @@ ) -class ForceFieldLibrary: - def __init__( - self, - bead_library: tuple[CgBead], - vdw_bond_cutoff: int, - prefix: str, - ) -> None: - self._bead_library = bead_library - self._vdw_bond_cutoff = vdw_bond_cutoff - self._prefix = prefix - self._bond_ranges: tuple = () - self._angle_ranges: tuple = () - self._torsion_ranges: tuple = () - self._nonbonded_ranges: tuple = () - - def add_bond_range(self, bond_range: tuple) -> None: - self._bond_ranges += (bond_range,) - - def add_angle_range(self, angle_range: tuple) -> None: - self._angle_ranges += (angle_range,) - - def add_torsion_range(self, torsion_range: tuple) -> None: - self._torsion_ranges += (torsion_range,) - - def add_nonbonded_range(self, nonbonded_range: tuple) -> None: - self._nonbonded_ranges += (nonbonded_range,) - - def _get_iterations(self) -> list: - iterations = [] - for bond_range in self._bond_ranges: - iterations.append(tuple(bond_range.yield_bonds())) - - for angle_range in self._angle_ranges: - iterations.append(tuple(angle_range.yield_angles())) - - for torsion_range in self._torsion_ranges: - iterations.append(tuple(torsion_range.yield_torsions())) - - for nonbonded_range in self._nonbonded_ranges: - iterations.append(tuple(nonbonded_range.yield_nonbondeds())) - return iterations - - def yield_forcefields(self): - iterations = self._get_iterations() - - for i, parameter_set in enumerate(itertools.product(*iterations)): - bond_terms = tuple( - i for i in parameter_set if "Bond" in i.__class__.__name__ - ) - angle_terms = tuple( - i - for i in parameter_set - if "Angle" in i.__class__.__name__ - # and "Pyramid" not in i.__class__.__name__ - ) - torsion_terms = tuple( - i - for i in parameter_set - if "Torsion" in i.__class__.__name__ - # if len(i.search_string) == 4 - ) - nonbonded_terms = tuple( - i for i in parameter_set if "Nonbonded" in i.__class__.__name__ - ) - - yield ForceField( - identifier=str(i), - prefix=self._prefix, - present_beads=self._bead_library, - bond_targets=bond_terms, - angle_targets=angle_terms, - torsion_targets=torsion_terms, - nonbonded_targets=nonbonded_terms, - vdw_bond_cutoff=self._vdw_bond_cutoff, - ) - - def __str__(self) -> str: - return ( - f"{self.__class__.__name__}(\n" - f" bead_library={self._bead_library},\n" - f" bond_ranges={self._bond_ranges},\n" - f" angle_ranges={self._angle_ranges},\n" - f" torsion_ranges={self._torsion_ranges},\n" - f" nonbonded_ranges={self._nonbonded_ranges}" - "\n)" - ) - - def __repr__(self) -> str: - return str(self) - - class ForceField: - def __init__( + """Define a CG ForceField.""" + + def __init__( # noqa: PLR0913: self, identifier: str, prefix: str, @@ -146,25 +57,29 @@ def __init__( nonbonded_targets: tuple[TargetNonbonded, ...], vdw_bond_cutoff: int, ) -> None: + """Initialize ForceField.""" # Check if you should use MartiniForceField. for bt in bond_targets: if isinstance(bt, TargetMartiniBond): - raise TypeError( + msg = ( f"{bt} is Martini type, probably use " "MartiniForceFieldLibrary/MartiniForceField" ) + raise TypeError(msg) for at in angle_targets: if isinstance(at, TargetMartiniAngle): - raise TypeError( + msg = ( f"{at} is Martini type, probably use " "MartiniForceFieldLibrary/MartiniForceField" ) + raise TypeError(msg) for tt in torsion_targets: if isinstance(tt, TargetMartiniTorsion): - raise TypeError( + msg = ( f"{tt} is Martini type, probably use " "MartiniForceFieldLibrary/MartiniForceField" ) + raise TypeError(msg) self._identifier = identifier self._prefix = prefix @@ -200,10 +115,10 @@ def _assign_bond_terms(self, molecule: stk.Molecule) -> tuple: continue assigned.add(cgbead_string) assigned.add(tuple(reversed(cgbead_string))) - try: - assert isinstance(target_term.bond_r, openmm.unit.Quantity) - assert isinstance(target_term.bond_k, openmm.unit.Quantity) - except AssertionError: + if not isinstance(target_term.bond_r, openmm.unit.Quantity): + msg = f"{target_term} in bonds does not have units" + raise ForceFieldUnitError(msg) + if not isinstance(target_term.bond_k, openmm.unit.Quantity): msg = f"{target_term} in bonds does not have units" raise ForceFieldUnitError(msg) @@ -229,13 +144,12 @@ def _assign_bond_terms(self, molecule: stk.Molecule) -> tuple: ) ) - logging.info( - "unassigned bond terms: " - f"{sorted((i for i in found if i not in assigned))}" - ) + unassigned = sorted(i for i in found if i not in assigned) + if len(unassigned) > 0: + logging.info(f"unassigned bond terms: {unassigned}") return tuple(bond_terms) - def _assign_angle_terms( + def _assign_angle_terms( # noqa: C901, PLR0912, PLR0915 self, molecule: stk.Molecule, ) -> tuple[Angle | CosineAngle, ...]: @@ -276,14 +190,17 @@ def _assign_angle_terms( assigned.add(cgbead_string) assigned.add(tuple(reversed(cgbead_string))) if isinstance(target_angle, TargetAngle): - try: - assert isinstance( - target_angle.angle, openmm.unit.Quantity - ) - assert isinstance( - target_angle.angle_k, openmm.unit.Quantity + if not isinstance( + target_angle.angle, openmm.unit.Quantity + ): + msg = ( + f"{target_angle} in angles does not have units for" + " parameters" ) - except AssertionError: + raise ForceFieldUnitError(msg) + if not isinstance( + target_angle.angle_k, openmm.unit.Quantity + ): msg = ( f"{target_angle} in angles does not have units for" " parameters" @@ -306,11 +223,11 @@ def _assign_angle_terms( angle_k=target_angle.angle_k, force="HarmonicAngleForce", ) - if central_bead.coordination == 4: + if central_bead.coordination == 4: # noqa: PLR2004 if central_name not in pyramid_angles: pyramid_angles[central_name] = [] pyramid_angles[central_name].append(actual_angle) - elif central_bead.coordination == 6: + elif central_bead.coordination == 6: # noqa: PLR2004 if central_name not in octahedral_angles: octahedral_angles[central_name] = [] octahedral_angles[central_name].append(actual_angle) @@ -318,11 +235,9 @@ def _assign_angle_terms( angle_terms.append(actual_angle) elif isinstance(target_angle, TargetCosineAngle): - try: - assert isinstance( - target_angle.angle_k, openmm.unit.Quantity - ) - except AssertionError: + if not isinstance( + target_angle.angle_k, openmm.unit.Quantity + ): msg = ( f"{target_angle} in angles does not have units for" " parameters" @@ -350,14 +265,17 @@ def _assign_angle_terms( ) elif isinstance(target_angle, TargetMartiniAngle): - try: - assert isinstance( - target_angle.angle, openmm.unit.Quantity - ) - assert isinstance( - target_angle.angle_k, openmm.unit.Quantity + if not isinstance( + target_angle.angle, openmm.unit.Quantity + ): + msg = ( + f"{target_angle} in angles does not have units for" + " parameters" ) - except AssertionError: + raise ForceFieldUnitError(msg) + if not isinstance( + target_angle.angle_k, openmm.unit.Quantity + ): msg = ( f"{target_angle} in angles does not have units for" " parameters" @@ -472,11 +390,9 @@ def _assign_angle_terms( force="HarmonicAngleForce", ), ) - - logging.info( - "unassigned angle terms: " - f"{sorted((i for i in found if i not in assigned))}" - ) + unassigned = sorted(i for i in found if i not in assigned) + if len(unassigned) > 0: + logging.info(f"unassigned angle terms: {unassigned}") return tuple(angle_terms) def _assign_torsion_terms( @@ -484,7 +400,6 @@ def _assign_torsion_terms( molecule: stk.Molecule, ) -> tuple: torsion_terms = [] - found = set() assigned = set() # Iterate over the different path lengths, and find all torsions @@ -500,8 +415,6 @@ def _assign_torsion_terms( for i in atom_estrings ] cgbead_string = tuple(i.bead_type for i in cgbeads) - found.add(cgbead_string) - found.add(tuple(reversed(cgbead_string))) for target_torsion in self._torsion_targets: if target_torsion.search_string not in ( cgbead_string, @@ -511,14 +424,16 @@ def _assign_torsion_terms( assigned.add(cgbead_string) assigned.add(tuple(reversed(cgbead_string))) - try: - assert isinstance( - target_torsion.phi0, openmm.unit.Quantity - ) - assert isinstance( - target_torsion.torsion_k, openmm.unit.Quantity + if not isinstance( + target_torsion.phi0, openmm.unit.Quantity + ): + msg = ( + f"{target_torsion} in torsions does not have units" ) - except AssertionError: + raise ForceFieldUnitError(msg) + if not isinstance( + target_torsion.torsion_k, openmm.unit.Quantity + ): msg = ( f"{target_torsion} in torsions does not have units" ) @@ -550,10 +465,11 @@ def _assign_torsion_terms( ) ) - logging.info( - "unassigned torsion terms: " - f"{sorted((i for i in found if i not in assigned))}" - ) + if len(assigned) > 0: + logging.info( + f"assigned torsion terms: {sorted(assigned)} " + f"({len(self._torsion_targets)} targets) " + ) return tuple(torsion_terms) def _assign_nonbonded_terms( @@ -561,20 +477,21 @@ def _assign_nonbonded_terms( molecule: stk.Molecule, ) -> tuple: nonbonded_terms = [] + found = set() + assigned = set() for atom in molecule.get_atoms(): atom_estring = atom.__class__.__name__ cgbead = get_cgbead_from_element(atom_estring, self.get_bead_set()) - + found.add(cgbead.bead_class) for target_term in self._nonbonded_targets: if target_term.bead_class != cgbead.bead_class: continue - try: - assert isinstance(target_term.sigma, openmm.unit.Quantity) - assert isinstance( - target_term.epsilon, openmm.unit.Quantity - ) - except AssertionError: + assigned.add(target_term.bead_class) + if not isinstance(target_term.sigma, openmm.unit.Quantity): + msg = f"{target_term} in nonbondeds does not have units" + raise ForceFieldUnitError(msg) + if not isinstance(target_term.epsilon, openmm.unit.Quantity): msg = f"{target_term} in nonbondeds does not have units" raise ForceFieldUnitError(msg) @@ -588,7 +505,9 @@ def _assign_nonbonded_terms( force=target_term.force, ) ) - + unassigned = sorted(i for i in found if i not in assigned) + if len(unassigned) > 0: + logging.info(f"unassigned nonbonded terms: {unassigned}") return tuple(nonbonded_terms) def assign_terms( @@ -597,6 +516,7 @@ def assign_terms( name: str, output_dir: pathlib.Path, ) -> ForcedSystem: + """Assign forcefield terms to molecule.""" assigned_terms = { "bond": self._assign_bond_terms(molecule), "angle": self._assign_angle_terms(molecule), @@ -612,7 +532,7 @@ def assign_terms( return AssignedSystem( molecule=molecule, - force_field_terms=assigned_terms, + forcefield_terms=assigned_terms, system_xml=( output_dir / f"{name}_{self._prefix}_{self._identifier}_syst.xml" @@ -626,21 +546,27 @@ def assign_terms( ) def get_bead_set(self) -> dict[str, CgBead]: + """Get beads in forcefield.""" return {i.bead_type: i for i in self._present_beads} def get_identifier(self) -> str: + """Get forcefield identifier.""" return self._identifier def get_prefix(self) -> str: + """Get forcefield prefix.""" return self._prefix def get_present_beads(self) -> tuple: + """Get beads present.""" return self._present_beads def get_vdw_bond_cutoff(self) -> int: + """Get vdW bond cutoff of forcefield.""" return self._vdw_bond_cutoff def get_targets(self) -> dict: + """Get targets of forcefield.""" return { "bonds": self._bond_targets, "angles": self._angle_targets, @@ -649,9 +575,11 @@ def get_targets(self) -> dict: } def get_hr_file_name(self) -> str: + """Get human-readable file name.""" return f"{self._hrprefix}_{self._prefix}_{self._identifier}.txt" - def write_human_readable(self, output_dir) -> None: + def write_human_readable(self, output_dir: pathlib.Path) -> None: + """Write forcefield definition to human-readable file.""" with open(output_dir / self.get_hr_file_name(), "w") as f: f.write(f"prefix: {self._prefix}\n") f.write(f"identifier: {self._identifier}\n") @@ -677,6 +605,7 @@ def write_human_readable(self, output_dir) -> None: f.write(f"{nt.human_readable()} \n") def __str__(self) -> str: + """Return a string representation of the Ensemble.""" return ( f"{self.__class__.__name__}(\n" f" present_beads={self._present_beads}, \n" @@ -688,41 +617,126 @@ def __str__(self) -> str: ) def __repr__(self) -> str: + """Return a string representation of the Ensemble.""" return str(self) -class MartiniForceFieldLibrary(ForceFieldLibrary): +class MartiniForceField(ForceField): + """Class defining a Martini Forcefield.""" + + def __init__( # noqa: PLR0913 + self, + identifier: str, + prefix: str, + present_beads: tuple[CgBead, ...], + bond_targets: tuple[TargetBond | TargetMartiniBond, ...], + angle_targets: tuple[TargetAngle | TargetMartiniAngle, ...], + torsion_targets: tuple[TargetTorsion | TargetMartiniTorsion, ...], + constraints: tuple[tuple], + vdw_bond_cutoff: int, + ) -> None: + """Initialize MartiniForceField.""" + self._identifier = identifier + self._prefix = prefix + self._present_beads = present_beads + self._bond_targets = bond_targets + self._angle_targets = angle_targets + self._torsion_targets = torsion_targets + self._vdw_bond_cutoff = vdw_bond_cutoff + self._constraints = constraints + self._hrprefix = "mffhr" + + def assign_terms( + self, + molecule: stk.Molecule, + name: str, + output_dir: pathlib.Path, + ) -> ForcedSystem: + """Assign forcefield terms to molecule.""" + assigned_terms = { + "bond": self._assign_bond_terms(molecule), + "angle": self._assign_angle_terms(molecule), + "torsion": self._assign_torsion_terms(molecule), + "nonbonded": (), + "constraints": self._constraints, + } + + return MartiniSystem( + molecule=molecule, + forcefield_terms=assigned_terms, + system_xml=( + output_dir + / f"{name}_{self._prefix}_{self._identifier}_syst.xml" + ), + topology_itp=( + output_dir + / f"{name}_{self._prefix}_{self._identifier}_topo.itp" + ), + bead_set=self.get_bead_set(), + vdw_bond_cutoff=self._vdw_bond_cutoff, + ) + + +class ForceFieldLibrary: + """Define a library of forcefields with varying parameters.""" + def __init__( self, bead_library: tuple[CgBead], vdw_bond_cutoff: int, prefix: str, ) -> None: + """Initialize ForceFieldLibrary.""" self._bead_library = bead_library self._vdw_bond_cutoff = vdw_bond_cutoff self._prefix = prefix self._bond_ranges: tuple = () self._angle_ranges: tuple = () self._torsion_ranges: tuple = () - self._constraints: tuple = () + self._nonbonded_ranges: tuple = () + + def add_bond_range(self, bond_range: tuple) -> None: + """Add a range of terms to library.""" + self._bond_ranges += (bond_range,) + + def add_angle_range(self, angle_range: tuple) -> None: + """Add a range of terms to library.""" + self._angle_ranges += (angle_range,) + + def add_torsion_range(self, torsion_range: tuple) -> None: + """Add a range of terms to library.""" + self._torsion_ranges += (torsion_range,) + + def add_nonbonded_range(self, nonbonded_range: tuple) -> None: + """Add a range of terms to library.""" + self._nonbonded_ranges += (nonbonded_range,) def _get_iterations(self) -> list: iterations = [] for bond_range in self._bond_ranges: - iterations.append(tuple(bond_range.yield_bonds())) + iterations.append(tuple(bond_range.yield_bonds())) # noqa: PERF401 for angle_range in self._angle_ranges: - iterations.append(tuple(angle_range.yield_angles())) + iterations.append( # noqa: PERF401 + tuple(angle_range.yield_angles()) + ) for torsion_range in self._torsion_ranges: - iterations.append(tuple(torsion_range.yield_torsions())) + iterations.append( # noqa: PERF401 + tuple(torsion_range.yield_torsions()) + ) + for nonbonded_range in self._nonbonded_ranges: + iterations.append( # noqa: PERF401 + tuple(nonbonded_range.yield_nonbondeds()) + ) return iterations - def yield_forcefields(self): + def yield_forcefields(self) -> abc.Iterable[ForceField]: + """Yield the forcefields in the library.""" iterations = self._get_iterations() - for i, parameter_set in enumerate(itertools.product(*iterations)): + for i, parameter_set in enumerate(it.product(*iterations)): bond_terms = tuple( i for i in parameter_set if "Bond" in i.__class__.__name__ ) @@ -738,78 +752,115 @@ def yield_forcefields(self): if "Torsion" in i.__class__.__name__ # if len(i.search_string) == 4 ) - yield MartiniForceField( + nonbonded_terms = tuple( + i for i in parameter_set if "Nonbonded" in i.__class__.__name__ + ) + + yield ForceField( identifier=str(i), prefix=self._prefix, present_beads=self._bead_library, bond_targets=bond_terms, angle_targets=angle_terms, torsion_targets=torsion_terms, - constraints=self._constraints, + nonbonded_targets=nonbonded_terms, vdw_bond_cutoff=self._vdw_bond_cutoff, ) def __str__(self) -> str: + """Return a string representation of the Ensemble.""" return ( f"{self.__class__.__name__}(\n" f" bead_library={self._bead_library},\n" f" bond_ranges={self._bond_ranges},\n" f" angle_ranges={self._angle_ranges},\n" f" torsion_ranges={self._torsion_ranges},\n" + f" nonbonded_ranges={self._nonbonded_ranges}" "\n)" ) def __repr__(self) -> str: + """Return a string representation of the Ensemble.""" return str(self) -class MartiniForceField(ForceField): +class MartiniForceFieldLibrary(ForceFieldLibrary): + """Define a library of forcefields with varying parameters.""" + def __init__( self, - identifier: str, - prefix: str, - present_beads: tuple[CgBead, ...], - bond_targets: tuple[TargetBond | TargetMartiniBond, ...], - angle_targets: tuple[TargetAngle | TargetMartiniAngle, ...], - torsion_targets: tuple[TargetTorsion | TargetMartiniTorsion, ...], - constraints: tuple[tuple], + bead_library: tuple[CgBead], vdw_bond_cutoff: int, + prefix: str, ) -> None: - self._identifier = identifier - self._prefix = prefix - self._present_beads = present_beads - self._bond_targets = bond_targets - self._angle_targets = angle_targets - self._torsion_targets = torsion_targets + """Initialize MartiniForceFieldLibrary.""" + self._bead_library = bead_library self._vdw_bond_cutoff = vdw_bond_cutoff - self._constraints = constraints - self._hrprefix = "mffhr" + self._prefix = prefix + self._bond_ranges: tuple = () + self._angle_ranges: tuple = () + self._torsion_ranges: tuple = () + self._constraints: tuple[tuple] = () # type: ignore[assignment] - def assign_terms( - self, - molecule: stk.Molecule, - name: str, - output_dir: pathlib.Path, - ) -> ForcedSystem: - assigned_terms = { - "bond": self._assign_bond_terms(molecule), - "angle": self._assign_angle_terms(molecule), - "torsion": self._assign_torsion_terms(molecule), - "nonbonded": (), - "constraints": self._constraints, - } + def _get_iterations(self) -> list: + iterations = [] + for bond_range in self._bond_ranges: + iterations.append(tuple(bond_range.yield_bonds())) # noqa: PERF401 - return MartiniSystem( - molecule=molecule, - force_field_terms=assigned_terms, - system_xml=( - output_dir - / f"{name}_{self._prefix}_{self._identifier}_syst.xml" - ), - topology_itp=( - output_dir - / f"{name}_{self._prefix}_{self._identifier}_topo.itp" - ), - bead_set=self.get_bead_set(), - vdw_bond_cutoff=self._vdw_bond_cutoff, + for angle_range in self._angle_ranges: + iterations.append( # noqa: PERF401 + tuple(angle_range.yield_angles()) + ) + + for torsion_range in self._torsion_ranges: + iterations.append( # noqa: PERF401 + tuple(torsion_range.yield_torsions()) + ) + + return iterations + + def yield_forcefields(self) -> abc.Iterable[MartiniForceField]: + """Yield forcefields from library.""" + iterations = self._get_iterations() + + for i, parameter_set in enumerate(it.product(*iterations)): + bond_terms = tuple( + i for i in parameter_set if "Bond" in i.__class__.__name__ + ) + angle_terms = tuple( + i + for i in parameter_set + if "Angle" in i.__class__.__name__ + # and "Pyramid" not in i.__class__.__name__ + ) + torsion_terms = tuple( + i + for i in parameter_set + if "Torsion" in i.__class__.__name__ + # if len(i.search_string) == 4 + ) + yield MartiniForceField( + identifier=str(i), + prefix=self._prefix, + present_beads=self._bead_library, + bond_targets=bond_terms, + angle_targets=angle_terms, + torsion_targets=torsion_terms, + constraints=self._constraints, + vdw_bond_cutoff=self._vdw_bond_cutoff, + ) + + def __str__(self) -> str: + """Return a string representation of the Ensemble.""" + return ( + f"{self.__class__.__name__}(\n" + f" bead_library={self._bead_library},\n" + f" bond_ranges={self._bond_ranges},\n" + f" angle_ranges={self._angle_ranges},\n" + f" torsion_ranges={self._torsion_ranges},\n" + "\n)" ) + + def __repr__(self) -> str: + """Return a string representation of the Ensemble.""" + return str(self) diff --git a/src/cgexplore/generation_utilities.py b/src/cgexplore/generation_utilities.py index cd1d827c..16937c6e 100644 --- a/src/cgexplore/generation_utilities.py +++ b/src/cgexplore/generation_utilities.py @@ -1,4 +1,3 @@ -#!/usr/bin/env python # Distributed under the terms of the MIT License. """Module for structure generation utilities. @@ -9,7 +8,6 @@ import logging -import os import pathlib from collections.abc import Iterator @@ -17,8 +15,8 @@ import stk from openmm import OpenMMException, openmm -from .angles import Angle -from .assigned_system import AssignedSystem +from .angles import Angle, CosineAngle +from .assigned_system import AssignedSystem, MartiniSystem from .beads import periodic_table from .bonds import Bond from .ensembles import Conformer, Ensemble @@ -35,7 +33,7 @@ def optimise_ligand( molecule: stk.Molecule, name: str, output_dir: pathlib.Path, - force_field: ForceField, + forcefield: ForceField, platform: str | None, ) -> stk.Molecule: """Optimise a building block. @@ -52,7 +50,7 @@ def optimise_ligand( output_dir: Directory to save outputs of optimisation process. - force_field: + forcefield: Define the forces used in the molecule. platform: @@ -64,13 +62,13 @@ def optimise_ligand( An stk molecule. """ - opt1_mol_file = os.path.join(output_dir, f"{name}_opted1.mol") + opt1_mol_file = pathlib.Path(output_dir) / f"{name}_opted1.mol" - if os.path.exists(opt1_mol_file): - molecule = molecule.with_structure_from_file(opt1_mol_file) + if opt1_mol_file.exists(): + molecule = molecule.with_structure_from_file(str(opt1_mol_file)) else: logging.info(f"optimising {name}, no max_iterations") - assigned_system = force_field.assign_terms( + assigned_system = forcefield.assign_terms( molecule=molecule, name=name, output_dir=output_dir, @@ -83,11 +81,16 @@ def optimise_ligand( molecule = opt.optimize(assigned_system) molecule = molecule.with_centroid(np.array((0, 0, 0))) molecule.write(opt1_mol_file) + energy_decomp = opt.read_final_energy_decomposition() + logging.info("optimised with energy:") + for i in energy_decomp: + e, u = energy_decomp[i] + logging.info(f"{i}: {round(e, 2)} [{u}]") return molecule -def soften_force_field( +def soften_forcefield( assigned_system: AssignedSystem, bond_ff_scale: float, angle_ff_scale: float, @@ -113,7 +116,40 @@ def soften_force_field( New assigned system. """ - new_force_field_terms = { + if isinstance(assigned_system, MartiniSystem): + msg = "soften not available for Martini yet." + raise TypeError(msg) + + angles = [] + for i in assigned_system.forcefield_terms["angle"]: + if isinstance(i, Angle): + angles.append( + Angle( + atom_names=i.atom_names, + atom_ids=i.atom_ids, + angle=i.angle, + angle_k=i.angle_k / angle_ff_scale, + atoms=i.atoms, + force=i.force, + ) + ) + elif isinstance(i, CosineAngle): + angles.append( + CosineAngle( # type: ignore[arg-type] + atom_names=i.atom_names, + atom_ids=i.atom_ids, + n=i.n, + b=i.b, + angle_k=i.angle_k / angle_ff_scale, + atoms=i.atoms, + force=i.force, + ) + ) + else: + msg = f"soften not available for {i} angle type." + raise TypeError(msg) + + new_forcefield_terms = { "bond": tuple( Bond( atom_names=i.atom_names, @@ -123,21 +159,11 @@ def soften_force_field( atoms=i.atoms, force=i.force, ) - for i in assigned_system.force_field_terms["bond"] - ), - "angle": tuple( - Angle( - atom_names=i.atom_names, - atom_ids=i.atom_ids, - angle=i.angle, - angle_k=i.angle_k / angle_ff_scale, - atoms=i.atoms, - force=i.force, - ) - for i in assigned_system.force_field_terms["angle"] + for i in assigned_system.forcefield_terms["bond"] ), + "angle": tuple(angles), "torsion": (), - "nonbonded": assigned_system.force_field_terms["nonbonded"], + "nonbonded": assigned_system.forcefield_terms["nonbonded"], } new_system_xml = pathlib.Path( @@ -148,7 +174,7 @@ def soften_force_field( return AssignedSystem( molecule=assigned_system.molecule, - force_field_terms=new_force_field_terms, + forcefield_terms=new_forcefield_terms, system_xml=new_system_xml, topology_xml=assigned_system.topology_xml, bead_set=assigned_system.bead_set, @@ -156,7 +182,7 @@ def soften_force_field( ) -def run_soft_md_cycle( +def run_soft_md_cycle( # noqa: PLR0913 name: str, assigned_system: AssignedSystem, output_dir: pathlib.Path, @@ -229,7 +255,7 @@ def run_soft_md_cycle( An OMMTrajectory containing the data and conformers. """ - soft_assigned_system = soften_force_field( + soft_assigned_system = soften_forcefield( assigned_system=assigned_system, bond_ff_scale=bond_ff_scale, angle_ff_scale=angle_ff_scale, @@ -256,7 +282,7 @@ def run_soft_md_cycle( return None -def run_constrained_optimisation( +def run_constrained_optimisation( # noqa: PLR0913 assigned_system: AssignedSystem, name: str, output_dir: pathlib.Path, @@ -298,7 +324,7 @@ def run_constrained_optimisation( An stk molecule. """ - soft_assigned_system = soften_force_field( + soft_assigned_system = soften_forcefield( assigned_system=assigned_system, bond_ff_scale=bond_ff_scale, angle_ff_scale=angle_ff_scale, @@ -325,7 +351,7 @@ def run_constrained_optimisation( return constrained_opt.optimize(soft_assigned_system) -def run_optimisation( +def run_optimisation( # noqa: PLR0913 assigned_system: AssignedSystem, name: str, file_suffix: str, @@ -416,14 +442,11 @@ def yield_near_models( ff_name = [i for i in name.split("_") if "f" in i][-1] for new_ff_id in neighbour_library: - if new_ff_id < 0: - continue - new_name = name.replace(ff_name, f"f{new_ff_id}") - new_fina_mol_file = os.path.join(output_dir, f"{new_name}_final.mol") - if os.path.exists(new_fina_mol_file): + new_fina_mol_file = pathlib.Path(output_dir) / f"{new_name}_final.mol" + if new_fina_mol_file.exists(): logging.info(f"found neigh: {new_fina_mol_file}") - yield molecule.with_structure_from_file(new_fina_mol_file) + yield molecule.with_structure_from_file(str(new_fina_mol_file)) def shift_beads( @@ -467,7 +490,7 @@ def shift_beads( def yield_shifted_models( molecule: stk.Molecule, - force_field: ForceField, + forcefield: ForceField, kicks: tuple[int], ) -> Iterator[stk.Molecule]: """Yield conformers with atom positions of particular beads shifted. @@ -477,7 +500,7 @@ def yield_shifted_models( molecule: The molecule to manipulate. - force_field: + forcefield: Defines the force field. kicks: @@ -487,7 +510,7 @@ def yield_shifted_models( An stk molecule. """ - for bead in force_field.get_present_beads(): + for bead in forcefield.get_present_beads(): atom_number = periodic_table()[bead.element_string] for kick in kicks: yield shift_beads(molecule, atom_number, kick) diff --git a/src/cgexplore/geom.py b/src/cgexplore/geom.py index 4136e5f2..89cf7049 100644 --- a/src/cgexplore/geom.py +++ b/src/cgexplore/geom.py @@ -6,14 +6,17 @@ import stk import stko -from rdkit.Chem import AllChem as rdkit +from rdkit.Chem import AllChem from .torsions import find_torsions from .utilities import get_dihedral class GeomMeasure: + """Class to perform geometry calculations.""" + def __init__(self, target_torsions: abc.Iterable | None = None) -> None: + """Initialize GeomMeasure.""" self._stko_analyser = stko.molecule_analysis.GeometryAnalyser() if target_torsions is None: self._target_torsions = None @@ -24,6 +27,7 @@ def calculate_min_distance( self, molecule: stk.Molecule, ) -> dict[str, float]: + """Calculate the minimum distance between beads and centroid.""" return { "min_distance": ( self._stko_analyser.get_min_centroid_distance(molecule) @@ -35,7 +39,7 @@ def _get_paths( molecule: stk.Molecule, path_length: int, ) -> tuple[tuple[int]]: - return rdkit.FindAllPathsOfLengthN( + return AllChem.FindAllPathsOfLengthN( mol=molecule.to_rdkit_mol(), length=path_length, useBonds=False, @@ -43,26 +47,36 @@ def _get_paths( ) def calculate_minb2b(self, molecule: stk.Molecule) -> float: + """Calculate the minimum distance between beads and beads.""" return self._stko_analyser.get_min_atom_atom_distance(molecule) def calculate_bonds( self, molecule: stk.Molecule, ) -> dict[tuple[str, ...], list[float]]: + """Calculate the bond lengths. + + Uses `stko..molecule_analysis.GeometryAnalyser` + """ return self._stko_analyser.calculate_bonds(molecule) def calculate_angles( self, molecule: stk.Molecule, ) -> dict[tuple[str, ...], list[float]]: + """Calculate the angle values. + + Uses `stko..molecule_analysis.GeometryAnalyser` + """ return self._stko_analyser.calculate_angles(molecule) def calculate_torsions( self, molecule: stk.Molecule, - absolute: bool, - as_search_string: bool = False, + absolute: bool, # noqa: FBT001 + as_search_string: bool = False, # noqa: FBT001, FBT002 ) -> dict[str, list[float]]: + """Calculate the value of target torsions.""" if self._target_torsions is None: return {} @@ -131,7 +145,15 @@ def calculate_torsions( return torsions def calculate_radius_gyration(self, molecule: stk.Molecule) -> float: + """Calculate the radius of gyration. + + Uses `stko..molecule_analysis.GeometryAnalyser` + """ return self._stko_analyser.get_radius_gyration(molecule) def calculate_max_diameter(self, molecule: stk.Molecule) -> float: + """Calculate the maximum diameter of the molecule. + + Uses `stko..molecule_analysis.GeometryAnalyser` + """ return self._stko_analyser.get_max_diameter(molecule) diff --git a/src/cgexplore/martini.py b/src/cgexplore/martini.py index 0064b28f..3fc3c9da 100644 --- a/src/cgexplore/martini.py +++ b/src/cgexplore/martini.py @@ -5,31 +5,28 @@ Contains class in https://github.com/maccallumlab/martini_openmm/tree/master """ +import contextlib import logging -import os import pathlib from openmm import app, openmm -try: +with contextlib.suppress(ModuleNotFoundError): import martini_openmm as martini -except ModuleNotFoundError: - pass logging.basicConfig( level=logging.INFO, format="%(asctime)s | %(levelname)s | %(message)s", ) -_martini_dir = ( - pathlib.Path(os.path.dirname(os.path.abspath(__file__))) / "data" -) +_martini_dir = pathlib.Path(__file__).resolve().parent / "data" class MartiniTopology: """Contains the :class:`MartiniTopFile`.""" def __init__(self, itp_file: pathlib.Path) -> None: + """Initialze MartiniTopology.""" self._itp_file = itp_file self._molecule_name = itp_file.name.replace(".itp", "") self._top_file = pathlib.Path(str(itp_file).replace(".itp", ".top")) @@ -61,13 +58,16 @@ def _write_top_file(self) -> None: f.write(string) def get_openmm_topology(self) -> app.topology.Topology: + """Return OpenMM.Topology object.""" return self._topology.topology def get_openmm_system(self) -> openmm.System: + """Return OpenMM.System object.""" return self._topology.create_system() -def get_martini_mass_by_type(bead_type) -> float: +def get_martini_mass_by_type(bead_type: str) -> float: + """Get the mass of martini types.""" bead_size = bead_type[0] return { "P": 72.0, diff --git a/src/cgexplore/molecule_construction.py b/src/cgexplore/molecule_construction.py index e16d9465..249fbdef 100644 --- a/src/cgexplore/molecule_construction.py +++ b/src/cgexplore/molecule_construction.py @@ -1,4 +1,3 @@ -#!/usr/bin/env python # Distributed under the terms of the MIT License. """Classes of topologies of precursors. @@ -15,24 +14,33 @@ class Precursor: + """Define model precursors.""" + def __init__(self) -> None: + """Initialize a precursor.""" self._bead_set: dict[str, CgBead] self._building_block: stk.BuildingBlock self._name: str raise NotImplementedError def get_bead_set(self) -> dict[str, CgBead]: + """Get beads in precursor.""" return self._bead_set def get_building_block(self) -> stk.BuildingBlock: + """Get building block defined by precursor.""" return self._building_block def get_name(self) -> str: + """Get name of precursor.""" return self._name class FourC0Arm(Precursor): + """A `FourC0Arm` Precursor.""" + def __init__(self, bead: CgBead) -> None: + """Initialize a precursor.""" self._bead = bead self._name = f"4C0{bead.bead_type}" self._bead_set = {bead.bead_type: bead} @@ -62,7 +70,10 @@ def __init__(self, bead: CgBead) -> None: class FourC1Arm(Precursor): + """A `FourC1Arm` Precursor.""" + def __init__(self, bead: CgBead, abead1: CgBead) -> None: + """Initialize a precursor.""" self._bead = bead self._abead1 = abead1 self._name = f"4C1{bead.bead_type}{abead1.bead_type}" @@ -111,7 +122,10 @@ def __init__(self, bead: CgBead, abead1: CgBead) -> None: class ThreeC0Arm(Precursor): + """A `ThreeC0Arm` Precursor.""" + def __init__(self, bead: CgBead) -> None: + """Initialize a precursor.""" self._bead = bead self._name = f"3C0{bead.bead_type}" self._bead_set = { @@ -142,7 +156,10 @@ def __init__(self, bead: CgBead) -> None: class ThreeC1Arm(Precursor): + """A `ThreeC1Arm` Precursor.""" + def __init__(self, bead: CgBead, abead1: CgBead) -> None: + """Initialize a precursor.""" self._bead = bead self._abead1 = abead1 self._name = f"3C1{bead.bead_type}{abead1.bead_type}" @@ -187,7 +204,10 @@ def __init__(self, bead: CgBead, abead1: CgBead) -> None: class ThreeC2Arm(Precursor): + """A `ThreeC2Arm` Precursor.""" + def __init__(self, bead: CgBead, abead1: CgBead, abead2: CgBead) -> None: + """Initialize a precursor.""" self._bead = bead self._abead1 = abead1 self._abead2 = abead2 @@ -240,7 +260,10 @@ def __init__(self, bead: CgBead, abead1: CgBead, abead2: CgBead) -> None: class TwoC0Arm(Precursor): + """A `TwoC0Arm` Precursor.""" + def __init__(self, bead: CgBead) -> None: + """Initialize a precursor.""" self._bead = bead self._name = f"2C0{bead.bead_type}" self._bead_set = {bead.bead_type: bead} @@ -261,7 +284,10 @@ def __init__(self, bead: CgBead) -> None: class TwoC1Arm(Precursor): + """A `TwoC1Arm` Precursor.""" + def __init__(self, bead: CgBead, abead1: CgBead) -> None: + """Initialize a precursor.""" self._bead = bead self._abead1 = abead1 self._name = f"2C1{bead.bead_type}{abead1.bead_type}" @@ -284,7 +310,10 @@ def __init__(self, bead: CgBead, abead1: CgBead) -> None: class TwoC2Arm(Precursor): + """A `TwoC2Arm` Precursor.""" + def __init__(self, bead: CgBead, abead1: CgBead, abead2: CgBead) -> None: + """Initialize a precursor.""" self._bead = bead self._abead1 = abead1 self._abead2 = abead2 @@ -321,6 +350,8 @@ def __init__(self, bead: CgBead, abead1: CgBead, abead2: CgBead) -> None: class TwoC3Arm(Precursor): + """A `TwoC3Arm` Precursor.""" + def __init__( self, bead: CgBead, @@ -328,6 +359,7 @@ def __init__( abead2: CgBead, abead3: CgBead, ) -> None: + """Initialize a precursor.""" self._bead = bead self._abead1 = abead1 self._abead2 = abead2 @@ -369,131 +401,3 @@ def __init__( ] ), ) - - -class UnsymmLigand(Precursor): - def __init__( - self, - centre_bead: CgBead, - lhs_bead: CgBead, - rhs_bead: CgBead, - binder_bead: CgBead, - ) -> None: - self._centre_bead = centre_bead - self._lhs_bead = lhs_bead - self._rhs_bead = rhs_bead - self._binder_bead = binder_bead - self._name = ( - f"UL{centre_bead.bead_type}{lhs_bead.bead_type}" - f"{rhs_bead.bead_type}{binder_bead.bead_type}" - ) - self._bead_set = { - centre_bead.bead_type: centre_bead, - lhs_bead.bead_type: lhs_bead, - rhs_bead.bead_type: rhs_bead, - binder_bead.bead_type: binder_bead, - } - - new_fgs = ( - stk.SmartsFunctionalGroupFactory( - smarts=( - f"[{binder_bead.element_string}X1]" - f"[{rhs_bead.element_string}]" - ), - bonders=(0,), - deleters=(), - placers=(0, 1), - ), - stk.SmartsFunctionalGroupFactory( - smarts=( - f"[{binder_bead.element_string}X1]" - f"[{lhs_bead.element_string}]" - ), - bonders=(0,), - deleters=(), - placers=(0, 1), - ), - ) - self._building_block = stk.BuildingBlock( - smiles=( - f"[{binder_bead.element_string}]" - f"[{lhs_bead.element_string}]" - f"[{centre_bead.element_string}]" - f"[{rhs_bead.element_string}]" - f"[{binder_bead.element_string}]" - ), - functional_groups=new_fgs, - position_matrix=np.array( - [ - [-10, 0, 0], - [-5, 3, 0], - [0, 5, 0], - [5, 3, 0], - [10, 0, 0], - ] - ), - ) - - -class UnsymmBiteLigand(Precursor): - def __init__( - self, - centre_bead: CgBead, - lhs_bead: CgBead, - rhs_bead: CgBead, - binder_bead: CgBead, - ) -> None: - self._centre_bead = centre_bead - self._lhs_bead = lhs_bead - self._rhs_bead = rhs_bead - self._binder_bead = binder_bead - self._name = ( - f"BL{centre_bead.bead_type}{lhs_bead.bead_type}" - f"{rhs_bead.bead_type}{binder_bead.bead_type}" - ) - self._bead_set = { - centre_bead.bead_type: centre_bead, - lhs_bead.bead_type: lhs_bead, - rhs_bead.bead_type: rhs_bead, - binder_bead.bead_type: binder_bead, - } - - new_fgs = ( - stk.SmartsFunctionalGroupFactory( - smarts=( - f"[{binder_bead.element_string}X1]" - f"[{rhs_bead.element_string}]" - ), - bonders=(0,), - deleters=(), - placers=(0, 1), - ), - stk.SmartsFunctionalGroupFactory( - smarts=( - f"[{binder_bead.element_string}X1]" - f"[{lhs_bead.element_string}]" - ), - bonders=(0,), - deleters=(), - placers=(0, 1), - ), - ) - self._building_block = stk.BuildingBlock( - smiles=( - f"[{binder_bead.element_string}]" - f"[{lhs_bead.element_string}]" - f"[{centre_bead.element_string}]" - f"[{rhs_bead.element_string}]" - f"[{binder_bead.element_string}]" - ), - functional_groups=new_fgs, - position_matrix=np.array( - [ - [-10, 0, 0], - [-5, 3, 0], - [0, 5, 0], - [5, 3, 0], - [10, 0, 0], - ] - ), - ) diff --git a/src/cgexplore/nonbonded.py b/src/cgexplore/nonbonded.py index 89088497..e366e6d3 100644 --- a/src/cgexplore/nonbonded.py +++ b/src/cgexplore/nonbonded.py @@ -1,4 +1,3 @@ -#!/usr/bin/env python # Distributed under the terms of the MIT License. """Module for handling nobonded interactions. @@ -7,8 +6,9 @@ """ -import itertools +import itertools as it import logging +from collections import abc from dataclasses import dataclass from openmm import openmm @@ -21,6 +21,8 @@ @dataclass class Nonbonded: + """Class containing term defintion.""" + atom_id: int bead_class: str bead_element: str @@ -31,13 +33,27 @@ class Nonbonded: @dataclass class TargetNonbonded: + """Defines a target term to search for in a molecule.""" + bead_class: str bead_element: str sigma: openmm.unit.Quantity epsilon: openmm.unit.Quantity force: str + def vector_key(self) -> str: + """Return key for vector defining this target term.""" + return self.bead_class + + def vector(self) -> tuple[float, float]: + """Return vector defining this target term.""" + return ( + self.sigma.value_in_unit(openmm.unit.angstrom), + self.epsilon.value_in_unit(openmm.unit.kilojoules_per_mole), + ) + def human_readable(self) -> str: + """Return human-readable definition of this target term.""" return ( f"{self.__class__.__name__}(" f"{self.bead_class}, " @@ -51,14 +67,17 @@ def human_readable(self) -> str: @dataclass class TargetNonbondedRange: + """Defines a target term and ranges in parameters to search for.""" + bead_class: str bead_element: str sigmas: tuple[openmm.unit.Quantity] epsilons: tuple[openmm.unit.Quantity] force: str - def yield_nonbondeds(self): - for sigma, epsilon in itertools.product(self.sigmas, self.epsilons): + def yield_nonbondeds(self) -> abc.Iterable[TargetNonbonded]: + """Find interactions matching target.""" + for sigma, epsilon in it.product(self.sigmas, self.epsilons): yield TargetNonbonded( bead_class=self.bead_class, bead_element=self.bead_element, diff --git a/src/cgexplore/openmm_optimizer.py b/src/cgexplore/openmm_optimizer.py index d0cf0711..74034347 100644 --- a/src/cgexplore/openmm_optimizer.py +++ b/src/cgexplore/openmm_optimizer.py @@ -27,7 +27,9 @@ class OMMTrajectory: - def __init__( + """Class for holding trajectory information from OpenMM.""" + + def __init__( # noqa: PLR0913 self, base_molecule: stk.Molecule, data_path: pathlib.Path, @@ -41,6 +43,7 @@ def __init__( reporting_freq: float, traj_freq: float, ) -> None: + """Initialize OMMTrajectory.""" self._base_molecule = base_molecule self._data_path = data_path self._traj_path = traj_path @@ -56,12 +59,15 @@ def __init__( self._friction = friction def get_data(self) -> pd.DataFrame: + """Get Trajectory data.""" return pd.read_csv(self._data_path) def get_base_molecule(self) -> stk.Molecule: + """Get the base molecule of this trajectory.""" return self._base_molecule def yield_conformers(self) -> abc.Iterator[Timestep]: + """Yield conformers from trajectory.""" num_atoms = self._base_molecule.get_num_atoms() start_trigger = "MODEL" triggered = False @@ -102,17 +108,21 @@ def yield_conformers(self) -> abc.Iterator[Timestep]: new_pos_mat.append([x, y, z]) def __str__(self) -> str: + """Return a string representation of the OMMTrajectory.""" return ( f"{self.__class__.__name__}(steps={self._num_steps}, " f"conformers={self._num_confs})" ) def __repr__(self) -> str: + """Return a string representation of the OMMTrajectory.""" return str(self) class CGOMMOptimizer: - def __init__( + """Optimiser of CG models using OpenMM.""" + + def __init__( # noqa: PLR0913 self, fileprefix: str, output_dir: pathlib.Path, @@ -120,6 +130,7 @@ def __init__( atom_constraints: abc.Iterable[tuple[int, int]] | None = None, platform: str | None = None, ) -> None: + """Initialize CGOMMOptimizer.""" self._fileprefix = fileprefix self._output_dir = output_dir self._output_path = output_dir / f"{fileprefix}_omm.out" @@ -131,7 +142,7 @@ def __init__( else: self._max_iterations = max_iterations - self._tolerance = 1e-6 * openmm.unit.kilojoules_per_mole + self._tolerance = 1e-6 if platform is not None: self._platform = openmm.Platform.getPlatformByName(platform) @@ -155,13 +166,14 @@ def _group_forces(self, system: openmm.System) -> dict: From https://github.com/XiaojuanHu/CTPOL_MD/pull/4/files Use these two methods as: - `fgrps=forcegroupify(system)` - ` - tt= getEnergyDecomposition(simulation.context, fgrps) - for idd in tt.keys(): - print(idd,tt[idd]) - print('\n') - ` + + ``` + fgrps=forcegroupify(system) + + tt= getEnergyDecomposition(simulation.context, fgrps) + for idd in tt.keys(): + print(idd,tt[idd]) + ``` """ forcegroups = {} @@ -321,10 +333,12 @@ def _update_stk_molecule( return molecule.with_position_matrix(positions * 10) def calculate_energy(self, assigned_system: ForcedSystem) -> float: + """Calculate energy of a system.""" simulation, _ = self._setup_simulation(assigned_system) return self._get_energy(simulation) def read_final_energy_decomposition(self) -> dict: + """Read the final energy decomposition in an output file.""" decomp_data = ( self._output_string.split("energy decomposition:")[-1] .split("end time:")[0] @@ -341,6 +355,7 @@ def read_final_energy_decomposition(self) -> dict: return decomposition def optimize(self, assigned_system: ForcedSystem) -> stk.Molecule: + """Optimize a molecule.""" start_time = time.time() self._output_string += f"start time: {start_time}\n" self._output_string += ( @@ -363,7 +378,9 @@ def optimize(self, assigned_system: ForcedSystem) -> stk.Molecule: class CGOMMDynamics(CGOMMOptimizer): - def __init__( + """Optimiser of CG models using OpenMM and Molecular Dynamics.""" + + def __init__( # noqa: PLR0913 self, fileprefix: str, output_dir: pathlib.Path, @@ -378,6 +395,7 @@ def __init__( atom_constraints: abc.Iterable[tuple[int, int]] | None = None, platform: str | None = None, ) -> None: + """Initialize CGOMMDynamics.""" self._fileprefix = fileprefix self._output_dir = output_dir self._output_path = output_dir / f"{fileprefix}_omm.out" @@ -390,7 +408,7 @@ def __init__( else: self._max_iterations = max_iterations - self._tolerance = 1e-6 * openmm.unit.kilojoules_per_mole + self._tolerance = 1e-6 self._atom_constraints = atom_constraints @@ -398,7 +416,8 @@ def __init__( if random_seed is None: logging.info("a random random seed is used") - self._random_seed = np.random.randint(1000) + rng = np.random.default_rng() + self._random_seed = rng.integers(low=1, high=4000) else: self._random_seed = random_seed @@ -515,6 +534,7 @@ def _get_trajectory(self, molecule: stk.Molecule) -> OMMTrajectory: ) def run_dynamics(self, assigned_system: ForcedSystem) -> OMMTrajectory: + """Run dynamics on an assigned system.""" start_time = time.time() self._output_string += f"start time: {start_time}\n" self._output_string += ( diff --git a/src/cgexplore/shape.py b/src/cgexplore/shape.py index a539accd..6cbe5b95 100644 --- a/src/cgexplore/shape.py +++ b/src/cgexplore/shape.py @@ -1,4 +1,3 @@ -#!/usr/bin/env python # Distributed under the terms of the MIT License. """Module for shape analysis. @@ -30,6 +29,11 @@ def test_shape_mol( name: str, topo_str: str, ) -> None: + """Test shape molecule. + + Requires hard-coded assumptions. + I suggest using the method in `ShapeMeasure`. + """ num_atoms = len(atoms) if num_atoms != topo_expected[topo_str]: msg = ( @@ -44,6 +48,7 @@ def fill_position_matrix_molecule( element: str, old_position_matrix: np.ndarray, ) -> tuple[list, list]: + """Make position matrix from filtering molecule based on element.""" position_matrix = [] atoms: list[stk.Atom] = [] @@ -68,6 +73,11 @@ def get_shape_molecule_byelement( element: str, topo_expected: dict[str, int], ) -> stk.Molecule | None: + """Get shape molecule. + + Requires hard-coded assumptions. + I suggest using the method in `ShapeMeasure`. + """ splits = name.split("_") topo_str = splits[0] if topo_str not in topo_expected: @@ -95,6 +105,7 @@ def fill_position_matrix( element: str, old_position_matrix: np.ndarray, ) -> tuple[list, list]: + """Make position matrix from filtering molecule based on target BB.""" position_matrix = [] atoms: list[stk.Atom] = [] @@ -125,6 +136,11 @@ def get_shape_molecule_nodes( element: str, topo_expected: dict[str, int], ) -> stk.Molecule | None: + """Get shape molecule. + + Requires hard-coded assumptions. + I suggest using the method in `ShapeMeasure`. + """ splits = name.split("_") topo_str = splits[0] if topo_str not in topo_expected: @@ -154,6 +170,11 @@ def get_shape_molecule_ligands( element: str, topo_expected: dict[str, int], ) -> stk.Molecule | None: + """Get shape molecule. + + Requires hard-coded assumptions. + I suggest using the method in `ShapeMeasure`. + """ splits = name.split("_") topo_str = splits[0] bbs = list(constructed_molecule.get_building_blocks()) @@ -192,6 +213,7 @@ def __init__( shape_path: pathlib.Path | str, shape_string: str | None = None, ) -> None: + """Initialize shape measure.""" self._output_dir = output_dir self._shape_path = shape_path self._shape_dict: dict[str, dict] @@ -205,7 +227,39 @@ def __init__( {int(self._shape_dict[i]["vertices"]) for i in self._shape_dict} ) + def _test_shape_mol(self, expected_points: int, atoms: list) -> None: + num_atoms = len(atoms) + if num_atoms != expected_points: + msg = ( + f"Expected num points not met ({expected_points}); has " + f"{num_atoms}" + ) + raise ValueError(msg) + + def get_shape_molecule_byelement( + self, + molecule: stk.Molecule, + element: str, + expected_points: int, + ) -> stk.Molecule | None: + """Get molecule to analyse by filtering by element.""" + old_position_matrix = molecule.get_position_matrix() + + position_matrix, atoms = fill_position_matrix_molecule( + molecule=molecule, + element=element, + old_position_matrix=old_position_matrix, + ) + + self._test_shape_mol(expected_points, atoms) + return stk.BuildingBlock.init( + atoms=atoms, + bonds=(), + position_matrix=np.array(position_matrix), + ) + def reference_shape_dict(self) -> dict[str, dict]: + """Reference shapes as dictionary.""" return { "L-2": { "code": "1", @@ -809,7 +863,9 @@ def _write_input_file( shape_numbers = tuple(i["code"] for i in possible_shapes) title = "$shape run by Andrew Tarzia - central atom=0 always.\n" - fix_perm = "%fixperm 0\\n" if num_vertices == 12 else "\n" + fix_perm = ( + "%fixperm 0\\n" if num_vertices == 12 else "\n" # noqa: PLR2004 + ) size_of_poly = f"{num_vertices} 0\n" codes = " ".join(shape_numbers) + "\n" @@ -829,18 +885,18 @@ def _run_calculation(self, structure_string: str) -> dict: structure_string=structure_string, ) - cmd = f"{self._shape_path} {input_file}" + # Note that sp.call will hold the program until completion + # of the calculation. + captured_output = sp.run( + [f"{self._shape_path}", f"{input_file}"], # noqa: S603 + stdin=sp.PIPE, + capture_output=True, + check=True, + # Shell is required to run complex arguments. + ) + with open(std_out, "w") as f: - # Note that sp.call will hold the program until completion - # of the calculation. - sp.call( - cmd, - stdin=sp.PIPE, - stdout=f, - stderr=sp.PIPE, - # Shell is required to run complex arguments. - shell=True, - ) + f.write(str(captured_output.stdout)) return self._collect_all_shape_values(output_file) @@ -857,9 +913,7 @@ def _get_centroids( bb_ids[aibbid] = [] bb_ids[aibbid].append(ai.get_atom().get_id()) - centroids = [] - for n in bb_ids: - centroids.append(molecule.get_centroid(atom_ids=bb_ids[n])) + centroids = [molecule.get_centroid(atom_ids=bb_ids[n]) for n in bb_ids] with open("cents.xyz", "w") as f: f.write(f"{len(centroids)}\n\n") @@ -877,13 +931,14 @@ def _get_possible_shapes(self, num_vertices: int) -> tuple[dict, ...]: ) def calculate(self, molecule: stk.Molecule) -> dict: - output_dir = os.path.abspath(self._output_dir) + """Calculate shape measures for a molecule.""" + output_dir = pathlib.Path(self._output_dir).resolve() - if os.path.exists(output_dir): + if output_dir.exists(): shutil.rmtree(output_dir) - os.mkdir(output_dir) + output_dir.mkdir() - init_dir = os.getcwd() + init_dir = pathlib.Path.cwd() try: os.chdir(output_dir) structure_string = "shape run by AT\n" @@ -915,13 +970,17 @@ def calculate_from_centroids( constructed_molecule: stk.ConstructedMolecule, target_atmnums: tuple[int], ) -> dict: - output_dir = os.path.abspath(self._output_dir) + """Calculate shape from the centroids of building blocks. + + Currently not implemented well. + """ + output_dir = pathlib.Path(self._output_dir).resolve() - if os.path.exists(output_dir): + if output_dir.exists(): shutil.rmtree(output_dir) - os.mkdir(output_dir) + output_dir.mkdir() - init_dir = os.getcwd() + init_dir = pathlib.Path.cwd() try: os.chdir(output_dir) centroids = self._get_centroids( diff --git a/src/cgexplore/torsions.py b/src/cgexplore/torsions.py index 5fa641b9..cae8f141 100644 --- a/src/cgexplore/torsions.py +++ b/src/cgexplore/torsions.py @@ -1,4 +1,3 @@ -#!/usr/bin/env python # Distributed under the terms of the MIT License. """Module for handling torsions. @@ -7,14 +6,14 @@ """ -import itertools +import itertools as it import logging from collections import abc from dataclasses import dataclass import stk from openmm import openmm -from rdkit.Chem import AllChem as rdkit +from rdkit.Chem import AllChem logging.basicConfig( level=logging.INFO, @@ -24,6 +23,8 @@ @dataclass class Torsion: + """Class containing torsion defintion.""" + atom_names: tuple[str, ...] atom_ids: tuple[int, ...] phi0: openmm.unit.Quantity @@ -35,6 +36,8 @@ class Torsion: @dataclass class TargetTorsion: + """Defines a target term to search for in a molecule.""" + search_string: tuple[str, ...] search_estring: tuple[str, ...] measured_atom_ids: tuple[int, int, int, int] @@ -43,7 +46,20 @@ class TargetTorsion: torsion_n: int funct: int = 0 + def vector_key(self) -> str: + """Return key for vector defining this target term.""" + return "".join(self.search_string) + + def vector(self) -> tuple[float, float, float]: + """Return vector defining this target term.""" + return ( + self.phi0.value_in_unit(openmm.unit.degrees), + self.torsion_k.value_in_unit(openmm.unit.kilojoules_per_mole), + self.torsion_n, + ) + def human_readable(self) -> str: + """Return human-readable definition of this target term.""" return ( f"{self.__class__.__name__}(" f'{"".join(self.search_string)}, ' @@ -58,6 +74,8 @@ def human_readable(self) -> str: @dataclass class TargetTorsionRange: + """Defines a target term and ranges in parameters to search for.""" + search_string: tuple[str, ...] search_estring: tuple[str, ...] measured_atom_ids: tuple[int, int, int, int] @@ -65,8 +83,9 @@ class TargetTorsionRange: torsion_ks: tuple[openmm.unit.Quantity] torsion_ns: tuple[int] - def yield_torsions(self): - for phi0, k, n in itertools.product( + def yield_torsions(self) -> abc.Iterable[TargetTorsion]: + """Find interactions matching target.""" + for phi0, k, n in it.product( self.phi0s, self.torsion_ks, self.torsion_ns ): yield TargetTorsion( @@ -81,6 +100,8 @@ def yield_torsions(self): @dataclass class FoundTorsion: + """Define a found forcefield term.""" + atoms: tuple[stk.Atom, ...] atom_ids: tuple[int, ...] @@ -89,7 +110,8 @@ def find_torsions( molecule: stk.Molecule, chain_length: int, ) -> abc.Iterator[FoundTorsion]: - paths = rdkit.FindAllPathsOfLengthN( + """Find torsions based on bonds in molecule.""" + paths = AllChem.FindAllPathsOfLengthN( mol=molecule.to_rdkit_mol(), length=chain_length, useBonds=False, @@ -105,6 +127,8 @@ def find_torsions( @dataclass class TargetMartiniTorsion: + """Defines a target angle to search for in a molecule.""" + search_string: tuple[str, ...] search_estring: tuple[str, ...] measured_atom_ids: tuple[int, int, int, int] @@ -114,6 +138,7 @@ class TargetMartiniTorsion: funct: int def human_readable(self) -> str: + """Return human-readable definition of this target term.""" return ( f"{self.__class__.__name__}(" f'{"".join(self.search_string)}, ' @@ -129,6 +154,8 @@ def human_readable(self) -> str: @dataclass class MartiniTorsionRange: + """Defines a target torsion and ranges in parameters to search for.""" + search_string: tuple[str, ...] search_estring: tuple[str, ...] measured_atom_ids: tuple[int, int, int, int] @@ -137,9 +164,11 @@ class MartiniTorsionRange: torsion_ns: tuple[int] funct: int - def yield_torsions(self): - raise NotImplementedError("handle torsions") - for phi0, k, n in itertools.product( + def yield_torsions(self) -> abc.Iterable[TargetMartiniTorsion]: + """Find torsions matching target.""" + msg = "handle torsions" + raise NotImplementedError(msg) + for phi0, k, n in it.product( # type: ignore[unreachable] self.phi0s, self.torsion_ks, self.torsion_ns ): yield TargetMartiniTorsion( diff --git a/src/cgexplore/utilities.py b/src/cgexplore/utilities.py index 9477c52b..362a877b 100644 --- a/src/cgexplore/utilities.py +++ b/src/cgexplore/utilities.py @@ -1,4 +1,3 @@ -#!/usr/bin/env python # Distributed under the terms of the MIT License. """Utilities module. @@ -9,10 +8,11 @@ import json import logging -import os import pathlib +import matplotlib.pyplot as plt import numpy as np +import numpy.typing as npt import stk from openmm import openmm from scipy.spatial.distance import euclidean @@ -35,8 +35,9 @@ def convert_pyramid_angle(outer_angle: float) -> float: def check_directory(path: pathlib.Path) -> None: - if not os.path.exists(path): - os.mkdir(path) + """Check if a directory exists, make if not.""" + if not path.exists(): + path.mkdir() def get_atom_distance( @@ -69,7 +70,7 @@ def get_atom_distance( return float(distance) -def unit_vector(vector): +def unit_vector(vector: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]: """Returns the unit vector of the vector. https://stackoverflow.com/questions/2827393/ @@ -80,8 +81,14 @@ def unit_vector(vector): return vector / np.linalg.norm(vector) -def angle_between(v1, v2, normal=None): - """Returns the angle in radians between vectors 'v1' and 'v2':: +def angle_between( + v1: npt.NDArray[np.float64], + v2: npt.NDArray[np.float64], + normal: npt.NDArray[np.float64] | None = None, +) -> float: + """Returns the angle in radians between vectors 'v1' and 'v2'. + + Defined as :: >>> angle_between((1, 0, 0), (0, 1, 0)) 1.5707963267948966 @@ -161,6 +168,7 @@ def get_dihedral( def custom_excluded_volume_force() -> openmm.CustomNonbondedForce: + """Define Custom Excluded Volume force.""" energy_expression = "epsilon*((sigma)/(r))^12;" energy_expression += "epsilon = sqrt(epsilon1*epsilon2);" energy_expression += "sigma = 0.5*(sigma1+sigma2);" @@ -171,6 +179,7 @@ def custom_excluded_volume_force() -> openmm.CustomNonbondedForce: def cosine_periodic_angle_force() -> openmm.CustomAngleForce: + """Define Custom Angle force.""" energy_expression = "F*C*(1-A*cos(n * theta));" energy_expression += "A = b*(min_n);" energy_expression += "C = (n^2 * k)/2;" @@ -181,3 +190,58 @@ def cosine_periodic_angle_force() -> openmm.CustomAngleForce: custom_force.addPerAngleParameter("b") custom_force.addPerAngleParameter("min_n") return custom_force + + +def draw_pie( + colours: list[str], + xpos: float, + ypos: float, + size: float, + ax: plt.Axes, +) -> None: + """Draw a pie chart at a specific point on ax. + + From: + https://stackoverflow.com/questions/56337732/how-to-plot-scatter- + pie-chart-using-matplotlib. + + """ + num_points = len(colours) + if num_points == 1: + ax.scatter( + xpos, + ypos, + c=colours[0], + edgecolors="k", + s=size, + ) + else: + ratios = [1 / num_points for i in range(num_points)] + if sum(ratios) > 1: + msg = ( + f"sum of ratios needs to be < 1 (np={num_points}, " + f"ratios={ratios})" + ) + raise AssertionError(msg) + + markers = [] + previous = 0.0 + # calculate the points of the pie pieces + for color, ratio in zip(colours, ratios, strict=True): + this = 2 * np.pi * ratio + previous + x = [0, *np.cos(np.linspace(previous, this, 100)).tolist(), 0] + y = [0, *np.sin(np.linspace(previous, this, 100)).tolist(), 0] + xy = np.column_stack([x, y]) + previous = this + markers.append( + { + "marker": xy, + "s": np.abs(xy).max() ** 2 * np.array(size), + "facecolor": color, + "edgecolors": "k", + } + ) + + # scatter each of the pie pieces to create pies + for marker in markers: + ax.scatter(xpos, ypos, **marker) diff --git a/src/cgexplore/visualisation.py b/src/cgexplore/visualisation.py index a78502ce..ec83f6c9 100644 --- a/src/cgexplore/visualisation.py +++ b/src/cgexplore/visualisation.py @@ -1,4 +1,3 @@ -#!/usr/bin/env python # Distributed under the terms of the MIT License. """Module for pymol visualisation. @@ -14,6 +13,8 @@ class Pymol: + """Pymol visualiser.""" + def __init__( self, output_dir: pathlib.Path, @@ -21,6 +22,7 @@ def __init__( pymol_path: pathlib.Path, settings: dict | None = None, ) -> None: + """Initialize Pymol visualiser.""" self._output_dir = output_dir self._file_prefix = file_prefix self._pymol = pymol_path @@ -49,11 +51,10 @@ def _get_zoom_string(self, structure_files: list) -> str: max_diam = stk.BuildingBlock.init_from_file( path=str(fi), ).get_maximum_diameter() - print(max_diam) max_max_diam = max((max_diam, max_max_diam)) return f"zoom center, {max_max_diam/2}" - def _write_pymol_script( + def _write_pymol_script( # noqa: PLR0913 self, structure_files: list, structure_colours: list | None, @@ -88,11 +89,11 @@ def _write_pymol_script( lstring += f"load {sf}\n" lname = str(sf.name).replace(".mol", "") lnames.append(lname) - col = col.replace("#", "0x") + tcol = col.replace("#", "0x") if structure_colours is None: cstring += "color orange, (name C*)\n" else: - cstring += f"color {col}, {lname}\n" + cstring += f"color {tcol}, {lname}\n" string = ( f"{lstring}\n" @@ -124,6 +125,7 @@ def visualise( orient_atoms: str | None = None, big_colour: str | None = None, ) -> None: + """Run pymol to visualise a molecule.""" pml_file = self._output_dir / f"{self._file_prefix}.pml" self._write_pymol_script( structure_files=structure_files, @@ -132,4 +134,4 @@ def visualise( orient_atoms=orient_atoms, big_colour=big_colour, ) - os.system(f"{self._pymol} -c -q {pml_file}") + os.system(f"{self._pymol} -c -q {pml_file}") # noqa: S605 diff --git a/tests/assigned/case_data.py b/tests/assigned/case_data.py index 1362f90c..d9c75904 100644 --- a/tests/assigned/case_data.py +++ b/tests/assigned/case_data.py @@ -1,3 +1,4 @@ +import cgexplore import stk @@ -7,11 +8,11 @@ class CaseData: def __init__( self, molecule: stk.Molecule, - force_field, - topology_xml_string, + forcefield: cgexplore.forcefield.ForceField, + topology_xml_string: str, name: str, ) -> None: self.molecule = molecule - self.force_field = force_field + self.forcefield = forcefield self.topology_xml_string = topology_xml_string self.name = name diff --git a/tests/assigned/conftest.py b/tests/assigned/conftest.py index 351a9b5d..a5717802 100644 --- a/tests/assigned/conftest.py +++ b/tests/assigned/conftest.py @@ -27,7 +27,7 @@ ) ), ), - force_field=ForceField( + forcefield=ForceField( identifier=name, prefix="assigned_tests", present_beads=( @@ -117,20 +117,6 @@ - - - - - - - - - - - - - - """ ), @@ -149,7 +135,7 @@ ) ), ), - force_field=ForceField( + forcefield=ForceField( identifier=name, prefix="assigned_tests", present_beads=( @@ -273,20 +259,6 @@ - - - - - - - - - - - - - - """ ), diff --git a/tests/assigned/molecule0_assigned_tests_molecule0_syst_saved.xml b/tests/assigned/molecule0_assigned_tests_molecule0_syst_saved.xml index c01d1f87..3e208b27 100644 --- a/tests/assigned/molecule0_assigned_tests_molecule0_syst_saved.xml +++ b/tests/assigned/molecule0_assigned_tests_molecule0_syst_saved.xml @@ -14,7 +14,6 @@ - diff --git a/tests/assigned/molecule1_assigned_tests_molecule1_syst_saved.xml b/tests/assigned/molecule1_assigned_tests_molecule1_syst_saved.xml index 8d6e5b92..ce83d3ea 100644 --- a/tests/assigned/molecule1_assigned_tests_molecule1_syst_saved.xml +++ b/tests/assigned/molecule1_assigned_tests_molecule1_syst_saved.xml @@ -14,7 +14,6 @@ - diff --git a/tests/assigned/test_system_xml_writer.py b/tests/assigned/test_system_xml_writer.py index 2441d315..b1192a0b 100644 --- a/tests/assigned/test_system_xml_writer.py +++ b/tests/assigned/test_system_xml_writer.py @@ -1,10 +1,11 @@ -import os import pathlib from cgexplore.errors import ForceFieldUnitError +from .case_data import CaseData -def test_system_xml_writer(molecule): + +def test_system_xml_writer(molecule: CaseData) -> None: """Test methods toward :meth:`.ForceField.get_xml_string`. Parameters: @@ -17,40 +18,33 @@ def test_system_xml_writer(molecule): """ try: - syst_xml_file = pathlib.Path( - os.path.dirname(os.path.realpath(__file__)) - ) / ( - f"{molecule.name}_{molecule.force_field.get_prefix()}_" + syst_xml_file = pathlib.Path(__file__).resolve().parent / ( + f"{molecule.name}_{molecule.forcefield.get_prefix()}_" f"{molecule.name}_syst.xml" ) - topo_xml_file = pathlib.Path( - os.path.dirname(os.path.realpath(__file__)) - ) / ( - f"{molecule.name}_{molecule.force_field.get_prefix()}_" - f"{molecule.name}_topo.xml" - ) - saved_syst_xml_file = pathlib.Path( - os.path.dirname(os.path.realpath(__file__)) - ) / ( - f"{molecule.name}_{molecule.force_field.get_prefix()}_" + saved_syst_xml_file = pathlib.Path(__file__).resolve().parent / ( + f"{molecule.name}_{molecule.forcefield.get_prefix()}_" f"{molecule.name}_syst_saved.xml" ) - assigned_system = molecule.force_field.assign_terms( + assigned_system = molecule.forcefield.assign_terms( molecule=molecule.molecule, - output_dir=pathlib.Path( - os.path.dirname(os.path.realpath(__file__)) - ), + output_dir=pathlib.Path(__file__).resolve().parent, name=molecule.name, ) assigned_system.get_openmm_system() with open(syst_xml_file) as f: - xml_string = f.read() + xml_string_list = f.readlines() with open(saved_syst_xml_file) as f: - test_xml_string = f.read() - print(xml_string) - assert xml_string == test_xml_string - os.system(f"rm {syst_xml_file}") - os.system(f"rm {topo_xml_file}") + test_xml_string_list = f.readlines() + + for xml, test in zip( + xml_string_list, test_xml_string_list, strict=True + ): + if "openmmVersion" in xml and "openmmVersion" in test: + continue + print(xml, test) + assert xml == test + syst_xml_file.unlink() except ForceFieldUnitError: assert molecule.num_forcefields == 0 diff --git a/tests/assigned/test_topology_xml_writer.py b/tests/assigned/test_topology_xml_writer.py index 364cb059..863b0469 100644 --- a/tests/assigned/test_topology_xml_writer.py +++ b/tests/assigned/test_topology_xml_writer.py @@ -1,8 +1,9 @@ -import os import pathlib +from .case_data import CaseData -def test_topology_xml_writer(molecule): + +def test_topology_xml_writer(molecule: CaseData) -> None: """Test methods toward :meth:`.AssignedSystem._get_topology_xml_string`. Parameters: @@ -14,11 +15,13 @@ def test_topology_xml_writer(molecule): None : :class:`NoneType` """ - assigned_system = molecule.force_field.assign_terms( + assigned_system = molecule.forcefield.assign_terms( molecule=molecule.molecule, - output_dir=pathlib.Path(os.path.dirname(os.path.realpath(__file__))), + output_dir=pathlib.Path(__file__).resolve().parent, name=molecule.name, ) - xml_string = assigned_system._get_topology_xml_string(molecule.molecule) + xml_string = assigned_system._get_topology_xml_string( # noqa: SLF001 + molecule.molecule + ) print(xml_string) assert xml_string == molecule.topology_xml_string diff --git a/tests/databases/test_databasing.py b/tests/databases/test_databasing.py index 8836ac0b..267139e6 100644 --- a/tests/databases/test_databasing.py +++ b/tests/databases/test_databasing.py @@ -1,24 +1,29 @@ -import os +import pathlib +import shutil import atomlite import pytest import stk from cgexplore.databases import AtomliteDatabase +from .case_data import CaseData -def test_databasing(molecule): + +def test_databasing(molecule: CaseData) -> None: """Test :class:`.AtomliteDatabase`.""" - path = os.path.join( - os.path.dirname(os.path.realpath(__file__)), - "test.db", - ) - if os.path.exists(path): - os.remove(path) + path = pathlib.Path(__file__).resolve().parent / "test.db" + + if path.exists(): + shutil.rmtree(path) database = AtomliteDatabase(path) assert database.get_num_entries() == 0 - for mol, prop in zip(molecule.molecules, molecule.property_dicts): + for mol, prop in zip( + molecule.molecules, + molecule.property_dicts, + strict=True, + ): print(mol, prop) key = stk.Smiles().get_key(mol) database.add_molecule(molecule=mol, key=key) @@ -50,4 +55,4 @@ def test_databasing(molecule): print(database.get_num_entries()) assert database.get_num_entries() == molecule.expected_count - os.remove(path) + path.unlink() diff --git a/tests/forcefield/case_data.py b/tests/forcefield/case_data.py index f37f01a3..6619d04f 100644 --- a/tests/forcefield/case_data.py +++ b/tests/forcefield/case_data.py @@ -1,39 +1,36 @@ +import cgexplore import stk class CaseData: - """A test case. - - Attributes: - - """ + """A test case.""" def __init__( self, molecule: stk.Molecule, - force_field_library, - bond_ranges, - angle_ranges, - torsion_ranges, - nonbonded_ranges, - present_bonds, - present_angles, - present_nonbondeds, - present_torsions, + forcefield_library: cgexplore.forcefield.ForceFieldLibrary, + bond_ranges: tuple[cgexplore.bonds.TargetBondRange], + angle_ranges: tuple[cgexplore.angles.TargetAngleRange], + torsion_ranges: tuple[cgexplore.torsions.TargetTorsionRange], + nonbonded_ranges: tuple[cgexplore.nonbonded.TargetNonbondedRange], + present_bonds: tuple[tuple], + present_angles: tuple[tuple], + present_nonbondeds: tuple[tuple], + present_torsions: tuple[tuple], num_forcefields: int, library_string: str, name: str, ) -> None: self.molecule = molecule - self.force_field_library = force_field_library + self.forcefield_library = forcefield_library for i in bond_ranges: - self.force_field_library.add_bond_range(i) + self.forcefield_library.add_bond_range(i) for i in angle_ranges: - self.force_field_library.add_angle_range(i) + self.forcefield_library.add_angle_range(i) for i in torsion_ranges: - self.force_field_library.add_torsion_range(i) + self.forcefield_library.add_torsion_range(i) for i in nonbonded_ranges: - self.force_field_library.add_nonbonded_range(i) + self.forcefield_library.add_nonbonded_range(i) self.num_forcefields = num_forcefields self.present_bonds = present_bonds self.present_angles = present_angles diff --git a/tests/forcefield/conftest.py b/tests/forcefield/conftest.py index 4b5504e8..d1b7ea70 100644 --- a/tests/forcefield/conftest.py +++ b/tests/forcefield/conftest.py @@ -79,7 +79,7 @@ ) ), ), - force_field_library=ForceFieldLibrary( + forcefield_library=ForceFieldLibrary( bead_library=(ag_bead, ba_bead, pb_bead), vdw_bond_cutoff=2, prefix="testff", @@ -368,7 +368,7 @@ ) ), ), - force_field_library=ForceFieldLibrary( + forcefield_library=ForceFieldLibrary( bead_library=(c_bead, n_bead, o_bead), vdw_bond_cutoff=2, prefix="testff", @@ -1395,7 +1395,7 @@ molecule=FourC1Arm( bead=pd_bead, abead1=pb_bead ).get_building_block(), - force_field_library=ForceFieldLibrary( + forcefield_library=ForceFieldLibrary( bead_library=(pd_bead, pb_bead), vdw_bond_cutoff=2, prefix="testff", @@ -1882,7 +1882,7 @@ ) ), ), - force_field_library=ForceFieldLibrary( + forcefield_library=ForceFieldLibrary( bead_library=(ag_bead, ba_bead, pb_bead), vdw_bond_cutoff=2, prefix="testff", @@ -1967,7 +1967,7 @@ ) ), ), - force_field_library=ForceFieldLibrary( + forcefield_library=ForceFieldLibrary( bead_library=(ag_bead, ba_bead, pb_bead), vdw_bond_cutoff=2, prefix="testff", @@ -2011,7 +2011,7 @@ smiles="[C][N][C]", position_matrix=np.array(([1, 0, 0], [0.5, 0, 0], [2, 0, 0])), ), - force_field_library=ForceFieldLibrary( + forcefield_library=ForceFieldLibrary( bead_library=(c_bead, n_bead, c_bead), vdw_bond_cutoff=2, prefix="testuff", diff --git a/tests/forcefield/test_fflibrary.py b/tests/forcefield/test_fflibrary.py index eb9c5f24..d4459f61 100644 --- a/tests/forcefield/test_fflibrary.py +++ b/tests/forcefield/test_fflibrary.py @@ -1,4 +1,7 @@ -def test_fflibrary(molecule): +from .case_data import CaseData + + +def test_fflibrary(molecule: CaseData) -> None: """Test methods toward :meth:`.ForceFieldLibrary`. Parameters: @@ -11,7 +14,7 @@ def test_fflibrary(molecule): """ if molecule.num_forcefields > 0: - force_fields = tuple(molecule.force_field_library.yield_forcefields()) - print(molecule.force_field_library) - assert molecule.num_forcefields == len(force_fields) - assert str(molecule.force_field_library) == molecule.library_string + forcefields = tuple(molecule.forcefield_library.yield_forcefields()) + print(molecule.forcefield_library) + assert molecule.num_forcefields == len(forcefields) + assert str(molecule.forcefield_library) == molecule.library_string diff --git a/tests/forcefield/test_present_angles.py b/tests/forcefield/test_present_angles.py index 2547ae19..f90345b4 100644 --- a/tests/forcefield/test_present_angles.py +++ b/tests/forcefield/test_present_angles.py @@ -1,13 +1,13 @@ -import os import pathlib from cgexplore.angles import Angle, CosineAngle from cgexplore.errors import ForceFieldUnitError +from .case_data import CaseData from .utilities import is_equivalent_atom -def test_present_angles(molecule): +def test_present_angles(molecule: CaseData) -> None: """Test methods toward :meth:`.ForceField._assign_angle_terms`. Parameters: @@ -20,17 +20,15 @@ def test_present_angles(molecule): """ try: - force_fields = tuple(molecule.force_field_library.yield_forcefields()) - for i, ff in enumerate(force_fields): + forcefields = tuple(molecule.forcefield_library.yield_forcefields()) + for i, ff in enumerate(forcefields): assigned_system = ff.assign_terms( molecule=molecule.molecule, - output_dir=pathlib.Path( - os.path.dirname(os.path.realpath(__file__)) - ), + output_dir=pathlib.Path(__file__).resolve().parent, name=molecule.name, ) - present_terms = assigned_system.force_field_terms["angle"] + present_terms = assigned_system.forcefield_terms["angle"] print(present_terms) assert len(present_terms) == len(molecule.present_angles[i]) for test, present in zip( diff --git a/tests/forcefield/test_present_bonds.py b/tests/forcefield/test_present_bonds.py index 1639204b..333fbe30 100644 --- a/tests/forcefield/test_present_bonds.py +++ b/tests/forcefield/test_present_bonds.py @@ -1,12 +1,12 @@ -import os import pathlib from cgexplore.errors import ForceFieldUnitError +from .case_data import CaseData from .utilities import is_equivalent_atom -def test_present_bonds(molecule): +def test_present_bonds(molecule: CaseData) -> None: """Test methods toward :meth:`.ForceField._assign_bond_terms`. Parameters: @@ -19,18 +19,16 @@ def test_present_bonds(molecule): """ try: - force_fields = tuple(molecule.force_field_library.yield_forcefields()) + forcefields = tuple(molecule.forcefield_library.yield_forcefields()) - for i, ff in enumerate(force_fields): + for i, ff in enumerate(forcefields): assigned_system = ff.assign_terms( molecule=molecule.molecule, - output_dir=pathlib.Path( - os.path.dirname(os.path.realpath(__file__)) - ), + output_dir=pathlib.Path(__file__).resolve().parent, name=molecule.name, ) - present_terms = assigned_system.force_field_terms["bond"] + present_terms = assigned_system.forcefield_terms["bond"] print(present_terms) assert len(present_terms) == len(molecule.present_bonds[i]) for test, present in zip( diff --git a/tests/forcefield/test_present_nonbondeds.py b/tests/forcefield/test_present_nonbondeds.py index 48bd0ca4..48160982 100644 --- a/tests/forcefield/test_present_nonbondeds.py +++ b/tests/forcefield/test_present_nonbondeds.py @@ -1,10 +1,11 @@ -import os import pathlib from cgexplore.errors import ForceFieldUnitError +from .case_data import CaseData -def test_present_nonbondeds(molecule): + +def test_present_nonbondeds(molecule: CaseData) -> None: """Test methods toward :meth:`.ForceField._assign_nonbonded_terms`. Parameters: @@ -17,17 +18,15 @@ def test_present_nonbondeds(molecule): """ try: - force_fields = tuple(molecule.force_field_library.yield_forcefields()) - for i, ff in enumerate(force_fields): + forcefields = tuple(molecule.forcefield_library.yield_forcefields()) + for i, ff in enumerate(forcefields): assigned_system = ff.assign_terms( molecule=molecule.molecule, - output_dir=pathlib.Path( - os.path.dirname(os.path.realpath(__file__)) - ), + output_dir=pathlib.Path(__file__).resolve().parent, name=molecule.name, ) - present_terms = assigned_system.force_field_terms["nonbonded"] + present_terms = assigned_system.forcefield_terms["nonbonded"] print(present_terms) assert len(present_terms) == len(molecule.present_nonbondeds[i]) for test, present in zip( diff --git a/tests/forcefield/test_present_torsions.py b/tests/forcefield/test_present_torsions.py index db49e5d7..ece3efbf 100644 --- a/tests/forcefield/test_present_torsions.py +++ b/tests/forcefield/test_present_torsions.py @@ -1,10 +1,11 @@ -import os import pathlib from cgexplore.errors import ForceFieldUnitError +from .case_data import CaseData -def test_present_torsions(molecule): + +def test_present_torsions(molecule: CaseData) -> None: """Test methods toward :meth:`.ForceField._assign_torsion_terms`. Parameters: @@ -17,17 +18,15 @@ def test_present_torsions(molecule): """ try: - force_fields = tuple(molecule.force_field_library.yield_forcefields()) - for i, ff in enumerate(force_fields): + forcefields = tuple(molecule.forcefield_library.yield_forcefields()) + for i, ff in enumerate(forcefields): assigned_system = ff.assign_terms( molecule=molecule.molecule, - output_dir=pathlib.Path( - os.path.dirname(os.path.realpath(__file__)) - ), + output_dir=pathlib.Path(__file__).resolve().parent, name=molecule.name, ) - present_terms = assigned_system.force_field_terms["torsion"] + present_terms = assigned_system.forcefield_terms["torsion"] print(present_terms) assert len(present_terms) == len(molecule.present_torsions[i]) for test, present in zip( diff --git a/tests/forcefield/utilities.py b/tests/forcefield/utilities.py index 40dfaf9b..4ef5f1ab 100644 --- a/tests/forcefield/utilities.py +++ b/tests/forcefield/utilities.py @@ -1,4 +1,8 @@ -def is_equivalent_atom(atom1, atom2): +import stk + + +def is_equivalent_atom(atom1: stk.Atom, atom2: stk.Atom) -> None: + """Check if two atoms are equivalent.""" assert atom1.get_id() == atom2.get_id() assert atom1.get_charge() == atom2.get_charge() assert atom1.get_atomic_number() == atom2.get_atomic_number() diff --git a/tests/geom/case_data.py b/tests/geom/case_data.py index c4c1d97d..1b669fa5 100644 --- a/tests/geom/case_data.py +++ b/tests/geom/case_data.py @@ -3,11 +3,7 @@ class CaseData: - """A test case. - - Attributes: - - """ + """A test case.""" def __init__( self, diff --git a/tests/geom/test_angles.py b/tests/geom/test_angles.py index 0e073d5a..44dec507 100644 --- a/tests/geom/test_angles.py +++ b/tests/geom/test_angles.py @@ -1,7 +1,9 @@ import numpy as np +from .case_data import CaseData -def test_angles(molecule): + +def test_angles(molecule: CaseData) -> None: """Test :meth:`.GeomMeasure.calculate_angles`. Parameters: diff --git a/tests/geom/test_bonds.py b/tests/geom/test_bonds.py index 8cee6dbf..06b7a87c 100644 --- a/tests/geom/test_bonds.py +++ b/tests/geom/test_bonds.py @@ -1,11 +1,13 @@ import numpy as np +from .case_data import CaseData -def test_bonds(molecule): + +def test_bonds(molecule: CaseData) -> None: """Test :meth:`.GeomMeasure.calculate_bonds`. - Parameters - ---------- + Parameters: + molecule: The molecule with bonds. diff --git a/tests/geom/test_max_diameter.py b/tests/geom/test_max_diameter.py index 201d6561..1f47e92e 100644 --- a/tests/geom/test_max_diameter.py +++ b/tests/geom/test_max_diameter.py @@ -1,7 +1,9 @@ import numpy as np +from .case_data import CaseData -def test_maxdiameter(molecule): + +def test_maxdiameter(molecule: CaseData) -> None: """Test :meth:`.GeomMeasure.calculate_max_diameter`. Parameters: diff --git a/tests/geom/test_minb2b.py b/tests/geom/test_minb2b.py index 42b64043..ed7b09c1 100644 --- a/tests/geom/test_minb2b.py +++ b/tests/geom/test_minb2b.py @@ -1,4 +1,7 @@ -def test_minb2b(molecule): +from .case_data import CaseData + + +def test_minb2b(molecule: CaseData) -> None: """Test :meth:`.GeomMeasure.calculate_minb2b`. Parameters: diff --git a/tests/geom/test_radius_gyration.py b/tests/geom/test_radius_gyration.py index 56715914..14d8876d 100644 --- a/tests/geom/test_radius_gyration.py +++ b/tests/geom/test_radius_gyration.py @@ -1,7 +1,9 @@ import numpy as np +from .case_data import CaseData -def test_radius_gyration(molecule): + +def test_radius_gyration(molecule: CaseData) -> None: """Test :meth:`.GeomMeasure.calculate_radius_gyration`. Parameters: diff --git a/tests/geom/test_torsions.py b/tests/geom/test_torsions.py index 775bf25f..2dfceaa9 100644 --- a/tests/geom/test_torsions.py +++ b/tests/geom/test_torsions.py @@ -1,7 +1,9 @@ import numpy as np +from .case_data import CaseData -def test_torsions(molecule): + +def test_torsions(molecule: CaseData) -> None: """Test :meth:`.GeomMeasure.calculate_torsions`. Parameters: diff --git a/tests/optimisers/case_data.py b/tests/optimisers/case_data.py index 6679209b..5ac834de 100644 --- a/tests/optimisers/case_data.py +++ b/tests/optimisers/case_data.py @@ -1,21 +1,18 @@ +import cgexplore import stk class CaseData: - """A test case. - - Attributes: - - """ + """A test case.""" def __init__( self, molecule: stk.Molecule, - force_field, - xml_string, + forcefield: cgexplore.forcefield.ForceField, + xml_string: str, name: str, ) -> None: self.molecule = molecule - self.force_field = force_field + self.forcefield = forcefield self.xml_string = xml_string self.name = name diff --git a/tests/optimisers/conftest.py b/tests/optimisers/conftest.py index a9497d45..1fde8af6 100644 --- a/tests/optimisers/conftest.py +++ b/tests/optimisers/conftest.py @@ -27,7 +27,7 @@ ) ), ), - force_field=ForceField( + forcefield=ForceField( identifier="testff", prefix="testffprefix", present_beads=( @@ -129,7 +129,7 @@ ) ), ), - force_field=ForceField( + forcefield=ForceField( identifier="testff", prefix="testffprefix", present_beads=( diff --git a/tests/precursors/case_data.py b/tests/precursors/case_data.py index 7e020362..3c954910 100644 --- a/tests/precursors/case_data.py +++ b/tests/precursors/case_data.py @@ -1,12 +1,10 @@ +import numpy as np +import numpy.typing as npt from cgexplore.molecule_construction import Precursor class CaseData: - """A test case. - - Attributes: - - """ + """A test case.""" def __init__( self, @@ -15,7 +13,7 @@ def __init__( num_fgs: int, bead_set: dict, smiles: str, - position_matrix, + position_matrix: npt.NDArray[np.float64], name: str, ) -> None: self.precursor = precursor diff --git a/tests/precursors/test_bead_set.py b/tests/precursors/test_bead_set.py index f5d778ff..6ac2b710 100644 --- a/tests/precursors/test_bead_set.py +++ b/tests/precursors/test_bead_set.py @@ -1,4 +1,7 @@ -def test_bead_set(precursor): +from .case_data import CaseData + + +def test_bead_set(precursor: CaseData) -> None: """Test :meth:`.Precursor.get_bead_set`. Parameters: diff --git a/tests/precursors/test_buildingblock.py b/tests/precursors/test_buildingblock.py index b190c8fb..74771d8a 100644 --- a/tests/precursors/test_buildingblock.py +++ b/tests/precursors/test_buildingblock.py @@ -1,8 +1,10 @@ import numpy as np import stk +from .case_data import CaseData -def test_building_block(precursor): + +def test_building_block(precursor: CaseData) -> None: """Test :meth:`.Precursor.get_building_block`. Parameters: diff --git a/tests/precursors/test_functional_groups.py b/tests/precursors/test_functional_groups.py index eec904a2..ffc24009 100644 --- a/tests/precursors/test_functional_groups.py +++ b/tests/precursors/test_functional_groups.py @@ -1,4 +1,7 @@ -def test_functional_groups(precursor): +from .case_data import CaseData + + +def test_functional_groups(precursor: CaseData) -> None: """Test :class:`.Precursor`. Parameters: diff --git a/tests/precursors/test_name.py b/tests/precursors/test_name.py index 43228293..5c545b78 100644 --- a/tests/precursors/test_name.py +++ b/tests/precursors/test_name.py @@ -1,4 +1,7 @@ -def test_name(precursor): +from .case_data import CaseData + + +def test_name(precursor: CaseData) -> None: """Test :meth:`.Precursor.get_name`. Parameters: diff --git a/tests/shape/__init__.py b/tests/shape/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/shape/case_data.py b/tests/shape/case_data.py new file mode 100644 index 00000000..2e28cae1 --- /dev/null +++ b/tests/shape/case_data.py @@ -0,0 +1,20 @@ +import stk + + +class CaseData: + """A test case.""" + + def __init__( + self, + molecule: stk.Molecule, + shape_dict: dict[str:float], + expected_points: int, + shape_string: str, + name: str, + ) -> None: + """Initialize CaseData.""" + self.molecule = molecule + self.shape_dict = shape_dict + self.expected_points = expected_points + self.shape_string = shape_string + self.name = name diff --git a/tests/shape/conftest.py b/tests/shape/conftest.py new file mode 100644 index 00000000..757b0749 --- /dev/null +++ b/tests/shape/conftest.py @@ -0,0 +1,143 @@ +import numpy as np +import pytest +import stk + +from .case_data import CaseData + +# Three tests cases with four atoms with known bond lengths, +# angles and torsions. + + +@pytest.fixture( + params=( + lambda name: CaseData( + molecule=stk.BuildingBlock( + smiles="C[Fe](C)(C)(C)(C)C", + position_matrix=np.array( + ( + [1, 0, 0], + [0, 0, 0], + [-1, 0, 0], + [0, 1, 0], + [0, -1, 0], + [0, 0, 1], + [0, 0, -1], + ) + ), + ), + shape_dict={ + "HP-6": 33.333, + "PPY-6": 30.153, + "OC-6": 0.0, + "TPR-6": 16.737, + "JPPY-6": 33.916, + }, + expected_points=6, + shape_string="[C].[C].[C].[C].[C].[C]", + name=name, + ), + lambda name: CaseData( + molecule=stk.BuildingBlock( + smiles="C[Fe](C)(C)(C)(C)C", + position_matrix=np.array( + ( + [1, 0.2, 0], + [0, 0, 0], + [-1.2, 0, 0], + [0.3, 1, 0], + [0, -1, 0], + [0, 0, 1], + [0, 0.5, -1], + ) + ), + ), + shape_dict={ + "HP-6": 25.184, + "PPY-6": 21.006, + "OC-6": 3.9, + "TPR-6": 11.582, + "JPPY-6": 23.9, + }, + expected_points=6, + shape_string="[C].[C].[C].[C].[C].[C]", + name=name, + ), + lambda name: CaseData( + molecule=stk.BuildingBlock( + smiles="C[Fe](C)(C)(C)(C)C", + position_matrix=np.array( + ( + [1, 0, 0], + [0, 0, 0], + [-1, 0, 0], + [0, 1, 0], + [0, -1, 0], + [0, 0, 1], + [0, 0, -1], + ) + ) + * 4, + ), + shape_dict={ + "HP-6": 33.333, + "PPY-6": 30.153, + "OC-6": 0.0, + "TPR-6": 16.737, + "JPPY-6": 33.916, + }, + expected_points=6, + shape_string="[C].[C].[C].[C].[C].[C]", + name=name, + ), + lambda name: CaseData( + molecule=stk.BuildingBlock( + smiles="C[Zn](C)(C)C", + position_matrix=np.array( + ( + [0.5541, 0.7996, 0.4965], + [0, 0, 0], + [0.6833, -0.8134, -0.2536], + [-0.7782, -0.3735, 0.6692], + [-0.4593, 0.3874, -0.9121], + ) + ), + ), + shape_dict={ + "SP-4": 33.332, + "T-4": 0.0, + "SS-4": 7.212, + "vTBPY-4": 2.287, + }, + expected_points=4, + shape_string="[C].[C].[C].[C]", + name=name, + ), + lambda name: CaseData( + molecule=stk.BuildingBlock( + smiles="C[Zn](C)(C)C", + position_matrix=np.array( + ( + [0.5541, 0.7996, 0.4965], + [0, 0, 0], + [0.6833, -0.2234, -0.2536], + [-0.7782, -0.4735, 0.6692], + [-0.4593, 0.3874, -0.9121], + ) + ), + ), + shape_dict={ + "SP-4": 25.71, + "T-4": 3.947, + "SS-4": 6.45, + "vTBPY-4": 1.731, + }, + expected_points=4, + shape_string="[C].[C].[C].[C]", + name=name, + ), + ) +) +def molecule(request: pytest.FixtureRequest) -> CaseData: + return request.param( + f"{request.fixturename}{request.param_index}", + ) diff --git a/tests/shape/test_shape.py b/tests/shape/test_shape.py new file mode 100644 index 00000000..618ad818 --- /dev/null +++ b/tests/shape/test_shape.py @@ -0,0 +1,48 @@ +import pathlib +import shutil + +import pytest +from cgexplore.shape import ShapeMeasure + +from .case_data import CaseData + + +def test_shape(molecule: CaseData) -> None: + """Test :meth:`.ShapeMeasure.calculate`. + + Parameters: + + molecule: + The molecule to test. + + Returns: + None : :class:`NoneType` + + """ + # Search for shape executable, if it exists, run the test. + expected_shape_path = pathlib.Path.home() / ( + "software/shape_2.1_linux_64/SHAPE_2.1_linux_64/shape_2.1_linux64" + ) + + if expected_shape_path.exists(): + # Do test. + output_dir = pathlib.Path(__file__).resolve().parent / ( + f"{molecule.name}_test_output" + ) + shape_calc = ShapeMeasure( + output_dir=output_dir, + shape_path=expected_shape_path, + ) + shape_mol = shape_calc.get_shape_molecule_byelement( + molecule=molecule.molecule, + element="C", + expected_points=molecule.expected_points, + ) + + shape_output = shape_calc.calculate(shape_mol) + print(shape_output, molecule.shape_dict) + assert shape_output == molecule.shape_dict + shutil.rmtree(output_dir) + + else: + pytest.skip(f"shape software not found at {expected_shape_path}") diff --git a/tests/shape/test_shape_molecule.py b/tests/shape/test_shape_molecule.py new file mode 100644 index 00000000..e7612506 --- /dev/null +++ b/tests/shape/test_shape_molecule.py @@ -0,0 +1,30 @@ +import stk +from cgexplore.shape import ShapeMeasure + +from .case_data import CaseData + + +def test_shape(molecule: CaseData) -> None: + """Test :meth:`.ShapeMeasure.get_shape_molecule_byelement`. + + Parameters: + + molecule: + The molecule to test. + + Returns: + None : :class:`NoneType` + + """ + shape_calc = ShapeMeasure( + output_dir="fake_output", + shape_path="fake_output", + ) + shape_mol = shape_calc.get_shape_molecule_byelement( + molecule=molecule.molecule, + element="C", + expected_points=molecule.expected_points, + ) + + print(stk.Smiles().get_key(shape_mol)) + assert stk.Smiles().get_key(shape_mol) == molecule.shape_string