From 9548d0cf0a4eca0abb31288f5f48697af923ffdd Mon Sep 17 00:00:00 2001 From: "Matthew W. Thompson" Date: Fri, 12 Jun 2026 09:08:10 -0500 Subject: [PATCH 1/6] Set up testing on Python 3.14 --- .github/workflows/base-ci.yaml | 4 +- .github/workflows/dev-ci.yaml | 110 ----------------------------- .github/workflows/examples-ci.yaml | 91 ------------------------ .github/workflows/gh-ci.yaml | 14 ++-- .github/workflows/gpu.yaml | 102 -------------------------- 5 files changed, 7 insertions(+), 314 deletions(-) delete mode 100644 .github/workflows/dev-ci.yaml delete mode 100644 .github/workflows/examples-ci.yaml delete mode 100644 .github/workflows/gpu.yaml diff --git a/.github/workflows/base-ci.yaml b/.github/workflows/base-ci.yaml index a14e5c2b..c12efe0d 100644 --- a/.github/workflows/base-ci.yaml +++ b/.github/workflows/base-ci.yaml @@ -31,8 +31,8 @@ jobs: fail-fast: false matrix: os: [macOS-latest, ubuntu-latest] - python-version: ["3.12"] - include-openeye: [false, true] + python-version: ["3.14"] + include-openeye: [false] steps: - uses: actions/checkout@v6 diff --git a/.github/workflows/dev-ci.yaml b/.github/workflows/dev-ci.yaml deleted file mode 100644 index ae0ed134..00000000 --- a/.github/workflows/dev-ci.yaml +++ /dev/null @@ -1,110 +0,0 @@ -name: Upstream nightly version CI -on: - schedule: - # weekly tests, Sundays at midnight - - cron: "0 0 * * 0" - workflow_dispatch: - -concurrency: - # Specific group naming so CI is only cancelled - # within same PR or on merge to main - group: ${{ github.ref }}-${{ github.head_ref }}-${{ github.workflow }} - cancel-in-progress: true - -defaults: - run: - shell: bash -l {0} - -env: - OE_LICENSE: ${{ github.workspace }}/oe_license.txt - DGL_HOME: ${{ github.workspace }}/dgl - -jobs: - nightly_check: - name: Nightly check - runs-on: ${{ matrix.os }} - strategy: - fail-fast: false - matrix: - os: [macOS-latest, ubuntu-latest] - python-version: ["3.12"] - - steps: - - uses: actions/checkout@v6 - - - name: Build information - run: | - uname -a - df -h - ulimit -a - - - name: Install environment - uses: mamba-org/setup-micromamba@v3 - with: - environment-file: devtools/conda-envs/nightly.yaml - create-args: >- - python=${{ matrix.python-version }} - pydantic=2 - - - name: Install nightly pytorch - run: | - python -m pip install https://github.com/Lightning-AI/lightning/archive/refs/heads/master.zip -U - - - name: Download DGL source - run: | - git clone --recurse-submodules https://github.com/dmlc/dgl.git - - - name: Set up DGL - if: matrix.os != 'ubuntu-latest' - run: | - # from https://docs.dgl.ai/en/latest/install/index.html#macos - cd dgl/ - mkdir build && cd build - cmake -DUSE_OPENMP=off -DUSE_LIBXSMM=OFF .. - make -j4 - cd ../.. - - - name: Set up DGL (Ubuntu) - if: matrix.os == 'ubuntu-latest' - run: | - cd dgl/ - mkdir build && cd build - cmake -DBUILD_TYPE=dev -DUSE_CUDA=OFF -DUSE_LIBXSMM=OFF .. - make - cd ../.. - - - name: Install DGL Python bindings - run: | - cd dgl/python - python setup.py install - python setup.py build_ext --inplace - cd ../.. - - - uses: ./.github/actions/include-openeye - with: - openeye-license-text: ${{ secrets.OE_LICENSE }} - openeye-license-file: ${{ env.OE_LICENSE }} - - - name: Check toolkit installations - shell: bash -l -c "python -u {0}" - run: | - from openff.toolkit.utils.toolkits import OPENEYE_AVAILABLE, RDKIT_AVAILABLE - assert str(OPENEYE_AVAILABLE).lower() == 'true' - assert str(RDKIT_AVAILABLE).lower() == 'true' - - - name: Check Python environment - run: | - pip list - - - name: Import DGL - run: | - python -c "import dgl; import dgl.nn" - - - name: Install package - run: | - python -m pip install . --no-deps - - - name: Run tests - run: | - python -m pytest -n 4 -v --cov=openff/nagl --cov-config=pyproject.toml --cov-append --cov-report=xml --color=yes openff/nagl/ - diff --git a/.github/workflows/examples-ci.yaml b/.github/workflows/examples-ci.yaml deleted file mode 100644 index 9a7c9130..00000000 --- a/.github/workflows/examples-ci.yaml +++ /dev/null @@ -1,91 +0,0 @@ -name: Examples CI -on: - push: - branches: - - main - pull_request: - branches: - - main - schedule: - # weekly tests, Sundays at midnight - - cron: "0 0 * * 0" - -concurrency: - # Specific group naming so CI is only cancelled - # within same PR or on merge to main - group: ${{ github.ref }}-${{ github.head_ref }}-${{ github.workflow }} - cancel-in-progress: true - -defaults: - run: - shell: bash -l {0} - -env: - OE_LICENSE: ${{ github.workspace }}/oe_license.txt - -jobs: - example_tests: - name: Examples CI (${{ matrix.os }}, py-${{ matrix.python-version }}, rdkit=${{ matrix.include-rdkit }}, openeye=${{ matrix.include-openeye }}, dgl=${{ matrix.include-dgl }}), pydantic=${{ matrix.pydantic-version }} - runs-on: ${{ matrix.os }} - strategy: - fail-fast: false - matrix: - os: [ubuntu-latest] - python-version: ["3.12"] - pydantic-version: ["2"] - include-rdkit: [true] - include-openeye: [false] - include-dgl: [true] - - - steps: - - uses: actions/checkout@v6 - - - name: Build information - run: | - uname -a - df -h - ulimit -a - - - name: Install environment - uses: mamba-org/setup-micromamba@v3 - with: - environment-file: devtools/conda-envs/examples_env.yaml - create-args: >- - python=${{ matrix.python-version }} - pydantic=${{ matrix.pydantic-version }} - - - name: Install package - run: | - python -m pip install . --no-deps - - - uses: ./.github/actions/include-openeye - if: matrix.include-openeye == true - with: - openeye-license-text: ${{ secrets.OE_LICENSE }} - openeye-license-file: ${{ env.OE_LICENSE }} - - - - name: Uninstall OpenEye - if: matrix.include-openeye == false - run: conda remove --force openeye-toolkits --yes || echo "openeye not installed" - - - name: Uninstall RDKit - if: matrix.include-rdkit == false - run: conda remove --force rdkit --yes || echo "rdkit not installed" - - - name: Python information - run: | - which python - conda info - conda list - - - name: Check toolkit installations - shell: bash -l -c "python -u {0}" - run: | - from openff.toolkit.utils.toolkits import OPENEYE_AVAILABLE, RDKIT_AVAILABLE - assert str(OPENEYE_AVAILABLE).lower() == '${{ matrix.include-openeye }}' - assert str(RDKIT_AVAILABLE).lower() == '${{ matrix.include-rdkit }}' - - - name: Run example notebooks - run: python -m pytest -r fE -v -x --tb=short -nauto --nbval-lax --nbval-cell-timeout=50000 --dist loadscope examples --ignore=examples/train-gnn-to-electric-field/ diff --git a/.github/workflows/gh-ci.yaml b/.github/workflows/gh-ci.yaml index cb1eb96f..469ff2fb 100644 --- a/.github/workflows/gh-ci.yaml +++ b/.github/workflows/gh-ci.yaml @@ -31,15 +31,11 @@ jobs: fail-fast: false matrix: os: [macOS-latest, ubuntu-latest] - python-version: ["3.12"] - pydantic-version: ["1", "2"] - include-rdkit: [false, true] - include-openeye: [false, true] - include-dgl: [false, true] - exclude: - - include-rdkit: false - include-openeye: false - + python-version: ["3.14"] + pydantic-version: ["2"] + include-rdkit: [true] + include-openeye: [false] + include-dgl: [false] steps: - uses: actions/checkout@v6 diff --git a/.github/workflows/gpu.yaml b/.github/workflows/gpu.yaml deleted file mode 100644 index f7661399..00000000 --- a/.github/workflows/gpu.yaml +++ /dev/null @@ -1,102 +0,0 @@ -name: GPU-enabled CI -on: - workflow_dispatch: - pull_request: - branches: - - "main" - -defaults: - run: - shell: bash -l {0} - -env: - OE_LICENSE: ${{ github.workspace }}/oe_license.txt - -jobs: - start-aws-runner: - runs-on: ubuntu-latest - permissions: - id-token: write - contents: read - outputs: - mapping: ${{ steps.aws-start.outputs.mapping }} - steps: - - name: Configure AWS credentials - uses: aws-actions/configure-aws-credentials@v6 - with: - role-to-assume: arn:aws:iam::649715411074:role/gh-actions-runner-role - aws-region: us-east-1 - - name: Create cloud runner - id: aws-start - uses: omsf/start-aws-gha-runner@v1.1.1 - with: - aws_image_id: ami-0d5079d9be06933e5 - aws_instance_type: g4dn.xlarge - # IAM default might be 5 GB? - aws_root_device_size: 125 - aws_region_name: us-east-1 - aws_home_dir: /home/ubuntu - env: - GH_PAT: ${{ secrets.GH_PAT }} - self-hosted-test: - runs-on: self-hosted - needs: - - start-aws-runner - steps: - - uses: actions/checkout@v6 - with: - fetch-depth: 0 - - - name: Print disk usage - run: "df -h" - - - name: Print Docker details - run: "docker version || true" - - - name: Setup Conda Environment - uses: mamba-org/setup-micromamba@v3 - with: - environment-file: devtools/conda-envs/test_cuda_env.yaml - - - name: Install Package - run: python -m pip install . --no-deps - - - name: Double-check local installation - run: python -c "from openff.nagl import __version__; print(__version__)" - - - name: Check that PyTorch can see CUDA - run: python -c "import torch; assert torch.cuda.is_available()" - - - name: Check we can see DGL - run: python -c "import dgl; print(dgl.__version__)" - - - uses: ./.github/actions/include-openeye - with: - openeye-license-text: ${{ secrets.OE_LICENSE }} - openeye-license-file: ${{ env.OE_LICENSE }} - - - name: Run tests - run: python -m pytest -n 4 -v --cov=openff/nagl --cov-config=pyproject.toml --cov-append --cov-report=xml --color=yes openff/nagl/ - - stop-aws-runner: - runs-on: ubuntu-latest - permissions: - id-token: write - contents: read - needs: - - start-aws-runner - - self-hosted-test - if: ${{ always() }} - steps: - - name: Configure AWS credentials - uses: aws-actions/configure-aws-credentials@v6 - with: - role-to-assume: arn:aws:iam::649715411074:role/gh-actions-runner-role - aws-region: us-east-1 - - name: Stop instances - uses: omsf/stop-aws-gha-runner@v1.0.0 - with: - instance_mapping: ${{ needs.start-aws-runner.outputs.mapping }} - aws_region_name: us-east-1 - env: - GH_PAT: ${{ secrets.GH_PAT }} From 1cd0fb22d91579a0829ea8b13f41faf7e236cfa2 Mon Sep 17 00:00:00 2001 From: "Matthew W. Thompson" Date: Fri, 12 Jun 2026 09:12:10 -0500 Subject: [PATCH 2/6] No AmberTools --- devtools/conda-envs/test_env_dgl_false.yaml | 1 - 1 file changed, 1 deletion(-) diff --git a/devtools/conda-envs/test_env_dgl_false.yaml b/devtools/conda-envs/test_env_dgl_false.yaml index 440d7347..3673af81 100644 --- a/devtools/conda-envs/test_env_dgl_false.yaml +++ b/devtools/conda-envs/test_env_dgl_false.yaml @@ -18,7 +18,6 @@ dependencies: - pydantic <3 - rdkit !=2024.03.6 - scipy - - ambertools # # database # - pyarrow From 7ae6fce389bb0c7927e05bddd89d050ec0c382b1 Mon Sep 17 00:00:00 2001 From: "Matthew W. Thompson" Date: Fri, 12 Jun 2026 10:12:07 -0500 Subject: [PATCH 3/6] Fix abc/base method sneaking in --- openff/nagl/_base/metaregistry.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/openff/nagl/_base/metaregistry.py b/openff/nagl/_base/metaregistry.py index b55f25e2..d90bf6b5 100644 --- a/openff/nagl/_base/metaregistry.py +++ b/openff/nagl/_base/metaregistry.py @@ -51,7 +51,7 @@ def __init__(cls, name, bases, namespace, **kwargs): super().__init__(name, bases, namespace, **kwargs) if hasattr(cls, cls._key_attribute): attr_val = getattr(cls, cls._key_attribute) - if not isinstance(attr_val, property) and attr_val: + if not isinstance(attr_val, property) and not callable(attr_val) and attr_val: cls.registry[cls._key_transform(attr_val)] = cls return RegistryMeta From 2ee9a2290d9846cdad6f8327d6d0fedee34f7cab Mon Sep 17 00:00:00 2001 From: "Matthew W. Thompson" Date: Fri, 12 Jun 2026 10:26:51 -0500 Subject: [PATCH 4/6] Debug --- .github/workflows/base-ci.yaml | 2 +- .github/workflows/gh-ci.yaml | 5 ++++- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/.github/workflows/base-ci.yaml b/.github/workflows/base-ci.yaml index c12efe0d..29f40854 100644 --- a/.github/workflows/base-ci.yaml +++ b/.github/workflows/base-ci.yaml @@ -18,7 +18,7 @@ concurrency: defaults: run: - shell: bash -l {0} + shell: bash -leo pipefail {0} env: OE_LICENSE: ${{ github.workspace }}/oe_license.txt diff --git a/.github/workflows/gh-ci.yaml b/.github/workflows/gh-ci.yaml index 469ff2fb..9a3539ec 100644 --- a/.github/workflows/gh-ci.yaml +++ b/.github/workflows/gh-ci.yaml @@ -18,7 +18,7 @@ concurrency: defaults: run: - shell: bash -l {0} + shell: bash -leo pipefail {0} env: OE_LICENSE: ${{ github.workspace }}/oe_license.txt @@ -79,6 +79,9 @@ jobs: mkdir -p ~/.dgl echo '{"backend": "pytorch"}' > ~/.dgl/config.json + - name: Target Pydantic error? + run: python -c "from openff.nagl.training.loss import TargetType" + - name: Check toolkit installations shell: bash -l -c "python -u {0}" run: | From 935369305d73515b856a35f0e140c728a75bbf8f Mon Sep 17 00:00:00 2001 From: "Matthew W. Thompson" Date: Fri, 12 Jun 2026 10:32:38 -0500 Subject: [PATCH 5/6] Minimal v2 fix --- openff/nagl/training/loss.py | 22 ++++++++++------------ 1 file changed, 10 insertions(+), 12 deletions(-) diff --git a/openff/nagl/training/loss.py b/openff/nagl/training/loss.py index 4b49b177..92dee12a 100644 --- a/openff/nagl/training/loss.py +++ b/openff/nagl/training/loss.py @@ -1,29 +1,27 @@ """Targets for calculating loss when training neural networks""" +from __future__ import annotations import abc import pathlib import typing import torch + +from openff.nagl._base.base import ImmutableModel from openff.nagl._base.metaregistry import create_registry_metaclass from openff.nagl.molecule._dgl import DGLMoleculeOrBatch -from openff.nagl.training.metrics import MetricType #MetricMeta, BaseMetric -from openff.nagl._base.base import ImmutableModel -from openff.nagl.nn._pooling import PoolingLayer from openff.nagl.nn._containers import ReadoutModule +from openff.nagl.nn._pooling import PoolingLayer +from openff.nagl.training.metrics import MetricType #MetricMeta, BaseMetric -try: - from pydantic.v1 import Field, validator - from pydantic.v1.main import ModelMetaclass -except ImportError: - from pydantic import Field, validator - from pydantic.main import ModelMetaclass +from pydantic import Field, field_validator if typing.TYPE_CHECKING: import torch - from openff.nagl.molecule._dgl import DGLMoleculeOrBatch from openff.toolkit import Molecule + from openff.nagl.molecule._dgl import DGLMoleculeOrBatch + __all__ = [ "ReadoutTarget", @@ -56,7 +54,7 @@ class _BaseTarget(ImmutableModel, abc.ABC): #, metaclass=_TargetMeta): ) ) - @validator("metric", pre=True) + @field_validator("metric") def _validate_metric(cls, v): if isinstance(v, str): v = {"name": v} @@ -477,4 +475,4 @@ def evaluate_target( SingleDipoleTarget, MultipleESPTarget, GeneralLinearFitTarget -] \ No newline at end of file +] From 5fd50c16d1b4b0c3a9187ee2131057fc37a053df Mon Sep 17 00:00:00 2001 From: "Matthew W. Thompson" Date: Mon, 15 Jun 2026 09:44:01 -0500 Subject: [PATCH 6/6] Add `from __future__ import annotations` --- docs/conf.py | 3 +- .../prepare-dataset/prepare-dataset.ipynb | 5 +- .../search_hyperparameters.py | 13 +++--- .../train-gnn-notebook.ipynb | 15 +++--- .../train-gnn-to-electric-field.ipynb | 46 ++++++++----------- examples/train-gnn/train_gnn.py | 2 + .../train-gnn-notebook.ipynb | 20 ++++---- openff/nagl/__init__.py | 3 ++ openff/nagl/_base/__init__.py | 2 + openff/nagl/_base/base.py | 7 +-- openff/nagl/_base/metaregistry.py | 5 +- openff/nagl/_version.py | 3 +- openff/nagl/config/__init__.py | 1 + openff/nagl/config/data.py | 4 +- openff/nagl/config/model.py | 6 +-- openff/nagl/config/optimizer.py | 1 + openff/nagl/config/training.py | 2 + openff/nagl/domains.py | 3 ++ openff/nagl/features/__init__.py | 2 + openff/nagl/features/_base.py | 6 ++- openff/nagl/features/_data.py | 2 + openff/nagl/features/_featurizers.py | 11 +++-- openff/nagl/features/_utils.py | 2 + openff/nagl/features/atoms.py | 10 ++-- openff/nagl/features/bonds.py | 4 +- openff/nagl/label/dataset.py | 12 +++-- openff/nagl/label/labels.py | 10 ++-- openff/nagl/label/utils.py | 2 + openff/nagl/lookups.py | 3 ++ openff/nagl/molecule/__init__.py | 5 +- openff/nagl/molecule/_base.py | 5 +- openff/nagl/molecule/_dgl/__init__.py | 2 + openff/nagl/molecule/_dgl/batch.py | 7 ++- openff/nagl/molecule/_dgl/molecule.py | 8 +++- openff/nagl/molecule/_dgl/utils.py | 15 +++--- openff/nagl/molecule/_graph/_batch.py | 5 +- openff/nagl/molecule/_graph/_graph.py | 15 +++--- openff/nagl/molecule/_graph/_update.py | 7 +-- openff/nagl/molecule/_graph/_utils.py | 10 ++-- openff/nagl/molecule/_graph/molecule.py | 7 ++- openff/nagl/molecule/_utils.py | 8 ++-- openff/nagl/nn/__init__.py | 12 ++--- openff/nagl/nn/_base.py | 2 + openff/nagl/nn/_containers.py | 11 +++-- openff/nagl/nn/_dataset.py | 14 +++--- openff/nagl/nn/_models.py | 13 ++++-- openff/nagl/nn/_pooling.py | 4 +- openff/nagl/nn/_sequential.py | 4 +- openff/nagl/nn/activation.py | 1 + openff/nagl/nn/gcn/__init__.py | 4 +- openff/nagl/nn/gcn/_base.py | 4 +- openff/nagl/nn/gcn/_function.py | 3 ++ openff/nagl/nn/gcn/_gin.py | 5 +- openff/nagl/nn/gcn/_sage.py | 11 +++-- openff/nagl/nn/postprocess.py | 1 + openff/nagl/tests/_base/test_base.py | 10 ++-- openff/nagl/tests/_base/test_metaregistry.py | 3 ++ openff/nagl/tests/config/test_model_config.py | 2 + openff/nagl/tests/conftest.py | 6 +-- .../data/example_charges/generate-new-sdfs.py | 4 +- openff/nagl/tests/data/files.py | 1 + openff/nagl/tests/features/test_features.py | 2 + .../nagl/tests/features/test_featurizers.py | 4 +- openff/nagl/tests/features/test_utils.py | 2 + openff/nagl/tests/label/conftest.py | 2 + .../nagl/tests/label/test_labelled_dataset.py | 7 ++- openff/nagl/tests/label/test_labels.py | 7 +-- openff/nagl/tests/molecule/dgl/test_batch.py | 3 +- .../nagl/tests/molecule/dgl/test_molecule.py | 4 +- openff/nagl/tests/molecule/dgl/test_utils.py | 6 ++- .../nagl/tests/molecule/graph/test_batch.py | 2 + .../nagl/tests/molecule/graph/test_graph.py | 22 +++++---- .../tests/molecule/graph/test_molecule.py | 6 ++- openff/nagl/tests/molecule/test_utils.py | 2 + openff/nagl/tests/nn/gcn/test_gin.py | 4 +- openff/nagl/tests/nn/gcn/test_sage.py | 6 ++- openff/nagl/tests/nn/test_activation.py | 6 ++- openff/nagl/tests/nn/test_containers.py | 8 ++-- openff/nagl/tests/nn/test_dataset.py | 16 ++++--- openff/nagl/tests/nn/test_model.py | 29 ++++++------ openff/nagl/tests/nn/test_pooling.py | 2 + openff/nagl/tests/nn/test_postprocess.py | 9 ++-- openff/nagl/tests/nn/test_sequential.py | 2 + openff/nagl/tests/test_domain.py | 3 ++ openff/nagl/tests/test_lookups.py | 10 ++-- openff/nagl/tests/testing/torch.py | 2 + openff/nagl/tests/testing/utils.py | 2 + openff/nagl/tests/training/test_loss.py | 20 ++++---- openff/nagl/tests/training/test_metrics.py | 7 ++- openff/nagl/tests/training/test_training.py | 37 ++++++--------- openff/nagl/tests/utils/test_openff.py | 35 +++++++------- openff/nagl/tests/utils/test_resonance.py | 5 +- openff/nagl/tests/utils/test_utils.py | 2 + openff/nagl/toolkits/__init__.py | 1 + openff/nagl/toolkits/_base.py | 5 +- openff/nagl/toolkits/openeye.py | 7 ++- openff/nagl/toolkits/openff.py | 20 ++++---- openff/nagl/toolkits/rdkit.py | 8 ++-- openff/nagl/toolkits/registry.py | 6 ++- openff/nagl/training/metrics.py | 3 +- openff/nagl/training/templates/files.py | 2 + openff/nagl/training/training.py | 16 +++---- openff/nagl/utils/_hash.py | 2 + openff/nagl/utils/_parallelization.py | 7 ++- openff/nagl/utils/_time.py | 2 + openff/nagl/utils/_types.py | 4 +- openff/nagl/utils/_utils.py | 2 + openff/nagl/utils/resonance.py | 8 ++-- pyproject.toml | 6 +++ versioneer.py | 7 +-- 110 files changed, 485 insertions(+), 309 deletions(-) diff --git a/docs/conf.py b/docs/conf.py index ac32bf95..c691dd03 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -13,13 +13,14 @@ # documentation root, use os.path.abspath to make it absolute, like shown here. # In case the project was not installed +from __future__ import annotations + import os import sys sys.path.insert(0, os.path.abspath("..")) import openff.nagl - # -- Project information ----------------------------------------------------- project = "OpenFF NAGL" diff --git a/examples/prepare-dataset/prepare-dataset.ipynb b/examples/prepare-dataset/prepare-dataset.ipynb index acd40cc0..48fa8e4d 100644 --- a/examples/prepare-dataset/prepare-dataset.ipynb +++ b/examples/prepare-dataset/prepare-dataset.ipynb @@ -41,11 +41,12 @@ } ], "source": [ - "from pathlib import Path\n", + "from __future__ import annotations\n", "\n", - "from tqdm import tqdm\n", + "from pathlib import Path\n", "\n", "from openff.toolkit.topology import Molecule\n", + "from tqdm import tqdm\n", "\n", "from openff.nagl.label.dataset import LabelledDataset\n", "from openff.nagl.label.labels import LabelCharges" diff --git a/examples/search-hyperparameters/search_hyperparameters.py b/examples/search-hyperparameters/search_hyperparameters.py index 6748baa5..0d20b7f5 100644 --- a/examples/search-hyperparameters/search_hyperparameters.py +++ b/examples/search-hyperparameters/search_hyperparameters.py @@ -1,11 +1,14 @@ +from __future__ import annotations + +import math import os import pathlib -import math import pickle -from typing import Dict, Any, Tuple +from typing import Any, Dict, Tuple import click + def train_model( config: Dict[str, Any] = {}, output_directory: str = ".", @@ -165,11 +168,9 @@ def tune_model( partial_charge_method: str = "am1", postprocess_layer: str = "compute_partial_charges" ): - import yaml - import ray - from ray import air - from ray import tune + import yaml + from ray import air, tune from ray.tune import CLIReporter from ray.tune.schedulers import ASHAScheduler diff --git a/examples/train-gnn-notebook/train-gnn-notebook.ipynb b/examples/train-gnn-notebook/train-gnn-notebook.ipynb index d838a2b3..be27aa44 100644 --- a/examples/train-gnn-notebook/train-gnn-notebook.ipynb +++ b/examples/train-gnn-notebook/train-gnn-notebook.ipynb @@ -34,10 +34,11 @@ "metadata": {}, "outputs": [], "source": [ + "from __future__ import annotations\n", + "\n", "from pathlib import Path\n", "\n", "import numpy as np\n", - "\n", "from openff.toolkit import Molecule\n", "from openff.units import unit" ] @@ -113,13 +114,14 @@ "metadata": {}, "outputs": [], "source": [ + "from torch.nn import ReLU\n", + "\n", + "from openff.nagl import GNNModel\n", "from openff.nagl.config.model import (\n", " ConvolutionLayer,\n", " ConvolutionModule,\n", ")\n", - "from openff.nagl import GNNModel\n", "from openff.nagl.nn.gcn import SAGEConvStack\n", - "from torch.nn import ReLU\n", "from openff.nagl.nn.postprocess import ComputePartialCharges\n", "\n", "single_convolution_layer = ConvolutionLayer(\n", @@ -547,10 +549,9 @@ "metadata": {}, "outputs": [], "source": [ - "from openff.nagl.config.data import DatasetConfig, DataConfig\n", - "from openff.nagl.training.metrics import RMSEMetric\n", + "from openff.nagl.config.data import DataConfig, DatasetConfig\n", "from openff.nagl.training.loss import ReadoutTarget\n", - "\n", + "from openff.nagl.training.metrics import RMSEMetric\n", "\n", "charge_rmse_target = ReadoutTarget(\n", " metric=RMSEMetric(), # use RMSE to calculate loss\n", @@ -624,7 +625,7 @@ "metadata": {}, "outputs": [], "source": [ - "from openff.nagl.training.training import TrainingGNNModel, DGLMoleculeDataModule\n", + "from openff.nagl.training.training import DGLMoleculeDataModule, TrainingGNNModel\n", "\n", "training_model = TrainingGNNModel(training_config)\n", "data_module = DGLMoleculeDataModule(training_config)" diff --git a/examples/train-gnn-to-electric-field/train-gnn-to-electric-field.ipynb b/examples/train-gnn-to-electric-field/train-gnn-to-electric-field.ipynb index 2e918cdd..42bf1a3d 100644 --- a/examples/train-gnn-to-electric-field/train-gnn-to-electric-field.ipynb +++ b/examples/train-gnn-to-electric-field/train-gnn-to-electric-field.ipynb @@ -24,23 +24,21 @@ "metadata": {}, "outputs": [], "source": [ - "import tqdm\n", - "\n", - "from qcportal import PortalClient\n", - "from openff.units import unit\n", + "from __future__ import annotations\n", "\n", - "from openff.toolkit import Molecule\n", + "import numpy as np\n", + "import pyarrow as pa\n", + "import pyarrow.parquet as pq\n", + "import torch\n", + "import tqdm\n", "from openff.qcsubmit.results import BasicResultCollection\n", - "from openff.recharge.esp.storage import MoleculeESPRecord\n", "from openff.recharge.esp.qcresults import from_qcportal_results\n", + "from openff.recharge.esp.storage import MoleculeESPRecord\n", "from openff.recharge.grids import MSKGridSettings\n", "from openff.recharge.utilities.geometry import compute_vector_field\n", - "\n", - "import pyarrow as pa\n", - "import pyarrow.parquet as pq\n", - "import numpy as np\n", - "\n", - "import torch" + "from openff.toolkit import Molecule\n", + "from openff.units import unit\n", + "from qcportal import PortalClient" ] }, { @@ -245,26 +243,22 @@ "metadata": {}, "outputs": [], "source": [ - "from openff.nagl.config import (\n", - " TrainingConfig,\n", - " OptimizerConfig,\n", - " ModelConfig,\n", - " DataConfig\n", - ")\n", + "from openff.nagl.config import DataConfig, ModelConfig, OptimizerConfig, TrainingConfig\n", + "from openff.nagl.config.data import DatasetConfig\n", "from openff.nagl.config.model import (\n", - " ConvolutionModule, ReadoutModule,\n", - " ConvolutionLayer, ForwardLayer,\n", + " ConvolutionLayer,\n", + " ConvolutionModule,\n", + " ForwardLayer,\n", + " ReadoutModule,\n", ")\n", - "from openff.nagl.config.data import DatasetConfig\n", - "from openff.nagl.training.training import TrainingGNNModel\n", "from openff.nagl.features.atoms import (\n", - " AtomicElement,\n", + " AtomAverageFormalCharge,\n", " AtomConnectivity,\n", + " AtomicElement,\n", " AtomInRingOfSize,\n", - " AtomAverageFormalCharge,\n", ")\n", - "\n", - "from openff.nagl.training.loss import GeneralLinearFitTarget" + "from openff.nagl.training.loss import GeneralLinearFitTarget\n", + "from openff.nagl.training.training import TrainingGNNModel" ] }, { diff --git a/examples/train-gnn/train_gnn.py b/examples/train-gnn/train_gnn.py index cad19abb..f1e1b5e1 100644 --- a/examples/train-gnn/train_gnn.py +++ b/examples/train-gnn/train_gnn.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import os import pathlib import pickle diff --git a/examples/train-multi-objective-gnn/train-gnn-notebook.ipynb b/examples/train-multi-objective-gnn/train-gnn-notebook.ipynb index e36e508c..d9af36d6 100644 --- a/examples/train-multi-objective-gnn/train-gnn-notebook.ipynb +++ b/examples/train-multi-objective-gnn/train-gnn-notebook.ipynb @@ -34,10 +34,11 @@ "metadata": {}, "outputs": [], "source": [ + "from __future__ import annotations\n", + "\n", "from pathlib import Path\n", "\n", "import numpy as np\n", - "\n", "from openff.toolkit import Molecule\n", "from openff.units import unit" ] @@ -94,13 +95,14 @@ "metadata": {}, "outputs": [], "source": [ + "from torch.nn import ReLU\n", + "\n", + "from openff.nagl import GNNModel\n", "from openff.nagl.config.model import (\n", " ConvolutionLayer,\n", " ConvolutionModule,\n", ")\n", - "from openff.nagl import GNNModel\n", "from openff.nagl.nn.gcn import SAGEConvStack\n", - "from torch.nn import ReLU\n", "from openff.nagl.nn.postprocess import ComputePartialCharges\n", "\n", "single_convolution_layer = ConvolutionLayer(\n", @@ -359,13 +361,14 @@ "metadata": {}, "outputs": [], "source": [ + "import openff.recharge\n", + "\n", "from openff.nagl.label.labels import (\n", - " LabelConformers,\n", " LabelCharges,\n", + " LabelConformers,\n", " LabelMultipleDipoles,\n", " LabelMultipleESPs,\n", ")\n", - "import openff.recharge\n", "\n", "# generate ELF conformers\n", "label_conformers = LabelConformers(\n", @@ -900,10 +903,9 @@ "metadata": {}, "outputs": [], "source": [ - "from openff.nagl.config.data import DatasetConfig, DataConfig\n", + "from openff.nagl.config.data import DataConfig, DatasetConfig\n", + "from openff.nagl.training.loss import MultipleDipoleTarget, MultipleESPTarget, ReadoutTarget\n", "from openff.nagl.training.metrics import RMSEMetric\n", - "from openff.nagl.training.loss import ReadoutTarget, MultipleDipoleTarget, MultipleESPTarget\n", - "\n", "\n", "am1_charge_rmse_target = ReadoutTarget(\n", " metric=RMSEMetric(), # use RMSE to calculate loss\n", @@ -1030,7 +1032,7 @@ "metadata": {}, "outputs": [], "source": [ - "from openff.nagl.training.training import TrainingGNNModel, DGLMoleculeDataModule\n", + "from openff.nagl.training.training import DGLMoleculeDataModule, TrainingGNNModel\n", "\n", "training_model = TrainingGNNModel(training_config)\n", "data_module = DGLMoleculeDataModule(training_config)" diff --git a/openff/nagl/__init__.py b/openff/nagl/__init__.py index 30ee756c..0ee30ba4 100644 --- a/openff/nagl/__init__.py +++ b/openff/nagl/__init__.py @@ -8,7 +8,10 @@ # importing any openff.nagl.toolkits submodule triggers a circular import: # toolkits._base → openff.toolkit.utils → GLOBAL_TOOLKIT_REGISTRY creation # → NAGLToolkitWrapper() → `from openff.nagl import __version__` → ImportError. +from __future__ import annotations + from . import _version + __version__ = _version.get_versions()['version'] from openff.nagl.nn._models import GNNModel diff --git a/openff/nagl/_base/__init__.py b/openff/nagl/_base/__init__.py index fb9f720f..83d5a596 100644 --- a/openff/nagl/_base/__init__.py +++ b/openff/nagl/_base/__init__.py @@ -1 +1,3 @@ +from __future__ import annotations + from .base import ImmutableModel diff --git a/openff/nagl/_base/base.py b/openff/nagl/_base/base.py index de15942f..2e5f577a 100644 --- a/openff/nagl/_base/base.py +++ b/openff/nagl/_base/base.py @@ -1,12 +1,13 @@ +from __future__ import annotations + import enum -import pathlib import json -import yaml +import pathlib import numpy as np +import yaml from openff.units import unit - try: from pydantic.v1 import BaseModel except ImportError: diff --git a/openff/nagl/_base/metaregistry.py b/openff/nagl/_base/metaregistry.py index d90bf6b5..a80367a6 100644 --- a/openff/nagl/_base/metaregistry.py +++ b/openff/nagl/_base/metaregistry.py @@ -37,7 +37,10 @@ def _get_object(cls, obj: Any): class MetaRegistryMixinAnyCase(MetaRegistryMixin): @classmethod def _key_transform(cls, key): - return key.lower() + try: + return key.lower() + except AttributeError as error: + raise AttributeError(f"Expected a string key, got {key} of type {type(key)}") from error def create_registry_metaclass(name_attribute: str = "name", ignore_case: bool = True): diff --git a/openff/nagl/_version.py b/openff/nagl/_version.py index c22fe7fe..6999e27c 100644 --- a/openff/nagl/_version.py +++ b/openff/nagl/_version.py @@ -10,14 +10,15 @@ # https://github.com/python-versioneer/python-versioneer """Git implementation of _version.py.""" +from __future__ import annotations import errno +import functools import os import re import subprocess import sys from typing import Any, Callable, Dict, List, Optional, Tuple -import functools def get_keywords() -> Dict[str, str]: diff --git a/openff/nagl/config/__init__.py b/openff/nagl/config/__init__.py index 25620eba..7ed179b6 100644 --- a/openff/nagl/config/__init__.py +++ b/openff/nagl/config/__init__.py @@ -1,6 +1,7 @@ """ Classes that define configuration options for training or using a GNN Model. """ +from __future__ import annotations from .data import DataConfig, DatasetConfig from .model import ModelConfig diff --git a/openff/nagl/config/data.py b/openff/nagl/config/data.py index 490319c2..a86848ba 100644 --- a/openff/nagl/config/data.py +++ b/openff/nagl/config/data.py @@ -1,15 +1,15 @@ """ Config classes for defining datasets. """ +from __future__ import annotations import pathlib import typing +from openff.nagl._base.base import ImmutableModel from openff.nagl.training.loss import ( TargetType, - ) -from openff.nagl._base.base import ImmutableModel from openff.nagl.utils._types import FromYamlMixin try: diff --git a/openff/nagl/config/model.py b/openff/nagl/config/model.py index ec535cb5..b8a48fd3 100644 --- a/openff/nagl/config/model.py +++ b/openff/nagl/config/model.py @@ -1,14 +1,14 @@ """Config classes for defining a GNNModel""" - +from __future__ import annotations import pathlib import typing from openff.nagl._base.base import ImmutableModel -from openff.nagl.nn.gcn._base import _GCNStackMeta -from openff.nagl.nn.activation import ActivationFunction from openff.nagl.features.atoms import DiscriminatedAtomFeatureType from openff.nagl.features.bonds import DiscriminatedBondFeatureType +from openff.nagl.nn.activation import ActivationFunction +from openff.nagl.nn.gcn._base import _GCNStackMeta from openff.nagl.utils._types import FromYamlMixin AggregatorType = typing.Literal["mean", "gcn", "pool", "lstm", "sum"] diff --git a/openff/nagl/config/optimizer.py b/openff/nagl/config/optimizer.py index 8baeb2f0..d9253647 100644 --- a/openff/nagl/config/optimizer.py +++ b/openff/nagl/config/optimizer.py @@ -1,4 +1,5 @@ """Config classes for defining optimizers""" +from __future__ import annotations import typing diff --git a/openff/nagl/config/training.py b/openff/nagl/config/training.py index db6c8946..f14d583f 100644 --- a/openff/nagl/config/training.py +++ b/openff/nagl/config/training.py @@ -1,6 +1,7 @@ """ The `TrainingConfig` class """ +from __future__ import annotations from openff.nagl._base.base import ImmutableModel from openff.nagl.config.data import DataConfig @@ -8,6 +9,7 @@ from openff.nagl.config.optimizer import OptimizerConfig from openff.nagl.utils._types import FromYamlMixin + class TrainingConfig(ImmutableModel, FromYamlMixin): """ Overall config class for training a model diff --git a/openff/nagl/domains.py b/openff/nagl/domains.py index 539406b8..1c929208 100644 --- a/openff/nagl/domains.py +++ b/openff/nagl/domains.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import typing from openff.nagl._base.base import ImmutableModel @@ -10,6 +12,7 @@ if typing.TYPE_CHECKING: from openff.toolkit.topology import Molecule + from openff.nagl.toolkits.registry import NAGLToolkitRegistry class ChemicalDomain(ImmutableModel): diff --git a/openff/nagl/features/__init__.py b/openff/nagl/features/__init__.py index 20b4984c..1afeb3b9 100644 --- a/openff/nagl/features/__init__.py +++ b/openff/nagl/features/__init__.py @@ -4,6 +4,8 @@ For an introduction to the theory behind this module, see :ref:`featurization_theory` """ +from __future__ import annotations + from ._base import Feature __all__ = [ diff --git a/openff/nagl/features/_base.py b/openff/nagl/features/_base.py index 05b68fe8..404bb5f6 100644 --- a/openff/nagl/features/_base.py +++ b/openff/nagl/features/_base.py @@ -1,12 +1,16 @@ +from __future__ import annotations + import abc import typing -from .._base.base import ImmutableModel from openff.nagl.toolkits.openff import ensure_toolkit_registry +from .._base.base import ImmutableModel + if typing.TYPE_CHECKING: import torch from openff.toolkit.topology import Molecule + from openff.nagl.toolkits.registry import NAGLToolkitRegistry diff --git a/openff/nagl/features/_data.py b/openff/nagl/features/_data.py index 4a31f68c..18dcb535 100644 --- a/openff/nagl/features/_data.py +++ b/openff/nagl/features/_data.py @@ -1,3 +1,5 @@ +from __future__ import annotations + ALLRED_ROCHOW_ELECTRONEGATIVITY = { # sourced from mendeleev 1: 0.0009765625, diff --git a/openff/nagl/features/_featurizers.py b/openff/nagl/features/_featurizers.py index 9f3e292c..42466ae8 100644 --- a/openff/nagl/features/_featurizers.py +++ b/openff/nagl/features/_featurizers.py @@ -1,15 +1,18 @@ -from typing import TYPE_CHECKING, Generic, List, TypeVar, Optional +from __future__ import annotations + +from typing import TYPE_CHECKING, Generic, List, Optional, TypeVar import torch -from .atoms import AtomFeature +from openff.nagl.toolkits.openff import ensure_toolkit_registry + from ._base import Feature +from .atoms import AtomFeature from .bonds import BondFeature -from openff.nagl.toolkits.openff import ensure_toolkit_registry - if TYPE_CHECKING: from openff.toolkit.topology import Molecule + from openff.nagl.toolkits.registry import NAGLToolkitRegistry diff --git a/openff/nagl/features/_utils.py b/openff/nagl/features/_utils.py index cc4cbf4b..5aa1b696 100644 --- a/openff/nagl/features/_utils.py +++ b/openff/nagl/features/_utils.py @@ -1,3 +1,5 @@ +from __future__ import annotations + from typing import List, TypeVar import torch diff --git a/openff/nagl/features/atoms.py b/openff/nagl/features/atoms.py index 73066e8f..79f01e86 100644 --- a/openff/nagl/features/atoms.py +++ b/openff/nagl/features/atoms.py @@ -20,26 +20,28 @@ to implement your own features. """ +from __future__ import annotations import copy import typing import numpy as np import torch +from openff.units import unit from openff.nagl.utils._types import HybridizationType -from openff.units import unit from ._base import CategoricalMixin, Feature from ._utils import one_hot_encode try: - from pydantic.v1 import validator, Field + from pydantic.v1 import Field, validator except ImportError: - from pydantic import validator, Field + from pydantic import Field, validator if typing.TYPE_CHECKING: from openff.toolkit.topology import Molecule + from openff.nagl.toolkits.registry import NAGLToolkitRegistry __all__ = [ @@ -262,8 +264,8 @@ class AtomAverageFormalCharge(AtomFeature): name: typing.Literal["atom_average_formal_charge"] = "atom_average_formal_charge" def _encode(self, molecule: "Molecule", toolkit_registry: typing.Optional["NAGLToolkitRegistry"] = None) -> torch.Tensor: - from openff.nagl.utils.resonance import enumerate_resonance_forms from openff.nagl.toolkits.openff import normalize_molecule + from openff.nagl.utils.resonance import enumerate_resonance_forms molecule = normalize_molecule(molecule, toolkit_registry=toolkit_registry) resonance_forms = enumerate_resonance_forms( diff --git a/openff/nagl/features/bonds.py b/openff/nagl/features/bonds.py index 366f8d1c..ddfd8e9f 100644 --- a/openff/nagl/features/bonds.py +++ b/openff/nagl/features/bonds.py @@ -19,11 +19,13 @@ """ # from typing import ClassVar, Dict, Type +from __future__ import annotations + import typing import torch -from ._base import CategoricalMixin, Feature #, FeatureMeta +from ._base import CategoricalMixin, Feature #, FeatureMeta from ._utils import one_hot_encode try: diff --git a/openff/nagl/label/dataset.py b/openff/nagl/label/dataset.py index aac3f7ca..b51c08a6 100644 --- a/openff/nagl/label/dataset.py +++ b/openff/nagl/label/dataset.py @@ -1,12 +1,15 @@ +from __future__ import annotations + import functools import pathlib -import tqdm import typing -from openff.nagl.utils._parallelization import get_mapper_to_processes -from openff.nagl.label.labels import LabellerType +import tqdm from openff.utilities import requires_package +from openff.nagl.label.labels import LabellerType +from openff.nagl.utils._parallelization import get_mapper_to_processes + if typing.TYPE_CHECKING: import pyarrow @@ -42,10 +45,9 @@ def from_smiles( verbose: bool = False, overwrite_existing: bool = False, ): - from openff.toolkit import Molecule - import pyarrow as pa import pyarrow.dataset as ds + from openff.toolkit import Molecule loader = functools.partial( Molecule.from_smiles, diff --git a/openff/nagl/label/labels.py b/openff/nagl/label/labels.py index 339ca653..dfb2b7a6 100644 --- a/openff/nagl/label/labels.py +++ b/openff/nagl/label/labels.py @@ -1,12 +1,14 @@ +from __future__ import annotations + import abc -from collections import defaultdict import functools import logging import pathlib -import tqdm import typing +from collections import defaultdict import numpy as np +import tqdm from openff.units import unit from openff.utilities import requires_package @@ -65,8 +67,8 @@ def apply( table: "pyarrow.Table", verbose: bool = False, ): - from openff.toolkit import Molecule import pyarrow as pa + from openff.toolkit import Molecule rms_cutoff = self.rms_cutoff if not isinstance(rms_cutoff, unit.Quantity): @@ -264,8 +266,8 @@ def _calculate_inverse_distance_grid( conformers: np.ndarray, n_conformers: int, ) -> typing.List[np.ndarray]: - from openff.toolkit import Molecule from openff.recharge.grids import GridGenerator, MSKGridSettings + from openff.toolkit import Molecule mol = Molecule.from_mapped_smiles( mapped_smiles, diff --git a/openff/nagl/label/utils.py b/openff/nagl/label/utils.py index 5b7fe081..d225c6b3 100644 --- a/openff/nagl/label/utils.py +++ b/openff/nagl/label/utils.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import logging import typing diff --git a/openff/nagl/lookups.py b/openff/nagl/lookups.py index 08a8abd3..a6ffa6d3 100644 --- a/openff/nagl/lookups.py +++ b/openff/nagl/lookups.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import types import typing @@ -14,6 +16,7 @@ if typing.TYPE_CHECKING: from openff.toolkit.topology import Molecule + from openff.nagl.toolkits.registry import NAGLToolkitRegistry diff --git a/openff/nagl/molecule/__init__.py b/openff/nagl/molecule/__init__.py index fd50bc86..6313934a 100644 --- a/openff/nagl/molecule/__init__.py +++ b/openff/nagl/molecule/__init__.py @@ -1,7 +1,8 @@ """ Structures for managing molecule graphs """ +from __future__ import annotations -from ._dgl.molecule import DGLMolecule from ._dgl.batch import DGLMoleculeBatch -from ._graph.molecule import GraphMolecule, GraphMoleculeBatch \ No newline at end of file +from ._dgl.molecule import DGLMolecule +from ._graph.molecule import GraphMolecule, GraphMoleculeBatch diff --git a/openff/nagl/molecule/_base.py b/openff/nagl/molecule/_base.py index 322a4d69..c5ea198e 100644 --- a/openff/nagl/molecule/_base.py +++ b/openff/nagl/molecule/_base.py @@ -1,9 +1,12 @@ -from typing import ClassVar, TYPE_CHECKING, Tuple, Optional +from __future__ import annotations + +from typing import TYPE_CHECKING, ClassVar, Optional, Tuple from openff.nagl.molecule._utils import FEATURE if TYPE_CHECKING: import torch + from openff.nagl.features.atoms import AtomFeature from openff.nagl.features.bonds import BondFeature diff --git a/openff/nagl/molecule/_dgl/__init__.py b/openff/nagl/molecule/_dgl/__init__.py index 741e6267..6f600feb 100644 --- a/openff/nagl/molecule/_dgl/__init__.py +++ b/openff/nagl/molecule/_dgl/__init__.py @@ -1,3 +1,5 @@ +from __future__ import annotations + from typing import Union from .batch import DGLMoleculeBatch diff --git a/openff/nagl/molecule/_dgl/batch.py b/openff/nagl/molecule/_dgl/batch.py index 87bbd2e6..7842d364 100644 --- a/openff/nagl/molecule/_dgl/batch.py +++ b/openff/nagl/molecule/_dgl/batch.py @@ -1,10 +1,13 @@ -from typing import List, Tuple +from __future__ import annotations +from typing import List, Tuple from openff.utilities import requires_package -from .molecule import DGLBase, DGLMolecule + from openff.nagl.molecule._base import BatchMixin +from .molecule import DGLBase, DGLMolecule + class DGLMoleculeBatch(BatchMixin, DGLBase): n_representations: Tuple[int, ...] diff --git a/openff/nagl/molecule/_dgl/molecule.py b/openff/nagl/molecule/_dgl/molecule.py index 9b19ec0a..89b99a74 100644 --- a/openff/nagl/molecule/_dgl/molecule.py +++ b/openff/nagl/molecule/_dgl/molecule.py @@ -1,11 +1,13 @@ -from typing import ClassVar, Optional, Tuple, TYPE_CHECKING +from __future__ import annotations + +from typing import TYPE_CHECKING, ClassVar, Optional, Tuple import torch from openff.utilities import requires_package from openff.nagl.features.atoms import AtomFeature from openff.nagl.features.bonds import BondFeature -from openff.nagl.molecule._base import NAGLMoleculeBase, MoleculeMixin +from openff.nagl.molecule._base import MoleculeMixin, NAGLMoleculeBase from openff.nagl.toolkits.openff import ensure_toolkit_registry from .utils import ( @@ -16,6 +18,7 @@ if TYPE_CHECKING: from openff.toolkit.topology import Molecule + from openff.nagl.toolkits.registry import NAGLToolkitRegistry @@ -60,6 +63,7 @@ def from_openff( ): toolkit_registry = ensure_toolkit_registry(toolkit_registry) import dgl + from openff.nagl.utils.resonance import ResonanceEnumerator if atom_features is None: diff --git a/openff/nagl/molecule/_dgl/utils.py b/openff/nagl/molecule/_dgl/utils.py index 8b8d519a..10ffae84 100644 --- a/openff/nagl/molecule/_dgl/utils.py +++ b/openff/nagl/molecule/_dgl/utils.py @@ -1,19 +1,21 @@ -from typing import Dict, List, TYPE_CHECKING, Optional +from __future__ import annotations + +from typing import TYPE_CHECKING, Dict, List, Optional -import torch import numpy as np +import torch from openff.utilities import requires_package +from openff.nagl.features._featurizers import AtomFeaturizer, BondFeaturizer from openff.nagl.features.atoms import AtomFeature from openff.nagl.features.bonds import BondFeature -from openff.nagl.features._featurizers import AtomFeaturizer, BondFeaturizer -from openff.nagl.molecule._utils import FORWARD, REVERSE, FEATURE +from openff.nagl.molecule._utils import FEATURE, FORWARD, REVERSE from openff.nagl.toolkits.openff import ensure_toolkit_registry - if TYPE_CHECKING: import dgl from openff.toolkit.topology.molecule import Molecule + from openff.nagl.toolkits.registry import NAGLToolkitRegistry @@ -28,6 +30,7 @@ def openff_molecule_to_base_dgl_graph( Convert an OpenFF Molecule to a DGL graph. """ import dgl + from openff.nagl.toolkits.openff import get_openff_molecule_bond_indices bonds = get_openff_molecule_bond_indices(molecule) @@ -128,7 +131,7 @@ def heterograph_to_homograph_no_edges(G: "dgl.DGLHeteroGraph", ndata=None, edata """ import dgl from dgl import backend as F - from dgl.base import EID, NID, ETYPE, NTYPE + from dgl.base import EID, ETYPE, NID, NTYPE from dgl.heterograph import combine_frames # TODO: revisit in case DGL accounts for this in the future diff --git a/openff/nagl/molecule/_graph/_batch.py b/openff/nagl/molecule/_graph/_batch.py index 9d442280..f7e12c66 100644 --- a/openff/nagl/molecule/_graph/_batch.py +++ b/openff/nagl/molecule/_graph/_batch.py @@ -1,10 +1,9 @@ -import copy +from __future__ import annotations +import copy from collections import UserDict - from typing import Dict - import torch diff --git a/openff/nagl/molecule/_graph/_graph.py b/openff/nagl/molecule/_graph/_graph.py index 33ba9b0d..8c54335b 100644 --- a/openff/nagl/molecule/_graph/_graph.py +++ b/openff/nagl/molecule/_graph/_graph.py @@ -1,19 +1,22 @@ +from __future__ import annotations + import contextlib import copy from collections import defaultdict -from typing import List, Tuple, TYPE_CHECKING +from typing import TYPE_CHECKING, List, Tuple + +import networkx as nx +import numpy as np +import torch +from openff.nagl.features._featurizers import AtomFeaturizer, BondFeaturizer from openff.nagl.features.atoms import AtomFeature from openff.nagl.features.bonds import BondFeature -from openff.nagl.features._featurizers import AtomFeaturizer, BondFeaturizer +from openff.nagl.molecule._utils import FEATURE, FORWARD, REVERSE from openff.nagl.toolkits.openff import ensure_toolkit_registry -import networkx as nx -import numpy as np -import torch from ._batch import FrameDict -from openff.nagl.molecule._utils import FORWARD, REVERSE, FEATURE if TYPE_CHECKING: from openff.toolkit import Molecule diff --git a/openff/nagl/molecule/_graph/_update.py b/openff/nagl/molecule/_graph/_update.py index 8cc0798b..a2870bfb 100644 --- a/openff/nagl/molecule/_graph/_update.py +++ b/openff/nagl/molecule/_graph/_update.py @@ -5,13 +5,14 @@ Functions here are not intended to be called directly and may be fragile. """ +from __future__ import annotations -from typing import TYPE_CHECKING, Dict, Optional, List, Callable +from typing import TYPE_CHECKING, Callable, Dict, List, Optional -import torch import numpy as np +import torch -from ._batch import EdgeBatch, NodeBatch, FrameDict +from ._batch import EdgeBatch, FrameDict, NodeBatch if TYPE_CHECKING: from ._graph import NXMolGraph diff --git a/openff/nagl/molecule/_graph/_utils.py b/openff/nagl/molecule/_graph/_utils.py index 6d3c3519..aa1cc24c 100644 --- a/openff/nagl/molecule/_graph/_utils.py +++ b/openff/nagl/molecule/_graph/_utils.py @@ -1,8 +1,10 @@ +from __future__ import annotations + from collections import defaultdict -from typing import List, Tuple, Iterable +from typing import Iterable, List, Tuple -import torch import networkx as nx +import torch def nonzero_1d(values): @@ -49,8 +51,9 @@ def bucketor(data) -> List[torch.Tensor]: def _batch_nx_graphs(graphs): - import torch import networkx as nx + import torch + from openff.nagl.molecule._graph._batch import FrameDict if not len(graphs): @@ -95,6 +98,7 @@ def _unbatch_nx_graphs( n_representations_per_molecule: Iterable[int], ) -> List[nx.Graph]: import networkx as nx + from openff.nagl.molecule._graph._batch import FrameDict unbatched_graphs = [] diff --git a/openff/nagl/molecule/_graph/molecule.py b/openff/nagl/molecule/_graph/molecule.py index 285de00a..03a5d675 100644 --- a/openff/nagl/molecule/_graph/molecule.py +++ b/openff/nagl/molecule/_graph/molecule.py @@ -1,11 +1,14 @@ -from typing import List, TYPE_CHECKING, Tuple, Optional +from __future__ import annotations -from openff.nagl.molecule._base import NAGLMoleculeBase, MoleculeMixin, BatchMixin +from typing import TYPE_CHECKING, List, Optional, Tuple + +from openff.nagl.molecule._base import BatchMixin, MoleculeMixin, NAGLMoleculeBase from openff.nagl.molecule._graph._graph import NXMolHeteroGraph from openff.nagl.toolkits.openff import ensure_toolkit_registry if TYPE_CHECKING: from openff.toolkit.topology import Molecule + from openff.nagl.features.atoms import AtomFeature from openff.nagl.features.bonds import BondFeature from openff.nagl.toolkits.registry import NAGLToolkitRegistry diff --git a/openff/nagl/molecule/_utils.py b/openff/nagl/molecule/_utils.py index 0e8ad907..458a5546 100644 --- a/openff/nagl/molecule/_utils.py +++ b/openff/nagl/molecule/_utils.py @@ -1,8 +1,10 @@ -from typing import Dict, TYPE_CHECKING +from __future__ import annotations + +from typing import TYPE_CHECKING, Dict if TYPE_CHECKING: - from openff.toolkit.topology.molecule import Molecule import torch + from openff.toolkit.topology.molecule import Molecule FORWARD = "forward" REVERSE = "reverse" @@ -12,8 +14,8 @@ def _get_openff_molecule_information( molecule: "Molecule", ) -> Dict[str, "torch.Tensor"]: - from openff.units import unit import torch + from openff.units import unit charges = [ atom.formal_charge.m_as(unit.elementary_charge) for atom in molecule.atoms diff --git a/openff/nagl/nn/__init__.py b/openff/nagl/nn/__init__.py index f922a69f..e57b8a4b 100644 --- a/openff/nagl/nn/__init__.py +++ b/openff/nagl/nn/__init__.py @@ -1,22 +1,18 @@ """ Components for constructing and processing GNN models """ +from __future__ import annotations - -from .activation import ActivationFunction from ._containers import ConvolutionModule, ReadoutModule from ._dataset import ( - DGLMoleculeDatasetEntry, DGLMoleculeDataLoader, DGLMoleculeDataset, + DGLMoleculeDatasetEntry, ) from ._pooling import PoolAtomFeatures, PoolBondFeatures from ._sequential import SequentialLayers -from .postprocess import ( - ComputePartialCharges, - RegularizedComputePartialCharges -) - +from .activation import ActivationFunction +from .postprocess import ComputePartialCharges, RegularizedComputePartialCharges __all__ = [ "ActivationFunction", diff --git a/openff/nagl/nn/_base.py b/openff/nagl/nn/_base.py index 3b8985ea..dbdd0a9a 100644 --- a/openff/nagl/nn/_base.py +++ b/openff/nagl/nn/_base.py @@ -1,3 +1,5 @@ +from __future__ import annotations + from typing import Any, List, Optional from openff.nagl.utils._utils import is_iterable diff --git a/openff/nagl/nn/_containers.py b/openff/nagl/nn/_containers.py index 792a0736..6c909dcd 100644 --- a/openff/nagl/nn/_containers.py +++ b/openff/nagl/nn/_containers.py @@ -1,14 +1,15 @@ +from __future__ import annotations + import copy -from typing import List, Optional, Union, Tuple, Callable +from typing import Callable, List, Optional, Tuple, Union import torch from openff.nagl.molecule._dgl import DGLMolecule, DGLMoleculeBatch - -from openff.nagl.nn.activation import ActivationFunction -from openff.nagl.nn.gcn._base import _GCNStackMeta, BaseConvModule -from openff.nagl.nn._sequential import SequentialLayers from openff.nagl.nn._pooling import PoolingLayer, get_pooling_layer +from openff.nagl.nn._sequential import SequentialLayers +from openff.nagl.nn.activation import ActivationFunction +from openff.nagl.nn.gcn._base import BaseConvModule, _GCNStackMeta from openff.nagl.nn.postprocess import PostprocessLayer, _PostprocessLayerMeta diff --git a/openff/nagl/nn/_dataset.py b/openff/nagl/nn/_dataset.py index 2fe03d3e..49bdf3c7 100644 --- a/openff/nagl/nn/_dataset.py +++ b/openff/nagl/nn/_dataset.py @@ -1,30 +1,30 @@ "Classes for handling featurized molecule data to train GNN models" +from __future__ import annotations -from collections import defaultdict import functools import glob import hashlib import io import logging +import pathlib import pickle import tempfile import typing +from collections import defaultdict -import tqdm +import numpy as np import torch +import tqdm from openff.utilities import requires_package -from torch.utils.data import Dataset, DataLoader, ConcatDataset +from torch.utils.data import ConcatDataset, DataLoader, Dataset from openff.nagl._base.base import ImmutableModel from openff.nagl.config.training import TrainingConfig from openff.nagl.features.atoms import AtomFeature from openff.nagl.features.bonds import BondFeature from openff.nagl.molecule._dgl import DGLMolecule, DGLMoleculeBatch, DGLMoleculeOrBatch -from openff.nagl.utils._parallelization import get_mapper_to_processes from openff.nagl.utils._hash import digest_file - -import pathlib -import numpy as np +from openff.nagl.utils._parallelization import get_mapper_to_processes if typing.TYPE_CHECKING: from openff.toolkit import Molecule diff --git a/openff/nagl/nn/_models.py b/openff/nagl/nn/_models.py index a0e60229..28298643 100644 --- a/openff/nagl/nn/_models.py +++ b/openff/nagl/nn/_models.py @@ -1,24 +1,27 @@ +from __future__ import annotations + import collections import logging import types -from typing import TYPE_CHECKING, Tuple, Dict, Union, Callable, Literal, Optional import warnings +from typing import TYPE_CHECKING, Callable, Dict, Literal, Optional, Tuple, Union -import torch import pytorch_lightning as pl - +import torch from openff.utilities.exceptions import MissingOptionalDependencyError -from openff.nagl.nn._containers import ConvolutionModule, ReadoutModule + from openff.nagl.config.model import ModelConfig from openff.nagl.domains import ChemicalDomain from openff.nagl.lookups import LookupTableType, _as_lookup_table -from openff.nagl.utils._utils import potential_dict_to_list +from openff.nagl.nn._containers import ConvolutionModule, ReadoutModule from openff.nagl.toolkits.openff import ensure_toolkit_registry +from openff.nagl.utils._utils import potential_dict_to_list logger = logging.getLogger(__name__) if TYPE_CHECKING: from openff.toolkit.topology import Molecule + from openff.nagl.molecule._dgl import DGLMoleculeOrBatch from openff.nagl.toolkits.registry import NAGLToolkitRegistry diff --git a/openff/nagl/nn/_pooling.py b/openff/nagl/nn/_pooling.py index 1be0b51d..a5196a60 100644 --- a/openff/nagl/nn/_pooling.py +++ b/openff/nagl/nn/_pooling.py @@ -1,6 +1,8 @@ +from __future__ import annotations + import abc import functools -from typing import ClassVar, Dict, Union, TYPE_CHECKING, Iterable +from typing import TYPE_CHECKING, ClassVar, Dict, Iterable, Union import torch.nn diff --git a/openff/nagl/nn/_sequential.py b/openff/nagl/nn/_sequential.py index a7285ca0..570fd6f9 100644 --- a/openff/nagl/nn/_sequential.py +++ b/openff/nagl/nn/_sequential.py @@ -1,9 +1,11 @@ +from __future__ import annotations + from typing import ClassVar, List, Optional import torch.nn -from .activation import ActivationFunction from ._base import ContainsLayersMixin +from .activation import ActivationFunction class SequentialLayers(torch.nn.Sequential, ContainsLayersMixin): diff --git a/openff/nagl/nn/activation.py b/openff/nagl/nn/activation.py index 771b2a8c..04b84468 100644 --- a/openff/nagl/nn/activation.py +++ b/openff/nagl/nn/activation.py @@ -1,4 +1,5 @@ "Activation functions" +from __future__ import annotations import enum from typing import Callable diff --git a/openff/nagl/nn/gcn/__init__.py b/openff/nagl/nn/gcn/__init__.py index cce7e642..08ea20ca 100644 --- a/openff/nagl/nn/gcn/__init__.py +++ b/openff/nagl/nn/gcn/__init__.py @@ -1,9 +1,9 @@ "Architectures for convolutional layers" +from __future__ import annotations from ._base import BaseGCNStack, _GCNStackMeta -from ._sage import SAGEConvStack from ._gin import GINConvStack - +from ._sage import SAGEConvStack __all__ = ["BaseGCNStack", "SAGEConvStack", "GINConvStack"] diff --git a/openff/nagl/nn/gcn/_base.py b/openff/nagl/nn/gcn/_base.py index a0c5a13c..aa0a42f0 100644 --- a/openff/nagl/nn/gcn/_base.py +++ b/openff/nagl/nn/gcn/_base.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import abc from typing import ClassVar, Dict, Generic, List, Optional, Type, TypeVar @@ -5,8 +7,8 @@ import torch.nn.functional from openff.nagl._base.metaregistry import create_registry_metaclass -from openff.nagl.nn.activation import ActivationFunction from openff.nagl.nn._base import ContainsLayersMixin +from openff.nagl.nn.activation import ActivationFunction GCNLayerType = TypeVar("GCNLayerType", bound=torch.nn.Module) diff --git a/openff/nagl/nn/gcn/_function.py b/openff/nagl/nn/gcn/_function.py index 9dc2bdc6..599235db 100644 --- a/openff/nagl/nn/gcn/_function.py +++ b/openff/nagl/nn/gcn/_function.py @@ -1,3 +1,6 @@ +from __future__ import annotations + + def copy_u(msg: str, out: str): def wrapped(edges): return {out: edges.src[msg]} diff --git a/openff/nagl/nn/gcn/_gin.py b/openff/nagl/nn/gcn/_gin.py index 8ea6b507..ed405485 100644 --- a/openff/nagl/nn/gcn/_gin.py +++ b/openff/nagl/nn/gcn/_gin.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import copy from typing import TYPE_CHECKING, Union @@ -5,9 +7,10 @@ from openff.utilities import requires_package from openff.utilities.exceptions import MissingOptionalDependencyError -from ._base import ActivationFunction, BaseGCNStack, BaseConvModule import openff.nagl.nn.gcn._function as _fn +from ._base import ActivationFunction, BaseConvModule, BaseGCNStack + if TYPE_CHECKING: import dgl diff --git a/openff/nagl/nn/gcn/_sage.py b/openff/nagl/nn/gcn/_sage.py index 1a33401b..b009a9ca 100644 --- a/openff/nagl/nn/gcn/_sage.py +++ b/openff/nagl/nn/gcn/_sage.py @@ -1,15 +1,18 @@ +from __future__ import annotations + import copy -from typing import Optional, Literal, Dict, TYPE_CHECKING, Union +from typing import TYPE_CHECKING, Dict, Literal, Optional, Union import torch -from ._base import ActivationFunction, BaseGCNStack, BaseConvModule -from openff.nagl.nn.gcn import _function as _fn - # import dgl.function as fn from openff.utilities import requires_package from openff.utilities.exceptions import MissingOptionalDependencyError +from openff.nagl.nn.gcn import _function as _fn + +from ._base import ActivationFunction, BaseConvModule, BaseGCNStack + if TYPE_CHECKING: import dgl diff --git a/openff/nagl/nn/postprocess.py b/openff/nagl/nn/postprocess.py index f125b132..a133eea1 100644 --- a/openff/nagl/nn/postprocess.py +++ b/openff/nagl/nn/postprocess.py @@ -1,6 +1,7 @@ """ Postprocessing functions to convert a graph representation to a predicted property """ +from __future__ import annotations import abc from typing import ClassVar, Dict, Type, Union diff --git a/openff/nagl/tests/_base/test_base.py b/openff/nagl/tests/_base/test_base.py index adf34e4c..18071afb 100644 --- a/openff/nagl/tests/_base/test_base.py +++ b/openff/nagl/tests/_base/test_base.py @@ -1,9 +1,13 @@ -from openff.nagl._base.base import MutableModel -from openff.units import unit -import numpy as np +from __future__ import annotations + import json import textwrap +import numpy as np +from openff.units import unit + +from openff.nagl._base.base import MutableModel + try: from pydantic.v1 import Field, validator except ImportError: diff --git a/openff/nagl/tests/_base/test_metaregistry.py b/openff/nagl/tests/_base/test_metaregistry.py index b1f6bf10..a973cab5 100644 --- a/openff/nagl/tests/_base/test_metaregistry.py +++ b/openff/nagl/tests/_base/test_metaregistry.py @@ -1,7 +1,10 @@ +from __future__ import annotations + import pytest from openff.nagl._base.metaregistry import create_registry_metaclass + class TestMetaRegistry: Registry = create_registry_metaclass(ignore_case=False) diff --git a/openff/nagl/tests/config/test_model_config.py b/openff/nagl/tests/config/test_model_config.py index 230ad147..e429bc9d 100644 --- a/openff/nagl/tests/config/test_model_config.py +++ b/openff/nagl/tests/config/test_model_config.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import textwrap from openff.nagl.config.model import ModelConfig diff --git a/openff/nagl/tests/conftest.py b/openff/nagl/tests/conftest.py index 66eac163..c7eab1ad 100644 --- a/openff/nagl/tests/conftest.py +++ b/openff/nagl/tests/conftest.py @@ -1,15 +1,15 @@ """ Global pytest fixtures """ - +from __future__ import annotations import numpy as np import pytest -from openff.nagl.molecule._dgl.molecule import DGLMolecule -from openff.nagl.molecule._dgl.batch import DGLMoleculeBatch from openff.nagl.features.atoms import AtomConnectivity, AtomicElement from openff.nagl.features.bonds import BondIsInRing +from openff.nagl.molecule._dgl.batch import DGLMoleculeBatch +from openff.nagl.molecule._dgl.molecule import DGLMolecule @pytest.fixture() diff --git a/openff/nagl/tests/data/example_charges/generate-new-sdfs.py b/openff/nagl/tests/data/example_charges/generate-new-sdfs.py index 60453ab9..0fa2b0ea 100644 --- a/openff/nagl/tests/data/example_charges/generate-new-sdfs.py +++ b/openff/nagl/tests/data/example_charges/generate-new-sdfs.py @@ -1,10 +1,12 @@ +from __future__ import annotations + import pathlib import click import tqdm - from openff.toolkit import Molecule from openff.units import unit + from openff.nagl import GNNModel diff --git a/openff/nagl/tests/data/files.py b/openff/nagl/tests/data/files.py index 2561f392..255f25c7 100644 --- a/openff/nagl/tests/data/files.py +++ b/openff/nagl/tests/data/files.py @@ -7,6 +7,7 @@ from openff.nagl.tests.data.files import * """ +from __future__ import annotations __all__ = [ # "EXAMPLE_MODEL_CONFIG", diff --git a/openff/nagl/tests/features/test_features.py b/openff/nagl/tests/features/test_features.py index 038a63d6..7f4fbb1e 100644 --- a/openff/nagl/tests/features/test_features.py +++ b/openff/nagl/tests/features/test_features.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import numpy as np import pytest from numpy.testing import assert_allclose, assert_equal diff --git a/openff/nagl/tests/features/test_featurizers.py b/openff/nagl/tests/features/test_featurizers.py index 52f3c016..3e3a3961 100644 --- a/openff/nagl/tests/features/test_featurizers.py +++ b/openff/nagl/tests/features/test_featurizers.py @@ -1,7 +1,10 @@ +from __future__ import annotations + import numpy as np import pytest from numpy.testing import assert_allclose, assert_equal +from openff.nagl.features._featurizers import AtomFeaturizer, BondFeaturizer from openff.nagl.features.atoms import ( AtomConnectivity, AtomIsAromatic, @@ -10,7 +13,6 @@ BondIsInRing, BondOrder, ) -from openff.nagl.features._featurizers import AtomFeaturizer, BondFeaturizer def test_atomfeaturizer(openff_methane_uncharged): diff --git a/openff/nagl/tests/features/test_utils.py b/openff/nagl/tests/features/test_utils.py index 48b31cc5..bb9d91c5 100644 --- a/openff/nagl/tests/features/test_utils.py +++ b/openff/nagl/tests/features/test_utils.py @@ -1,3 +1,5 @@ +from __future__ import annotations + from numpy.testing import assert_allclose from openff.nagl.features._utils import one_hot_encode diff --git a/openff/nagl/tests/label/conftest.py b/openff/nagl/tests/label/conftest.py index 4ee14da3..cc3faf50 100644 --- a/openff/nagl/tests/label/conftest.py +++ b/openff/nagl/tests/label/conftest.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import pytest from openff.nagl.label.dataset import LabelledDataset diff --git a/openff/nagl/tests/label/test_labelled_dataset.py b/openff/nagl/tests/label/test_labelled_dataset.py index 010f6dfa..21b5bb55 100644 --- a/openff/nagl/tests/label/test_labelled_dataset.py +++ b/openff/nagl/tests/label/test_labelled_dataset.py @@ -1,10 +1,9 @@ +from __future__ import annotations + import pytest from openff.nagl.label.dataset import LabelledDataset -from openff.nagl.label.labels import ( - LabelConformers, - LabelCharges -) +from openff.nagl.label.labels import LabelCharges, LabelConformers pa = pytest.importorskip("pyarrow") class TestLabelledDataset: diff --git a/openff/nagl/tests/label/test_labels.py b/openff/nagl/tests/label/test_labels.py index 3937080d..c791e6fe 100644 --- a/openff/nagl/tests/label/test_labels.py +++ b/openff/nagl/tests/label/test_labels.py @@ -1,15 +1,16 @@ -import pytest +from __future__ import annotations import numpy as np +import pytest +from openff.toolkit.utils import OPENEYE_AVAILABLE from openff.nagl.label.dataset import LabelledDataset from openff.nagl.label.labels import ( - LabelConformers, LabelCharges, + LabelConformers, LabelMultipleDipoles, LabelMultipleESPs, ) -from openff.toolkit.utils import OPENEYE_AVAILABLE pytest.importorskip("pyarrow") diff --git a/openff/nagl/tests/molecule/dgl/test_batch.py b/openff/nagl/tests/molecule/dgl/test_batch.py index 94731673..8c4477b4 100644 --- a/openff/nagl/tests/molecule/dgl/test_batch.py +++ b/openff/nagl/tests/molecule/dgl/test_batch.py @@ -1,9 +1,10 @@ +from __future__ import annotations + import pytest from numpy.testing import assert_allclose, assert_array_almost_equal from openff.nagl.molecule._dgl.batch import DGLMolecule, DGLMoleculeBatch - dgl = pytest.importorskip("dgl") diff --git a/openff/nagl/tests/molecule/dgl/test_molecule.py b/openff/nagl/tests/molecule/dgl/test_molecule.py index cdd89a39..785981a5 100644 --- a/openff/nagl/tests/molecule/dgl/test_molecule.py +++ b/openff/nagl/tests/molecule/dgl/test_molecule.py @@ -1,9 +1,11 @@ +from __future__ import annotations + import pytest from numpy.testing import assert_allclose, assert_array_almost_equal -from openff.nagl.molecule._dgl.molecule import DGLMolecule from openff.nagl.features.atoms import AtomConnectivity from openff.nagl.features.bonds import BondIsInRing +from openff.nagl.molecule._dgl.molecule import DGLMolecule dgl = pytest.importorskip("dgl") diff --git a/openff/nagl/tests/molecule/dgl/test_utils.py b/openff/nagl/tests/molecule/dgl/test_utils.py index 01319419..ecbee6dc 100644 --- a/openff/nagl/tests/molecule/dgl/test_utils.py +++ b/openff/nagl/tests/molecule/dgl/test_utils.py @@ -1,15 +1,17 @@ +from __future__ import annotations + import pytest import torch from openff.toolkit.topology.molecule import Molecule from torch.testing import assert_close +from openff.nagl.features.atoms import AtomConnectivity +from openff.nagl.features.bonds import BondIsInRing from openff.nagl.molecule._dgl.utils import ( dgl_heterograph_to_homograph, openff_molecule_to_base_dgl_graph, openff_molecule_to_dgl_graph, ) -from openff.nagl.features.atoms import AtomConnectivity -from openff.nagl.features.bonds import BondIsInRing pytest.importorskip("dgl") diff --git a/openff/nagl/tests/molecule/graph/test_batch.py b/openff/nagl/tests/molecule/graph/test_batch.py index ab141725..cfc6594f 100644 --- a/openff/nagl/tests/molecule/graph/test_batch.py +++ b/openff/nagl/tests/molecule/graph/test_batch.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import pytest from numpy.testing import assert_allclose, assert_array_almost_equal diff --git a/openff/nagl/tests/molecule/graph/test_graph.py b/openff/nagl/tests/molecule/graph/test_graph.py index 7cca5919..32ad088c 100644 --- a/openff/nagl/tests/molecule/graph/test_graph.py +++ b/openff/nagl/tests/molecule/graph/test_graph.py @@ -1,24 +1,26 @@ -import pytest +from __future__ import annotations from collections import defaultdict + import networkx as nx import numpy as np -from numpy.testing import assert_allclose, assert_equal +import pytest import torch +from numpy.testing import assert_allclose, assert_equal -from openff.nagl.molecule._graph._batch import FrameDict -from openff.nagl.molecule._graph._graph import ( - NXMolGraph, - NXMolHomoGraph, - NXMolHeteroGraph, -) from openff.nagl.features.atoms import ( - AtomicElement, - AtomConnectivity, AtomAverageFormalCharge, + AtomConnectivity, + AtomicElement, AtomInRingOfSize, ) from openff.nagl.features.bonds import BondOrder +from openff.nagl.molecule._graph._batch import FrameDict +from openff.nagl.molecule._graph._graph import ( + NXMolGraph, + NXMolHeteroGraph, + NXMolHomoGraph, +) from openff.nagl.molecule._graph.molecule import GraphMolecule from openff.nagl.nn.gcn import _function as _fn diff --git a/openff/nagl/tests/molecule/graph/test_molecule.py b/openff/nagl/tests/molecule/graph/test_molecule.py index 9fb7e4c7..d57fdd0a 100644 --- a/openff/nagl/tests/molecule/graph/test_molecule.py +++ b/openff/nagl/tests/molecule/graph/test_molecule.py @@ -1,10 +1,12 @@ +from __future__ import annotations + import pytest from numpy.testing import assert_allclose, assert_array_almost_equal -from openff.nagl.molecule._graph.molecule import GraphMolecule -from openff.nagl.molecule._graph._graph import NXMolHeteroGraph, NXMolHomoGraph from openff.nagl.features.atoms import AtomConnectivity from openff.nagl.features.bonds import BondIsInRing +from openff.nagl.molecule._graph._graph import NXMolHeteroGraph, NXMolHomoGraph +from openff.nagl.molecule._graph.molecule import GraphMolecule class TestNXMolecule: diff --git a/openff/nagl/tests/molecule/test_utils.py b/openff/nagl/tests/molecule/test_utils.py index 1b8959f2..1dc23972 100644 --- a/openff/nagl/tests/molecule/test_utils.py +++ b/openff/nagl/tests/molecule/test_utils.py @@ -1,3 +1,5 @@ +from __future__ import annotations + from openff.nagl.molecule._utils import _get_openff_molecule_information diff --git a/openff/nagl/tests/nn/gcn/test_gin.py b/openff/nagl/tests/nn/gcn/test_gin.py index 171210bb..6861b2e8 100644 --- a/openff/nagl/tests/nn/gcn/test_gin.py +++ b/openff/nagl/tests/nn/gcn/test_gin.py @@ -1,9 +1,11 @@ +from __future__ import annotations + import numpy as np import pytest import torch from openff.nagl.nn.activation import ActivationFunction -from openff.nagl.nn.gcn._gin import GINConvStack, GINConv, DGLGINConv +from openff.nagl.nn.gcn._gin import DGLGINConv, GINConv, GINConvStack try: import dgl diff --git a/openff/nagl/tests/nn/gcn/test_sage.py b/openff/nagl/tests/nn/gcn/test_sage.py index 40ce85c9..1e8fa8df 100644 --- a/openff/nagl/tests/nn/gcn/test_sage.py +++ b/openff/nagl/tests/nn/gcn/test_sage.py @@ -1,9 +1,11 @@ +from __future__ import annotations + import numpy as np -from numpy.testing import assert_allclose, assert_array_almost_equal import pytest import torch - +from numpy.testing import assert_allclose, assert_array_almost_equal from openff.utilities import requires_package + from openff.nagl.nn.activation import ActivationFunction from openff.nagl.nn.gcn import SAGEConvStack from openff.nagl.nn.gcn._sage import SAGEConv diff --git a/openff/nagl/tests/nn/test_activation.py b/openff/nagl/tests/nn/test_activation.py index 4e5e52e7..ac574fc0 100644 --- a/openff/nagl/tests/nn/test_activation.py +++ b/openff/nagl/tests/nn/test_activation.py @@ -1,10 +1,12 @@ -import torch -import torch.nn.functional as F +from __future__ import annotations import pytest +import torch +import torch.nn.functional as F from openff.nagl.nn.activation import ActivationFunction + @pytest.mark.parametrize("name", ["relu", "ReLU"]) def test_get(name): assert ActivationFunction.get(name) == ActivationFunction.ReLU diff --git a/openff/nagl/tests/nn/test_containers.py b/openff/nagl/tests/nn/test_containers.py index 551bba52..57f60a62 100644 --- a/openff/nagl/tests/nn/test_containers.py +++ b/openff/nagl/tests/nn/test_containers.py @@ -1,14 +1,16 @@ +from __future__ import annotations + import torch -from openff.nagl.nn.activation import ActivationFunction -from openff.nagl.nn.gcn import SAGEConvStack from openff.nagl.nn._containers import ( ConvolutionModule, ReadoutModule, ) from openff.nagl.nn._pooling import PoolAtomFeatures -from openff.nagl.nn.postprocess import ComputePartialCharges from openff.nagl.nn._sequential import SequentialLayers +from openff.nagl.nn.activation import ActivationFunction +from openff.nagl.nn.gcn import SAGEConvStack +from openff.nagl.nn.postprocess import ComputePartialCharges class TestConvolutionModule: diff --git a/openff/nagl/tests/nn/test_dataset.py b/openff/nagl/tests/nn/test_dataset.py index ec9128bf..d9f9b57a 100644 --- a/openff/nagl/tests/nn/test_dataset.py +++ b/openff/nagl/tests/nn/test_dataset.py @@ -1,23 +1,25 @@ +from __future__ import annotations + import typing import numpy as np import pytest import torch -from torch.utils.data import ConcatDataset from openff.units import unit +from torch.utils.data import ConcatDataset -from openff.nagl.molecule._dgl import DGLMolecule, DGLMoleculeBatch from openff.nagl.features.atoms import AtomConnectivity, AtomFormalCharge, AtomicElement from openff.nagl.features.bonds import BondIsInRing, BondOrder +from openff.nagl.molecule._dgl import DGLMolecule, DGLMoleculeBatch from openff.nagl.nn._dataset import ( - DGLMoleculeDatasetEntry, + DataHash, + DGLMoleculeDataLoader, DGLMoleculeDataset, + DGLMoleculeDatasetEntry, + _get_hashed_arrow_dataset_path, _LazyDGLMoleculeDataset, - DGLMoleculeDataLoader, - DataHash, - _get_hashed_arrow_dataset_path ) -from openff.nagl.tests.data.files import EXAMPLE_UNFEATURIZED_PARQUET_DATASET, EXAMPLE_FEATURIZED_PARQUET_DATASET +from openff.nagl.tests.data.files import EXAMPLE_FEATURIZED_PARQUET_DATASET, EXAMPLE_UNFEATURIZED_PARQUET_DATASET pytest.importorskip("dgl") diff --git a/openff/nagl/tests/nn/test_model.py b/openff/nagl/tests/nn/test_model.py index 901b48f5..41406c97 100644 --- a/openff/nagl/tests/nn/test_model.py +++ b/openff/nagl/tests/nn/test_model.py @@ -1,31 +1,30 @@ +from __future__ import annotations + import importlib.resources + import numpy as np import pytest import torch from numpy.testing import assert_allclose - -from openff.units import unit from openff.toolkit.topology import Molecule from openff.toolkit.utils.toolkits import RDKIT_AVAILABLE +from openff.units import unit -from openff.nagl.nn.gcn._sage import SAGEConvStack -from openff.nagl.nn._containers import ConvolutionModule, ReadoutModule -from openff.nagl.nn._models import BaseGNNModel, GNNModel -from openff.nagl.nn._pooling import PoolAtomFeatures, PoolBondFeatures -from openff.nagl.nn.postprocess import ComputePartialCharges -from openff.nagl.nn._sequential import SequentialLayers from openff.nagl.domains import ChemicalDomain -from openff.nagl.lookups import AtomPropertiesLookupTable, AtomPropertiesLookupTableEntry -from openff.nagl.tests.data.files import ( - EXAMPLE_AM1BCC_MODEL, - data_directory -) from openff.nagl.features.atoms import ( - AtomicElement, - AtomConnectivity, AtomAverageFormalCharge, + AtomConnectivity, + AtomicElement, AtomInRingOfSize, ) +from openff.nagl.lookups import AtomPropertiesLookupTable, AtomPropertiesLookupTableEntry +from openff.nagl.nn._containers import ConvolutionModule, ReadoutModule +from openff.nagl.nn._models import BaseGNNModel, GNNModel +from openff.nagl.nn._pooling import PoolAtomFeatures, PoolBondFeatures +from openff.nagl.nn._sequential import SequentialLayers +from openff.nagl.nn.gcn._sage import SAGEConvStack +from openff.nagl.nn.postprocess import ComputePartialCharges +from openff.nagl.tests.data.files import EXAMPLE_AM1BCC_MODEL, data_directory @pytest.fixture() diff --git a/openff/nagl/tests/nn/test_pooling.py b/openff/nagl/tests/nn/test_pooling.py index 768d057c..193028ed 100644 --- a/openff/nagl/tests/nn/test_pooling.py +++ b/openff/nagl/tests/nn/test_pooling.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import numpy as np import pytest import torch diff --git a/openff/nagl/tests/nn/test_postprocess.py b/openff/nagl/tests/nn/test_postprocess.py index dd1fb371..6b0bba4f 100644 --- a/openff/nagl/tests/nn/test_postprocess.py +++ b/openff/nagl/tests/nn/test_postprocess.py @@ -1,12 +1,11 @@ -import pytest +from __future__ import annotations + import numpy as np +import pytest import torch from openff.nagl.molecule._dgl import DGLMolecule, DGLMoleculeBatch -from openff.nagl.nn.postprocess import ( - ComputePartialCharges, - RegularizedComputePartialCharges -) +from openff.nagl.nn.postprocess import ComputePartialCharges, RegularizedComputePartialCharges # @pytest.fixture # def dgl_carboxylate diff --git a/openff/nagl/tests/nn/test_sequential.py b/openff/nagl/tests/nn/test_sequential.py index f41d65e4..f64512c5 100644 --- a/openff/nagl/tests/nn/test_sequential.py +++ b/openff/nagl/tests/nn/test_sequential.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import numpy as np import pytest import torch diff --git a/openff/nagl/tests/test_domain.py b/openff/nagl/tests/test_domain.py index b06c004c..89d57953 100644 --- a/openff/nagl/tests/test_domain.py +++ b/openff/nagl/tests/test_domain.py @@ -1,7 +1,10 @@ +from __future__ import annotations + import pytest from openff.nagl.domains import ChemicalDomain + class TestChemicalDomain: @pytest.mark.parametrize( diff --git a/openff/nagl/tests/test_lookups.py b/openff/nagl/tests/test_lookups.py index 7c858d23..d0f8a604 100644 --- a/openff/nagl/tests/test_lookups.py +++ b/openff/nagl/tests/test_lookups.py @@ -1,11 +1,13 @@ +from __future__ import annotations + import types -import pytest -import torch import numpy as np +import pytest +import torch from numpy.testing import assert_allclose - from openff.toolkit import Molecule + from openff.nagl.lookups import ( AtomPropertiesLookupTable, AtomPropertiesLookupTableEntry, @@ -99,9 +101,9 @@ def test_lookup_with_different_connectivity(self, lookup_table): assert_allclose(properties.numpy(), np.array([-0.103, 0.234, -0.209, -0.209, 0.096, 0.096, 0.096])) def test_lookup_long(self, lookup_table): - from packaging.version import Version from openff.toolkit import __version__ as toolkit_version from openff.utilities import has_package + from packaging.version import Version mol = Molecule.from_smiles(341 * "C") diff --git a/openff/nagl/tests/testing/torch.py b/openff/nagl/tests/testing/torch.py index b3a14760..f298e8b2 100644 --- a/openff/nagl/tests/testing/torch.py +++ b/openff/nagl/tests/testing/torch.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import functools from torch.testing import assert_close diff --git a/openff/nagl/tests/testing/utils.py b/openff/nagl/tests/testing/utils.py index 5c64a4ca..2601a2f0 100644 --- a/openff/nagl/tests/testing/utils.py +++ b/openff/nagl/tests/testing/utils.py @@ -1,3 +1,5 @@ +from __future__ import annotations + def rdkit_molecule_to_smiles(rdkit_molecule): from openff.toolkit.topology.molecule import Molecule diff --git a/openff/nagl/tests/training/test_loss.py b/openff/nagl/tests/training/test_loss.py index dbd5aa93..e4dca058 100644 --- a/openff/nagl/tests/training/test_loss.py +++ b/openff/nagl/tests/training/test_loss.py @@ -1,20 +1,24 @@ +from __future__ import annotations + import typing from typing import Dict, List + +import numpy as np import pytest -from openff.nagl.nn._containers import ReadoutModule import torch -import numpy as np -from openff.nagl.training.metrics import RMSEMetric +from openff.nagl.nn._containers import ReadoutModule from openff.nagl.training.loss import ( - _BaseTarget, - MultipleDipoleTarget, - SingleDipoleTarget, + GeneralLinearFitTarget, HeavyAtomReadoutTarget, - ReadoutTarget, + MultipleDipoleTarget, MultipleESPTarget, - GeneralLinearFitTarget + ReadoutTarget, + SingleDipoleTarget, + _BaseTarget, ) +from openff.nagl.training.metrics import RMSEMetric + class TestBaseTarget: diff --git a/openff/nagl/tests/training/test_metrics.py b/openff/nagl/tests/training/test_metrics.py index 6a8ba6e3..26f0b196 100644 --- a/openff/nagl/tests/training/test_metrics.py +++ b/openff/nagl/tests/training/test_metrics.py @@ -1,11 +1,14 @@ +from __future__ import annotations + import torch from openff.nagl.training.metrics import ( - RMSEMetric, - MSEMetric, MAEMetric, + MSEMetric, + RMSEMetric, ) + def test_rmse(): metric = RMSEMetric() predicted = torch.tensor([-1., 3.]) diff --git a/openff/nagl/tests/training/test_training.py b/openff/nagl/tests/training/test_training.py index 44dd493e..13ea62f3 100644 --- a/openff/nagl/tests/training/test_training.py +++ b/openff/nagl/tests/training/test_training.py @@ -1,42 +1,33 @@ +from __future__ import annotations + import itertools import pathlib -import pytest import shutil -import torch import numpy as np +import pytest import pytorch_lightning as pl +import torch -from openff.nagl.config.data import DatasetConfig, DataConfig +from openff.nagl.config.data import DataConfig, DatasetConfig +from openff.nagl.config.model import ConvolutionLayer, ConvolutionModule, ForwardLayer, ModelConfig, ReadoutModule from openff.nagl.config.optimizer import OptimizerConfig from openff.nagl.config.training import TrainingConfig -from openff.nagl.config.model import ( - ForwardLayer, - ReadoutModule, - ModelConfig, - ConvolutionLayer, - ConvolutionModule -) from openff.nagl.features import atoms -from openff.nagl.training.metrics import RMSEMetric -from openff.nagl.training.loss import ReadoutTarget -from openff.nagl.training.training import DGLMoleculeDataModule, DataHash, TrainingGNNModel +from openff.nagl.nn._dataset import DGLMoleculeDataLoader, DGLMoleculeDataset, _LazyDGLMoleculeDataset from openff.nagl.nn._models import GNNModel -from openff.nagl.nn._dataset import ( - DGLMoleculeDataLoader, - DGLMoleculeDataset, - _LazyDGLMoleculeDataset -) -from openff.nagl.config.training import TrainingConfig from openff.nagl.tests.data.files import ( - EXAMPLE_UNFEATURIZED_PARQUET_DATASET, + EXAMPLE_FEATURIZED_LAZY_DATA, + EXAMPLE_FEATURIZED_LAZY_DATA_SHORT, EXAMPLE_FEATURIZED_PARQUET_DATASET, - EXAMPLE_UNFEATURIZED_PARQUET_DATASET_SHORT, EXAMPLE_TRAINING_CONFIG, EXAMPLE_TRAINING_CONFIG_LAZY, - EXAMPLE_FEATURIZED_LAZY_DATA, - EXAMPLE_FEATURIZED_LAZY_DATA_SHORT, + EXAMPLE_UNFEATURIZED_PARQUET_DATASET, + EXAMPLE_UNFEATURIZED_PARQUET_DATASET_SHORT, ) +from openff.nagl.training.loss import ReadoutTarget +from openff.nagl.training.metrics import RMSEMetric +from openff.nagl.training.training import DataHash, DGLMoleculeDataModule, TrainingGNNModel dgl = pytest.importorskip("dgl") diff --git a/openff/nagl/tests/utils/test_openff.py b/openff/nagl/tests/utils/test_openff.py index 70a97523..5d6829a6 100644 --- a/openff/nagl/tests/utils/test_openff.py +++ b/openff/nagl/tests/utils/test_openff.py @@ -1,35 +1,37 @@ +from __future__ import annotations + import gzip import warnings import numpy as np import pytest from numpy.testing import assert_allclose +from openff.toolkit import AmberToolsToolkitWrapper, Molecule, OpenEyeToolkitWrapper, RDKitToolkitWrapper from openff.toolkit.topology import Molecule -from openff.toolkit import Molecule, RDKitToolkitWrapper, AmberToolsToolkitWrapper, OpenEyeToolkitWrapper -from openff.nagl.toolkits.rdkit import NAGLRDKitToolkitWrapper -from openff.nagl.toolkits.registry import NAGLToolkitRegistry -from openff.toolkit.utils.toolkit_registry import toolkit_registry_manager, ToolkitRegistry -from openff.toolkit.utils.toolkits import RDKIT_AVAILABLE, OPENEYE_AVAILABLE from openff.toolkit.utils.exceptions import MultipleComponentsInMoleculeWarning +from openff.toolkit.utils.toolkit_registry import ToolkitRegistry, toolkit_registry_manager +from openff.toolkit.utils.toolkits import OPENEYE_AVAILABLE, RDKIT_AVAILABLE from openff.units import unit +from openff.nagl.tests.data.files import COFACTOR_SDF_GZ, EXAMPLE_MODEL_RC4 from openff.nagl.toolkits.openff import ( + _molecule_from_dict, + _molecule_to_dict, + calculate_circular_fingerprint_similarity, + capture_toolkit_warnings, get_best_rmsd, get_openff_molecule_bond_indices, is_conformer_identical, map_indexed_smiles, + molecule_from_networkx, normalize_molecule, smiles_to_inchi_key, - calculate_circular_fingerprint_similarity, - capture_toolkit_warnings, - molecule_from_networkx, - _molecule_from_dict, - _molecule_to_dict, split_up_molecule, ) +from openff.nagl.toolkits.rdkit import NAGLRDKitToolkitWrapper +from openff.nagl.toolkits.registry import NAGLToolkitRegistry from openff.nagl.utils._utils import transform_coordinates -from openff.nagl.tests.data.files import COFACTOR_SDF_GZ, EXAMPLE_MODEL_RC4 def _load_rdkit_molecule_exactly(mapped_smiles: str): """ @@ -278,8 +280,8 @@ def test_calculate_circular_fingerprint_similarity( @pytest.mark.skipif(not RDKIT_AVAILABLE, reason="requires rdkit") def test_get_best_rmsd(): - from rdkit.Chem import rdMolAlign from openff.toolkit.topology.molecule import Molecule + from rdkit.Chem import rdMolAlign offmol = Molecule.from_smiles("CCC") offmol._conformers = [ @@ -468,8 +470,8 @@ def test_toolkit_registry_passes_through_nagl(toolkit_combinations): Tests issue #177: OpenEye being called when disallowed by the native toolkit registry manager """ - from rdkit.Chem import ForwardSDMolSupplier from openff.toolkit.utils.nagl_wrapper import NAGLToolkitWrapper + from rdkit.Chem import ForwardSDMolSupplier suppl = ForwardSDMolSupplier(gzip.open(COFACTOR_SDF_GZ), removeHs=False) @@ -498,8 +500,9 @@ def test_compute_partial_charges_with_toolkit_registry(toolkit_combinations): Tests issue #177: OpenEye being called when disallowed by the native toolkit registry manager """ - from rdkit.Chem import ForwardSDMolSupplier from openff.toolkit.utils.nagl_wrapper import NAGLToolkitWrapper + from rdkit.Chem import ForwardSDMolSupplier + from openff.nagl import GNNModel suppl = ForwardSDMolSupplier(gzip.open(COFACTOR_SDF_GZ), removeHs=False) @@ -524,10 +527,10 @@ def test_toolkit_registry_passes_through_nagl_and_fails(): Tests issue #177: OpenEye being called when disallowed by the native toolkit registry manager """ - from rdkit.Chem import ForwardSDMolSupplier from openff.toolkit.utils import GLOBAL_TOOLKIT_REGISTRY - from openff.toolkit.utils.nagl_wrapper import NAGLToolkitWrapper from openff.toolkit.utils.exceptions import InconsistentStereochemistryError + from openff.toolkit.utils.nagl_wrapper import NAGLToolkitWrapper + from rdkit.Chem import ForwardSDMolSupplier suppl = ForwardSDMolSupplier(gzip.open(COFACTOR_SDF_GZ), removeHs=False) rdmol = list(suppl)[0] diff --git a/openff/nagl/tests/utils/test_resonance.py b/openff/nagl/tests/utils/test_resonance.py index 7cf714a5..eb17e0d6 100644 --- a/openff/nagl/tests/utils/test_resonance.py +++ b/openff/nagl/tests/utils/test_resonance.py @@ -1,9 +1,10 @@ -import pytest +from __future__ import annotations +import pytest from openff.units import unit -from openff.nagl.utils.resonance import ResonanceEnumerator, FragmentEnumerator from openff.nagl.tests.testing.utils import assert_smiles_equal +from openff.nagl.utils.resonance import FragmentEnumerator, ResonanceEnumerator @pytest.fixture diff --git a/openff/nagl/tests/utils/test_utils.py b/openff/nagl/tests/utils/test_utils.py index b0929b92..d43ac702 100644 --- a/openff/nagl/tests/utils/test_utils.py +++ b/openff/nagl/tests/utils/test_utils.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import numpy as np import pytest import torch diff --git a/openff/nagl/toolkits/__init__.py b/openff/nagl/toolkits/__init__.py index 280900e1..7429e762 100644 --- a/openff/nagl/toolkits/__init__.py +++ b/openff/nagl/toolkits/__init__.py @@ -1,4 +1,5 @@ "Extensions of the OpenFF Toolkit ToolkitWrappers for use with NAGL." +from __future__ import annotations from openff.nagl.toolkits.openeye import NAGLOpenEyeToolkitWrapper from openff.nagl.toolkits.rdkit import NAGLRDKitToolkitWrapper diff --git a/openff/nagl/toolkits/_base.py b/openff/nagl/toolkits/_base.py index 455a6030..1ffabbaa 100644 --- a/openff/nagl/toolkits/_base.py +++ b/openff/nagl/toolkits/_base.py @@ -1,8 +1,11 @@ +from __future__ import annotations + from typing import Union -from openff.nagl._base.metaregistry import create_registry_metaclass from openff.toolkit.utils.base_wrapper import ToolkitWrapper +from openff.nagl._base.metaregistry import create_registry_metaclass + class NAGLToolkitWrapperMeta(create_registry_metaclass("name", ignore_case=True)): pass diff --git a/openff/nagl/toolkits/openeye.py b/openff/nagl/toolkits/openeye.py index b8b29e1c..1d4b3580 100644 --- a/openff/nagl/toolkits/openeye.py +++ b/openff/nagl/toolkits/openeye.py @@ -1,17 +1,16 @@ """Functions using the OpenEye toolkit""" +from __future__ import annotations import copy -from typing import Tuple, TYPE_CHECKING, List, Union +from typing import TYPE_CHECKING, List, Tuple, Union import numpy as np - +from openff.toolkit.utils.openeye_wrapper import OpenEyeToolkitWrapper from openff.units import unit from openff.nagl.toolkits._base import NAGLToolkitWrapperBase -from openff.toolkit.utils.openeye_wrapper import OpenEyeToolkitWrapper from openff.nagl.utils._types import HybridizationType - if TYPE_CHECKING: from openff.toolkit.topology import Molecule diff --git a/openff/nagl/toolkits/openff.py b/openff/nagl/toolkits/openff.py index cc2807b4..49ebec42 100644 --- a/openff/nagl/toolkits/openff.py +++ b/openff/nagl/toolkits/openff.py @@ -1,18 +1,19 @@ -import warnings +from __future__ import annotations + import contextlib import copy import functools -from typing import TYPE_CHECKING, Tuple, List, Union, Dict, NamedTuple, Any, Optional +import warnings +from typing import TYPE_CHECKING, Any, Dict, List, NamedTuple, Optional, Tuple, Union import numpy as np - from openff.units import unit - from openff.utilities import requires_package from openff.utilities.exceptions import MissingOptionalDependencyError if TYPE_CHECKING: from openff.toolkit.topology import Molecule + from openff.nagl.utils._types import HybridizationType @@ -33,14 +34,15 @@ def call_toolkit_function(function_name, toolkit_registry, *args, **kwargs): The keyword arguments to pass to the function. """ from openff.toolkit import GLOBAL_TOOLKIT_REGISTRY, ToolkitRegistry - from openff.nagl.toolkits.registry import NAGLToolkitRegistry + from openff.toolkit.utils.exceptions import InvalidToolkitRegistryError + from openff.nagl.toolkits._base import ( - NAGLToolkitWrapperMeta, NAGLToolkitWrapperBase, + NAGLToolkitWrapperMeta, ) - from openff.toolkit.utils.exceptions import InvalidToolkitRegistryError - from openff.nagl.toolkits.rdkit import NAGLRDKitToolkitWrapper from openff.nagl.toolkits.openeye import NAGLOpenEyeToolkitWrapper + from openff.nagl.toolkits.rdkit import NAGLRDKitToolkitWrapper + from openff.nagl.toolkits.registry import NAGLToolkitRegistry if isinstance(toolkit_registry, NAGLToolkitWrapperMeta): # case of NAGLRDKitToolkitWrapper (not instantiated) @@ -617,8 +619,8 @@ def stream_molecules_from_smiles_file( ------- A generator of openff.toolkit.topology.Molecule objects or SMILES strings """ - from openff.toolkit.topology.molecule import SmilesParsingError from openff.toolkit.topology import Molecule + from openff.toolkit.topology.molecule import SmilesParsingError with open(file, "r") as f: smiles = [x.strip() for x in f.readlines()] diff --git a/openff/nagl/toolkits/rdkit.py b/openff/nagl/toolkits/rdkit.py index 85d1285d..a4d4cd6f 100644 --- a/openff/nagl/toolkits/rdkit.py +++ b/openff/nagl/toolkits/rdkit.py @@ -1,17 +1,15 @@ """Functions using the RDKit toolkit""" - +from __future__ import annotations import copy import functools -from typing import Tuple, TYPE_CHECKING, List, Union +from typing import TYPE_CHECKING, List, Tuple, Union import numpy as np - +from openff.toolkit.utils.rdkit_wrapper import RDKitToolkitWrapper from openff.units import unit - from openff.nagl.toolkits._base import NAGLToolkitWrapperBase -from openff.toolkit.utils.rdkit_wrapper import RDKitToolkitWrapper from openff.nagl.utils._types import HybridizationType if TYPE_CHECKING: diff --git a/openff/nagl/toolkits/registry.py b/openff/nagl/toolkits/registry.py index 3cf47cfe..47f35c28 100644 --- a/openff/nagl/toolkits/registry.py +++ b/openff/nagl/toolkits/registry.py @@ -1,15 +1,18 @@ """Defining the NAGLToolkitRegistry class.""" +from __future__ import annotations from typing import List, Optional from openff.toolkit.utils.toolkit_registry import ( ToolkitRegistry as _ToolkitRegistry, +) +from openff.toolkit.utils.toolkit_registry import ( ToolkitUnavailableException, ) from openff.nagl.toolkits._base import ( - NAGLToolkitWrapperMeta, NAGLToolkitWrapperBase, + NAGLToolkitWrapperMeta, ToolkitWrapperType, ) @@ -87,6 +90,7 @@ def from_openff_toolkit_registry(cls, toolkit_registry: _ToolkitRegistry) -> "NA A NAGLToolkitRegistry with the same toolkits as the input registry """ from openff.toolkit.utils import OpenEyeToolkitWrapper, RDKitToolkitWrapper + from openff.nagl.toolkits.openeye import NAGLOpenEyeToolkitWrapper from openff.nagl.toolkits.rdkit import NAGLRDKitToolkitWrapper diff --git a/openff/nagl/training/metrics.py b/openff/nagl/training/metrics.py index c0d2e722..9b6a3a77 100644 --- a/openff/nagl/training/metrics.py +++ b/openff/nagl/training/metrics.py @@ -1,6 +1,8 @@ """ Metrics for evaluating loss """ +from __future__ import annotations + import abc import typing @@ -8,7 +10,6 @@ from openff.nagl._base.base import ImmutableModel - if typing.TYPE_CHECKING: import torch diff --git a/openff/nagl/training/templates/files.py b/openff/nagl/training/templates/files.py index 71f17ee6..d7e04e6f 100644 --- a/openff/nagl/training/templates/files.py +++ b/openff/nagl/training/templates/files.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import importlib.resources data_directory = importlib.resources.files("openff.nagl") / "training" / "templates" diff --git a/openff/nagl/training/training.py b/openff/nagl/training/training.py index 2b110d0c..3366e26e 100644 --- a/openff/nagl/training/training.py +++ b/openff/nagl/training/training.py @@ -1,25 +1,23 @@ """Classes for training a GNN""" +from __future__ import annotations -from collections import defaultdict import functools import logging import pathlib import pickle import typing +from collections import defaultdict -import torch import pytorch_lightning as pl +import torch -from openff.nagl.config.training import TrainingConfig -from openff.nagl.config.data import DatasetConfig -from openff.nagl.nn._models import GNNModel from openff.nagl._base.base import ImmutableModel +from openff.nagl.config.data import DatasetConfig +from openff.nagl.config.training import TrainingConfig from openff.nagl.features.atoms import AtomFeature from openff.nagl.features.bonds import BondFeature -from openff.nagl.nn._dataset import ( - DGLMoleculeDataset, DataHash, - _LazyDGLMoleculeDataset -) +from openff.nagl.nn._dataset import DataHash, DGLMoleculeDataset, _LazyDGLMoleculeDataset +from openff.nagl.nn._models import GNNModel if typing.TYPE_CHECKING: from openff.nagl.molecule._dgl import DGLMoleculeOrBatch diff --git a/openff/nagl/utils/_hash.py b/openff/nagl/utils/_hash.py index 71cae626..aad0aac2 100644 --- a/openff/nagl/utils/_hash.py +++ b/openff/nagl/utils/_hash.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import enum import hashlib import json diff --git a/openff/nagl/utils/_parallelization.py b/openff/nagl/utils/_parallelization.py index d74e29a0..2c5b444d 100644 --- a/openff/nagl/utils/_parallelization.py +++ b/openff/nagl/utils/_parallelization.py @@ -1,12 +1,15 @@ +from __future__ import annotations + import contextlib import functools import itertools import logging import math import multiprocessing -import tqdm import typing +import tqdm + _S = typing.TypeVar("_S") _T = typing.TypeVar("_T") @@ -106,9 +109,9 @@ def batch_distributed( **kwargs ): import dask - from distributed import LocalCluster from dask import distributed from dask_jobqueue import LSFCluster, SLURMCluster + from distributed import LocalCluster n_workers, n_batches, batch_size = reconcile_batch_workers( entries, n_entries, batch_size, n_workers diff --git a/openff/nagl/utils/_time.py b/openff/nagl/utils/_time.py index e300d8a4..fd2207f1 100644 --- a/openff/nagl/utils/_time.py +++ b/openff/nagl/utils/_time.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import time diff --git a/openff/nagl/utils/_types.py b/openff/nagl/utils/_types.py index 0d6a9623..12b64891 100644 --- a/openff/nagl/utils/_types.py +++ b/openff/nagl/utils/_types.py @@ -1,6 +1,8 @@ +from __future__ import annotations + import enum import pathlib -from typing import Union, Tuple, NamedTuple, Dict, Literal +from typing import Dict, Literal, NamedTuple, Tuple, Union Pathlike = Union[str, pathlib.Path] diff --git a/openff/nagl/utils/_utils.py b/openff/nagl/utils/_utils.py index 67e7133e..23f93a4f 100644 --- a/openff/nagl/utils/_utils.py +++ b/openff/nagl/utils/_utils.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import typing import numpy as np diff --git a/openff/nagl/utils/resonance.py b/openff/nagl/utils/resonance.py index d396dff1..7a393d07 100644 --- a/openff/nagl/utils/resonance.py +++ b/openff/nagl/utils/resonance.py @@ -1,18 +1,18 @@ """Enumerate resonance forms of a molecule""" +from __future__ import annotations import copy import itertools import json -from typing import Dict, Optional, List, Generator, Tuple, Any, Union +from typing import Any, Dict, Generator, List, Optional, Tuple, Union import networkx as nx import numpy as np - -from openff.units import unit from openff.toolkit.topology import Molecule +from openff.units import unit -from openff.nagl.utils._types import ResonanceType, ResonanceAtomType from openff.nagl.toolkits.openff import _molecule_from_dict, _molecule_to_dict +from openff.nagl.utils._types import ResonanceAtomType, ResonanceType __all__ = ["ResonanceEnumerator", "enumerate_resonance_forms"] diff --git a/pyproject.toml b/pyproject.toml index 9b25d447..adaff330 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -71,3 +71,9 @@ versionfile_source = "openff/nagl/_version.py" versionfile_build = "openff/nagl/_version.py" tag_prefix = "v" parentdir_prefix = "openff-nagl-" + +[tool.ruff] +target-version = "py314" +line-length = 119 +lint.select = [ "I" ] +lint.isort.required-imports = [ "from __future__ import annotations" ] \ No newline at end of file diff --git a/versioneer.py b/versioneer.py index 1e3753e6..842c9112 100644 --- a/versioneer.py +++ b/versioneer.py @@ -308,17 +308,18 @@ # pylint:disable=too-few-public-methods,redefined-outer-name,consider-using-with # pylint:disable=attribute-defined-outside-init,too-many-arguments +from __future__ import annotations + import configparser import errno +import functools import json import os import re import subprocess import sys from pathlib import Path -from typing import Any, Callable, cast, Dict, List, Optional, Tuple, Union -from typing import NoReturn -import functools +from typing import Any, Callable, Dict, List, NoReturn, Optional, Tuple, Union, cast have_tomllib = True if sys.version_info >= (3, 11):