rc-foundry 0.1.1__py3-none-any.whl

foundry/callbacks/metrics_logging.py

@@ -0,0 +1,211 @@
+import os
+from copy import deepcopy
+from pathlib import Path
+
+import pandas as pd
+from atomworks.ml.utils import nested_dict
+from beartype.typing import Any, Literal
+from omegaconf import ListConfig
+
+from foundry.callbacks.callback import BaseCallback
+from foundry.utils.ddp import RankedLogger
+
+ranked_logger = RankedLogger(__name__, rank_zero_only=True)
+
+
+class StoreValidationMetricsInDFCallback(BaseCallback):
+    """Saves the validation outputs in a DataFrame for each rank and concatenates them at the end of the validation epoch."""
+
+    def __init__(
+        self,
+        save_dir: os.PathLike,
+        metrics_to_save: list[str] | Literal["all"] = "all",
+    ):
+        self.save_dir = Path(save_dir)
+        self.metrics_to_save = metrics_to_save
+
+    def _save_dataframe_for_rank(self, rank: int, epoch: int):
+        """Saves per-GPU output dataframe of metrics to a rank-specific CSV."""
+        self.save_dir.mkdir(parents=True, exist_ok=True)
+        file_path = self.save_dir / f"validation_output_rank_{rank}_epoch_{epoch}.csv"
+
+        # Flush explicitly to ensure the file is written to disk
+        with open(file_path, "w") as f:
+            self.per_gpu_outputs_df.to_csv(f, index=False)
+            f.flush()
+            os.fsync(f.fileno())
+
+        ranked_logger.info(
+            f"Saved validation outputs to {file_path} for rank {rank}, epoch {epoch}"
+        )
+
+    def on_validation_epoch_start(self, trainer):
+        self.per_gpu_outputs_df = pd.DataFrame()
+
+    def on_validation_batch_end(
+        self,
+        trainer,
+        outputs: dict,
+        batch: Any,
+        batch_idx: int,
+        num_batches: int,
+        dataset_name: str | None = None,
+    ):
+        """Build a flattened DataFrame from the metrics output and accumulate with the prior batches"""
+        assert "metrics_output" in outputs, "Validation outputs must contain metrics."
+        metrics_output = deepcopy(outputs["metrics_output"])
+
+        # ... assemble a flat DataFrame from the metrics output
+        example_id = metrics_output.pop("example_id")
+        metrics_as_list_of_dicts = []
+
+        # ... remove metrics that are not in the save list
+        if self.metrics_to_save != "all" and isinstance(
+            self.metrics_to_save, list | ListConfig
+        ):
+            metrics_output = {
+                k: v
+                for k, v in metrics_output.items()
+                if any(k.startswith(prefix) for prefix in self.metrics_to_save)
+            }
+
+        def _build_row_from_flattened_dict(
+            dict_to_flatten: dict, prefix: str, example_id: str
+        ):
+            """Helper function to build a DataFrame row"""
+            flattened_dict = nested_dict.flatten(dict_to_flatten, fuse_keys=".")
+            row_data = {"example_id": example_id}
+            for sub_k, sub_v in flattened_dict.items():
+                # Convert lists to tuples so that they are hashable
+                if isinstance(sub_v, list):
+                    sub_v = tuple(sub_v)
+                row_data[f"{prefix}.{sub_k}"] = sub_v
+            return row_data
+
+        scalar_metrics = {"example_id": example_id}
+        for key, value in metrics_output.items():
+            if isinstance(value, dict):
+                # Flatten once for this dict => 1 row.
+                metrics_as_list_of_dicts.append(
+                    _build_row_from_flattened_dict(value, key, example_id)
+                )
+            elif isinstance(value, list) and all(isinstance(x, dict) for x in value):
+                # Flatten each dict in the list => multiple rows.
+                for subdict in value:
+                    metrics_as_list_of_dicts.append(
+                        _build_row_from_flattened_dict(subdict, key, example_id)
+                    )
+            else:
+                # Scalar (string, float, int, or list that isn't list-of-dicts)
+                assert key not in scalar_metrics, f"Duplicate key: {key}"
+                scalar_metrics[key] = value
+
+        metrics_as_list_of_dicts.append(scalar_metrics)
+
+        # ... convert the list of dicts to a DataFrame and add epoch and dataset columns
+        batch_df = pd.DataFrame(metrics_as_list_of_dicts)
+        batch_df["epoch"] = trainer.state["current_epoch"]
+        batch_df["dataset"] = dataset_name
+
+        # Assert no duplicate rows
+        assert (
+            batch_df.duplicated().sum() == 0
+        ), "Duplicate rows found in the metrics DataFrame!"
+
+        # Accumulate into the per-rank DataFrame
+        self.per_gpu_outputs_df = pd.concat(
+            [self.per_gpu_outputs_df, batch_df], ignore_index=True
+        )
+
+        ranked_logger.info(
+            f"Validation Progress: {100 * (batch_idx + 1) / num_batches:.0f}% for {dataset_name}"
+        )
+
+    def on_validation_epoch_end(self, trainer):
+        """Aggregate and log the validation metrics at the end of the epoch.
+
+        Each rank writes out its partial CSV. Then rank 0 aggregates them, logs grouped metrics by dataset,
+        and appends them to a master file containing data from all epochs.
+        """
+
+        #  ... write out partial CSV for this rank
+        rank = trainer.fabric.global_rank
+        epoch = trainer.state["current_epoch"]
+        self._save_dataframe_for_rank(rank, epoch)
+
+        # Synchronize all processes
+        ranked_logger.info(
+            "Synchronizing all processes before concatenating DataFrames..."
+        )
+        trainer.fabric.barrier()
+
+        # Only rank 0 loads and concatenates the DataFrames
+        ranked_logger.info("Loading and concatenating DataFrames...")
+        if trainer.fabric.is_global_zero:
+            # ... load all partial CSVs
+            merged_df = self._load_and_concatenate_csvs(epoch)
+
+            # ... append to master CSV for all epochs
+            master_path = self.save_dir / "validation_output_all_epochs.csv"
+            if master_path.exists():
+                old_df = pd.read_csv(master_path)
+                merged_df = pd.concat(
+                    [old_df, merged_df], ignore_index=True, sort=False
+                )
+            merged_df.to_csv(master_path, index=False)
+            ranked_logger.info(f"Appended epoch={epoch} results to {master_path}")
+
+            # Store the path to the master CSV in the Trainer
+            trainer.validation_results_path = master_path
+
+            # Cleanup
+            self._cleanup_temp_files()
+
+    def _load_and_concatenate_csvs(self, epoch: int) -> pd.DataFrame:
+        """Load rank-specific CSVs for the given epoch and concatenate them without duplicating examples."""
+        pattern = f"validation_output_rank_*_epoch_{epoch}.csv"
+        files = list(self.save_dir.glob(pattern))
+
+        # Track which example_id + dataset combinations we've already seen
+        seen_examples = set()
+        final_dataframes = []
+
+        for f in files:
+            try:
+                df = pd.read_csv(f)
+
+                # Create a filter for rows with new example_id + dataset combinations
+                if not df.empty:
+                    # Create a unique identifier for each example_id + dataset combination
+                    df["_example_key"] = (
+                        df["example_id"].astype(str) + "|" + df["dataset"].astype(str)
+                    )
+
+                    # Filter out rows with example_id + dataset combinations we've already seen
+                    new_examples_mask = ~df["_example_key"].isin(seen_examples)
+
+                    # If there are any new examples, add them to our final list
+                    if new_examples_mask.any():
+                        new_examples_df = df[new_examples_mask].copy()
+
+                        # Update our set of seen examples
+                        seen_examples.update(new_examples_df["_example_key"].tolist())
+
+                        # Remove the temporary column before adding to final list
+                        new_examples_df.drop("_example_key", axis=1, inplace=True)
+                        final_dataframes.append(new_examples_df)
+
+            except pd.errors.EmptyDataError:
+                ranked_logger.warning(f"Skipping empty CSV: {f}")
+
+        # Concatenate dataframes, filling missing columns with NaN
+        return pd.concat(final_dataframes, axis=0, ignore_index=True, sort=False)
+
+    def _cleanup_temp_files(self):
+        """Remove temporary files used to store individual rank outputs."""
+        all_files = list(self.save_dir.rglob("validation_output_rank_*_epoch_*.csv"))
+        for file in all_files:
+            try:
+                file.unlink()  # Remove the file
+            except Exception as e:
+                ranked_logger.warning(f"Failed to delete file {file}: {e}")

foundry/callbacks/timing_logging.py

@@ -0,0 +1,67 @@
+import pandas as pd
+from lightning.fabric.utilities.rank_zero import rank_zero_only
+
+from foundry.callbacks.callback import BaseCallback
+from foundry.utils.logging import print_df_as_table
+from foundry.utils.torch import Timers
+
+
+class TimingCallback(BaseCallback):
+    """Fabric callback to print timing metrics."""
+
+    def __init__(self, log_every_n: int = 100):
+        super().__init__()
+        self.log_every_n = log_every_n
+        self.timers = Timers()
+        self.n_steps_since_last_log = 0
+
+    @rank_zero_only
+    def on_train_epoch_start(self, trainer, **kwargs):
+        self.timers.start("train_loader_iter")
+
+    @rank_zero_only
+    def on_after_train_loader_iter(self, trainer, **kwargs):
+        self.timers.stop("train_loader_iter")
+
+    @rank_zero_only
+    def on_before_train_loader_next(self, trainer, **kwargs):
+        self.timers.start("train_step", "train_loader_next")
+
+    @rank_zero_only
+    def on_train_batch_start(self, trainer, **kwargs):
+        self.timers.start("forward_loss_backward")
+        self.timers.stop("train_loader_next")
+
+    @rank_zero_only
+    def on_train_batch_end(self, trainer, **kwargs):
+        self.timers.stop("forward_loss_backward")
+        self.timers.stop("train_step")
+
+    @rank_zero_only
+    def on_before_optimizer_step(self, trainer, **kwargs):
+        self.timers.start("optimizer_step")
+
+    @rank_zero_only
+    def on_after_optimizer_step(self, optimizer, **kwargs):
+        self.timers.stop("optimizer_step")
+
+    @rank_zero_only
+    def optimizer_step(self, trainer, optimizer):
+        step = trainer.state["global_step"]
+        self.n_steps_since_last_log += 1
+        if step % self.log_every_n == 0:
+            timings = self.timers.elapsed(*self.timers.timers.keys(), reset=True)
+            timings = {
+                f"timings/{k}": v / self.n_steps_since_last_log
+                for k, v in timings.items()
+            }
+            trainer.fabric.log_dict(timings, step=step)
+            if trainer.fabric.is_global_zero:
+                self._print_timings(timings)
+
+    def _print_timings(self, timings: dict[str, float]):
+        df = pd.DataFrame(timings.items(), columns=["Step", "Time (s)"])
+        print_df_as_table(
+            df, title=f"Timing stats (over {self.n_steps_since_last_log} steps)"
+        )
+        self.n_steps_since_last_log = 0

foundry/callbacks/train_logging.py

@@ -0,0 +1,278 @@
+import time
+from collections import defaultdict
+
+import pandas as pd
+from atomworks.ml.example_id import parse_example_id
+from beartype.typing import Any
+from lightning.fabric.wrappers import (
+    _FabricOptimizer,
+)
+from rf3.utils.loss import convert_batched_losses_to_list_of_dicts, mean_losses
+from rich.console import Group
+from rich.panel import Panel
+from rich.table import Table
+from torchmetrics.aggregation import MeanMetric
+
+from foundry.callbacks.callback import BaseCallback
+from foundry.utils.ddp import RankedLogger
+from foundry.utils.logging import (
+    print_df_as_table,
+    print_model_parameters,
+    safe_print,
+    table_from_df,
+)
+
+
+class LogModelParametersCallback(BaseCallback):
+    """Print a table of the total and trainable parameters of the model at the start of training."""
+
+    def on_fit_start(self, trainer):
+        print_model_parameters(trainer.state["model"])
+
+
+class PrintExampleIDBeforeForwardPassCallback(BaseCallback):
+    """Print the example ID for each rank at the start of the forward pass for each batch.
+
+    WARNING: Spams the console. Use only for debugging purposes.
+    """
+
+    def __init__(self, rank_zero_only: bool = True):
+        self.logger = RankedLogger(__name__, rank_zero_only=rank_zero_only)
+
+    def on_train_batch_start(self, trainer, batch: Any, batch_idx: int):
+        example_id = batch[0]["example_id"]
+
+        # Prepare the formatted strings with colors
+        rank_info = f"[grey]<Rank {trainer.fabric.global_rank}>[/grey]"
+        epoch_batch_info = (
+            f"[blue]Epoch {trainer.state['current_epoch']} Batch {batch_idx}[/blue]"
+        )
+        example_id_info = f"[bold yellow]Example ID: {example_id}[/bold yellow]"
+
+        safe_print(
+            f"{rank_info} {epoch_batch_info} - {example_id_info}",
+            logger=self.logger,
+        )
+
+
+class LogDatasetSamplingRatiosCallback(BaseCallback):
+    """Monitor the sampling ratios of the datasets and log after each epoch."""
+
+    def on_fit_start(self, trainer):
+        self.dataset_sampling_counts = defaultdict(int)
+
+    def on_train_batch_start(self, trainer, batch, batch_idx):
+        example_id = batch[0]["example_id"]
+
+        if trainer.fabric.is_global_zero:
+            dataset_string = "/".join(parse_example_id(example_id)["datasets"])
+            self.dataset_sampling_counts[dataset_string] += 1
+
+    def on_train_epoch_end(self, trainer):
+        if trainer.fabric.is_global_zero:
+            total_samples = sum(self.dataset_sampling_counts.values())
+
+            data = {
+                "Dataset": list(self.dataset_sampling_counts.keys()),
+                "Count": list(self.dataset_sampling_counts.values()),
+                "Percentage": [
+                    f"{(count / total_samples) * 100:.2f}%"
+                    for count in self.dataset_sampling_counts.values()
+                ],
+            }
+
+            print_df_as_table(
+                df=pd.DataFrame(data),
+                title=f"Epoch {trainer.state['current_epoch']}: Dataset Sampling Ratios",
+            )
+
+            # Reset the counts for the next epoch
+            self.dataset_sampling_counts.clear()
+
+
+class LogLearningRateCallback(BaseCallback):
+    """Monitor the learning rate of the optimizer
+
+    Args:
+        log_every_n: Log the learning rate every n optimizer steps.
+    """
+
+    def __init__(self, log_every_n: int):
+        self.log_every_n = log_every_n
+
+    def optimizer_step(self, trainer, optimizer: _FabricOptimizer):
+        # Get the current global step
+        current_step = trainer.state["global_step"]
+
+        # Log the learning rate only every `log_every_n` steps
+        if current_step % self.log_every_n == 0:
+            trainer.fabric.log(
+                "train/learning_rate",
+                optimizer.param_groups[0]["lr"],
+                step=current_step,
+            )
+
+
+class LogAF3TrainingLossesCallback(BaseCallback):
+    """Log the primary model losses for AF3.
+
+    Includes:
+        - The mean training losses every `log_every_n` batches
+        - The mean training losses at the end of each epoch
+        - The time taken to complete each epoch
+        - (Optionally) The full batch losses for each structure in the diffusion batch
+
+    Args:
+        log_every_n (int): Print the training loss after every n batches.
+    """
+
+    def __init__(
+        self,
+        log_full_batch_losses: bool = False,
+        log_every_n: int = 10,
+    ):
+        """
+        Args:
+            log_full_batch_losses(bool): Log losses for every structure within the diffusion batch.
+            log_every_n (int): Print the training loss after every n batches.
+            console_width (int): Width of the console for printing.
+        """
+        self.log_every_n = log_every_n
+        self.log_full_batch_losses = log_full_batch_losses
+
+        self.start_time = None
+        self.logger = RankedLogger(__name__, rank_zero_only=True)
+
+        # This dict will store key -> MeanMetric() for each loss
+        self.loss_trackers = {}
+
+    def on_train_epoch_start(self, trainer):
+        # Record the start time of the epoch
+        self.start_time = time.time()
+
+    def on_train_batch_end(self, trainer, outputs: Any, batch: Any, batch_idx: int):
+        mean_loss_dict = {}
+        if "loss_dict" in outputs:
+            mean_loss_dict.update(mean_losses(outputs["loss_dict"]))
+
+        for key, val in mean_loss_dict.items():
+            if key not in self.loss_trackers:
+                self.loss_trackers[key] = trainer.fabric.to_device(MeanMetric())
+            self.loss_trackers[key].update(val)
+
+        if trainer.fabric.is_global_zero and batch_idx % self.log_every_n == 0:
+            # ... log losses for each structure in the batch
+            if self.log_full_batch_losses:
+                full_batch_loss_dicts = convert_batched_losses_to_list_of_dicts(
+                    outputs["loss_dict"]
+                )
+                for loss_dict in full_batch_loss_dicts:
+                    loss_dict = {
+                        f"train/per_structure/{k}": v for k, v in loss_dict.items()
+                    }
+                    trainer.fabric.log_dict(
+                        loss_dict, step=trainer.state["global_step"]
+                    )
+
+            # ... log losses meaned across the batch
+            # (Prepend "train/batch_mean" to the keys in the loss dictionary)
+            mean_loss_dict_for_logging = {
+                f"train/batch_mean/{k}": v for k, v in mean_loss_dict.items()
+            }
+            trainer.fabric.log_dict(
+                mean_loss_dict_for_logging, step=trainer.state["global_step"]
+            )
+
+            # ... print the mean losses in a table
+            df_losses = pd.DataFrame(
+                {
+                    "Train Loss Name": [
+                        k.replace("_", " ").title() for k in mean_loss_dict.keys()
+                    ],
+                    "Value": [v for v in mean_loss_dict.values()],
+                }
+            )
+            table = table_from_df(df_losses, title="Training Losses")
+
+            # (percentage of batch count)
+            percentage_complete = (batch_idx / trainer.n_batches_per_epoch) * 100
+
+            # Simple progress bar using Unicode blocks
+            progress_bar_length = 10  # Length of the progress bar
+            filled_length = int(progress_bar_length * percentage_complete // 100)
+            progress_bar = "█" * filled_length + "░" * (
+                progress_bar_length - filled_length
+            )
+            percentage_str = f"[bold magenta]{percentage_complete:.2f}%[/bold magenta]"
+
+            # Create a panel for the epoch and batch info with a progress bar
+            epoch_batch_info = (
+                f"[grey]<Rank {trainer.fabric.global_rank}>[/grey] "
+                f"Epoch {trainer.state['current_epoch']} Batch {batch_idx} "
+                f"[{progress_bar}] {percentage_str}"
+            )
+
+            epoch_batch_panel = Panel(
+                epoch_batch_info,
+                border_style="bold blue",
+            )
+
+            # Create a panel for the example ID
+            example_id = batch[0]["example_id"]
+            example_id_str = f"[bold yellow]{example_id}[/bold yellow]"
+            example_id_panel = Panel(
+                example_id_str,
+                border_style="bold green",
+            )
+
+            # Combine all components vertically
+            combined_content = Group(epoch_batch_panel, example_id_panel, table)
+
+            safe_print(combined_content)
+
+    def on_train_epoch_end(self, trainer):
+        # Gather final epoch means (must be run on all ranks)
+        final_means = {
+            k: tracker.compute().item() for k, tracker in self.loss_trackers.items()
+        }
+
+        # Calculate elapsed time and number of batches (from the total_loss tracker, if available)
+        elapsed_time = time.time() - self.start_time
+        num_batches = (
+            self.loss_trackers["total_loss"].update_count
+            if "total_loss" in self.loss_trackers
+            else trainer.n_batches_per_epoch
+        )
+
+        if trainer.fabric.is_global_zero:
+            # Create a summary table
+            table = Table(
+                title=f"Epoch {trainer.state['current_epoch']} Summary",
+                show_header=False,
+                header_style="bold magenta",
+            )
+            table.add_column("Loss Name", style="bold cyan", justify="left")
+            table.add_column("Value", style="green", justify="right")
+
+            for k, v in final_means.items():
+                table.add_row(f"<Train> Mean {k}", f"{v:.4f}")
+
+            table.add_section()
+            table.add_row("Total Optimizer Steps", str(trainer.state["global_step"]))
+            table.add_row("Number of Batches", str(num_batches))
+            table.add_row("Elapsed Time (s)", f"{elapsed_time:.2f}")
+            table.add_row(
+                "Mean Time per Batch (s)", f"{elapsed_time / num_batches:.2f}"
+            )
+
+            safe_print(table)
+
+        # Log these final epoch means (prepend "train/per_epoch_" to each key)
+        trainer.fabric.log_dict(
+            {f"train/per_epoch_{k}": v for k, v in final_means.items()},
+            step=trainer.state["current_epoch"],
+        )
+
+        # Reset the trackers for the next epoch
+        for metric in self.loss_trackers.values():
+            metric.reset()

foundry/common.py

@@ -0,0 +1,108 @@
+from functools import wraps
+
+import torch
+from beartype.typing import Any, Callable, Iterable
+from toolz import merge_with
+
+
+def run_once(fn: Callable) -> Callable:
+    """Decorator to ensure a function is only executed once per process.
+
+    Args:
+        fn (Callable): The function to decorate.
+
+    Returns:
+        Callable: A wrapped function that only executes once.
+    """
+
+    @wraps(fn)
+    def wrapper(*args: Any, **kwargs: Any) -> Any:
+        if getattr(wrapper, "_has_run", False):
+            return
+        wrapper._has_run = True
+        return fn(*args, **kwargs)
+
+    return wrapper
+
+
+def do_nothing(*args: Any, **kwargs: Any) -> None:
+    """Does nothing, just returns None"""
+    pass
+
+
+def exists(obj: Any) -> bool:
+    """True iff object is not None"""
+    return obj is not None
+
+
+def default(obj: Any, default: Any) -> Any:
+    """Return obj if it exists, otherwise return default"""
+    return obj if exists(obj) else default
+
+
+def exactly_one_exists(*args: object) -> bool:
+    """True iff exactly one of the arguments exists"""
+    return sum(exists(arg) for arg in args) == 1
+
+
+def at_least_one_exists(*args: object) -> bool:
+    """True iff at least one of the arguments exists"""
+    return any(exists(arg) for arg in args)
+
+
+def concat_dicts(*dicts: dict) -> dict:
+    """
+    Concatenate a list of dicts with the same keys into a single dict.
+
+    Example:
+        >>> d1 = {"a": 1, "b": 2}
+        >>> d2 = {"a": 3, "b": 4}
+        >>> concat_dicts(d1, d2)
+        {'a': [1, 3], 'b': [2, 4]}
+    """
+    return merge_with(list, *dicts)
+
+
+def listmap(fn: Callable, lst: Iterable[Any]) -> list:
+    """
+    Apply a function to each element of a single list.
+
+    Args:
+        - fn (Callable): Function to apply to each element
+        - lst (list): Input list
+
+    Returns:
+        - list: Result of applying fn to each element
+
+    Example:
+        >>> listmap(lambda x: x + 1, [1, 2, 3])
+        [2, 3, 4]
+    """
+    return [fn(x) for x in lst]
+
+
+def listmap_with_idx(fn: Callable[[int, Any], Any], lst: Iterable[Any]) -> list:
+    """Maps a function over a list while providing both index and value to the function.
+
+    A convenience wrapper around listmap that allows the mapping function to access both the index and value
+    of each element in the input list.
+
+    Args:
+        - fn (Callable[[int, Any], Any]): Function that takes two arguments (index, value) and returns a transformed value.
+        - lst (list): Input list to map over.
+
+    Returns:
+        - list: New list containing the results of applying fn to each (index, value) pair.
+
+    Example:
+        >>> def add_index(i, x):
+        ...     return f"{i}_{x}"
+        >>> listmap_with_idx(add_index, ["a", "b", "c"])
+        ['0_a', '1_b', '2_c']
+    """
+    return [fn(idx, x) for idx, x in enumerate(lst)]
+
+
+def ensure_dtype(tensor: torch.Tensor, dtype: torch.dtype) -> torch.Tensor:
+    """Convert tensor to target dtype if it's not already that dtype."""
+    return tensor if tensor.dtype == dtype else tensor.to(dtype)

foundry/constants.py

@@ -0,0 +1,28 @@
+# fmt: off
+# ... For convenience, define BKBN, or TIP to be used as a shortcut | TIP is the largest set of fixed atom given at least 2 tip atoms
+TIP_BY_RESTYPE = {
+    "TRP": ["CG","CD1","CD2","NE1","CE2","CE3","CZ2","CZ3","CH2"],  # fix both rings
+    "HIS": ["CG","ND1","CD2","CE1","NE2"],  # fixed ring
+    "TYR": ["CZ","OH"],  # keeps ring dihedral flexible
+    "PHE": ["CG","CD1","CD2","CE1","CE2","CZ"],
+    "ASN": ["CB", "CG","OD1","ND2"],
+    "ASP": ["CB", "CG","OD1","OD2"],
+    "GLN": ["CG", "CD","OE1","NE2"],
+    "GLU": ["CG", "CD","OE1","OE2"],
+    "CYS": ["CB", "SG"],
+    "SER": ["CB", "OG"],
+    "THR": ["CB", "OG1"],
+    "LEU": ["CB", "CG", "CD1", "CD2"],
+    "VAL": ["CG1", "CG2"],
+    "ILE": ["CB", "CG2"],
+    "MET": ["SD", "CE"],
+    "LYS": ["CE","NZ"],
+    "ARG": ["CD","NE","CZ","NH1","NH2"],
+    "PRO": None,
+    "ALA": None,
+    "GLY": None,
+    "UNK": None,
+    "MSK": None
+}
+
+# fmt: on