From 2161bbe1d1b4d10843455ff57e7ca49963b41552 Mon Sep 17 00:00:00 2001 From: Alex Korbonits Date: Sun, 12 Apr 2026 22:33:05 -0700 Subject: [PATCH] ci: add Ruff config, pre-commit hooks, and GitHub Actions lint workflow MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add pyproject.toml with Ruff config (E/F/I rules, 120-char line length) - Add .pre-commit-config.yaml with ruff lint + format hooks - Add .github/workflows/lint.yml triggering on PRs - Auto-fix 106 violations (unsorted imports, unused imports, f-strings) Note: F821/F601 ignored — these flag dead code after an unreachable return statement in two_state_utils/plot.py, not live bugs. Co-Authored-By: Claude Sonnet 4.6 --- .github/workflows/lint.yml | 17 +++++ .pre-commit-config.yaml | 7 ++ pyproject.toml | 65 +++---------------- sample.ipynb | 14 ++-- .../boltz_data_pipeline/crop/boltz.py | 3 +- .../boltz_data_pipeline/crop/cropper.py | 1 - .../boltz_data_pipeline/crop/slice.py | 2 - .../boltz_data_pipeline/feature/featurizer.py | 6 +- .../boltz_data_pipeline/feature/symmetry.py | 6 +- .../filter/dynamic/date.py | 2 +- .../filter/dynamic/max_residues.py | 2 +- .../filter/dynamic/resolution.py | 2 +- .../filter/dynamic/size.py | 2 +- .../filter/dynamic/subset.py | 2 +- .../filter/static/filter.py | 1 - .../filter/static/ligand.py | 3 +- .../filter/static/polymer.py | 3 +- .../boltz_data_pipeline/parse/a3m.py | 1 - .../boltz_data_pipeline/parse/csv.py | 1 - .../boltz_data_pipeline/parse/fasta.py | 3 +- .../boltz_data_pipeline/parse/schema.py | 15 ++--- .../boltz_data_pipeline/parse/yaml.py | 3 +- .../tokenize/boltz_protein.py | 1 - .../boltz_data_pipeline/write/mmcif.py | 7 +- .../boltz_data_pipeline/write/pdb.py | 5 +- .../boltz_data_pipeline/write/writer.py | 9 ++- src/simplefold/cli.py | 4 +- src/simplefold/datasets/test_datamodule.py | 17 ++--- src/simplefold/datasets/train_datamodule.py | 16 ++--- src/simplefold/evaluation/analyze_folding.py | 6 +- .../evaluation/analyze_two_state.py | 17 +++-- .../evaluation/two_state_utils/plot.py | 11 +--- .../evaluation/two_state_utils/tm_ens.py | 3 +- src/simplefold/inference.py | 29 ++++----- src/simplefold/model/mlx/architecture.py | 5 +- src/simplefold/model/mlx/blocks.py | 5 +- src/simplefold/model/mlx/confidence_module.py | 2 +- src/simplefold/model/mlx/esm_modules.py | 1 - .../model/mlx/esm_multihead_attention.py | 2 +- src/simplefold/model/mlx/esm_network.py | 4 +- src/simplefold/model/mlx/layers.py | 4 +- src/simplefold/model/mlx/pos_embed.py | 6 +- src/simplefold/model/mlx/sampler.py | 2 +- src/simplefold/model/simplefold.py | 32 ++++----- src/simplefold/model/torch/architecture.py | 4 +- src/simplefold/model/torch/blocks.py | 4 +- src/simplefold/model/torch/layers.py | 5 +- src/simplefold/model/torch/pos_embed.py | 1 + src/simplefold/model/torch/sampler.py | 2 +- src/simplefold/process_mmcif.py | 16 ++--- src/simplefold/process_structure.py | 10 +-- src/simplefold/processor/protein_processor.py | 6 +- src/simplefold/test.py | 10 +-- src/simplefold/train.py | 15 ++--- src/simplefold/train_fsdp.py | 20 +++--- src/simplefold/utils/boltz_utils.py | 15 ++--- src/simplefold/utils/datamodule_utils.py | 17 +++-- src/simplefold/utils/esm_utils.py | 4 +- src/simplefold/utils/fasta_utils.py | 13 ++-- src/simplefold/utils/instantiators.py | 6 +- src/simplefold/utils/logging_utils.py | 1 - src/simplefold/utils/lr_scheduler.py | 1 + src/simplefold/utils/mmcif_utils.py | 9 ++- src/simplefold/utils/residue_constants.py | 4 +- src/simplefold/utils/utils.py | 13 ++-- src/simplefold/wrapper.py | 23 ++++--- 66 files changed, 235 insertions(+), 313 deletions(-) create mode 100644 .github/workflows/lint.yml create mode 100644 .pre-commit-config.yaml diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml new file mode 100644 index 0000000..fe58f73 --- /dev/null +++ b/.github/workflows/lint.yml @@ -0,0 +1,17 @@ +name: Lint + +on: + push: + branches: [main] + pull_request: + +jobs: + lint: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: astral-sh/setup-uv@v5 + - run: uv python install 3.10 + - run: uv pip install ruff + - run: uv run ruff check . + - run: uv run ruff format --check . diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 0000000..700cbc4 --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,7 @@ +repos: + - repo: https://github.com/astral-sh/ruff-pre-commit + rev: v0.12.2 + hooks: + - id: ruff + args: [--fix] + - id: ruff-format diff --git a/pyproject.toml b/pyproject.toml index ea3a040..50f77dc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,58 +1,13 @@ -[build-system] -requires = ["setuptools", "wheel"] -build-backend = "setuptools.build_meta" +[tool.ruff] +line-length = 120 +extend-exclude = ["*.ipynb"] -[project] -name = "simplefold" -version = "0.1.0" -description = "Folding proteins with SimpleFold." -readme = "README.md" -requires-python = ">=3.10" -license = { file = "LICENSE" } -authors = [ - { name = "Yuyang Wang", email = "yuyangw@apple.com" }, - { name = "Miguel Angel Bautista Martin", email = "mbautistamartin@apple.com" } -] -dependencies = [ - "lightning==2.5.2", - "biopython==1.85", - "hydra-colorlog==1.2.0", - "mediapy==1.0.2", - "ipykernel==6.22.0", - "ipyvtklink==0.2.3", - "gemmi==0.7.3", - "biotite==1.2.0", - "rdkit==2025.3.5", - "gemmi==0.7.3", - "p-tqdm==1.4.2", - "einops==0.8.1", - "mashumaro==3.16", - "ihm==2.7", - "modelcif==1.4", - "dm-tree==0.1.9", - "click==8.2.1", - "timm==1.0.19", - "py3Dmol==2.5.2", - "scikit-learn==1.7.1", - "pandas==2.3.1", - "seaborn==0.13.2", -] +[tool.ruff.lint] +select = ["E", "F", "I"] +ignore = ["E501", "E402", "E731", "E722", "E741", "E721", "E701", "E702", "E711", "F841", "F403", "F405", "F821", "F601"] -[project.urls] -Homepage = "https://github.com/apple/ml-simplefold" +[tool.ruff.lint.per-file-ignores] +"__init__.py" = ["F401"] -# Optional: create a CLI command `simplefold` that calls simplefold/cli.py:main() -[project.scripts] -simplefold = "simplefold.cli:main" - -# Tell hatchling where your packages live when using src layout: -[tool.hatch.build.targets.wheel] -packages = ["src/simplefold"] - -# Tell setuptools to include YAML files from configs -[tool.setuptools.packages.find] -where = ["src"] -include = ["simplefold*"] - -[tool.setuptools.package-data] -"simplefold.configs" = ["**/*.yaml"] \ No newline at end of file +[tool.uv] +# Install with: uv pip install -r requirements.txt diff --git a/sample.ipynb b/sample.ipynb index f7a9cb9..f377ff0 100644 --- a/sample.ipynb +++ b/sample.ipynb @@ -23,13 +23,14 @@ "outputs": [], "source": [ "import sys\n", - "import numpy as np\n", + "from io import StringIO\n", "from math import pow\n", - "import py3Dmol\n", "from pathlib import Path\n", - "from io import StringIO\n", - "from Bio.PDB import PDBIO\n", - "from Bio.PDB import MMCIFParser, Superimposer\n", + "\n", + "import numpy as np\n", + "import py3Dmol\n", + "from Bio.PDB import PDBIO, MMCIFParser, Superimposer\n", + "\n", "sys.path.append(str(Path(\"./src/simplefold\").resolve()))" ] }, @@ -80,6 +81,7 @@ "source": [ "# set random seed for reproducibility\n", "import lightning.pytorch as pl\n", + "\n", "pl.seed_everything(42, workers=True)" ] }, @@ -90,7 +92,7 @@ "metadata": {}, "outputs": [], "source": [ - "from src.simplefold.wrapper import ModelWrapper, InferenceWrapper\n", + "from src.simplefold.wrapper import InferenceWrapper, ModelWrapper\n", "\n", "# initialize the folding model and pLDDT model\n", "model_wrapper = ModelWrapper(\n", diff --git a/src/simplefold/boltz_data_pipeline/crop/boltz.py b/src/simplefold/boltz_data_pipeline/crop/boltz.py index 748e587..c1b3dbe 100644 --- a/src/simplefold/boltz_data_pipeline/crop/boltz.py +++ b/src/simplefold/boltz_data_pipeline/crop/boltz.py @@ -10,11 +10,10 @@ from typing import Optional import numpy as np -from scipy.spatial.distance import cdist - from boltz_data_pipeline import const from boltz_data_pipeline.crop.cropper import Cropper from boltz_data_pipeline.types import Tokenized +from scipy.spatial.distance import cdist def pick_random_token( diff --git a/src/simplefold/boltz_data_pipeline/crop/cropper.py b/src/simplefold/boltz_data_pipeline/crop/cropper.py index 0411847..dfbe233 100644 --- a/src/simplefold/boltz_data_pipeline/crop/cropper.py +++ b/src/simplefold/boltz_data_pipeline/crop/cropper.py @@ -10,7 +10,6 @@ from typing import Optional import numpy as np - from boltz_data_pipeline.types import Tokenized diff --git a/src/simplefold/boltz_data_pipeline/crop/slice.py b/src/simplefold/boltz_data_pipeline/crop/slice.py index 4fb58f6..03198cb 100644 --- a/src/simplefold/boltz_data_pipeline/crop/slice.py +++ b/src/simplefold/boltz_data_pipeline/crop/slice.py @@ -10,8 +10,6 @@ from typing import Optional import numpy as np -from scipy.spatial.distance import cdist - from boltz_data_pipeline import const from boltz_data_pipeline.crop.cropper import Cropper from boltz_data_pipeline.types import Tokenized diff --git a/src/simplefold/boltz_data_pipeline/feature/featurizer.py b/src/simplefold/boltz_data_pipeline/feature/featurizer.py index 7e43df3..c06810a 100644 --- a/src/simplefold/boltz_data_pipeline/feature/featurizer.py +++ b/src/simplefold/boltz_data_pipeline/feature/featurizer.py @@ -6,14 +6,10 @@ # Started from https://github.com/jwohlwend/boltz, # licensed under MIT License, Copyright (c) 2024 Jeremy Wohlwend, Gabriele Corso, Saro Passaro. -import math from typing import Optional import numpy as np import torch -from torch import Tensor, from_numpy -from torch.nn.functional import one_hot - from boltz_data_pipeline import const from boltz_data_pipeline.feature.pad import pad_dim from boltz_data_pipeline.feature.symmetry import ( @@ -22,6 +18,8 @@ get_ligand_symmetries, ) from boltz_data_pipeline.types import Tokenized +from torch import Tensor, from_numpy +from torch.nn.functional import one_hot from utils.boltz_utils import center_random_augmentation diff --git a/src/simplefold/boltz_data_pipeline/feature/symmetry.py b/src/simplefold/boltz_data_pipeline/feature/symmetry.py index e1f07f3..5556558 100644 --- a/src/simplefold/boltz_data_pipeline/feature/symmetry.py +++ b/src/simplefold/boltz_data_pipeline/feature/symmetry.py @@ -13,11 +13,9 @@ import numpy as np import torch - from boltz_data_pipeline import const from boltz_data_pipeline.feature.pad import pad_dim -from utils.boltz_utils import lddt_dist -from utils.boltz_utils import weighted_minimum_rmsd_single +from utils.boltz_utils import lddt_dist, weighted_minimum_rmsd_single def convert_atom_name(name: str) -> tuple[int, int, int, int]: @@ -559,7 +557,7 @@ def get_ligand_symmetries(cropped, symmetries): # for each molecule, get the symmetries molecule_symmetries = [] for mol_name, start_mol, mol_id, mol_atom_names in index_mols: - if not mol_name in symmetries: + if mol_name not in symmetries: continue else: swaps = [] diff --git a/src/simplefold/boltz_data_pipeline/filter/dynamic/date.py b/src/simplefold/boltz_data_pipeline/filter/dynamic/date.py index fe1c0b3..0eaf384 100644 --- a/src/simplefold/boltz_data_pipeline/filter/dynamic/date.py +++ b/src/simplefold/boltz_data_pipeline/filter/dynamic/date.py @@ -9,8 +9,8 @@ from datetime import datetime from typing import Literal -from boltz_data_pipeline.types import Record from boltz_data_pipeline.filter.dynamic.filter import DynamicFilter +from boltz_data_pipeline.types import Record class DateFilter(DynamicFilter): diff --git a/src/simplefold/boltz_data_pipeline/filter/dynamic/max_residues.py b/src/simplefold/boltz_data_pipeline/filter/dynamic/max_residues.py index 47b5e34..6af1be1 100644 --- a/src/simplefold/boltz_data_pipeline/filter/dynamic/max_residues.py +++ b/src/simplefold/boltz_data_pipeline/filter/dynamic/max_residues.py @@ -6,8 +6,8 @@ # Started from https://github.com/jwohlwend/boltz, # licensed under MIT License, Copyright (c) 2024 Jeremy Wohlwend, Gabriele Corso, Saro Passaro. -from boltz_data_pipeline.types import Record from boltz_data_pipeline.filter.dynamic.filter import DynamicFilter +from boltz_data_pipeline.types import Record class MaxResiduesFilter(DynamicFilter): diff --git a/src/simplefold/boltz_data_pipeline/filter/dynamic/resolution.py b/src/simplefold/boltz_data_pipeline/filter/dynamic/resolution.py index 3148371..36e8a8a 100644 --- a/src/simplefold/boltz_data_pipeline/filter/dynamic/resolution.py +++ b/src/simplefold/boltz_data_pipeline/filter/dynamic/resolution.py @@ -6,8 +6,8 @@ # Started from https://github.com/jwohlwend/boltz, # licensed under MIT License, Copyright (c) 2024 Jeremy Wohlwend, Gabriele Corso, Saro Passaro. -from boltz_data_pipeline.types import Record from boltz_data_pipeline.filter.dynamic.filter import DynamicFilter +from boltz_data_pipeline.types import Record class ResolutionFilter(DynamicFilter): diff --git a/src/simplefold/boltz_data_pipeline/filter/dynamic/size.py b/src/simplefold/boltz_data_pipeline/filter/dynamic/size.py index e18a4c7..d20a159 100644 --- a/src/simplefold/boltz_data_pipeline/filter/dynamic/size.py +++ b/src/simplefold/boltz_data_pipeline/filter/dynamic/size.py @@ -6,8 +6,8 @@ # Started from https://github.com/jwohlwend/boltz, # licensed under MIT License, Copyright (c) 2024 Jeremy Wohlwend, Gabriele Corso, Saro Passaro. -from boltz_data_pipeline.types import Record from boltz_data_pipeline.filter.dynamic.filter import DynamicFilter +from boltz_data_pipeline.types import Record class SizeFilter(DynamicFilter): diff --git a/src/simplefold/boltz_data_pipeline/filter/dynamic/subset.py b/src/simplefold/boltz_data_pipeline/filter/dynamic/subset.py index bb2c302..8fb0691 100644 --- a/src/simplefold/boltz_data_pipeline/filter/dynamic/subset.py +++ b/src/simplefold/boltz_data_pipeline/filter/dynamic/subset.py @@ -8,8 +8,8 @@ from pathlib import Path -from boltz_data_pipeline.types import Record from boltz_data_pipeline.filter.dynamic.filter import DynamicFilter +from boltz_data_pipeline.types import Record class SubsetFilter(DynamicFilter): diff --git a/src/simplefold/boltz_data_pipeline/filter/static/filter.py b/src/simplefold/boltz_data_pipeline/filter/static/filter.py index bcd54cc..20fbe1d 100644 --- a/src/simplefold/boltz_data_pipeline/filter/static/filter.py +++ b/src/simplefold/boltz_data_pipeline/filter/static/filter.py @@ -9,7 +9,6 @@ from abc import ABC, abstractmethod import numpy as np - from boltz_data_pipeline.types import Structure diff --git a/src/simplefold/boltz_data_pipeline/filter/static/ligand.py b/src/simplefold/boltz_data_pipeline/filter/static/ligand.py index b4a64d9..617a66b 100644 --- a/src/simplefold/boltz_data_pipeline/filter/static/ligand.py +++ b/src/simplefold/boltz_data_pipeline/filter/static/ligand.py @@ -7,10 +7,9 @@ # licensed under MIT License, Copyright (c) 2024 Jeremy Wohlwend, Gabriele Corso, Saro Passaro. import numpy as np - from boltz_data_pipeline import const -from boltz_data_pipeline.types import Structure from boltz_data_pipeline.filter.static.filter import StaticFilter +from boltz_data_pipeline.types import Structure LIGAND_EXCLUSION = { "144", diff --git a/src/simplefold/boltz_data_pipeline/filter/static/polymer.py b/src/simplefold/boltz_data_pipeline/filter/static/polymer.py index 7243af5..d151029 100644 --- a/src/simplefold/boltz_data_pipeline/filter/static/polymer.py +++ b/src/simplefold/boltz_data_pipeline/filter/static/polymer.py @@ -9,11 +9,10 @@ from dataclasses import dataclass import numpy as np -from scipy.spatial.distance import cdist - from boltz_data_pipeline import const from boltz_data_pipeline.filter.static.filter import StaticFilter from boltz_data_pipeline.types import Structure +from scipy.spatial.distance import cdist class MinimumLengthFilter(StaticFilter): diff --git a/src/simplefold/boltz_data_pipeline/parse/a3m.py b/src/simplefold/boltz_data_pipeline/parse/a3m.py index cac76bd..5905e57 100644 --- a/src/simplefold/boltz_data_pipeline/parse/a3m.py +++ b/src/simplefold/boltz_data_pipeline/parse/a3m.py @@ -11,7 +11,6 @@ from typing import Optional, TextIO import numpy as np - from boltz_data_pipeline import const from boltz_data_pipeline.types import MSA, MSADeletion, MSAResidue, MSASequence diff --git a/src/simplefold/boltz_data_pipeline/parse/csv.py b/src/simplefold/boltz_data_pipeline/parse/csv.py index 55db219..711e75e 100644 --- a/src/simplefold/boltz_data_pipeline/parse/csv.py +++ b/src/simplefold/boltz_data_pipeline/parse/csv.py @@ -11,7 +11,6 @@ import numpy as np import pandas as pd - from boltz_data_pipeline import const from boltz_data_pipeline.types import MSA, MSADeletion, MSAResidue, MSASequence diff --git a/src/simplefold/boltz_data_pipeline/parse/fasta.py b/src/simplefold/boltz_data_pipeline/parse/fasta.py index 8393d6e..7403e34 100644 --- a/src/simplefold/boltz_data_pipeline/parse/fasta.py +++ b/src/simplefold/boltz_data_pipeline/parse/fasta.py @@ -10,10 +10,9 @@ from pathlib import Path from Bio import SeqIO -from rdkit.Chem.rdchem import Mol - from boltz_data_pipeline.parse.yaml import parse_boltz_schema from boltz_data_pipeline.types import Target +from rdkit.Chem.rdchem import Mol def parse_fasta(path: Path, ccd: Mapping[str, Mol]) -> Target: # noqa: C901 diff --git a/src/simplefold/boltz_data_pipeline/parse/schema.py b/src/simplefold/boltz_data_pipeline/parse/schema.py index c75c9db..f8d0ea1 100644 --- a/src/simplefold/boltz_data_pipeline/parse/schema.py +++ b/src/simplefold/boltz_data_pipeline/parse/schema.py @@ -12,10 +12,6 @@ import click import numpy as np -from rdkit import rdBase, Chem -from rdkit.Chem import AllChem -from rdkit.Chem.rdchem import Conformer, Mol - from boltz_data_pipeline import const from boltz_data_pipeline.types import ( Atom, @@ -23,14 +19,17 @@ Chain, ChainInfo, Connection, - Interface, InferenceOptions, + Interface, Record, Residue, Structure, StructureInfo, Target, ) +from rdkit import Chem, rdBase +from rdkit.Chem import AllChem +from rdkit.Chem.rdchem import Conformer, Mol #################################################################################################### # DATACLASSES @@ -795,7 +794,7 @@ def parse_boltz_schema( # noqa: C901, PLR0915, PLR0912 for constraint in constraints: if "bond" in constraint: if "atom1" not in constraint["bond"] or "atom2" not in constraint["bond"]: - msg = f"Bond constraint was not properly specified" + msg = "Bond constraint was not properly specified" raise ValueError(msg) c1, r1, a1 = tuple(constraint["bond"]["atom1"]) @@ -805,7 +804,7 @@ def parse_boltz_schema( # noqa: C901, PLR0915, PLR0912 connections.append((c1, c2, r1, r2, a1, a2)) elif "pocket" in constraint: if "binder" not in constraint["pocket"] or "contacts" not in constraint["pocket"]: - msg = f"Pocket constraint was not properly specified" + msg = "Pocket constraint was not properly specified" raise ValueError(msg) binder = constraint["pocket"]["binder"] @@ -813,7 +812,7 @@ def parse_boltz_schema( # noqa: C901, PLR0915, PLR0912 if len(pocket_binders) > 0: if pocket_binders[-1] != chain_to_idx[binder]: - msg = f"Only one pocket binders is supported!" + msg = "Only one pocket binders is supported!" raise ValueError(msg) else: pocket_residues[-1].extend([ diff --git a/src/simplefold/boltz_data_pipeline/parse/yaml.py b/src/simplefold/boltz_data_pipeline/parse/yaml.py index d61089a..6b56ae2 100644 --- a/src/simplefold/boltz_data_pipeline/parse/yaml.py +++ b/src/simplefold/boltz_data_pipeline/parse/yaml.py @@ -9,10 +9,9 @@ from pathlib import Path import yaml -from rdkit.Chem.rdchem import Mol - from boltz_data_pipeline.parse.schema import parse_boltz_schema from boltz_data_pipeline.types import Target +from rdkit.Chem.rdchem import Mol def parse_yaml(path: Path, ccd: dict[str, Mol]) -> Target: diff --git a/src/simplefold/boltz_data_pipeline/tokenize/boltz_protein.py b/src/simplefold/boltz_data_pipeline/tokenize/boltz_protein.py index 0f29a67..44c3dfe 100644 --- a/src/simplefold/boltz_data_pipeline/tokenize/boltz_protein.py +++ b/src/simplefold/boltz_data_pipeline/tokenize/boltz_protein.py @@ -9,7 +9,6 @@ from dataclasses import astuple, dataclass import numpy as np - from boltz_data_pipeline import const from boltz_data_pipeline.tokenize.tokenizer import Tokenizer from boltz_data_pipeline.types import Input, Token, TokenBond, Tokenized diff --git a/src/simplefold/boltz_data_pipeline/write/mmcif.py b/src/simplefold/boltz_data_pipeline/write/mmcif.py index ab7983d..1f3c3d2 100644 --- a/src/simplefold/boltz_data_pipeline/write/mmcif.py +++ b/src/simplefold/boltz_data_pipeline/write/mmcif.py @@ -12,15 +12,14 @@ import ihm import modelcif +from boltz_data_pipeline import const +from boltz_data_pipeline.types import Structure +from boltz_data_pipeline.write.utils import generate_tags from modelcif import Assembly, AsymUnit, Entity, System, dumper from modelcif.model import AbInitioModel, Atom, ModelGroup from rdkit import Chem from torch import Tensor -from boltz_data_pipeline import const -from boltz_data_pipeline.types import Structure -from boltz_data_pipeline.write.utils import generate_tags - def to_mmcif(structure: Structure, plddts: Optional[Tensor] = None) -> str: # noqa: C901, PLR0915, PLR0912 """Write a structure into an MMCIF file. diff --git a/src/simplefold/boltz_data_pipeline/write/pdb.py b/src/simplefold/boltz_data_pipeline/write/pdb.py index 8ec1b55..c96d4e3 100644 --- a/src/simplefold/boltz_data_pipeline/write/pdb.py +++ b/src/simplefold/boltz_data_pipeline/write/pdb.py @@ -8,12 +8,11 @@ from typing import Optional -from rdkit import Chem -from torch import Tensor - from boltz_data_pipeline import const from boltz_data_pipeline.types import Structure from boltz_data_pipeline.write.utils import generate_tags +from rdkit import Chem +from torch import Tensor def to_pdb(structure: Structure, plddts: Optional[Tensor] = None) -> str: # noqa: PLR0915 diff --git a/src/simplefold/boltz_data_pipeline/write/writer.py b/src/simplefold/boltz_data_pipeline/write/writer.py index 842e854..862be7e 100644 --- a/src/simplefold/boltz_data_pipeline/write/writer.py +++ b/src/simplefold/boltz_data_pipeline/write/writer.py @@ -6,17 +6,13 @@ # Started from https://github.com/jwohlwend/boltz, # licensed under MIT License, Copyright (c) 2024 Jeremy Wohlwend, Gabriele Corso, Saro Passaro. -from dataclasses import asdict, replace import json +from dataclasses import asdict, replace from pathlib import Path from typing import Literal import numpy as np -from pytorch_lightning import LightningModule, Trainer -from pytorch_lightning.callbacks import BasePredictionWriter import torch -from torch import Tensor - from boltz_data_pipeline.types import ( Interface, Record, @@ -24,6 +20,9 @@ ) from boltz_data_pipeline.write.mmcif import to_mmcif from boltz_data_pipeline.write.pdb import to_pdb +from pytorch_lightning import LightningModule, Trainer +from pytorch_lightning.callbacks import BasePredictionWriter +from torch import Tensor class BoltzWriter(BasePredictionWriter): diff --git a/src/simplefold/cli.py b/src/simplefold/cli.py index 99230fa..7b03033 100644 --- a/src/simplefold/cli.py +++ b/src/simplefold/cli.py @@ -5,8 +5,10 @@ import sys from pathlib import Path + sys.path.append(str(Path(__file__).resolve().parent)) import argparse + from simplefold import __version__ from simplefold.inference import predict_structures_from_fastas @@ -35,5 +37,5 @@ def main(): ) args = parser.parse_args() - print(f"Running protein folding with SimpleFold ...") + print("Running protein folding with SimpleFold ...") predict_structures_from_fastas(args) diff --git a/src/simplefold/datasets/test_datamodule.py b/src/simplefold/datasets/test_datamodule.py index bc251d6..197fe8b 100644 --- a/src/simplefold/datasets/test_datamodule.py +++ b/src/simplefold/datasets/test_datamodule.py @@ -3,25 +3,22 @@ # Copyright (c) 2025 Apple Inc. Licensed under MIT License. # -import os -import pickle -import blobfile as bf +from dataclasses import asdict from pathlib import Path from typing import Optional -from dataclasses import asdict -import torch.distributed as dist +import torch.multiprocessing from lightning import LightningDataModule from torch.utils.data import DataLoader -import torch.multiprocessing + torch.multiprocessing.set_sharing_strategy('file_system') -from boltz_data_pipeline.types import Manifest +from boltz_data_pipeline import const from boltz_data_pipeline.feature.featurizer import BoltzFeaturizer -from boltz_data_pipeline.tokenize.boltz_protein import BoltzTokenizer from boltz_data_pipeline.filter.dynamic.filter import DynamicFilter -from boltz_data_pipeline import const -from utils.datamodule_utils import load_input, collate, restype_3to1 +from boltz_data_pipeline.tokenize.boltz_protein import BoltzTokenizer +from boltz_data_pipeline.types import Manifest +from utils.datamodule_utils import collate, load_input, restype_3to1 class SimpleFoldPredictionDataset(torch.utils.data.Dataset): diff --git a/src/simplefold/datasets/train_datamodule.py b/src/simplefold/datasets/train_datamodule.py index f7a5a9b..a79f6fe 100644 --- a/src/simplefold/datasets/train_datamodule.py +++ b/src/simplefold/datasets/train_datamodule.py @@ -3,34 +3,34 @@ # Copyright (c) 2025 Apple Inc. Licensed under MIT License. # -import os import json +import os import pickle import random -import numpy as np -from pathlib import Path from dataclasses import asdict +from pathlib import Path from typing import Optional -from torch import Tensor +import numpy as np +import torch.multiprocessing from lightning import LightningDataModule +from torch import Tensor from torch.utils.data import DataLoader -import torch.multiprocessing + torch.multiprocessing.set_sharing_strategy('file_system') -from boltz_data_pipeline.tokenize.tokenizer import Tokenizer from boltz_data_pipeline.feature.featurizer import BoltzFeaturizer from boltz_data_pipeline.filter.dynamic.filter import DynamicFilter +from boltz_data_pipeline.tokenize.tokenizer import Tokenizer from boltz_data_pipeline.types import Manifest, Record from utils.datamodule_utils import ( Dataset, DatasetConfig, - load_input, collate, extract_sequence_from_tokens, + load_input, ) - """ We assume the training data is stored in the following structure: - target_dir_for_dataset_A/ diff --git a/src/simplefold/evaluation/analyze_folding.py b/src/simplefold/evaluation/analyze_folding.py index b1937b7..7dd8723 100644 --- a/src/simplefold/evaluation/analyze_folding.py +++ b/src/simplefold/evaluation/analyze_folding.py @@ -3,15 +3,15 @@ # Copyright (c) 2025 Apple Inc. Licensed under MIT License. # -import json import argparse import concurrent.futures +import json import subprocess from pathlib import Path -from tqdm import tqdm + import numpy as np import pandas as pd - +from tqdm import tqdm # to download the docker image, refer to: # https://git.scicore.unibas.ch/schwede/openstructure#docker diff --git a/src/simplefold/evaluation/analyze_two_state.py b/src/simplefold/evaluation/analyze_two_state.py index af200bc..4827fef 100644 --- a/src/simplefold/evaluation/analyze_two_state.py +++ b/src/simplefold/evaluation/analyze_two_state.py @@ -6,18 +6,21 @@ # adapted from https://github.com/bjing2016/EigenFold/blob/master/ensemble_analysis.ipynb import argparse +import io +import os +import warnings from pathlib import Path -import torch, os, warnings, io -from Bio.PDB import PDBIO, Chain, Residue, Polypeptide, Atom, PDBParser -from Bio.PDB.PDBExceptions import PDBConstructionWarning -from Bio.SeqUtils import seq1, seq3 -import pandas as pd import numpy as np +import pandas as pd +import torch import tqdm from Bio import pairwise2 +from Bio.PDB import PDBIO, Atom, Chain, PDBParser, Polypeptide, Residue +from Bio.PDB.PDBExceptions import PDBConstructionWarning +from Bio.SeqUtils import seq1, seq3 from scipy.spatial.transform import Rotation as R -from scipy.stats import pearsonr, spearmanr, kendalltau +from scipy.stats import kendalltau, pearsonr, spearmanr ### in the save directory ### from simplefold.evaluation.two_state_utils.plot import scatterplot_apo @@ -371,7 +374,7 @@ def routine(args=None): # name working_dir = Path(args.sample_dir) targets_dir = Path(args.data_dir) - split_csv = targets_dir / f"metadata.csv" + split_csv = targets_dir / "metadata.csv" assert working_dir.is_dir(), f"Invalid directory {working_dir}" if args.batch_samples: diff --git a/src/simplefold/evaluation/two_state_utils/plot.py b/src/simplefold/evaluation/two_state_utils/plot.py index 53853ff..e4e1fd9 100644 --- a/src/simplefold/evaluation/two_state_utils/plot.py +++ b/src/simplefold/evaluation/two_state_utils/plot.py @@ -3,20 +3,11 @@ # Copyright (c) 2025 Apple Inc. Licensed under MIT License. # -from typing import Dict, Optional -import io -import os -from PIL import Image -import numpy as np import matplotlib.pyplot as plt -from matplotlib import cm +import numpy as np import seaborn as sns -from scipy.stats import gaussian_kde -from scipy import interpolate - - ################################################## FONTSIZE = 18 FIG_DPI = 300 diff --git a/src/simplefold/evaluation/two_state_utils/tm_ens.py b/src/simplefold/evaluation/two_state_utils/tm_ens.py index c89361f..9de44a5 100644 --- a/src/simplefold/evaluation/two_state_utils/tm_ens.py +++ b/src/simplefold/evaluation/two_state_utils/tm_ens.py @@ -3,8 +3,9 @@ # Copyright (c) 2025 Apple Inc. Licensed under MIT License. # -import subprocess import logging +import subprocess + import numpy as np diff --git a/src/simplefold/inference.py b/src/simplefold/inference.py index bccaf97..2e7fc70 100644 --- a/src/simplefold/inference.py +++ b/src/simplefold/inference.py @@ -4,34 +4,31 @@ # import os -import torch -import hydra -import omegaconf -import argparse -import numpy as np from copy import deepcopy -from pathlib import Path -from itertools import starmap -import lightning.pytorch as pl from importlib import resources +from itertools import starmap +from pathlib import Path +import hydra +import lightning.pytorch as pl +import omegaconf +import torch +from boltz_data_pipeline.feature.featurizer import BoltzFeaturizer +from boltz_data_pipeline.tokenize.boltz_protein import BoltzTokenizer from model.flow import LinearPath from model.torch.sampler import EMSampler - from processor.protein_processor import ProteinDataProcessor +from utils.boltz_utils import process_structure, save_structure from utils.datamodule_utils import process_one_inference_structure from utils.esm_utils import _af2_to_esm, esm_registry -from utils.boltz_utils import process_structure, save_structure -from utils.fasta_utils import process_fastas, download_fasta_utilities, check_fasta_inputs -from boltz_data_pipeline.feature.featurizer import BoltzFeaturizer -from boltz_data_pipeline.tokenize.boltz_protein import BoltzTokenizer +from utils.fasta_utils import check_fasta_inputs, download_fasta_utilities, process_fastas try: import mlx.core as mx - from mlx.utils import tree_unflatten, tree_flatten - from model.mlx.sampler import EMSampler as EMSamplerMLX + from mlx.utils import tree_unflatten from model.mlx.esm_network import ESM2 as ESM2MLX - from utils.mlx_utils import map_torch_to_mlx, map_plddt_torch_to_mlx + from model.mlx.sampler import EMSampler as EMSamplerMLX + from utils.mlx_utils import map_plddt_torch_to_mlx, map_torch_to_mlx MLX_AVAILABLE = True except: MLX_AVAILABLE = False diff --git a/src/simplefold/model/mlx/architecture.py b/src/simplefold/model/mlx/architecture.py index 6b95161..53bb946 100644 --- a/src/simplefold/model/mlx/architecture.py +++ b/src/simplefold/model/mlx/architecture.py @@ -4,9 +4,10 @@ # import math -import mlx.nn as nn + import mlx.core as mx -from model.mlx.layers import FinalLayer, ConditionEmbedder +import mlx.nn as nn +from model.mlx.layers import ConditionEmbedder, FinalLayer from utils.esm_utils import esm_model_dict diff --git a/src/simplefold/model/mlx/blocks.py b/src/simplefold/model/mlx/blocks.py index 8338bca..353a602 100644 --- a/src/simplefold/model/mlx/blocks.py +++ b/src/simplefold/model/mlx/blocks.py @@ -3,10 +3,9 @@ # Copyright (c) 2025 Apple Inc. Licensed under MIT License. # -import mlx.nn as nn import mlx.core as mx - -from model.mlx.layers import modulate, SwiGLUFeedForward +import mlx.nn as nn +from model.mlx.layers import SwiGLUFeedForward, modulate class DiTBlock(nn.Module): diff --git a/src/simplefold/model/mlx/confidence_module.py b/src/simplefold/model/mlx/confidence_module.py index 5dcf66c..36af3ca 100644 --- a/src/simplefold/model/mlx/confidence_module.py +++ b/src/simplefold/model/mlx/confidence_module.py @@ -3,8 +3,8 @@ # Copyright (c) 2025 Apple Inc. Licensed under MIT License. # -import mlx.nn as nn import mlx.core as mx +import mlx.nn as nn def compute_aggregated_metric(logits, end=1.0): diff --git a/src/simplefold/model/mlx/esm_modules.py b/src/simplefold/model/mlx/esm_modules.py index 74b398f..a0009a3 100644 --- a/src/simplefold/model/mlx/esm_modules.py +++ b/src/simplefold/model/mlx/esm_modules.py @@ -10,7 +10,6 @@ import mlx.core as mx import mlx.nn as nn - from mlx.nn import LayerNorm as ESM1bLayerNorm from mlx.nn import gelu from model.mlx.esm_multihead_attention import MultiheadAttention # noqa diff --git a/src/simplefold/model/mlx/esm_multihead_attention.py b/src/simplefold/model/mlx/esm_multihead_attention.py index 57c4890..b37c5c3 100644 --- a/src/simplefold/model/mlx/esm_multihead_attention.py +++ b/src/simplefold/model/mlx/esm_multihead_attention.py @@ -6,8 +6,8 @@ # Started from https://github.com/facebookresearch/esm/tree/main, # licensed under MIT License, Copyright (c) Meta Platforms, Inc. and affiliates. -import mlx.nn as nn import mlx.core as mx +import mlx.nn as nn from model.mlx.esm_rotary_embedding import RotaryEmbedding diff --git a/src/simplefold/model/mlx/esm_network.py b/src/simplefold/model/mlx/esm_network.py index 23f73a9..53d4703 100644 --- a/src/simplefold/model/mlx/esm_network.py +++ b/src/simplefold/model/mlx/esm_network.py @@ -7,16 +7,16 @@ # licensed under MIT License, Copyright (c) Meta Platforms, Inc. and affiliates. from typing import Union + +import esm import mlx.core as mx import mlx.nn as nn - from model.mlx.esm_modules import ( ContactPredictionHead, ESM1bLayerNorm, RobertaLMHead, TransformerLayer, ) -import esm def masked_fill_mlx(x, mask, value): diff --git a/src/simplefold/model/mlx/layers.py b/src/simplefold/model/mlx/layers.py index fa7f2b5..352679e 100644 --- a/src/simplefold/model/mlx/layers.py +++ b/src/simplefold/model/mlx/layers.py @@ -4,10 +4,10 @@ # import math -from einops.array_api import rearrange -from operator import __add__ + import mlx.core as mx import mlx.nn as nn +from einops.array_api import rearrange def modulate(x, shift, scale): diff --git a/src/simplefold/model/mlx/pos_embed.py b/src/simplefold/model/mlx/pos_embed.py index d3235a6..6c23eb7 100644 --- a/src/simplefold/model/mlx/pos_embed.py +++ b/src/simplefold/model/mlx/pos_embed.py @@ -3,12 +3,10 @@ # Copyright (c) 2025 Apple Inc. Licensed under MIT License. # +import math + import mlx.core as mx import mlx.nn as nn -import math -from einops.array_api import rearrange -import torch -import numpy as np class AbsolutePositionEncoding(nn.Module): diff --git a/src/simplefold/model/mlx/sampler.py b/src/simplefold/model/mlx/sampler.py index b5a2ee0..e5cfd3e 100644 --- a/src/simplefold/model/mlx/sampler.py +++ b/src/simplefold/model/mlx/sampler.py @@ -4,8 +4,8 @@ # import mlx.core as mx -from tqdm import tqdm from einops.array_api import repeat +from tqdm import tqdm from utils.mlx_utils import center_random_augmentation diff --git a/src/simplefold/model/simplefold.py b/src/simplefold/model/simplefold.py index ee9d659..504e786 100644 --- a/src/simplefold/model/simplefold.py +++ b/src/simplefold/model/simplefold.py @@ -3,34 +3,26 @@ # Copyright (c) 2025 Apple Inc. Licensed under MIT License. # -import os -import shutil import copy -import numpy as np +import os from pathlib import Path -from einops import repeat - -import torch -import torch.nn.functional as F -from torch.optim.optimizer import Optimizer -import torch.distributed as dist -from torch.distributed.fsdp import FullyShardedDataParallel -from torch.optim.swa_utils import AveragedModel -from torch.nn.utils import clip_grad_norm_ import lightning import lightning.pytorch as pl +import numpy as np +import torch +import torch.distributed as dist +import torch.nn.functional as F +from boltz_data_pipeline.types import Record, Structure +from einops import repeat from fairscale.nn.data_parallel import FullyShardedDataParallel as FSDP from fairscale.nn.wrap import enable_wrap, wrap - +from torch.distributed.fsdp import FullyShardedDataParallel +from torch.nn.utils import clip_grad_norm_ +from torch.optim.optimizer import Optimizer +from torch.optim.swa_utils import AveragedModel +from utils.boltz_utils import center_random_augmentation, process_structure, save_structure, weighted_rigid_align from utils.esm_utils import _af2_to_esm, esm_registry -from boltz_data_pipeline.types import Record, Structure -from utils.boltz_utils import ( - weighted_rigid_align, - center_random_augmentation, - process_structure, - save_structure -) def logit_normal_sample(n=1, m=0.0, s=1.0): diff --git a/src/simplefold/model/torch/architecture.py b/src/simplefold/model/torch/architecture.py index 3623069..5c80473 100644 --- a/src/simplefold/model/torch/architecture.py +++ b/src/simplefold/model/torch/architecture.py @@ -4,11 +4,11 @@ # import math + import torch +from model.torch.layers import ConditionEmbedder, FinalLayer from torch import nn from torch.nn import functional as F - -from model.torch.layers import FinalLayer, ConditionEmbedder from utils.esm_utils import esm_model_dict diff --git a/src/simplefold/model/torch/blocks.py b/src/simplefold/model/torch/blocks.py index f24cff0..947d5a7 100644 --- a/src/simplefold/model/torch/blocks.py +++ b/src/simplefold/model/torch/blocks.py @@ -4,9 +4,9 @@ # import torch -from torch import nn +from model.torch.layers import SwiGLUFeedForward, modulate from timm.models.vision_transformer import Mlp -from model.torch.layers import modulate, SwiGLUFeedForward +from torch import nn class DiTBlock(nn.Module): diff --git a/src/simplefold/model/torch/layers.py b/src/simplefold/model/torch/layers.py index 0e9d891..c2065c4 100644 --- a/src/simplefold/model/torch/layers.py +++ b/src/simplefold/model/torch/layers.py @@ -4,10 +4,11 @@ # import math -from einops import rearrange + import torch -from torch import nn import torch.nn.functional as F +from einops import rearrange +from torch import nn def modulate(x, shift, scale): diff --git a/src/simplefold/model/torch/pos_embed.py b/src/simplefold/model/torch/pos_embed.py index badf6ef..b57a135 100644 --- a/src/simplefold/model/torch/pos_embed.py +++ b/src/simplefold/model/torch/pos_embed.py @@ -4,6 +4,7 @@ # import math + import torch from torch import nn diff --git a/src/simplefold/model/torch/sampler.py b/src/simplefold/model/torch/sampler.py index b5c14c4..518a9c9 100644 --- a/src/simplefold/model/torch/sampler.py +++ b/src/simplefold/model/torch/sampler.py @@ -4,8 +4,8 @@ # import torch -from tqdm import tqdm from einops import repeat +from tqdm import tqdm from utils.boltz_utils import center_random_augmentation diff --git a/src/simplefold/process_mmcif.py b/src/simplefold/process_mmcif.py index 739921d..044a810 100644 --- a/src/simplefold/process_mmcif.py +++ b/src/simplefold/process_mmcif.py @@ -18,22 +18,20 @@ import numpy as np import rdkit -from redis import Redis -from tqdm import tqdm -from p_tqdm import p_umap - from boltz_data_pipeline.filter.static.filter import StaticFilter from boltz_data_pipeline.filter.static.ligand import ExcludedLigands from boltz_data_pipeline.filter.static.polymer import ( - ClashingChainsFilter, - ConsecutiveCA, - MinimumLengthFilter, - UnknownFilter + ClashingChainsFilter, + ConsecutiveCA, + MinimumLengthFilter, + UnknownFilter, ) from boltz_data_pipeline.types import ChainInfo, Record, Target +from p_tqdm import p_umap +from redis import Redis +from tqdm import tqdm from utils.mmcif_utils import parse_mmcif - """ run: wget https://boltz1.s3.us-east-2.amazonaws.com/ccd.rdb diff --git a/src/simplefold/process_structure.py b/src/simplefold/process_structure.py index 1016f22..a0aef9a 100644 --- a/src/simplefold/process_structure.py +++ b/src/simplefold/process_structure.py @@ -6,18 +6,18 @@ # Started from https://github.com/jwohlwend/boltz, # licensed under MIT License, Copyright (c) 2024 Jeremy Wohlwend, Gabriele Corso, Saro Passaro. -import os +import argparse import json +import os import pickle -import argparse -from tqdm import tqdm -import numpy as np +from dataclasses import asdict, dataclass, replace from pathlib import Path from typing import Optional -from dataclasses import asdict, replace, dataclass +import numpy as np from boltz_data_pipeline.tokenize.boltz_protein import BoltzTokenizer from boltz_data_pipeline.types import Connection, Input, Manifest, Record, Structure +from tqdm import tqdm @dataclass diff --git a/src/simplefold/processor/protein_processor.py b/src/simplefold/processor/protein_processor.py index 2de580a..31a3566 100644 --- a/src/simplefold/processor/protein_processor.py +++ b/src/simplefold/processor/protein_processor.py @@ -3,14 +3,14 @@ # Copyright (c) 2025 Apple Inc. Licensed under MIT License. # -import torch import numpy as np +import torch +from utils.boltz_utils import center_random_augmentation as torch_center_random from utils.esm_utils import ( af2_idx_to_esm_idx, - compute_language_model_representations, batch_encode_sequences, + compute_language_model_representations, ) -from utils.boltz_utils import center_random_augmentation as torch_center_random try: import mlx.core as mx diff --git a/src/simplefold/test.py b/src/simplefold/test.py index c884560..2d1dd2f 100644 --- a/src/simplefold/test.py +++ b/src/simplefold/test.py @@ -4,18 +4,14 @@ # import hydra +import lightning.pytorch as pl import torch +from lightning.pytorch import LightningDataModule, LightningModule from omegaconf import OmegaConf -import lightning.pytorch as pl -from lightning.pytorch import ( - LightningDataModule, - LightningModule -) - -from utils.utils import extras, create_folders, task_wrapper from utils.instantiators import instantiate_callbacks from utils.logging_utils import log_hyperparameters from utils.pylogger import RankedLogger +from utils.utils import create_folders, extras, task_wrapper torch.set_float32_matmul_precision("medium") log = RankedLogger(__name__, rank_zero_only=True) diff --git a/src/simplefold/train.py b/src/simplefold/train.py index 782399e..65305f6 100644 --- a/src/simplefold/train.py +++ b/src/simplefold/train.py @@ -3,17 +3,11 @@ # Copyright (c) 2025 Apple Inc. Licensed under MIT License. # -import torch +import hydra import lightning.pytorch as pl +import torch from lightning.pytorch import LightningDataModule, LightningModule -import hydra from omegaconf import OmegaConf - -from utils.utils import ( - extras, - create_folders, - task_wrapper, -) from utils.instantiators import ( instantiate_callbacks, instantiate_loggers, @@ -21,6 +15,11 @@ ) from utils.logging_utils import log_hyperparameters from utils.pylogger import RankedLogger +from utils.utils import ( + create_folders, + extras, + task_wrapper, +) log = RankedLogger(__name__, rank_zero_only=True) diff --git a/src/simplefold/train_fsdp.py b/src/simplefold/train_fsdp.py index 0b72a9b..868f881 100644 --- a/src/simplefold/train_fsdp.py +++ b/src/simplefold/train_fsdp.py @@ -3,29 +3,27 @@ # Copyright (c) 2025 Apple Inc. Licensed under MIT License. # -import os -import hydra -import torch import functools -from omegaconf import OmegaConf +import hydra import lightning.pytorch as pl +import torch from lightning.pytorch import LightningDataModule, LightningModule from lightning.pytorch.strategies import FSDPStrategy -from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy - from model.torch.blocks import DiTBlock -from utils.utils import ( - extras, - create_folders, - task_wrapper, -) +from omegaconf import OmegaConf +from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy from utils.instantiators import ( instantiate_callbacks, instantiate_loggers, ) from utils.logging_utils import log_hyperparameters from utils.pylogger import RankedLogger +from utils.utils import ( + create_folders, + extras, + task_wrapper, +) log = RankedLogger(__name__, rank_zero_only=True) torch.set_float32_matmul_precision("medium") diff --git a/src/simplefold/utils/boltz_utils.py b/src/simplefold/utils/boltz_utils.py index 53d6628..54781a4 100644 --- a/src/simplefold/utils/boltz_utils.py +++ b/src/simplefold/utils/boltz_utils.py @@ -9,22 +9,19 @@ # Started from code from https://github.com/lucidrains/alphafold3-pytorch, # licensed under MIT License, Copyright (c) 2024 Phil Wang -import numpy as np +from dataclasses import replace from functools import partial from typing import Optional -from einops import einsum -from dataclasses import replace +import numpy as np import torch -from torch.nn import Linear -import torch.nn.functional as F -from torch.types import Device - from boltz_data_pipeline import const -from boltz_data_pipeline.types import Interface, Structure, ChainInfo +from boltz_data_pipeline.types import ChainInfo, Interface, Structure from boltz_data_pipeline.write.mmcif import to_mmcif from boltz_data_pipeline.write.pdb import to_pdb - +from einops import einsum +from torch.nn import Linear +from torch.types import Device LinearNoBias = partial(Linear, bias=False) diff --git a/src/simplefold/utils/datamodule_utils.py b/src/simplefold/utils/datamodule_utils.py index ef7c3be..48a752f 100644 --- a/src/simplefold/utils/datamodule_utils.py +++ b/src/simplefold/utils/datamodule_utils.py @@ -3,22 +3,21 @@ # Copyright (c) 2025 Apple Inc. Licensed under MIT License. # -import torch -from torch import Tensor import json -import numpy as np -from typing import Optional from dataclasses import dataclass from pathlib import Path +from typing import Optional -from boltz_data_pipeline.feature.pad import pad_to_max +import numpy as np +import torch +from boltz_data_pipeline import const from boltz_data_pipeline.crop.cropper import Cropper -from boltz_data_pipeline.tokenize.boltz_protein import BoltzTokenizer from boltz_data_pipeline.feature.featurizer import BoltzFeaturizer -from processor.protein_processor import ProteinDataProcessor +from boltz_data_pipeline.feature.pad import pad_to_max +from boltz_data_pipeline.tokenize.boltz_protein import BoltzTokenizer from boltz_data_pipeline.types import Connection, Input, Manifest, Record, Structure -from boltz_data_pipeline import const - +from processor.protein_processor import ProteinDataProcessor +from torch import Tensor restype_1to3 = { "A": "ALA", diff --git a/src/simplefold/utils/esm_utils.py b/src/simplefold/utils/esm_utils.py index be0080c..a592d3f 100644 --- a/src/simplefold/utils/esm_utils.py +++ b/src/simplefold/utils/esm_utils.py @@ -6,11 +6,11 @@ # Started from https://github.com/facebookresearch/esm/tree/main, # licensed under MIT License, Copyright (c) Meta Platforms, Inc. and affiliates. -import torch import typing as T -import numpy as np from functools import partial +import numpy as np +import torch from utils import residue_constants try: diff --git a/src/simplefold/utils/fasta_utils.py b/src/simplefold/utils/fasta_utils.py index a64612a..88ade52 100644 --- a/src/simplefold/utils/fasta_utils.py +++ b/src/simplefold/utils/fasta_utils.py @@ -6,16 +6,15 @@ # Started from https://github.com/jwohlwend/boltz, # licensed under MIT License, Copyright (c) 2024 Jeremy Wohlwend, Gabriele Corso, Saro Passaro. -import click -import pickle import json +import pickle import urllib.request -from pathlib import Path -from tqdm import tqdm from dataclasses import asdict +from pathlib import Path +import click from boltz_data_pipeline.types import Manifest, Record - +from tqdm import tqdm CCD_URL = "https://huggingface.co/boltz-community/boltz-1/resolve/main/ccd.pkl" MODEL_URL = ( @@ -24,13 +23,11 @@ from collections.abc import Mapping -from pathlib import Path from Bio import SeqIO -from rdkit.Chem.rdchem import Mol - from boltz_data_pipeline.parse.yaml import parse_boltz_schema from boltz_data_pipeline.types import Target +from rdkit.Chem.rdchem import Mol def parse_fasta(path: Path, ccd: Mapping[str, Mol]) -> Target: # noqa: C901 diff --git a/src/simplefold/utils/instantiators.py b/src/simplefold/utils/instantiators.py index 271fbb5..302ca65 100644 --- a/src/simplefold/utils/instantiators.py +++ b/src/simplefold/utils/instantiators.py @@ -6,13 +6,13 @@ import os import shlex import subprocess -import torch -import hydra from typing import List + +import hydra +import torch from lightning import Callback from lightning.pytorch.loggers import Logger from omegaconf import DictConfig - from utils import pylogger log = pylogger.RankedLogger(__name__, rank_zero_only=True) diff --git a/src/simplefold/utils/logging_utils.py b/src/simplefold/utils/logging_utils.py index 0070607..992060f 100644 --- a/src/simplefold/utils/logging_utils.py +++ b/src/simplefold/utils/logging_utils.py @@ -7,7 +7,6 @@ from lightning_utilities.core.rank_zero import rank_zero_only from omegaconf import OmegaConf - from utils import pylogger log = pylogger.RankedLogger(__name__, rank_zero_only=True) diff --git a/src/simplefold/utils/lr_scheduler.py b/src/simplefold/utils/lr_scheduler.py index b99feb2..61888d6 100644 --- a/src/simplefold/utils/lr_scheduler.py +++ b/src/simplefold/utils/lr_scheduler.py @@ -4,6 +4,7 @@ # import math + import torch from torch.optim.lr_scheduler import _LRScheduler diff --git a/src/simplefold/utils/mmcif_utils.py b/src/simplefold/utils/mmcif_utils.py index 8e75714..e493ef3 100644 --- a/src/simplefold/utils/mmcif_utils.py +++ b/src/simplefold/utils/mmcif_utils.py @@ -12,11 +12,6 @@ import gemmi import numpy as np -from rdkit import rdBase -from rdkit.Chem import AllChem -from rdkit.Chem.rdchem import Conformer, Mol -from sklearn.neighbors import KDTree - from boltz_data_pipeline import const from boltz_data_pipeline.types import ( Atom, @@ -28,6 +23,10 @@ Structure, StructureInfo, ) +from rdkit import rdBase +from rdkit.Chem import AllChem +from rdkit.Chem.rdchem import Conformer, Mol +from sklearn.neighbors import KDTree #################################################################################################### # DATACLASSES diff --git a/src/simplefold/utils/residue_constants.py b/src/simplefold/utils/residue_constants.py index 42df036..95ec3df 100644 --- a/src/simplefold/utils/residue_constants.py +++ b/src/simplefold/utils/residue_constants.py @@ -22,14 +22,12 @@ import collections import functools -import os -from typing import Mapping, List, Tuple from importlib import resources +from typing import List, Mapping, Tuple import numpy as np import tree - # Distance from one CA to next CA [trans configuration: omega = 180]. ca_ca = 3.80209737096 diff --git a/src/simplefold/utils/utils.py b/src/simplefold/utils/utils.py index ee01e0b..decdb95 100644 --- a/src/simplefold/utils/utils.py +++ b/src/simplefold/utils/utils.py @@ -4,20 +4,19 @@ # import os +import warnings +from pathlib import Path +from typing import Any, Callable, Dict, Sequence, Tuple + import rich import rich.syntax import rich.tree -from rich.prompt import Prompt -import warnings -from typing import Any, Callable, Dict, Tuple -from omegaconf import DictConfig -from pathlib import Path -from typing import Sequence from hydra.core.hydra_config import HydraConfig from lightning_utilities.core.rank_zero import rank_zero_only from omegaconf import DictConfig, OmegaConf, open_dict - +from rich.prompt import Prompt from utils.pylogger import RankedLogger + log = RankedLogger(__name__, rank_zero_only=True) diff --git a/src/simplefold/wrapper.py b/src/simplefold/wrapper.py index add3734..68c51c3 100644 --- a/src/simplefold/wrapper.py +++ b/src/simplefold/wrapper.py @@ -4,29 +4,28 @@ # import os -import torch -import hydra -import omegaconf -from pathlib import Path from itertools import starmap +from pathlib import Path +import hydra +import omegaconf +import torch +from boltz_data_pipeline.feature.featurizer import BoltzFeaturizer +from boltz_data_pipeline.tokenize.boltz_protein import BoltzTokenizer from model.flow import LinearPath from model.torch.sampler import EMSampler - from processor.protein_processor import ProteinDataProcessor +from utils.boltz_utils import process_structure, save_structure from utils.datamodule_utils import process_one_inference_structure from utils.esm_utils import _af2_to_esm, esm_registry -from utils.boltz_utils import process_structure, save_structure -from utils.fasta_utils import process_fastas, download_fasta_utilities -from boltz_data_pipeline.feature.featurizer import BoltzFeaturizer -from boltz_data_pipeline.tokenize.boltz_protein import BoltzTokenizer +from utils.fasta_utils import download_fasta_utilities, process_fastas try: import mlx.core as mx - from mlx.utils import tree_unflatten, tree_flatten - from model.mlx.sampler import EMSampler as EMSamplerMLX + from mlx.utils import tree_unflatten from model.mlx.esm_network import ESM2 as ESM2MLX - from utils.mlx_utils import map_torch_to_mlx, map_plddt_torch_to_mlx + from model.mlx.sampler import EMSampler as EMSamplerMLX + from utils.mlx_utils import map_plddt_torch_to_mlx, map_torch_to_mlx MLX_AVAILABLE = True except: MLX_AVAILABLE = False