rc-foundry 0.1.1__py3-none-any.whl

rfd3/testing/testing_utils.py

@@ -0,0 +1,356 @@
+import copy
+import getpass
+import json
+import logging
+import os
+import sys
+import tempfile
+from pathlib import Path
+from unittest.mock import patch
+
+import hydra
+import numpy as np
+import rootutils
+import torch
+from hydra import compose, initialize
+from hydra.core.global_hydra import GlobalHydra
+from lightning.fabric import seed_everything
+from omegaconf import DictConfig
+
+sys.path.append(os.path.dirname(os.path.abspath(__file__)) + "/../../..")
+sys.path.append(os.path.dirname(os.path.abspath(__file__)) + "/../../../src")
+
+import atomworks
+from atomworks import parse
+from atomworks.io.parser import STANDARD_PARSER_ARGS
+from atomworks.io.utils.io_utils import to_cif_file
+from hydra.core.hydra_config import HydraConfig
+from omegaconf import OmegaConf
+from rfd3.inference.input_parsing import (
+    DesignInputSpecification,
+    create_atom_array_from_design_specification,
+)
+from rfd3.transforms.pipelines import (
+    build_atom14_base_pipeline,
+)
+
+from foundry.hydra.resolvers import register_resolvers
+
+register_resolvers()
+
+print("Atomworks version:", atomworks.__version__)
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Same as train.py
+rootutils.setup_root(__file__ + "/../..", indicator=".project-root", pythonpath=True)
+_config_path = "pkg://configs"
+
+print("_config_path:", _config_path)
+print(os.getcwd())
+
+# Set globals
+DEBUG_MODE_USERS = ["jbutch"]
+DEBUG_MODE = getpass.getuser() in DEBUG_MODE_USERS
+seed = 42
+np.random.seed(seed)
+torch.manual_seed(seed)
+
+BASE_HYDRA_OVERRIDES = ["hydra.job.num=0", "hydra.runtime.output_dir=/dev/null"]
+
+PATH_TO_REPO_ROOT = os.path.dirname(os.path.abspath(__file__)) + "/../../.."
+
+
+def load_inference_cfg():
+    GlobalHydra.instance().clear()  # Clear existing context
+    with patch.object(sys, "argv", ["compile_cfg"]):
+        initialize(config_path=_config_path, version_base="1.3", job_name="engine_cfg")
+        cfg = compose(
+            config_name="inference",
+            overrides=["ckpt_path=null", "out_dir=null", "inputs=null"]
+            + BASE_HYDRA_OVERRIDES,
+            return_hydra_config=True,
+        )
+        HydraConfig.instance().set_config(cfg)
+    return cfg
+
+
+def load_train_or_val_cfg(name="pretrain", is_val_cfg=False):
+    GlobalHydra.instance().clear()  # Clear existing context
+    with patch.object(sys, "argv", ["compile_cfg"]):
+        initialize(config_path=_config_path, version_base="1.3", job_name="trainer_cfg")
+        cfg = compose(
+            config_name="validate" if is_val_cfg else "train",
+            overrides=["experiment={}".format(name)] + BASE_HYDRA_OVERRIDES,
+            return_hydra_config=True,
+        )
+        HydraConfig.instance().set_config(cfg)
+    return cfg
+
+
+TEST_CFG_INFERENCE = load_inference_cfg()
+TEST_CFG_TRAIN = load_train_or_val_cfg()
+
+
+##########################################################################################
+# Additional utilities
+##########################################################################################
+
+DIRS = [
+    os.path.join(os.path.dirname(os.path.abspath(__file__)), "../../../tests"),
+    os.path.join(os.path.dirname(os.path.abspath(__file__))),
+    TEST_CFG_TRAIN.paths.data.design_benchmark_data_dir,
+]
+
+
+def load_test_json():
+    test_files = ["demo.json", "demo_extended.json", "tests.json"]
+    test_files += [
+        "mcsa_41.json",
+        "rfd_unindexed.json",
+        "sym_tests.json",
+        "brk_regression.json",
+    ]
+    test_json_data = {}
+    for dir in DIRS:
+        test_data_dir = Path(dir, "test_data")
+
+        for test_file in test_files:
+            test_json = test_data_dir / test_file
+            if not test_json.exists():
+                print(f"Test json file {test_json} does not exist.")
+                continue
+
+            print(f"Loading test json file {test_json}...")
+            with open(test_json, "r") as f:
+                test_json_data_ = json.load(f)
+
+            print("Found test json data:", test_json_data.keys())
+            assert test_json_data_, "Test json data is empty."
+
+            nonexistent_keys = []
+            for k, d in test_json_data_.items():
+                if "input" in d:
+                    file = test_data_dir / d["input"]
+                    test_json_data_[k]["input"] = file
+                    if not file.exists():
+                        nonexistent_keys.append(k)
+
+            if nonexistent_keys:
+                test_json_data_ = {
+                    k: v
+                    for k, v in test_json_data_.items()
+                    if k not in nonexistent_keys
+                }
+                print(f"Test input files {nonexistent_keys} do not exist for test.")
+                print("Skipping the examples.")
+
+            test_json_data = test_json_data | test_json_data_
+    return test_json_data
+
+
+TEST_JSON_DATA = load_test_json()
+assert TEST_JSON_DATA, "No test json data loaded!"
+
+
+def filter_inference_args(args):
+    return {k: v for k, v in args.items() if k in valid_keys_}
+
+
+def instantiate_example(args, is_inference=True):
+    """Saves and reloads atom array for pipeline forwarding"""
+    args = copy.deepcopy(args)
+
+    if is_inference:
+        input = DesignInputSpecification.safe_init(**args).to_pipeline_input(
+            example_id=args.get("example_id", "example")
+        )
+    else:
+        file = args.get("input")
+        if file is None:
+            # args = filter_inference_args(args)
+            atom_array, spec = create_atom_array_from_design_specification(**args)
+            outdir = tempfile.TemporaryDirectory()
+            outdir = Path(outdir.name)
+            tmpfile = outdir / "tmp.cif"
+            to_cif_file(
+                atom_array,
+                tmpfile,
+                id="example",
+                extra_fields="all",
+            )
+            file = tmpfile
+        result_dict = parse(
+            filename=file,
+            build_assembly=("1",),
+            **STANDARD_PARSER_ARGS,
+        )
+        input = {
+            "atom_array": result_dict["assemblies"]["1"][0],  # First model
+            "chain_info": result_dict["chain_info"],
+            "ligand_info": result_dict["ligand_info"],
+            "metadata": result_dict["metadata"],
+        }
+    return input
+
+
+# Pre-build pipeline | in debug mode, exclude random augmentation so that the noise scales are more clear
+
+
+def build_pipelines(
+    cfg_name="pretrain",
+    train_dataset_nested_keys=["pdb", "sub_datasets", "interface"],
+    val_dataset_nested_keys=["unconditional"],
+    composed_config: DictConfig | None = None,
+    standardize_crop_size: bool = True,
+    **transform_kwargs,
+):
+    pipes = {}
+    for is_validation in [True, False]:
+        if composed_config is None:
+            config = load_train_or_val_cfg(name=cfg_name, is_val_cfg=is_validation)
+        else:
+            config = composed_config
+        dataset_config = config["datasets"]["val" if is_validation else "train"]
+        for key in (
+            val_dataset_nested_keys if is_validation else train_dataset_nested_keys
+        ):
+            dataset_config = dataset_config[key]
+        transforms_config = dataset_config["dataset"]["transform"]
+
+        OmegaConf.set_struct(transforms_config, False)
+        transforms_config.return_atom_array = True
+        transforms_config.seed = seed
+        transforms_config.pop("_target_")
+
+        if standardize_crop_size:
+            transforms_config.crop_size = 9999
+            transforms_config.max_atoms_in_crop = 99999
+
+        for k, v in transform_kwargs.items():
+            transforms_config[k] = v
+
+        pipes |= {
+            is_validation: build_atom14_base_pipeline(
+                **(transforms_config),
+            )
+        }
+
+    if DEBUG_MODE:
+        from rfd3.testing.debug_utils import to_debug_pipe
+
+        for is_inference in [True, False]:
+            pipes[is_inference] = to_debug_pipe(pipes[is_inference])
+    return pipes
+
+
+def get_train_dataloader(cfg_name="pretrain", seed=42):
+    print("This function is deprecated!")
+    cfg = load_train_or_val_cfg(name=cfg_name, is_val_cfg=False)
+    from foundry.utils.datasets import (
+        assemble_distributed_loader,
+        recursively_instantiate_datasets_and_samplers,
+    )
+
+    seed_everything(seed, workers=True, verbose=True)
+
+    # Number of examples per epoch (accross all GPUs)
+    # (We must sample this many indices from our sampler)
+    n_examples_per_epoch = cfg.trainer.n_examples_per_epoch
+
+    # HACK: I'm getting OOMS using the full parquets, so overwrite with these
+    # NOTE: If using this test, check that your desired use-case is represented in these mini-datasets!
+    try:
+        PATH_TO_DATADIR_ROOT = "/projects/ml/datahub/dfs/af3_splits/2024_12_16/"
+        if "interface" in cfg.datasets.train.pdb.sub_datasets:
+            path = PATH_TO_DATADIR_ROOT + "/interfaces_df_train.parquet"
+            assert os.path.exists(path), f"Path {path} does not exist!"
+            cfg.datasets.train.pdb.sub_datasets.interface.dataset.dataset.data = path
+        if "pn_unit" in cfg.datasets.train.pdb.sub_datasets:
+            path = PATH_TO_DATADIR_ROOT + "/pn_units_df_train.parquet"
+            assert os.path.exists(path), f"Path {path} does not exist!"
+            cfg.datasets.train.pdb.sub_datasets.pn_unit.dataset.dataset.data = path
+    except Exception as e:
+        print("Could not overwrite dataset paths for train dataloader:", e)
+
+    # ... build the train dataset
+    # assert (
+    #     "train" in cfg.datasets and cfg.datasets.train
+    # ), "No 'train' dataloader configuration provided! If only performing validation, use `validate.py` instead."
+    dataset_and_sampler = recursively_instantiate_datasets_and_samplers(
+        cfg.datasets.train
+    )
+
+    train_dataset = dataset_and_sampler["dataset"]
+    train_sampler = dataset_and_sampler["sampler"]
+
+    train_loader = assemble_distributed_loader(
+        dataset=train_dataset,
+        sampler=train_sampler,
+        rank=0,
+        world_size=1,
+        n_examples_per_epoch=n_examples_per_epoch,
+        loader_cfg=cfg.dataloader["train"],
+    )
+
+    return train_loader
+
+
+def assert_tensor_or_array_equal(actual, expected, error_msg: str):
+    """Compare tensors/arrays with dtype-aware tolerances."""
+    if torch.is_tensor(actual):
+        if actual.dtype == torch.bool or actual.dtype in (torch.int32, torch.int64):
+            assert torch.equal(actual, expected), error_msg
+        else:
+            assert torch.allclose(
+                actual, expected, atol=1e-5, rtol=1e-5, equal_nan=True
+            ), error_msg
+    elif isinstance(actual, np.ndarray):
+        if actual.dtype.kind in {"U", "S"}:
+            assert np.array_equal(actual, expected), error_msg
+        elif actual.dtype == bool or np.issubdtype(actual.dtype, np.integer):
+            assert np.array_equal(actual, expected), error_msg
+        else:
+            assert np.allclose(
+                actual, expected, atol=1e-5, rtol=1e-5, equal_nan=True
+            ), error_msg
+    else:
+        assert actual == expected, error_msg
+
+
+def _assert_tensor_or_array_equal(actual, expected, error_msg: str):
+    """Backwards-compatible alias for legacy imports."""
+    return assert_tensor_or_array_equal(actual, expected, error_msg)
+
+
+def build_pipes_constant():
+    cfgs = {
+        True: copy.deepcopy(
+            TEST_CFG_TRAIN.datasets.val.unconditional.dataset.transform
+        ),
+        False: copy.deepcopy(
+            TEST_CFG_TRAIN.datasets.train.pdb.sub_datasets.pn_unit.dataset.transform
+        ),
+    }
+    # Standardize crop size and atom array return
+    pipes = {}
+    for is_inference, cfg in cfgs.items():
+        OmegaConf.set_struct(cfg, False)
+        cfg.seed = seed
+        cfg.return_atom_array = True
+        cfg.crop_size = 9999
+        cfg.max_atoms_in_crop = 99999
+        pipes[is_inference] = hydra.utils.instantiate(cfg)
+        # pipes[is_inference] = build_atom14_base_pipeline(
+        #     **cfg,
+        # )
+    if DEBUG_MODE:
+        from rfd3.testing.debug_utils import to_debug_pipe
+
+        for is_inference in [True, False]:
+            pipes[is_inference] = to_debug_pipe(pipes[is_inference])
+    return pipes
+
+
+PIPES = build_pipes_constant()

rfd3/train.py

@@ -0,0 +1,194 @@
+#!/usr/bin/env -S /bin/sh -c '"$(dirname "$0")/../../../../.ipd/shebang/rfd3_exec.sh" "$0" "$@"'
+
+import logging
+import os
+
+import hydra
+import rootutils
+from dotenv import load_dotenv
+from omegaconf import DictConfig
+
+from foundry.utils.logging import suppress_warnings
+from foundry.utils.weights import CheckpointConfig
+
+# Setup root dir and environment variables (more info: https://github.com/ashleve/rootutils)
+# NOTE: Sets the `PROJECT_ROOT` environment variable to the root directory of the project (where `.project-root` is located)
+rootutils.setup_root(__file__, indicator=".project-root", pythonpath=True)
+
+load_dotenv(override=True)
+
+_config_path = os.path.join(os.environ["PROJECT_ROOT"], "models/rfd3/configs")
+
+_spawning_process_logger = logging.getLogger(__name__)
+
+
+@hydra.main(config_path=_config_path, config_name="train", version_base="1.3")
+def train(cfg: DictConfig) -> None:
+    # ==============================================================================
+    # Import dependencies and resolve Hydra configuration
+    # ==============================================================================
+
+    _spawning_process_logger.info("Importing dependencies...")
+
+    # Lazy imports to make config generation fast
+    import torch
+    from lightning.fabric import seed_everything
+    from lightning.fabric.loggers import Logger
+
+    # If training on DIGS L40, set precision of matrix multiplication to balance speed and accuracy
+    # Reference: https://pytorch.org/docs/stable/generated/torch.set_float32_matmul_precision.html#torch.set_float32_matmul_precision
+    torch.set_float32_matmul_precision("medium")
+
+    from foundry.callbacks.callback import BaseCallback  # noqa
+    from foundry.utils.instantiators import instantiate_loggers, instantiate_callbacks  # noqa
+    from foundry.utils.logging import (
+        print_config_tree,
+        log_hyperparameters_with_all_loggers,
+    )  # noqa
+    from foundry.utils.ddp import RankedLogger  # noqa
+    from foundry.utils.ddp import is_rank_zero, set_accelerator_based_on_availability  # noqa
+    from foundry.utils.datasets import (
+        recursively_instantiate_datasets_and_samplers,
+        assemble_distributed_loader,
+        subset_dataset_to_example_ids,
+        assemble_val_loader_dict,
+    )  # noqa
+
+    set_accelerator_based_on_availability(cfg)
+
+    ranked_logger = RankedLogger(__name__, rank_zero_only=True)
+    _spawning_process_logger.info("Completed dependency imports ...")
+
+    # ... print the configuration tree (NOTE: Only prints for rank 0)
+    print_config_tree(cfg, resolve=True)
+
+    # ==============================================================================
+    # Logging and Callback instantiation
+    # ==============================================================================
+
+    # Reduce the logging level for all dataset and sampler loggers (unless rank 0)
+    # We will still see messages from Rank 0; they are identical, since all ranks load and sample from the same datasets
+    if not is_rank_zero():
+        dataset_logger = logging.getLogger("datasets")
+        sampler_logger = logging.getLogger("atomworks.ml.samplers")
+        dataset_logger.setLevel(logging.WARNING)
+        sampler_logger.setLevel(logging.ERROR)
+
+    # ... seed everything (NOTE: By setting `workers=True`, we ensure that the dataloaders are seeded as well)
+    # (`PL_GLOBAL_SEED` environment varaible will be passed to the spawned subprocessed; e.g., through `ddp_spawn` backend)
+    if cfg.get("seed"):
+        ranked_logger.info(f"Seeding everything with seed={cfg.seed}...")
+        seed_everything(cfg.seed, workers=True, verbose=True)
+    else:
+        ranked_logger.warning("No seed provided - Not seeding anything!")
+
+    ranked_logger.info("Instantiating loggers...")
+    loggers: list[Logger] = instantiate_loggers(cfg.get("logger"))
+
+    ranked_logger.info("Instantiating callbacks...")
+    callbacks: list[BaseCallback] = instantiate_callbacks(cfg.get("callbacks"))
+
+    # ==============================================================================
+    # Trainer and model instantiation
+    # ==============================================================================
+
+    # ... instantiate the trainer
+    ranked_logger.info("Instantiating trainer...")
+    trainer = hydra.utils.instantiate(
+        cfg.trainer,
+        loggers=loggers or None,
+        callbacks=callbacks or None,
+        _convert_="partial",
+        _recursive_=False,
+    )
+    # (Store the Hydra configuration in the trainer state)
+    trainer.initialize_or_update_trainer_state({"train_cfg": cfg})
+
+    # ... spawn processes for distributed training
+    # (We spawn here, rather than within `fit`, so we can use Fabric's `init_module` to efficiently initialize the model on the appropriate device)
+    ranked_logger.info(
+        f"Spawning {trainer.fabric.world_size} processes from {trainer.fabric.global_rank}..."
+    )
+    trainer.fabric.launch()
+
+    # ... construct the model
+    trainer.construct_model()
+
+    # ... construct the optimizer and schedule (which requires the model to be constructed)
+    trainer.construct_optimizer()
+    trainer.construct_scheduler()
+
+    # ==============================================================================
+    # Dataset instantiation
+    # ==============================================================================
+
+    # Number of examples per epoch (accross all GPUs)
+    # (We must sample this many indices from our sampler)
+    n_examples_per_epoch = cfg.trainer.n_examples_per_epoch
+
+    # ... build the train dataset
+    assert (
+        "train" in cfg.datasets and cfg.datasets.train
+    ), "No 'train' dataloader configuration provided! If only performing validation, use `validate.py` instead."
+    dataset_and_sampler = recursively_instantiate_datasets_and_samplers(
+        cfg.datasets.train
+    )
+    train_dataset, train_sampler = (
+        dataset_and_sampler["dataset"],
+        dataset_and_sampler["sampler"],
+    )
+
+    # ... compose the train loader
+    if "subset_to_example_ids" in cfg.datasets:
+        # Backdoor for debugging and overfitting: subset the dataset to a specific set of example IDs
+        train_dataset = subset_dataset_to_example_ids(
+            train_dataset, cfg.datasets.subset_to_example_ids
+        )
+        train_sampler = None  # Sampler is no longer valid, since we are using a subset of the dataset
+
+    train_loader = assemble_distributed_loader(
+        dataset=train_dataset,
+        sampler=train_sampler,
+        rank=trainer.fabric.global_rank,
+        world_size=trainer.fabric.world_size,
+        n_examples_per_epoch=n_examples_per_epoch,
+        loader_cfg=cfg.dataloader["train"],
+    )
+
+    # ... compose the validation loader(s)
+    if "val" in cfg.datasets and cfg.datasets.val:
+        val_loaders = assemble_val_loader_dict(
+            cfg=cfg.datasets.val,
+            rank=trainer.fabric.global_rank,
+            world_size=trainer.fabric.world_size,
+            loader_cfg=cfg.dataloader["val"],
+        )
+    else:
+        ranked_logger.warning("No validation datasets provided! Skipping validation...")
+        val_loaders = None
+
+    ranked_logger.info("Logging hyperparameters...")
+    log_hyperparameters_with_all_loggers(
+        trainer=trainer, cfg=cfg, model=trainer.state["model"]
+    )
+
+    # ... load the checkpoint configuration
+    ckpt_config = None
+    if "ckpt_config" in cfg and cfg.ckpt_config:
+        ckpt_config = hydra.utils.instantiate(cfg.ckpt_config)
+    elif "ckpt_path" in cfg and cfg.ckpt_path:
+        # Just a checkpoint path
+        if cfg.ckpt_path is not None:
+            ckpt_config = CheckpointConfig(path=cfg.ckpt_path)
+
+    # ... train the model
+    ranked_logger.info("Training model...")
+
+    with suppress_warnings():
+        trainer.fit(
+            train_loader=train_loader, val_loaders=val_loaders, ckpt_config=ckpt_config
+        )
+
+
+if __name__ == "__main__":
+    train()

rfd3/trainer/dump_validation_structures.py

@@ -0,0 +1,154 @@
+import functools
+from os import PathLike
+from pathlib import Path
+
+from atomworks.ml.example_id import parse_example_id
+from beartype.typing import Any, Literal
+from rfd3.utils.io import (
+    dump_metadata,
+    dump_structures,
+    dump_trajectories,
+)
+
+from foundry.callbacks.callback import BaseCallback
+from foundry.utils.ddp import RankedLogger
+
+ranked_logger = RankedLogger(__name__, rank_zero_only=True)
+
+
+class DumpValidationStructuresCallback(BaseCallback):
+    """Dump predicted structures and/or diffusion trajectories during validation"""
+
+    def __init__(
+        self,
+        save_dir: PathLike,
+        dump_predictions: bool = False,
+        dump_prediction_metadata_json: bool = True,
+        one_model_per_file: bool = False,
+        dump_trajectories: bool = False,
+        dump_denoised_trajectories_only: bool = False,
+        dump_every_n: int = 1,
+        align_trajectories: bool = True,
+        verbose: bool = True,
+        extra_fields: list[str] | Literal["all"] = [],
+    ):
+        """
+        Args:
+            dump_predictions: Whether to dump structures (CIF files) after validation batches.
+            one_model_per_file: If True, write each structure within a diffusion batch to its own CIF files. If False,
+                include each structure within a diffusion batch as a separate model within one CIF file.
+            dump_trajectories: Whether to dump denoising trajectories after validation batches.
+            dump_every_n: Dump validation structures every n epochs.
+            extra_fields: List of extra annotations to include in the saved CIF file.
+        """
+        super().__init__()
+        self.save_dir = Path(save_dir)
+        self.dump_predictions = dump_predictions
+        self.dump_prediction_metadata_json = dump_prediction_metadata_json
+        self.dump_trajectories = dump_trajectories
+        self.one_model_per_file = one_model_per_file
+        self.dump_denoised_trajectories_only = dump_denoised_trajectories_only
+        self.dump_every_n = dump_every_n
+        self.align_trajectories = align_trajectories
+        self.verbose = verbose
+        self.extra_fields = extra_fields
+
+    def _build_path_from_example_id(
+        self,
+        example_id,
+        dir: str,
+        extra: str = "",
+        epoch: str = None,
+        dataset_name: str = "",
+    ) -> Path:
+        """Helper function to build a path from a training or validation example_id."""
+        try:
+            # ... try to extract the PDB ID and assembly ID from the example ID
+            parsed_id = parse_example_id(example_id)
+            identifier = f"{parsed_id['pdb_id']}_{parsed_id['assembly_id']}"
+        except (KeyError, ValueError):
+            # ... if parsing fails, fall back to the original example ID
+            identifier = example_id
+
+        # ... parse the example_id into a dictionary of components
+        epoch_str = "epoch_{}".format(epoch) if epoch else ""
+        return (
+            self.save_dir / dir / f"{epoch_str}" / dataset_name / f"{identifier}{extra}"
+        )
+
+    def on_validation_batch_end(
+        self,
+        *,
+        trainer,
+        outputs: dict,
+        batch: Any,
+        dataset_name: str,
+        **_,
+    ):
+        if (not self.dump_predictions) and (not self.dump_trajectories):
+            return  # Nothing to do
+        current_epoch = trainer.state["current_epoch"]
+        if current_epoch % self.dump_every_n != 0:
+            ranked_logger.debug(
+                f"Skipping validation batch dump at step {current_epoch} (not every {self.dump_every_n} epochs)"
+            )
+            return
+
+        assert (
+            "network_output" in outputs
+        ), "Validation outputs must contain `network_output` to dump structures!"
+
+        network_output = outputs["network_output"]
+        atom_array_stack = outputs["predicted_atom_array_stack"]
+        prediction_metadata = outputs["prediction_metadata"]
+        example = batch[0]  # Assume batch size = 1
+
+        _build_path_from_example_id = functools.partial(
+            self._build_path_from_example_id,
+            example_id=example["example_id"],
+            epoch=current_epoch,
+            dataset_name=dataset_name,
+        )
+
+        if self.dump_predictions:
+            dump_structures(
+                atom_arrays=atom_array_stack,
+                base_path=(f := _build_path_from_example_id(dir="predictions")),
+                one_model_per_file=self.one_model_per_file,
+                extra_fields=self.extra_fields,
+            )
+            if self.verbose:
+                ranked_logger.info(f"Dumped validation predictions to {f}")
+
+        if self.dump_prediction_metadata_json and prediction_metadata:
+            dump_metadata(
+                prediction_metadata=prediction_metadata,
+                base_path=(f := _build_path_from_example_id(dir="predictions")),
+                one_model_per_file=self.one_model_per_file,
+            )
+            if self.verbose:
+                ranked_logger.info(f"Dumped validation metadata to {f}")
+
+        if self.dump_trajectories:
+            dump_trajectories(
+                trajectory_list=network_output["X_denoised_L_traj"],
+                atom_array=example["atom_array"],
+                base_path=(
+                    f := _build_path_from_example_id(
+                        dir="trajectories", extra="_denoised"
+                    )
+                ),
+                align_structures=self.align_trajectories,
+            )
+            if not self.dump_denoised_trajectories_only:
+                dump_trajectories(
+                    trajectory_list=network_output["X_noisy_L_traj"],
+                    atom_array=example["atom_array"],
+                    base_path=_build_path_from_example_id(
+                        dir="trajectories", extra="_noisy"
+                    ),
+                    align_structures=self.align_trajectories,
+                )
+
+            if self.verbose:
+                ranked_logger.info(f"Dumped validation trajectories to {f}")