Upload 36 files

src/data/components/pinder_dataset.py

@@ -0,0 +1,64 @@
+from typing import List
+
+import __main__
+import rootutils
+import torch
+from torch_geometric.data import Dataset
+
+# setup root dir and pythonpath
+rootutils.setup_root(__file__, indicator=".project-root", pythonpath=True)
+from src.data.components.prepare_data import CropPairedPDB
+
+setattr(__main__, "CropPairedPDB", CropPairedPDB)
+
+
+class PinderDataset(Dataset):
+    """Pinder dataset.
+
+    Args:
+        Dataset: PyTorch Geometric Dataset.
+    """
+
+    def __init__(self, file_paths: List[str]) -> None:
+        """Initialize the PinderDataset.
+
+        Args:
+            file_paths: List of file paths.
+        """
+        super().__init__()
+        self.file_paths = file_paths
+
+    @property
+    def processed_file_names(self) -> List[str]:
+        """Return the processed file names.
+
+        Returns:
+            List[str]: List of processed
+        """
+        return self.file_paths
+
+    def len(self) -> int:
+        """Return the length of the dataset.
+
+        Returns:
+            int: Length of the dataset
+        """
+        return len(self.processed_file_names)
+
+    def get(self, idx) -> CropPairedPDB:
+        """Get the data at the given index.
+
+        Args:
+            idx: Index of the data.
+
+        Returns:
+            CropPairedPDB: CropPairedPDB object.
+        """
+        data = torch.load(self.processed_file_names[idx], weights_only=False)
+        return data
+
+
+if __name__ == "__main__":
+    file_paths = ["./data/processed/apo/test/1a19__A1_P11540--1a19__B1_P11540.pt"]
+    dataset = PinderDataset(file_paths=file_paths)
+    print(dataset[0])

src/data/components/prepare_data.py

@@ -0,0 +1,175 @@
+import multiprocessing
+import os
+from argparse import ArgumentParser
+from pathlib import Path
+from typing import Optional
+
+import rootutils
+import torch
+from loguru import logger
+from pinder.core import PinderSystem, get_index
+from pinder.core.loader.geodata import PairedPDB, structure2tensor
+from pinder.core.loader.structure import Structure
+from tqdm.auto import tqdm
+
+# setup root dir and pythonpath
+rootutils.setup_root(__file__, indicator=".project-root", pythonpath=True)
+
+try:
+    from torch_cluster import knn_graph
+
+    torch_cluster_installed = True
+except ImportError:
+    logger.warning(
+        "torch-cluster is not installed!"
+        "Please install the appropriate library for your pytorch installation."
+        "See https://github.com/rusty1s/pytorch_cluster/issues/185 for background."
+    )
+    torch_cluster_installed = False
+
+
+def create_lr_files(system_id: str, apo_complex_path: str, save_path: str):
+    apo_r_path = os.path.join(save_path, f"apo_r_{system_id}.pdb")
+    apo_l_path = os.path.join(save_path, f"apo_l_{system_id}.pdb")
+    native_path = apo_complex_path.with_name(apo_complex_path.stem + f"{system_id}.pdb")
+    with open(native_path) as infile, open(apo_r_path, "w") as output_r, open(
+        apo_l_path, "w"
+    ) as output_l:
+
+        for line in infile:
+            # Check if the line is an ATOM or HETATM line and has a chain ID at position 21
+            if line.startswith("ATOM") or line.startswith("HETATM"):
+                chain_id = line[21]
+                if chain_id == "R":
+                    output_r.write(line)
+                elif chain_id == "L":
+                    output_l.write(line)
+            else:
+                # Write other lines (e.g., HEADER, REMARK) to both files
+                output_r.write(line)
+                output_l.write(line)
+    return apo_r_path, apo_l_path
+
+
+class CropPairedPDB(PairedPDB):
+    @classmethod
+    def from_crop_system(
+        cls,
+        system_id: str,
+        root: str = "./data/",
+        k: int = 10,
+        add_edges: bool = True,
+        predicted_structures: bool = True,
+        split: str = "train",
+    ) -> None:
+        system = PinderSystem(system_id)
+        # Create directories if they do not exist
+        for subdir in ["apo", "holo", "predicted"]:
+            os.makedirs(Path(root) / "raw" / subdir / split, exist_ok=True)
+
+        try:
+            holo_complex, apo_complex, pred_complex = system.create_masked_bound_unbound_complexes(
+                renumber_residues=True
+            )
+            for complex_type, complex_obj in zip(
+                ["apo", "holo", "predicted"], [apo_complex, holo_complex, pred_complex]
+            ):
+                complex_obj.to_pdb(
+                    Path(root) / "raw" / complex_type / split / f"{system_id}_complex.pdb"
+                )
+        except Exception as e:
+            logger.error(f"Error in writing PDB files: {e}, {system_id}")
+            return None
+
+        if predicted_structures:
+            apo_complex = pred_complex
+            save_path = os.path.join(root, "processed", "predicted", split)
+        else:
+            save_path = os.path.join(root, "processed", "apo", split)
+
+        # create the directory if it does not exist
+        os.makedirs(save_path, exist_ok=True)
+
+        graph = cls.from_structure_pair(
+            holo_complex=holo_complex,
+            apo_complex=apo_complex,
+            add_edges=add_edges,
+            k=k,
+        )
+        torch.save(graph, os.path.join(save_path, f"{system_id}.pt"))
+
+    @classmethod
+    def from_structure_pair(
+        cls,
+        holo_complex: Structure,
+        apo_complex: Structure,
+        add_edges: bool = True,
+        k: int = 10,
+    ) -> PairedPDB:
+        def get_structure_props(structure: Structure, start: int, end: Optional[int]):
+            calpha = structure.filter("atom_name", mask=["CA"])
+            return structure2tensor(
+                atom_coordinates=structure.coords[start:end],
+                atom_types=structure.atom_array.atom_name[start:end],
+                element_types=structure.atom_array.element[start:end],
+                residue_coordinates=calpha.coords[start:end],
+                residue_types=calpha.atom_array.res_name[start:end],
+                residue_ids=calpha.atom_array.res_id[start:end],
+            )
+
+        graph = cls()
+        r_h = (holo_complex.dataframe["chain_id"] == "R").sum()
+        r_a = (apo_complex.dataframe["chain_id"] == "R").sum()
+
+        holo_r_props = get_structure_props(holo_complex, 0, r_h)
+        holo_l_props = get_structure_props(holo_complex, r_h, None)
+        apo_r_props = get_structure_props(apo_complex, 0, r_a)
+        apo_l_props = get_structure_props(apo_complex, r_a, None)
+
+        graph["ligand"].x = apo_l_props["atom_types"]
+        graph["ligand"].pos = apo_l_props["atom_coordinates"]
+        graph["receptor"].x = apo_r_props["atom_types"]
+        graph["receptor"].pos = apo_r_props["atom_coordinates"]
+        graph["ligand"].y = holo_l_props["atom_coordinates"]
+        graph["receptor"].y = holo_r_props["atom_coordinates"]
+
+        if add_edges and torch_cluster_installed:
+            graph["ligand", "ligand"].edge_index = knn_graph(graph["ligand"].pos, k=k)
+            graph["receptor", "receptor"].edge_index = knn_graph(graph["receptor"].pos, k=k)
+
+        return graph
+
+
+if __name__ == "__main__":
+    parser = ArgumentParser()
+    parser.add_argument("--n_jobs", type=int, default=20)
+    parser.add_argument("--k", type=int, default=10)
+    parser.add_argument("--predicted_structures", action="store_true")
+    parser.add_argument("--split", type=str, default="train")
+    args = parser.parse_args()
+
+    predicted_structures = args.predicted_structures
+
+    # get indices for train, validation, and test splits
+    indices = get_index()
+
+    if predicted_structures:
+        query = '(split == "{split}") and ((apo_R == False and apo_L == False) and (predicted_R==True and predicted_L==True))'
+    else:
+        query = '(split == "{split}") and (apo_R == True and apo_L == True)'
+
+    system_idx = indices.query(query.format(split=args.split)).reset_index(drop=True)
+
+    system_ids = system_idx.id.tolist()
+
+    def process_system_id(system_id: str):
+        graph = CropPairedPDB.from_crop_system(
+            system_id,
+            predicted_structures=predicted_structures,
+            k=args.k,
+            split=args.split,
+        )
+        return graph
+
+    with multiprocessing.Pool(args.n_jobs) as pool:
+        results = list(tqdm(pool.imap(process_system_id, system_ids), total=len(system_ids)))

src/data/pinder_datamodule.py

@@ -0,0 +1,167 @@
+import os
+from typing import Any, Dict, Optional
+
+import pandas as pd
+import rootutils
+from lightning import LightningDataModule
+from torch_geometric.data import Dataset
+from torch_geometric.loader import DataLoader
+
+rootutils.setup_root(__file__, indicator=".project-root", pythonpath=True)
+
+from src.data.components.pinder_dataset import PinderDataset
+
+
+class PINDERDataModule(LightningDataModule):
+    """`LightningDataModule` for the PINDER dataset."""
+
+    def __init__(
+        self,
+        data_dir: str = "data/processed",
+        predicted_structures: bool = False,
+        high_quality: bool = False,
+        batch_size: int = 1,
+        num_workers: int = 0,
+        pin_memory: bool = True,
+    ) -> None:
+        """Initialize the `PINDERDataModule`.
+
+        Args:
+            data_dir: Data for pinder. Defaults to "data/processed".
+            predicted_structures: Whether to use predicted structures. Defaults to True.
+            batch_size: Batch size. Defaults to 64.
+            num_workers: Number of workers for parallel processing. Defaults to 0.
+            pin_memory: Whether to pin memory. Defaults to True.
+        """
+        super().__init__()
+
+        # this line allows to access init params with 'self.hparams' attribute
+        # also ensures init params will be stored in ckpt
+        self.save_hyperparameters(logger=False)
+
+        # get metadata
+        metadata = pd.read_csv(os.path.join(self.hparams.data_dir, "metadata.csv"))
+
+        def get_files(split: str, complex_types: list) -> list:
+            file_df = metadata[
+                (metadata["split"] == split) & (metadata["complex"].isin(complex_types))
+            ]
+            file_df["file_paths"] = file_df.apply(
+                lambda row: os.path.join(
+                    "./data/processed", row["complex"], row["split"], row["file_paths"]
+                ),
+                axis=1,
+            )
+            return file_df["file_paths"].tolist()
+
+        complex_types = ["apo", "predicted"] if self.hparams.predicted_structures else ["apo"]
+        self.train_files = get_files("train", complex_types)
+        self.val_files = get_files("val", complex_types)
+        self.test_files = get_files("test", complex_types)
+
+        self.data_train: Optional[Dataset] = None
+        self.data_val: Optional[Dataset] = None
+        self.data_test: Optional[Dataset] = None
+
+        self.batch_size_per_device = batch_size
+
+    def setup(self, stage: Optional[str] = None) -> None:
+        """Load data. Set variables: `self.data_train`, `self.data_val`, `self.data_test`.
+
+        This method is called by Lightning before `trainer.fit()`, `trainer.validate()`, `trainer.test()`, and
+        `trainer.predict()`, so be careful not to execute things like random split twice! Also, it is called after
+        `self.prepare_data()` and there is a barrier in between which ensures that all the processes proceed to
+        `self.setup()` once the data is prepared and available for use.
+
+        :param stage: The stage to setup. Either `"fit"`, `"validate"`, `"test"`, or `"predict"`. Defaults to ``None``.
+        """
+        # Divide batch size by the number of devices.
+        if self.trainer is not None:
+            if self.hparams.batch_size % self.trainer.world_size != 0:
+                raise RuntimeError(
+                    f"Batch size ({self.hparams.batch_size}) is not divisible by the number of devices ({self.trainer.world_size})."
+                )
+            self.batch_size_per_device = self.hparams.batch_size // self.trainer.world_size
+
+        # load and split datasets only if not loaded already
+        if not self.data_train and not self.data_val and not self.data_test:
+            self.data_train = PinderDataset(self.train_files)
+            self.data_val = PinderDataset(self.val_files)
+            self.data_test = PinderDataset(self.test_files)
+
+    def train_dataloader(self) -> DataLoader:
+        """Create and return the train dataloader.
+
+        :return: The train dataloader.
+        """
+        return DataLoader(
+            dataset=self.data_train,
+            batch_size=self.batch_size_per_device,
+            num_workers=self.hparams.num_workers,
+            pin_memory=self.hparams.pin_memory,
+            shuffle=True,
+            drop_last=True,
+        )
+
+    def val_dataloader(self) -> DataLoader:
+        """Create and return the validation dataloader.
+
+        :return: The validation dataloader.
+        """
+        return DataLoader(
+            dataset=self.data_val,
+            batch_size=self.batch_size_per_device,
+            num_workers=self.hparams.num_workers,
+            pin_memory=self.hparams.pin_memory,
+            shuffle=False,
+        )
+
+    def test_dataloader(self) -> DataLoader:
+        """Create and return the test dataloader.
+
+        :return: The test dataloader.
+        """
+        return DataLoader(
+            dataset=self.data_test,
+            batch_size=self.batch_size_per_device,
+            num_workers=self.hparams.num_workers,
+            pin_memory=self.hparams.pin_memory,
+            shuffle=False,
+        )
+
+    def teardown(self, stage: Optional[str] = None) -> None:
+        """Lightning hook for cleaning up after `trainer.fit()`, `trainer.validate()`,
+        `trainer.test()`, and `trainer.predict()`.
+
+        :param stage: The stage being torn down. Either `"fit"`, `"validate"`, `"test"`, or `"predict"`.
+            Defaults to ``None``.
+        """
+        pass
+
+    def state_dict(self) -> Dict[Any, Any]:
+        """Called when saving a checkpoint. Implement to generate and save the datamodule state.
+
+        :return: A dictionary containing the datamodule state that you want to save.
+        """
+        return {}
+
+    def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
+        """Called when loading a checkpoint. Implement to reload datamodule state given datamodule
+        `state_dict()`.
+
+        :param state_dict: The datamodule state returned by `self.state_dict()`.
+        """
+        pass
+
+
+if __name__ == "__main__":
+    datamodule = PINDERDataModule()
+    datamodule.setup()
+    # print(datamodule.train_files[64])
+    train_loader = datamodule.train_dataloader()
+    val_loader = datamodule.val_dataloader()
+    test_loader = datamodule.test_dataloader()
+    print(f"Number of training batches: {len(train_loader)}")
+    print(f"Number of validation batches: {len(val_loader)}")
+    print(f"Number of test batches: {len(test_loader)}")
+    print(next(iter(train_loader)))

src/eval.py

@@ -0,0 +1,99 @@
+from typing import Any, Dict, List, Tuple
+
+import hydra
+import rootutils
+from lightning import LightningDataModule, LightningModule, Trainer
+from lightning.pytorch.loggers import Logger
+from omegaconf import DictConfig
+
+rootutils.setup_root(__file__, indicator=".project-root", pythonpath=True)
+# ------------------------------------------------------------------------------------ #
+# the setup_root above is equivalent to:
+# - adding project root dir to PYTHONPATH
+#       (so you don't need to force user to install project as a package)
+#       (necessary before importing any local modules e.g. `from src import utils`)
+# - setting up PROJECT_ROOT environment variable
+#       (which is used as a base for paths in "configs/paths/default.yaml")
+#       (this way all filepaths are the same no matter where you run the code)
+# - loading environment variables from ".env" in root dir
+#
+# you can remove it if you:
+# 1. either install project as a package or move entry files to project root dir
+# 2. set `root_dir` to "." in "configs/paths/default.yaml"
+#
+# more info: https://github.com/ashleve/rootutils
+# ------------------------------------------------------------------------------------ #
+
+from src.utils import (
+    RankedLogger,
+    extras,
+    instantiate_loggers,
+    log_hyperparameters,
+    task_wrapper,
+)
+
+log = RankedLogger(__name__, rank_zero_only=True)
+
+
+@task_wrapper
+def evaluate(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]:
+    """Evaluates given checkpoint on a datamodule testset.
+
+    This method is wrapped in optional @task_wrapper decorator, that controls the behavior during
+    failure. Useful for multiruns, saving info about the crash, etc.
+
+    :param cfg: DictConfig configuration composed by Hydra.
+    :return: Tuple[dict, dict] with metrics and dict with all instantiated objects.
+    """
+    assert cfg.ckpt_path
+
+    log.info(f"Instantiating datamodule <{cfg.data._target_}>")
+    datamodule: LightningDataModule = hydra.utils.instantiate(cfg.data)
+
+    log.info(f"Instantiating model <{cfg.model._target_}>")
+    model: LightningModule = hydra.utils.instantiate(cfg.model)
+
+    log.info("Instantiating loggers...")
+    logger: List[Logger] = instantiate_loggers(cfg.get("logger"))
+
+    log.info(f"Instantiating trainer <{cfg.trainer._target_}>")
+    trainer: Trainer = hydra.utils.instantiate(cfg.trainer, logger=logger)
+
+    object_dict = {
+        "cfg": cfg,
+        "datamodule": datamodule,
+        "model": model,
+        "logger": logger,
+        "trainer": trainer,
+    }
+
+    if logger:
+        log.info("Logging hyperparameters!")
+        log_hyperparameters(object_dict)
+
+    log.info("Starting testing!")
+    trainer.test(model=model, datamodule=datamodule, ckpt_path=cfg.ckpt_path)
+
+    # for predictions use trainer.predict(...)
+    # predictions = trainer.predict(model=model, dataloaders=dataloaders, ckpt_path=cfg.ckpt_path)
+
+    metric_dict = trainer.callback_metrics
+
+    return metric_dict, object_dict
+
+
+@hydra.main(version_base="1.3", config_path="../configs", config_name="eval.yaml")
+def main(cfg: DictConfig) -> None:
+    """Main entry point for evaluation.
+
+    :param cfg: DictConfig configuration composed by Hydra.
+    """
+    # apply extra utilities
+    # (e.g. ask for tags if none are provided in cfg, print cfg tree, etc.)
+    extras(cfg)
+
+    evaluate(cfg)
+
+
+if __name__ == "__main__":
+    main()

src/models/components/equivariant_mpnn.py

@@ -0,0 +1,231 @@
+import rootutils
+import torch
+from torch import nn
+from torch.nn import BatchNorm1d, Linear, Module, ReLU, Sequential
+from torch_geometric.loader import DataLoader
+from torch_geometric.nn import MessagePassing
+from torch_scatter import scatter
+
+# setup root dir and pythonpath
+rootutils.setup_root(__file__, indicator=".project-root", pythonpath=True)
+
+from src.data.components.pinder_dataset import PinderDataset
+from src.models.components.utils import (
+    compute_euler_angles_from_rotation_matrices,
+    compute_rotation_matrix_from_ortho6d,
+)
+
+
+class EquivariantMPNNLayer(MessagePassing):
+    def __init__(self, emb_dim=64, out_dim=128, aggr="add"):
+        r"""Message Passing Neural Network Layer
+
+        This layer is equivariant to 3D rotations and translations.
+
+        Args:
+            emb_dim: (int) - hidden dimension d
+            edge_dim: (int) - edge feature dimension d_e
+            aggr: (str) - aggregation function \oplus (sum/mean/max)
+        """
+        # Set the aggregation function
+        super().__init__(aggr=aggr)
+
+        self.emb_dim = emb_dim
+
+        #
+        self.mlp_msg = Sequential(
+            Linear(2 * emb_dim + 1, emb_dim),
+            BatchNorm1d(emb_dim),
+            ReLU(),
+            Linear(emb_dim, emb_dim),
+            BatchNorm1d(emb_dim),
+            ReLU(),
+        )
+
+        self.mlp_pos = Sequential(
+            Linear(emb_dim, emb_dim), BatchNorm1d(emb_dim), ReLU(), Linear(emb_dim, 1)
+        )  # MLP \psi
+        self.mlp_upd = Sequential(
+            Linear(2 * emb_dim, emb_dim),
+            BatchNorm1d(emb_dim),
+            ReLU(),
+            Linear(emb_dim, emb_dim),
+            BatchNorm1d(emb_dim),
+            ReLU(),
+        )  # MLP \phi
+        # ===========================================
+
+        self.lin_out = Linear(emb_dim, out_dim)
+
+    def forward(self, data):
+        """
+        The forward pass updates node features h via one round of message passing.
+
+        Args:
+            h: (n, d) - initial node features
+            pos: (n, 3) - initial node coordinates
+            edge_index: (e, 2) - pairs of edges (i, j)
+            edge_attr: (e, d_e) - edge features
+
+        Returns:
+            out: [(n, d),(n,3)] - updated node features
+        """
+
+        #
+        h, pos, edge_index = data
+        h_out, pos_out = self.propagate(edge_index=edge_index, h=h, pos=pos)
+        h_out = self.lin_out(h_out)
+        return h_out, pos_out, edge_index
+        # ==========================================
+
+    #
+    def message(self, h_i, h_j, pos_i, pos_j):
+        # Compute distance between nodes i and j (Euclidean distance)
+        # distance_ij = torch.norm(pos_i - pos_j, dim=-1, keepdim=True)  # (e, 1)
+        pos_diff = pos_i - pos_j
+        dists = torch.norm(pos_diff, dim=-1).unsqueeze(1)
+
+        # Concatenate node features, edge features, and distance
+        msg = torch.cat([h_i, h_j, dists], dim=-1)
+        msg = self.mlp_msg(msg)
+        pos_diff = pos_diff * self.mlp_pos(msg)  # (e, 2d + d_e + 1)
+
+        # (e, d)
+        return msg, pos_diff
+
+    #   ...
+    #
+    def aggregate(self, inputs, index):
+        """The aggregate function aggregates the messages from neighboring nodes,
+        according to the chosen aggregation function ('sum' by default).
+
+        Args:
+            inputs: (e, d) - messages m_ij from destination to source nodes
+            index: (e, 1) - list of source nodes for each edge/message in input
+
+        Returns:
+            aggr_out: (n, d) - aggregated messages m_i
+        """
+        msgs, pos_diffs = inputs
+
+        msg_aggr = scatter(msgs, index, dim=self.node_dim, reduce=self.aggr)
+
+        pos_aggr = scatter(pos_diffs, index, dim=self.node_dim, reduce="mean")
+
+        return msg_aggr, pos_aggr
+
+    def update(self, aggr_out, h, pos):
+        msg_aggr, pos_aggr = aggr_out
+
+        upd_out = self.mlp_upd(torch.cat((h, msg_aggr), dim=-1))
+
+        upd_pos = pos + pos_aggr
+
+        return upd_out, upd_pos
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}(emb_dim={self.emb_dim}, aggr={self.aggr})"
+
+
+class PinderMPNNModel(Module):
+    def __init__(self, input_dim=1, emb_dim=64, num_heads=5):
+        """Message Passing Neural Network model for graph property prediction
+
+        This model uses both node features and coordinates as inputs, and
+        is invariant to 3D rotations and translations (the constituent MPNN layers
+        are equivariant to 3D rotations and translations).
+
+        Args:
+            emb_dim: (int) - hidden dimension d
+            input_dim: (int) - initial node feature dimension d_n
+            edge_dim: (int) - edge feature dimension d_e
+            out_dim: (int) - output dimension (fixed to 1)
+        """
+        super().__init__()
+
+        # Linear projection for initial node features
+        self.lin_in_rec = Linear(input_dim, emb_dim)
+        self.lin_in_lig = Linear(input_dim, emb_dim)
+
+        # Stack of MPNN layers
+        self.receptor_mpnn = Sequential(
+            EquivariantMPNNLayer(emb_dim, 128, aggr="mean"),
+            EquivariantMPNNLayer(128, 256, aggr="mean"),
+            # EquivariantMPNNLayer(256, 512, aggr="mean"),
+            # EquivariantMPNNLayer(512, 512, aggr="mean"),
+        )
+        self.ligand_mpnn = Sequential(
+            EquivariantMPNNLayer(64, 128, aggr="mean"),
+            EquivariantMPNNLayer(128, 256, aggr="mean"),
+            # EquivariantMPNNLayer(256, 512, aggr="mean"),
+            # EquivariantMPNNLayer(512, 512, aggr="mean"),
+        )
+
+        # Cross-attention layer
+        self.rec_cross_attention = nn.MultiheadAttention(256, num_heads, batch_first=True)
+        self.lig_cross_attention = nn.MultiheadAttention(256, num_heads, batch_first=True)
+
+        # MLPs for translation prediction
+        self.fc_translation_rec = nn.Linear(256 + 3, 3)
+        self.fc_translation_lig = nn.Linear(256 + 3, 3)
+
+    def forward(self, batch):
+        """
+        The main forward pass of the model.
+
+        Args:
+            batch: Same as in forward_rot_trans.
+
+        Returns:
+            transformed_ligands: List of tensors, each of shape (1, num_ligand_atoms, 3)
+            representing the transformed ligand coordinates after applying the predicted
+            rotation and translation.
+        """
+        h_receptor = self.lin_in_rec(batch["receptor"].x)
+        h_ligand = self.lin_in_lig(batch["ligand"].x)
+
+        pos_receptor = batch["receptor"].pos
+        pos_ligand = batch["ligand"].pos
+
+        h_receptor, pos_receptor, _ = self.receptor_mpnn(
+            (h_receptor, pos_receptor, batch["receptor", "receptor"].edge_index)
+        )
+
+        h_ligand, pos_ligand, _ = self.ligand_mpnn(
+            (h_ligand, pos_ligand, batch["ligand", "ligand"].edge_index)
+        )
+
+        attn_output_rec, _ = self.rec_cross_attention(h_receptor, h_ligand, h_ligand)
+
+        attn_output_lig, _ = self.lig_cross_attention(h_ligand, h_receptor, h_receptor)
+
+        emb_features_receptor = torch.cat((attn_output_rec, pos_receptor), dim=-1)
+        emb_features_ligand = torch.cat((attn_output_lig, pos_ligand), dim=-1)
+
+        translation_vector_r = self.fc_translation_rec(emb_features_receptor)
+        translation_vector_l = self.fc_translation_lig(emb_features_ligand)
+
+        ortho_6d_rec = compute_rotation_matrix_from_ortho6d(attn_output_rec)
+        ortho_6d_lig = compute_rotation_matrix_from_ortho6d(attn_output_lig)
+
+        receptor_coords = (
+            compute_euler_angles_from_rotation_matrices(ortho_6d_rec) * 180 / torch.pi
+        )
+        ligand_coords = compute_euler_angles_from_rotation_matrices(ortho_6d_lig) * 180 / torch.pi
+
+        receptor_coords = receptor_coords + translation_vector_r
+        ligand_coords = ligand_coords + translation_vector_l
+
+        return receptor_coords, ligand_coords
+
+
+if __name__ == "__main__":
+    file_paths = ["./data/processed/apo/test/1a19__A1_P11540--1a19__B1_P11540.pt"]
+    dataset = PinderDataset(file_paths=file_paths * 3)
+    loader = DataLoader(dataset, batch_size=3, shuffle=False)
+    batch = next(iter(loader))
+    model = PinderMPNNModel()
+    print("Number of parameters:", sum(p.numel() for p in model.parameters()))
+    receptor_coords, ligand_coords = model(batch)
+    print(receptor_coords.shape)
+    print(ligand_coords.shape)

src/models/components/utils.py

@@ -0,0 +1,100 @@
+import torch
+
+
+# batch*n
+def normalize_vector(v):
+    batch = v.shape[0]
+    v_mag = torch.sqrt(v.pow(2).sum(1))  # batch
+    eps = torch.tensor(1e-8, device=v.device)
+    v_mag = torch.max(v_mag, eps)
+    v_mag = v_mag.view(batch, 1).expand(batch, v.shape[1])
+    v = v / v_mag
+    return v
+
+
+# u, v batch*n
+def cross_product(u, v):
+    batch = u.shape[0]
+    # print (u.shape)
+    # print (v.shape)
+    i = u[:, 1] * v[:, 2] - u[:, 2] * v[:, 1]
+    j = u[:, 2] * v[:, 0] - u[:, 0] * v[:, 2]
+    k = u[:, 0] * v[:, 1] - u[:, 1] * v[:, 0]
+
+    out = torch.cat((i.view(batch, 1), j.view(batch, 1), k.view(batch, 1)), 1)  # batch*3
+
+    return out
+
+
+# poses batch*6
+# poses
+def compute_rotation_matrix_from_ortho6d(poses):
+    x_raw = poses[:, 0:3]  # batch*3
+    y_raw = poses[:, 3:6]  # batch*3
+
+    x = normalize_vector(x_raw)  # batch*3
+    z = cross_product(x, y_raw)  # batch*3
+    z = normalize_vector(z)  # batch*3
+    y = cross_product(z, x)  # batch*3
+
+    x = x.view(-1, 3, 1)
+    y = y.view(-1, 3, 1)
+    z = z.view(-1, 3, 1)
+    matrix = torch.cat((x, y, z), 2)  # batch*3*3
+    return matrix
+
+
+# input batch*4*4 or batch*3*3
+# output torch batch*3 x, y, z in radiant
+# the rotation is in the sequence of x,y,z
+def compute_euler_angles_from_rotation_matrices(rotation_matrices):
+    batch = rotation_matrices.shape[0]
+    R = rotation_matrices
+    sy = torch.sqrt(R[:, 0, 0] * R[:, 0, 0] + R[:, 1, 0] * R[:, 1, 0])
+    singular = sy < 1e-6
+    singular = singular.float()
+
+    x = torch.atan2(R[:, 2, 1], R[:, 2, 2])
+    y = torch.atan2(-R[:, 2, 0], sy)
+    z = torch.atan2(R[:, 1, 0], R[:, 0, 0])
+
+    xs = torch.atan2(-R[:, 1, 2], R[:, 1, 1])
+    ys = torch.atan2(-R[:, 2, 0], sy)
+    zs = R[:, 1, 0] * 0
+
+    out_euler = torch.zeros(batch, 3, device=rotation_matrices.device)
+
+    out_euler[:, 0] = x * (1 - singular) + xs * singular
+    out_euler[:, 1] = y * (1 - singular) + ys * singular
+    out_euler[:, 2] = z * (1 - singular) + zs * singular
+
+    return out_euler
+
+
+def get_R(x, y, z):
+    """Get rotation matrix from three rotation angles (radians). right-handed.
+    Args:
+        x: rotation angle around x-axis
+        y: rotation angle around y-axis
+        z: rotation angle around z-axis
+    Returns:
+        R: [3, 3]. rotation matrix.
+    """
+    # x
+    Rx = torch.tensor(
+        [[1, 0, 0], [0, torch.cos(x), -torch.sin(x)], [0, torch.sin(x), torch.cos(x)]],
+        device=x.device,
+    )
+    # y
+    Ry = torch.tensor(
+        [[torch.cos(y), 0, torch.sin(y)], [0, 1, 0], [-torch.sin(y), 0, torch.cos(y)]],
+        device=y.device,
+    )
+    # z
+    Rz = torch.tensor(
+        [[torch.cos(z), -torch.sin(z), 0], [torch.sin(z), torch.cos(z), 0], [0, 0, 1]],
+        device=z.device,
+    )
+
+    R = torch.mm(Rz, torch.mm(Ry, Rx))
+    return R

src/models/pinder_module.py

@@ -0,0 +1,297 @@
+from typing import Any, Dict, Tuple
+
+import torch
+from lightning import LightningModule
+from torchmetrics import MeanMetric, MinMetric
+from torchmetrics.regression import MeanAbsoluteError, MeanSquaredError
+
+
+class PinderLitModule(LightningModule):
+    """Example of a `LightningModule` for MNIST classification.
+
+    A `LightningModule` implements 8 key methods:
+
+    ```python
+    def __init__(self):
+    # Define initialization code here.
+
+    def setup(self, stage):
+    # Things to setup before each stage, 'fit', 'validate', 'test', 'predict'.
+    # This hook is called on every process when using DDP.
+
+    def training_step(self, batch, batch_idx):
+    # The complete training step.
+
+    def validation_step(self, batch, batch_idx):
+    # The complete validation step.
+
+    def test_step(self, batch, batch_idx):
+    # The complete test step.
+
+    def predict_step(self, batch, batch_idx):
+    # The complete predict step.
+
+    def configure_optimizers(self):
+    # Define and configure optimizers and LR schedulers.
+    ```
+
+    Docs:
+        https://lightning.ai/docs/pytorch/latest/common/lightning_module.html
+    """
+
+    def __init__(
+        self,
+        net: torch.nn.Module,
+        optimizer: torch.optim.Optimizer,
+        scheduler: torch.optim.lr_scheduler,
+        compile: bool,
+    ) -> None:
+        """Initialize a `MNISTLitModule`.
+
+        :param net: The model to train.
+        :param optimizer: The optimizer to use for training.
+        :param scheduler: The learning rate scheduler to use for training.
+        """
+        super().__init__()
+
+        # this line allows to access init params with 'self.hparams' attribute
+        # also ensures init params will be stored in ckpt
+        self.save_hyperparameters(logger=False)
+
+        self.net = net
+
+        # loss function
+        self.criterion = torch.nn.MSELoss()
+
+        # metric objects for calculating and averaging accuracy across batches
+        self.train_mse_ligand = MeanSquaredError()
+        self.val_mse_ligand = MeanSquaredError()
+        self.test_mse_ligand = MeanSquaredError()
+
+        self.train_mse_receptor = MeanSquaredError()
+        self.val_mse_receptor = MeanSquaredError()
+        self.test_mse_receptor = MeanSquaredError()
+
+        self.train_mae_receptor = MeanAbsoluteError()
+        self.val_mae_receptor = MeanAbsoluteError()
+        self.test_mae_receptor = MeanAbsoluteError()
+
+        self.train_mae_ligand = MeanAbsoluteError()
+        self.val_mae_ligand = MeanAbsoluteError()
+        self.test_mae_ligand = MeanAbsoluteError()
+
+        # for averaging loss across batches
+        self.train_loss = MeanMetric()
+        self.val_loss = MeanMetric()
+        self.test_loss = MeanMetric()
+
+        # for tracking best so far validation mse
+        self.val_mse_best = MinMetric()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Perform a forward pass through the model `self.net`.
+
+        :param x: A tensor of images.
+        :return: A tensor of logits.
+        """
+        return self.net(x)
+
+    def on_train_start(self) -> None:
+        """Lightning hook that is called when training begins."""
+        # by default lightning executes validation step sanity checks before training starts,
+        # so it's worth to make sure validation metrics don't store results from these checks
+        self.val_loss.reset()
+        self.val_mse_ligand.reset()
+        self.val_mse_receptor.reset()
+        self.val_mae_receptor.reset()
+        self.val_mae_ligand.reset()
+        self.val_mse_best.reset()
+
+    def model_step(
+        self, batch: Tuple[torch.Tensor, torch.Tensor]
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Perform a single model step on a batch of data.
+
+        :param batch: A batch of data (a tuple) containing the input tensor of images and target labels.
+
+        :return: A tuple containing (in order):
+            - A tensor of losses.
+            - A tensor of predictions.
+            - A tensor of target labels.
+        """
+
+        receptor_coords, ligand_coords = self.forward(batch)
+        loss_receptor = self.criterion(receptor_coords, batch["receptor"].y)
+        loss_ligand = self.criterion(ligand_coords, batch["ligand"].y)
+        loss = loss_receptor + loss_ligand
+        return loss, receptor_coords, ligand_coords, batch["receptor"].y, batch["ligand"].y
+
+    def training_step(
+        self, batch: Tuple[torch.Tensor, torch.Tensor], batch_idx: int
+    ) -> torch.Tensor:
+        """Perform a single training step on a batch of data from the training set.
+
+        :param batch: A batch of data (a tuple) containing the input tensor of images and target
+            labels.
+        :param batch_idx: The index of the current batch.
+        :return: A tensor of losses between model predictions and targets.
+        """
+        loss, receptor_coords, ligand_coords, receptor_targets, ligand_targets = self.model_step(
+            batch
+        )
+
+        # update and log metrics
+        self.train_loss(loss)
+        self.train_mse_ligand(ligand_coords, ligand_targets)
+        self.train_mse_receptor(receptor_coords, receptor_targets)
+        self.train_mae_ligand(ligand_coords, ligand_targets)
+        self.train_mae_receptor(receptor_coords, receptor_targets)
+        self.log("train/loss", self.train_loss, on_step=True, on_epoch=False, prog_bar=True)
+        self.log(
+            "train/mse_ligand", self.train_mse_ligand, on_step=True, on_epoch=False, prog_bar=True
+        )
+        self.log(
+            "train/mse_receptor",
+            self.train_mse_receptor,
+            on_step=True,
+            on_epoch=False,
+            prog_bar=True,
+        )
+        self.log(
+            "train/mae_ligand", self.train_mae_ligand, on_step=True, on_epoch=False, prog_bar=True
+        )
+        self.log(
+            "train/mae_receptor",
+            self.train_mae_receptor,
+            on_step=True,
+            on_epoch=False,
+            prog_bar=True,
+        )
+
+        # return loss or backpropagation will fail
+        return loss
+
+    def on_train_epoch_end(self) -> None:
+        "Lightning hook that is called when a training epoch ends."
+        pass
+
+    def validation_step(self, batch: Tuple[torch.Tensor, torch.Tensor], batch_idx: int) -> None:
+        """Perform a single validation step on a batch of data from the validation set.
+
+        :param batch: A batch of data (a tuple) containing the input tensor of images and target
+            labels.
+        :param batch_idx: The index of the current batch.
+        """
+        loss, receptor_coords, ligand_coords, receptor_targets, ligand_targets = self.model_step(
+            batch
+        )
+
+        # update and log metrics
+        self.val_loss(loss)
+        self.val_mse_ligand(ligand_coords, ligand_targets)
+        self.val_mse_receptor(receptor_coords, receptor_targets)
+        self.val_mae_ligand(ligand_coords, ligand_targets)
+        self.val_mae_receptor(receptor_coords, receptor_targets)
+        self.log("val/loss", self.val_loss, on_step=False, on_epoch=True, prog_bar=True)
+        self.log(
+            "val/mse_ligand", self.val_mse_ligand, on_step=False, on_epoch=True, prog_bar=True
+        )
+        self.log(
+            "val/mse_receptor", self.val_mse_receptor, on_step=False, on_epoch=True, prog_bar=True
+        )
+        self.log(
+            "val/mae_ligand", self.val_mae_ligand, on_step=False, on_epoch=True, prog_bar=True
+        )
+        self.log(
+            "val/mae_receptor", self.val_mae_receptor, on_step=False, on_epoch=True, prog_bar=True
+        )
+
+    def on_validation_epoch_end(self) -> None:
+        "Lightning hook that is called when a validation epoch ends."
+        acc = self.val_mse_ligand.compute()  # get current val acc
+        self.val_mse_best(acc)  # update best so far val acc
+        # log `val_acc_best` as a value through `.compute()` method, instead of as a metric object
+        # otherwise metric would be reset by lightning after each epoch
+        self.log("val/acc_best", self.val_mse_best.compute(), sync_dist=True, prog_bar=True)
+
+    def test_step(self, batch: Tuple[torch.Tensor, torch.Tensor], batch_idx: int) -> None:
+        """Perform a single test step on a batch of data from the test set.
+
+        :param batch: A batch of data (a tuple) containing the input tensor of images and target
+            labels.
+        :param batch_idx: The index of the current batch.
+        """
+        loss, receptor_coords, ligand_coords, receptor_targets, ligand_targets = self.model_step(
+            batch
+        )
+
+        # update and log metrics
+        self.test_loss(loss)
+        self.test_mse_ligand(ligand_coords, ligand_targets)
+        self.test_mse_receptor(receptor_coords, receptor_targets)
+        self.test_mae_ligand(ligand_coords, ligand_targets)
+        self.test_mae_receptor(receptor_coords, receptor_targets)
+        self.log("test/loss", self.test_loss, on_step=False, on_epoch=True, prog_bar=True)
+        self.log(
+            "test/mse_ligand", self.test_mse_ligand, on_step=False, on_epoch=True, prog_bar=True
+        )
+        self.log(
+            "test/mse_receptor",
+            self.test_mse_receptor,
+            on_step=False,
+            on_epoch=True,
+            prog_bar=True,
+        )
+        self.log(
+            "test/mae_ligand", self.test_mae_ligand, on_step=False, on_epoch=True, prog_bar=True
+        )
+        self.log(
+            "test/mae_receptor",
+            self.test_mae_receptor,
+            on_step=False,
+            on_epoch=True,
+            prog_bar=True,
+        )
+
+    def on_test_epoch_end(self) -> None:
+        """Lightning hook that is called when a test epoch ends."""
+        pass
+
+    def setup(self, stage: str) -> None:
+        """Lightning hook that is called at the beginning of fit (train + validate), validate,
+        test, or predict.
+
+        This is a good hook when you need to build models dynamically or adjust something about
+        them. This hook is called on every process when using DDP.
+
+        :param stage: Either `"fit"`, `"validate"`, `"test"`, or `"predict"`.
+        """
+        if self.hparams.compile and stage == "fit":
+            self.net = torch.compile(self.net)
+
+    def configure_optimizers(self) -> Dict[str, Any]:
+        """Choose what optimizers and learning-rate schedulers to use in your optimization.
+        Normally you'd need one. But in the case of GANs or similar you might have multiple.
+
+        Examples:
+            https://lightning.ai/docs/pytorch/latest/common/lightning_module.html#configure-optimizers
+
+        :return: A dict containing the configured optimizers and learning-rate schedulers to be used for training.
+        """
+        optimizer = self.hparams.optimizer(params=self.trainer.model.parameters())
+        if self.hparams.scheduler is not None:
+            scheduler = self.hparams.scheduler(optimizer=optimizer)
+            return {
+                "optimizer": optimizer,
+                "lr_scheduler": {
+                    "scheduler": scheduler,
+                    "monitor": "val/loss",
+                    "interval": "epoch",
+                    "frequency": 1,
+                },
+            }
+        return {"optimizer": optimizer}
+
+
+if __name__ == "__main__":
+    _ = PinderLitModule(None, None, None, None)

src/train.py

@@ -0,0 +1,133 @@
+from typing import Any, Dict, List, Optional, Tuple
+
+import hydra
+import lightning as L
+import rootutils
+import torch
+from lightning import Callback, LightningDataModule, LightningModule, Trainer
+from lightning.pytorch.loggers import Logger
+from omegaconf import DictConfig
+
+rootutils.setup_root(__file__, indicator=".project-root", pythonpath=True)
+# ------------------------------------------------------------------------------------ #
+# the setup_root above is equivalent to:
+# - adding project root dir to PYTHONPATH
+#       (so you don't need to force user to install project as a package)
+#       (necessary before importing any local modules e.g. `from src import utils`)
+# - setting up PROJECT_ROOT environment variable
+#       (which is used as a base for paths in "configs/paths/default.yaml")
+#       (this way all filepaths are the same no matter where you run the code)
+# - loading environment variables from ".env" in root dir
+#
+# you can remove it if you:
+# 1. either install project as a package or move entry files to project root dir
+# 2. set `root_dir` to "." in "configs/paths/default.yaml"
+#
+# more info: https://github.com/ashleve/rootutils
+# ------------------------------------------------------------------------------------ #
+
+from src.utils import (
+    RankedLogger,
+    extras,
+    get_metric_value,
+    instantiate_callbacks,
+    instantiate_loggers,
+    log_hyperparameters,
+    task_wrapper,
+)
+
+log = RankedLogger(__name__, rank_zero_only=True)
+
+
+@task_wrapper
+def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]:
+    """Trains the model. Can additionally evaluate on a testset, using best weights obtained during
+    training.
+
+    This method is wrapped in optional @task_wrapper decorator, that controls the behavior during
+    failure. Useful for multiruns, saving info about the crash, etc.
+
+    :param cfg: A DictConfig configuration composed by Hydra.
+    :return: A tuple with metrics and dict with all instantiated objects.
+    """
+    # set seed for random number generators in pytorch, numpy and python.random
+    if cfg.get("seed"):
+        L.seed_everything(cfg.seed, workers=True)
+
+    log.info(f"Instantiating datamodule <{cfg.data._target_}>")
+    datamodule: LightningDataModule = hydra.utils.instantiate(cfg.data)
+
+    log.info(f"Instantiating model <{cfg.model._target_}>")
+    model: LightningModule = hydra.utils.instantiate(cfg.model)
+
+    log.info("Instantiating callbacks...")
+    callbacks: List[Callback] = instantiate_callbacks(cfg.get("callbacks"))
+
+    log.info("Instantiating loggers...")
+    logger: List[Logger] = instantiate_loggers(cfg.get("logger"))
+
+    log.info(f"Instantiating trainer <{cfg.trainer._target_}>")
+    trainer: Trainer = hydra.utils.instantiate(cfg.trainer, callbacks=callbacks, logger=logger)
+
+    object_dict = {
+        "cfg": cfg,
+        "datamodule": datamodule,
+        "model": model,
+        "callbacks": callbacks,
+        "logger": logger,
+        "trainer": trainer,
+    }
+
+    if logger:
+        log.info("Logging hyperparameters!")
+        log_hyperparameters(object_dict)
+
+    if cfg.get("train"):
+        log.info("Starting training!")
+        trainer.fit(model=model, datamodule=datamodule, ckpt_path=cfg.get("ckpt_path"))
+
+    train_metrics = trainer.callback_metrics
+
+    if cfg.get("test"):
+        log.info("Starting testing!")
+        ckpt_path = trainer.checkpoint_callback.best_model_path
+        if ckpt_path == "":
+            log.warning("Best ckpt not found! Using current weights for testing...")
+            ckpt_path = None
+        trainer.test(model=model, datamodule=datamodule, ckpt_path=ckpt_path)
+        log.info(f"Best ckpt path: {ckpt_path}")
+
+    test_metrics = trainer.callback_metrics
+
+    # merge train and test metrics
+    metric_dict = {**train_metrics, **test_metrics}
+
+    return metric_dict, object_dict
+
+
+@hydra.main(version_base="1.3", config_path="../configs", config_name="train.yaml")
+def main(cfg: DictConfig) -> Optional[float]:
+    """Main entry point for training.
+
+    :param cfg: DictConfig configuration composed by Hydra.
+    :return: Optional[float] with optimized metric value.
+    """
+    # apply extra utilities
+    # (e.g. ask for tags if none are provided in cfg, print cfg tree, etc.)
+    extras(cfg)
+
+    # train the model
+    metric_dict, _ = train(cfg)
+
+    # safely retrieve metric value for hydra-based hyperparameter optimization
+    metric_value = get_metric_value(
+        metric_dict=metric_dict, metric_name=cfg.get("optimized_metric")
+    )
+
+    # return optimized metric
+    return metric_value
+
+
+if __name__ == "__main__":
+    torch.set_float32_matmul_precision("high")
+    main()

src/utils/__init__.py

@@ -0,0 +1,5 @@
+from src.utils.instantiators import instantiate_callbacks, instantiate_loggers
+from src.utils.logging_utils import log_hyperparameters
+from src.utils.pylogger import RankedLogger
+from src.utils.rich_utils import enforce_tags, print_config_tree
+from src.utils.utils import extras, get_metric_value, task_wrapper

src/utils/instantiators.py

@@ -0,0 +1,56 @@
+from typing import List
+
+import hydra
+from lightning import Callback
+from lightning.pytorch.loggers import Logger
+from omegaconf import DictConfig
+
+from src.utils import pylogger
+
+log = pylogger.RankedLogger(__name__, rank_zero_only=True)
+
+
+def instantiate_callbacks(callbacks_cfg: DictConfig) -> List[Callback]:
+    """Instantiates callbacks from config.
+
+    :param callbacks_cfg: A DictConfig object containing callback configurations.
+    :return: A list of instantiated callbacks.
+    """
+    callbacks: List[Callback] = []
+
+    if not callbacks_cfg:
+        log.warning("No callback configs found! Skipping..")
+        return callbacks
+
+    if not isinstance(callbacks_cfg, DictConfig):
+        raise TypeError("Callbacks config must be a DictConfig!")
+
+    for _, cb_conf in callbacks_cfg.items():
+        if isinstance(cb_conf, DictConfig) and "_target_" in cb_conf:
+            log.info(f"Instantiating callback <{cb_conf._target_}>")
+            callbacks.append(hydra.utils.instantiate(cb_conf))
+
+    return callbacks
+
+
+def instantiate_loggers(logger_cfg: DictConfig) -> List[Logger]:
+    """Instantiates loggers from config.
+
+    :param logger_cfg: A DictConfig object containing logger configurations.
+    :return: A list of instantiated loggers.
+    """
+    logger: List[Logger] = []
+
+    if not logger_cfg:
+        log.warning("No logger configs found! Skipping...")
+        return logger
+
+    if not isinstance(logger_cfg, DictConfig):
+        raise TypeError("Logger config must be a DictConfig!")
+
+    for _, lg_conf in logger_cfg.items():
+        if isinstance(lg_conf, DictConfig) and "_target_" in lg_conf:
+            log.info(f"Instantiating logger <{lg_conf._target_}>")
+            logger.append(hydra.utils.instantiate(lg_conf))
+
+    return logger

src/utils/logging_utils.py

@@ -0,0 +1,57 @@
+from typing import Any, Dict
+
+from lightning_utilities.core.rank_zero import rank_zero_only
+from omegaconf import OmegaConf
+
+from src.utils import pylogger
+
+log = pylogger.RankedLogger(__name__, rank_zero_only=True)
+
+
+@rank_zero_only
+def log_hyperparameters(object_dict: Dict[str, Any]) -> None:
+    """Controls which config parts are saved by Lightning loggers.
+
+    Additionally saves:
+        - Number of model parameters
+
+    :param object_dict: A dictionary containing the following objects:
+        - `"cfg"`: A DictConfig object containing the main config.
+        - `"model"`: The Lightning model.
+        - `"trainer"`: The Lightning trainer.
+    """
+    hparams = {}
+
+    cfg = OmegaConf.to_container(object_dict["cfg"])
+    model = object_dict["model"]
+    trainer = object_dict["trainer"]
+
+    if not trainer.logger:
+        log.warning("Logger not found! Skipping hyperparameter logging...")
+        return
+
+    hparams["model"] = cfg["model"]
+
+    # save number of model parameters
+    hparams["model/params/total"] = sum(p.numel() for p in model.parameters())
+    hparams["model/params/trainable"] = sum(
+        p.numel() for p in model.parameters() if p.requires_grad
+    )
+    hparams["model/params/non_trainable"] = sum(
+        p.numel() for p in model.parameters() if not p.requires_grad
+    )
+
+    hparams["data"] = cfg["data"]
+    hparams["trainer"] = cfg["trainer"]
+
+    hparams["callbacks"] = cfg.get("callbacks")
+    hparams["extras"] = cfg.get("extras")
+
+    hparams["task_name"] = cfg.get("task_name")
+    hparams["tags"] = cfg.get("tags")
+    hparams["ckpt_path"] = cfg.get("ckpt_path")
+    hparams["seed"] = cfg.get("seed")
+
+    # send hparams to all loggers
+    for logger in trainer.loggers:
+        logger.log_hyperparams(hparams)

src/utils/pylogger.py

@@ -0,0 +1,51 @@
+import logging
+from typing import Mapping, Optional
+
+from lightning_utilities.core.rank_zero import rank_prefixed_message, rank_zero_only
+
+
+class RankedLogger(logging.LoggerAdapter):
+    """A multi-GPU-friendly python command line logger."""
+
+    def __init__(
+        self,
+        name: str = __name__,
+        rank_zero_only: bool = False,
+        extra: Optional[Mapping[str, object]] = None,
+    ) -> None:
+        """Initializes a multi-GPU-friendly python command line logger that logs on all processes
+        with their rank prefixed in the log message.
+
+        :param name: The name of the logger. Default is ``__name__``.
+        :param rank_zero_only: Whether to force all logs to only occur on the rank zero process. Default is `False`.
+        :param extra: (Optional) A dict-like object which provides contextual information. See `logging.LoggerAdapter`.
+        """
+        logger = logging.getLogger(name)
+        super().__init__(logger=logger, extra=extra)
+        self.rank_zero_only = rank_zero_only
+
+    def log(self, level: int, msg: str, rank: Optional[int] = None, *args, **kwargs) -> None:
+        """Delegate a log call to the underlying logger, after prefixing its message with the rank
+        of the process it's being logged from. If `'rank'` is provided, then the log will only
+        occur on that rank/process.
+
+        :param level: The level to log at. Look at `logging.__init__.py` for more information.
+        :param msg: The message to log.
+        :param rank: The rank to log at.
+        :param args: Additional args to pass to the underlying logging function.
+        :param kwargs: Any additional keyword args to pass to the underlying logging function.
+        """
+        if self.isEnabledFor(level):
+            msg, kwargs = self.process(msg, kwargs)
+            current_rank = getattr(rank_zero_only, "rank", None)
+            if current_rank is None:
+                raise RuntimeError("The `rank_zero_only.rank` needs to be set before use")
+            msg = rank_prefixed_message(msg, current_rank)
+            if self.rank_zero_only:
+                if current_rank == 0:
+                    self.logger.log(level, msg, *args, **kwargs)
+            else:
+                if rank is None:
+                    self.logger.log(level, msg, *args, **kwargs)
+                elif current_rank == rank:
+                    self.logger.log(level, msg, *args, **kwargs)

src/utils/rich_utils.py

@@ -0,0 +1,103 @@
+from pathlib import Path
+from typing import Sequence
+
+import rich
+import rich.syntax
+import rich.tree
+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 src.utils import pylogger
+
+log = pylogger.RankedLogger(__name__, rank_zero_only=True)
+
+
+@rank_zero_only
+def print_config_tree(
+    cfg: DictConfig,
+    print_order: Sequence[str] = (
+        "data",
+        "model",
+        "callbacks",
+        "logger",
+        "trainer",
+        "paths",
+        "extras",
+    ),
+    resolve: bool = False,
+    save_to_file: bool = False,
+) -> None:
+    """Prints the contents of a DictConfig as a tree structure using the Rich library.
+
+    :param cfg: A DictConfig composed by Hydra.
+    :param print_order: Determines in what order config components are printed. Default is ``("data", "model",
+    "callbacks", "logger", "trainer", "paths", "extras")``.
+    :param resolve: Whether to resolve reference fields of DictConfig. Default is ``False``.
+    :param save_to_file: Whether to export config to the hydra output folder. Default is ``False``.
+    """
+    style = "dim"
+    tree = rich.tree.Tree("CONFIG", style=style, guide_style=style)
+
+    queue = []
+
+    # add fields from `print_order` to queue
+    for field in print_order:
+        (
+            queue.append(field)
+            if field in cfg
+            else log.warning(
+                f"Field '{field}' not found in config. Skipping '{field}' config printing..."
+            )
+        )
+
+    # add all the other fields to queue (not specified in `print_order`)
+    for field in cfg:
+        if field not in queue:
+            queue.append(field)
+
+    # generate config tree from queue
+    for field in queue:
+        branch = tree.add(field, style=style, guide_style=style)
+
+        config_group = cfg[field]
+        if isinstance(config_group, DictConfig):
+            branch_content = OmegaConf.to_yaml(config_group, resolve=resolve)
+        else:
+            branch_content = str(config_group)
+
+        branch.add(rich.syntax.Syntax(branch_content, "yaml"))
+
+    # print config tree
+    rich.print(tree)
+
+    # save config tree to file
+    if save_to_file:
+        with open(Path(cfg.paths.output_dir, "config_tree.log"), "w") as file:
+            rich.print(tree, file=file)
+
+
+@rank_zero_only
+def enforce_tags(cfg: DictConfig, save_to_file: bool = False) -> None:
+    """Prompts user to input tags from command line if no tags are provided in config.
+
+    :param cfg: A DictConfig composed by Hydra.
+    :param save_to_file: Whether to export tags to the hydra output folder. Default is ``False``.
+    """
+    if not cfg.get("tags"):
+        if "id" in HydraConfig().cfg.hydra.job:
+            raise ValueError("Specify tags before launching a multirun!")
+
+        log.warning("No tags provided in config. Prompting user to input tags...")
+        tags = Prompt.ask("Enter a list of comma separated tags", default="dev")
+        tags = [t.strip() for t in tags.split(",") if t != ""]
+
+        with open_dict(cfg):
+            cfg.tags = tags
+
+        log.info(f"Tags: {cfg.tags}")
+
+    if save_to_file:
+        with open(Path(cfg.paths.output_dir, "tags.log"), "w") as file:
+            rich.print(cfg.tags, file=file)

src/utils/utils.py

@@ -0,0 +1,119 @@
+import warnings
+from importlib.util import find_spec
+from typing import Any, Callable, Dict, Optional, Tuple
+
+from omegaconf import DictConfig
+
+from src.utils import pylogger, rich_utils
+
+log = pylogger.RankedLogger(__name__, rank_zero_only=True)
+
+
+def extras(cfg: DictConfig) -> None:
+    """Applies optional utilities before the task is started.
+
+    Utilities:
+        - Ignoring python warnings
+        - Setting tags from command line
+        - Rich config printing
+
+    :param cfg: A DictConfig object containing the config tree.
+    """
+    # return if no `extras` config
+    if not cfg.get("extras"):
+        log.warning("Extras config not found! <cfg.extras=null>")
+        return
+
+    # disable python warnings
+    if cfg.extras.get("ignore_warnings"):
+        log.info("Disabling python warnings! <cfg.extras.ignore_warnings=True>")
+        warnings.filterwarnings("ignore")
+
+    # prompt user to input tags from command line if none are provided in the config
+    if cfg.extras.get("enforce_tags"):
+        log.info("Enforcing tags! <cfg.extras.enforce_tags=True>")
+        rich_utils.enforce_tags(cfg, save_to_file=True)
+
+    # pretty print config tree using Rich library
+    if cfg.extras.get("print_config"):
+        log.info("Printing config tree with Rich! <cfg.extras.print_config=True>")
+        rich_utils.print_config_tree(cfg, resolve=True, save_to_file=True)
+
+
+def task_wrapper(task_func: Callable) -> Callable:
+    """Optional decorator that controls the failure behavior when executing the task function.
+
+    This wrapper can be used to:
+        - make sure loggers are closed even if the task function raises an exception (prevents multirun failure)
+        - save the exception to a `.log` file
+        - mark the run as failed with a dedicated file in the `logs/` folder (so we can find and rerun it later)
+        - etc. (adjust depending on your needs)
+
+    Example:
+    ```
+    @utils.task_wrapper
+    def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]:
+        ...
+        return metric_dict, object_dict
+    ```
+
+    :param task_func: The task function to be wrapped.
+
+    :return: The wrapped task function.
+    """
+
+    def wrap(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]:
+        # execute the task
+        try:
+            metric_dict, object_dict = task_func(cfg=cfg)
+
+        # things to do if exception occurs
+        except Exception as ex:
+            # save exception to `.log` file
+            log.exception("")
+
+            # some hyperparameter combinations might be invalid or cause out-of-memory errors
+            # so when using hparam search plugins like Optuna, you might want to disable
+            # raising the below exception to avoid multirun failure
+            raise ex
+
+        # things to always do after either success or exception
+        finally:
+            # display output dir path in terminal
+            log.info(f"Output dir: {cfg.paths.output_dir}")
+
+            # always close wandb run (even if exception occurs so multirun won't fail)
+            if find_spec("wandb"):  # check if wandb is installed
+                import wandb
+
+                if wandb.run:
+                    log.info("Closing wandb!")
+                    wandb.finish()
+
+        return metric_dict, object_dict
+
+    return wrap
+
+
+def get_metric_value(metric_dict: Dict[str, Any], metric_name: Optional[str]) -> Optional[float]:
+    """Safely retrieves value of the metric logged in LightningModule.
+
+    :param metric_dict: A dict containing metric values.
+    :param metric_name: If provided, the name of the metric to retrieve.
+    :return: If a metric name was provided, the value of the metric.
+    """
+    if not metric_name:
+        log.info("Metric name is None! Skipping metric value retrieval...")
+        return None
+
+    if metric_name not in metric_dict:
+        raise Exception(
+            f"Metric value not found! <metric_name={metric_name}>\n"
+            "Make sure metric name logged in LightningModule is correct!\n"
+            "Make sure `optimized_metric` name in `hparams_search` config is correct!"
+        )
+
+    metric_value = metric_dict[metric_name].item()
+    log.info(f"Retrieved metric value! <{metric_name}={metric_value}>")
+
+    return metric_value