rc-foundry 0.1.1__py3-none-any.whl

foundry/hydra/resolvers.py

@@ -0,0 +1,77 @@
+"""Resolvers for Hydra configuration files.
+
+Documentation on custom resolvers:
+- https://omegaconf.readthedocs.io/en/latest/custom_resolvers.html
+"""
+
+import importlib
+
+from atomworks.enums import ChainType, ChainTypeInfo
+from beartype.typing import Any
+from omegaconf import OmegaConf
+
+from ..common import run_once
+
+
+#  (Custom resolvers)
+@run_once
+def register_resolvers():
+    resolvers = {
+        "resolve_import": resolve_import,
+        "chain_type_info_to_regex": chain_type_info_to_regex,
+    }
+
+    for name, resolver in resolvers.items():
+        OmegaConf.register_new_resolver(name, resolver)
+
+
+def resolve_import(module_path: str, attribute_path: str = None) -> Any:
+    """
+    Import a module and access a specific attribute from it.
+
+    Args:
+        module_path (str): The path to the module.
+        attribute_path (str): The path to the attribute within the module.
+
+    Returns:
+        The imported attribute.
+    """
+    module = importlib.import_module(module_path)
+    if attribute_path is not None:
+        # Split the attribute path to navigate through nested attributes
+        attributes = attribute_path.split(".")
+        attr = module
+        for attr_name in attributes:
+            attr = getattr(attr, attr_name)
+        return attr
+    else:
+        return module
+
+
+def chain_type_info_to_regex(*args) -> Any:
+    """Convert a combination of ChainType or ChainTypeInfo attributes to a regex string.
+
+    Primarily used for filtering a dataset by chain type prior to training/validation.
+
+    Example filter:
+    - "pn_unit_1_type.astype('str').str.match('${chain_type_info_to_regex:PROTEINS}')"
+
+    """
+    regex_str = ""
+
+    for arg in args:
+        if hasattr(ChainType, arg):
+            regex_str += f"{getattr(ChainType, arg).value}|"
+        elif hasattr(ChainTypeInfo, arg):
+            chain_types_list = getattr(ChainTypeInfo, arg)
+            for ct in chain_types_list:
+                regex_str += f"{ct.value}|"
+        else:
+            raise ValueError(
+                f"Attribute not found for ChainType or ChainTypeInfo: {arg}."
+            )
+
+    # Remove the trailing '|'
+    regex_str = regex_str[:-1]
+
+    return regex_str

foundry/inference_engines/base.py

@@ -0,0 +1,235 @@
+import logging
+import os
+from os import PathLike
+from pathlib import Path
+from typing import Any, Dict
+
+import hydra
+import torch
+from biotite.structure import AtomArray
+from lightning.fabric import seed_everything
+from omegaconf import OmegaConf
+
+from foundry.inference_engines.checkpoint_registry import REGISTERED_CHECKPOINTS
+from foundry.utils.ddp import RankedLogger, set_accelerator_based_on_availability
+from foundry.utils.logging import (
+    configure_minimal_inference_logging,
+    print_config_tree,
+)
+
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s %(levelname)s %(name)s: %(message)s",
+    datefmt="%H:%M:%S",
+)
+ranked_logger = RankedLogger(__name__, rank_zero_only=True)
+
+
+def merge(cfg, overrides: dict):
+    return OmegaConf.merge(cfg, OmegaConf.create(overrides))
+
+
+class BaseInferenceEngine:
+    """
+    Base inference engine.
+    Separates model setup (expensive, once) from inference (can run multiple times).
+    """
+
+    def __init__(
+        self,
+        ckpt_path: PathLike,
+        num_nodes: int = 1,
+        devices_per_node: int = 1,
+        # Config overrides
+        transform_overrides={},
+        inference_sampler_overrides={},
+        trainer_overrides={},
+        # Debug
+        verbose: bool = False,
+        seed: int | None = None,
+    ):
+        """Initialize inference engine and load model.
+
+        Model config is loaded from checkpoint and overridden with parameters provided here.
+
+        Args:
+          ckpt_path: Path to model checkpoint.
+          seed: Random seed. If None, uses external RNG state. Defaults to ``None``.
+          num_nodes: Number of nodes for distributed inference. Defaults to ``1``.
+          devices_per_node: Number of devices per node. Defaults to ``1``.
+          verbose: If True, show detailed logging and config trees. Defaults to ``False``.
+        """
+        if not verbose:
+            configure_minimal_inference_logging()
+
+        # Set attrs
+        self.initialized_ = False
+        self.trainer = None
+        self.pipeline = None
+        self.verbose = verbose
+
+        # Resolve checkpoint path
+        if '.' not in str(ckpt_path):
+            # Assume registered model
+            name = str(ckpt_path)
+            assert name in REGISTERED_CHECKPOINTS, 'Checkpoint provided not and not in registered checkpoints'
+            ckpt = REGISTERED_CHECKPOINTS[name]
+            
+            ckpt_path = ckpt.get_default_path()
+            ranked_logger.info("Using checkpoint from default installation directory, got: {}".format(str(ckpt_path)))
+            assert os.path.exists(ckpt_path), 'Invalid checkpoint: {}. And could not find checkpoint in default installation location: {}'.format(name, ckpt_path)
+        self.ckpt_path = Path(ckpt_path).resolve()
+
+        # Set random seed (only if seed is not None)
+        if seed is not None:
+            ranked_logger.info(f"Seeding everything with seed={seed}...")
+            seed_everything(seed, workers=True, verbose=True)
+        else:
+            ranked_logger.info("Seed is None - using external RNG state")
+        self.seed = seed
+
+        # Stored for later;
+        self.transform_overrides = transform_overrides
+        self.overrides: dict[str, Any] = {}
+
+        base_overrides = {
+            "trainer.seed": seed,
+            "trainer.metrics": {},
+            "trainer.loss": None,
+            "trainer.num_nodes": num_nodes,
+            "trainer.devices_per_node": devices_per_node,
+        }
+        for key, value in base_overrides.items():
+            self._assign_override(key, value)
+
+        for key, value in trainer_overrides.items():
+            self._assign_override(f"trainer.{key}", value)
+
+        for key, value in inference_sampler_overrides.items():
+            self._assign_override(f"model.net.inference_sampler.{key}", value)
+
+    ###################################################################################
+    # Required subclasss methods
+    ###################################################################################
+
+    def initialize(self):
+        if self.initialized_:
+            return getattr(self, "cfg", None)
+
+        # Load checkpoint and config
+        ranked_logger.info(
+            f"Loading checkpoint from {Path(self.ckpt_path).resolve()}..."
+        )
+        checkpoint = torch.load(self.ckpt_path, "cpu", weights_only=False)
+        cfg = self._override_checkpoint_config(checkpoint["train_cfg"])
+
+        # Load pipeline first before trainer/model
+        self._construct_pipeline(cfg)
+        self._construct_trainer(cfg, checkpoint=checkpoint)
+
+        ranked_logger.info("Model loaded and ready for inference.")
+        self.initialized_ = True
+        return cfg
+
+    def run(
+        self,
+        inputs: (
+            Dict[str, dict] | AtomArray | list[AtomArray] | PathLike | list[PathLike]
+        ),
+        *_,
+    ) -> dict[str, dict] | None:
+        self.initialize()
+        raise NotImplementedError(
+            "Subclasses must implement inference logic in `run` method."
+        )
+
+    ###################################################################################
+    # Util methods
+    ###################################################################################
+
+    def _override_checkpoint_config(self, cfg):
+        cfg = merge(cfg, self.overrides)
+        cfg = set_accelerator_based_on_availability(cfg)
+        return cfg
+
+    def _construct_trainer(self, cfg, checkpoint=None):
+        """
+        Sets attr self.trainer
+        """
+        # Instantiate trainer
+        ranked_logger.info("Instantiating trainer...")
+        if self.verbose:
+            print_config_tree(
+                cfg.trainer, resolve=True, title="INFERENCE TRAINER CONFIGURATION"
+            )
+        trainer = hydra.utils.instantiate(
+            cfg.trainer,
+            _convert_="partial",
+            _recursive_=False,
+        )
+
+        # Setup model
+        ranked_logger.info("Setting up model...")
+        trainer.fabric.launch()
+        trainer.initialize_or_update_trainer_state(
+            {"train_cfg": cfg}
+        )  # config from training stores net params
+        trainer.construct_model()
+
+        ranked_logger.info("Loading model weights from checkpoint...")
+        trainer.load_checkpoint(checkpoint=checkpoint or self.ckpt_path)
+
+        # Ensure optimizer isn't loaded
+        trainer.state["optimizer"] = None
+        trainer.state["train_cfg"].model.optimizer = None
+        trainer.setup_model_optimizers_and_schedulers()
+        trainer.state["model"].eval()
+        self.trainer = trainer
+
+    def _assign_override(self, dotted_key: str, value: Any) -> None:
+        """Assign ``value`` into ``self.overrides`` using a dotted path."""
+        target = self.overrides
+        keys = dotted_key.split(".")
+        for key in keys[:-1]:
+            if key not in target or not isinstance(target[key], dict):
+                target[key] = {}
+            target = target[key]
+        target[keys[-1]] = value
+
+    def _construct_pipeline(self, cfg):
+        """
+        Sets attr self.pipeline
+        """
+        # Construct pipeline
+        ranked_logger.info("Building Transform pipeline...")
+        first_val_dataset_key, first_val_dataset = next(iter(cfg.datasets.val.items()))
+        ranked_logger.info(
+            f"Using settings from validation dataset: {first_val_dataset_key}."
+        )
+        transform = first_val_dataset.dataset.transform
+        transform = merge(transform, self.transform_overrides)
+
+        if self.verbose:
+            print_config_tree(
+                transform,
+                resolve=True,
+                title="INFERENCE TRANSFORM PIPELINE",
+            )
+
+        self.pipeline = hydra.utils.instantiate(transform)
+
+    # aliases for run
+    def forward(self, *args, **kwargs):
+        return self.run(*args, **kwargs)
+
+    def __call__(self, *args, **kwargs):
+        return self.run(*args, **kwargs)
+
+    # for use as a context manager: e.g. `with BaseInferenceEngine(...) as engine:` to automatically cleanup
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc, tb):
+        self.trainer = None
+        self.pipeline = None
+        self.initialized_ = False

foundry/inference_engines/checkpoint_registry.py

@@ -0,0 +1,66 @@
+'''Management of checkpoints'''
+import os
+from dataclasses import dataclass
+from pathlib import Path
+
+
+def get_default_checkpoint_dir() -> Path:
+    """Get the default checkpoint directory.
+
+    Priority:
+    1. FOUNDRY_CHECKPOINTS_DIR environment variable
+    2. ~/.foundry/checkpoints
+    """
+    if "FOUNDRY_CHECKPOINTS_DIR" in os.environ and os.environ.get("FOUNDRY_CHECKPOINTS_DIR"):
+        return Path(os.environ["FOUNDRY_CHECKPOINTS_DIR"]).absolute()
+    return Path.home() / ".foundry" / "checkpoints"
+
+@dataclass
+class RegisteredCheckpoint:
+    url: str
+    filename: str
+    description: str
+    sha256: None = None  # Optional: add checksum for verification
+
+    def get_default_path(self):
+        return get_default_checkpoint_dir() / self.filename
+
+
+REGISTERED_CHECKPOINTS = {
+    "rfd3": RegisteredCheckpoint(
+        url = "https://files.ipd.uw.edu/pub/rfd3/rfd3_foundry_2025_12_01_remapped.ckpt",
+        filename = "rfd3_latest.ckpt",
+        description = "RFdiffusion3 checkpoint",
+    ),
+    "rf3":  RegisteredCheckpoint(
+        url = "https://files.ipd.uw.edu/pub/rf3/rf3_foundry_01_24_latest_remapped.ckpt",
+        filename= "rf3_foundry_01_24_latest_remapped.ckpt",
+        description= "latest RF3 checkpoint trained with data until 1/2024 (expect best performance)",
+    ),
+    "proteinmpnn":  RegisteredCheckpoint(
+        url = "https://files.ipd.uw.edu/pub/ligandmpnn/proteinmpnn_v_48_020.pt",
+        filename= "proteinmpnn_v_48_020.pt",
+        description= "ProteinMPNN checkpoint",
+    ),
+    "ligandmpnn": RegisteredCheckpoint(
+        url = "https://files.ipd.uw.edu/pub/ligandmpnn/ligandmpnn_v_32_010_25.pt",
+        filename= "ligandmpnn_v_32_010_25.pt",
+        description= "LigandMPNN checkpoint",
+    ),
+    # Other models
+    "rf3_preprint_921": RegisteredCheckpoint(
+        url = "https://files.ipd.uw.edu/pub/rf3/rf3_foundry_09_21_preprint_remapped.ckpt",
+        filename = "rf3_foundry_09_21_preprint_remapped.ckpt",
+        description = "RF3 preprint checkpoint trained with data until 9/2021",
+    ),
+    "rf3_preprint_124": RegisteredCheckpoint(
+        url = "https://files.ipd.uw.edu/pub/rf3/rf3_foundry_01_24_preprint_remapped.ckpt",
+        filename = "rf3_foundry_01_24_preprint_remapped.ckpt",
+        description= "RF3 preprint checkpoint trained with data until 1/2024",
+    ),
+    "solublempnn":  RegisteredCheckpoint(
+        url =  "https://files.ipd.uw.edu/pub/ligandmpnn/solublempnn_v_48_020.pt",
+        filename= "solublempnn_v_48_020.pt",
+        description= "SolubleMPNN checkpoint"
+    )
+}

foundry/metrics/__init__.py

@@ -0,0 +1,12 @@
+"""Metrics for model evaluation.
+
+This module provides the base metric framework.
+"""
+
+from foundry.metrics.metric import Metric, MetricInputError, MetricManager
+
+__all__ = [
+    "Metric",
+    "MetricManager",
+    "MetricInputError",
+]

foundry/metrics/losses.py

@@ -0,0 +1,30 @@
+import hydra
+import torch.nn as nn
+from omegaconf import DictConfig
+
+
+class Loss(nn.Module):
+    def __init__(self, **losses):
+        super().__init__()
+        self.to_compute = []
+        for loss_name, loss in losses.items():
+            loss_fn = hydra.utils.instantiate(loss)
+            self.to_compute.append(loss_fn)
+            assert not isinstance(
+                loss_fn, DictConfig
+            ), f"Loss {loss_name} was instantiated as a DictConfig. Is _target_ present?."
+
+    def forward(
+        self,
+        network_input,
+        network_output,
+        loss_input,
+    ):
+        loss_dict = {}
+        loss = 0
+        for loss_fn in self.to_compute:
+            loss_, loss_dict_ = loss_fn(network_input, network_output, loss_input)
+            loss += loss_
+            loss_dict.update(loss_dict_)
+        loss_dict["total_loss"] = loss.detach()
+        return loss, loss_dict