models/partial_autoencoder/autoencoder.py

import os
from typing import Dict, List, Literal

import lightning as L
import matplotlib.pyplot as plt
import torch
import wandb
from jaxtyping import Bool, Float, Int
from lightning.pytorch.utilities.rank_zero import rank_zero_only
from loguru import logger
from sklearn.decomposition import PCA
from torch import Tensor

from models.partial_autoencoder.decoder import DecoderTransformer
from models.partial_autoencoder.decoder_ff import DecoderFFLocal
from models.partial_autoencoder.encoder import EncoderTransformer
from models.utils.coors_utils import nm_to_ang
from models.utils.pdb_utils import write_prot_to_pdb

COLORS_RT = [
    "#FF0000",  # Red
    "#008000",  # Green
    "#0000FF",  # Blue
    "#FFFF00",  # Yellow
    "#FFA500",  # Orange
    "#800080",  # Purple
    "#00FFFF",  # Cyan
    "#FF00FF",  # Magenta
    "#00FF00",  # Lime
    "#FFC0CB",  # Pink
    "#008080",  # Teal
    "#E6E6FA",  # Lavender
    "#A52A2A",  # Brown
    "#F5F5DC",  # Beige
    "#800000",  # Maroon
    "#808000",  # Olive
    "#FF7F50",  # Coral
    "#000080",  # Navy
    "#AAF0D1",  # Mint
    "#FFDB58",  # Mustard
]


@rank_zero_only
def create_dir(dir):
    if not os.path.exists(dir):
        os.makedirs(dir, exist_ok=True)


class AutoEncoder(L.LightningModule):
    def __init__(self, cfg_ae, store_dir=None):
        super().__init__()
        self.save_hyperparameters()
        self.cfg_ae = cfg_ae

        self.store_dir = store_dir if store_dir is not None else "./tmp_ae"
        self.val_path_tmp = os.path.join(self.store_dir, "val_stuff")
        create_dir(self.val_path_tmp)

        self.encoder = EncoderTransformer(**self.cfg_ae.nn_ae)

        decoder_type = self.cfg_ae.nn_ae.decoder.get("type", "transformer")
        if decoder_type == "transformer":
            self.decoder = DecoderTransformer(**self.cfg_ae.nn_ae)
        elif decoder_type == "ff_local":
            self.decoder = DecoderFFLocal(**self.cfg_ae.nn_ae)
        else:
            raise IOError(f"Invalid decoder_type {decoder_type}")

        self.nsamples_processed = 0
        self.nparams_enc = sum(
            p.numel() for p in self.encoder.parameters() if p.requires_grad
        )
        self.nparams_dec = sum(
            p.numel() for p in self.decoder.parameters() if p.requires_grad
        )

        self.validation_output = []
        self.validation_data_samples = []
        self.validation_rec_samples = []

        self.latent_dim = self.cfg_ae.nn_ae["latent_z_dim"]

    def log_histogram(self, id_log: str, v: Float[torch.Tensor, "r"]):
        """
        Logs histogram, v must be a flat tensor.
        """
        assert v.ndim == 1, f"Tensor v has shape {v.shape}, cannot log histogram"
        try:
            self.logger.experiment.log(
                {id_log: wandb.Histogram(v.cpu().detach().numpy())}
            )
        except:
            pass

    def configure_optimizers(self):
        optimizer = torch.optim.AdamW(
            [p for p in self.parameters() if p.requires_grad],
            lr=self.cfg_ae.opt.lr,
            amsgrad=True,
            weight_decay=1e-2,
        )
        return optimizer

    def on_save_checkpoint(self, checkpoint):
        """Adds additional variables to checkpoint."""
        checkpoint["nsamples_processed"] = self.nsamples_processed

    def on_load_checkpoint(self, checkpoint):
        """Loads additional variables from checkpoint."""
        try:
            self.nsamples_processed = checkpoint["nsamples_processed"]
        except:
            logger.info("Failed to load nsamples_processed from checkpoint")
            self.nsamples_processed = 0

    def encode(self, batch: Dict) -> Float[torch.Tensor, "b n d"]:
        """
        Runs the encoder and returns only the latent variables.
        """
        mask = batch["mask_dict"]["coords"][..., 0, 0]  # [b, n] boolean
        batch["mask"] = mask
        return self.encoder(batch)  # z_latent, mean, log_scale

    def decode(
        self,
        z_latent: Float[torch.Tensor, "b n d"],
        ca_coors_nm: Float[torch.Tensor, "b n 3"],
        mask: Bool[torch.Tensor, "b n"],
    ) -> Dict:
        """
        Runs the decoder and returns a dictionary with all necessary decoding information.
        """
        input_decoder = {
            "z_latent": z_latent,
            "ca_coors_nm": ca_coors_nm,
            "residue_mask": mask,
            "mask": mask,
        }
        output_dec = self.decoder(input_decoder)
        mask = output_dec["residue_mask"]  # [b, n]
        atom_mask = output_dec["atom_mask"]  # [b, n, 37]
        coors_nm = (
            output_dec["coors_nm"] * mask[..., None, None] * atom_mask[..., None]
        )  # [b, n, 37, 3]
        return {
            "coors_nm": coors_nm,
            "residue_type": output_dec["aatype_max"] * mask,
            "residue_mask": mask,
            "atom_mask": atom_mask,
        }

    def training_step(self, batch: Dict, batch_idx: int):
        """
        Computes training loss for batch of samples.

        Args:
            batch: Data batch.

        Returns:
            Training loss averaged over batch dimension.
        """
        val_step = batch_idx == -1  # validation step is indicated with batch_idx -1
        log_prefix = "validation_loss" if val_step else "train"
        histogram_every_n = 5000
        pca_every_n = 5000
        per_aatype_kl = True

        mask = batch["mask_dict"]["coords"][..., 0, 0]  # [b, n] boolean
        batch["mask"] = mask
        ca_coors_nm = batch["coords_nm"][..., 1, :]  # [b, n, 3]
        ca_coors_nm = ca_coors_nm * mask[..., None]  # [b, n, 3]
        bs, n = mask.shape[0], mask.shape[1]

        output_enc = self.encoder(batch)
        # {
        #   "z_latent": latent_sample, shape [b, n, d]
        #   "mean": mean of latent (diag) Gaussian dist, shape [b, n, d]
        #   "log_scale": log standard deviation of latent (diag) Gaussian dist, shape [b, n, d]
        # }
        log_prefix_stats = (
            "train_stats_latent" if "train" in log_prefix else "val_stats_latent"
        )
        self.log_tensor_statistics(
            bs=bs,
            v=output_enc["mean"],
            log_prefix=log_prefix_stats + "_mean_LS",
            mask=mask,
            histogram_every_n=histogram_every_n,
        )
        self.log_tensor_statistics(
            bs=bs,
            v=torch.exp(output_enc["log_scale"]),
            log_prefix=log_prefix_stats + "_scale_LS",
            mask=mask,
            histogram_every_n=histogram_every_n,
        )
        self.log_tensor_statistics(
            bs=bs,
            v=output_enc["z_latent"],
            log_prefix=log_prefix_stats + "_z_LS",
            mask=mask,
            histogram_every_n=histogram_every_n,
        )
        self.log_pca(
            v=output_enc["z_latent"],
            log_prefix=log_prefix_stats + "_z_PCA",
            mask=mask,
            every_n=pca_every_n,
        )
        self.log_pca_per_residue_type(
            v=output_enc["z_latent"],
            log_prefix=log_prefix_stats + "_z_PCA_pre_aatype",
            mask=mask,
            every_n=pca_every_n,
            res_ty=batch["residue_type"],
        )

        input_decoder = {
            "z_latent": output_enc["z_latent"],
            "ca_coors_nm": ca_coors_nm,
            "residue_mask": mask,
            "mask": mask,
        }
        output_dec = self.decoder(input_decoder)
        # {
        #   "coors_nm": all atom coordinates, shape [b, n, 37, 3], in nm
        #   "seq_logits": logits for the residue types, shape [b, n, 20]
        #   "residue_mask": boolean [b, n]
        #   "aatype_max": residue type by taking the most likely logit, shape [b, n], with integer values {0, ..., 19}
        #   "atom_mask": boolean [b, n, 37, 3], atom37 mask corresponding to aatype_max
        # }

        losses = (
            {}
        )  # Will be a Dict[str, tensor[b]]. If "_justlog" in name just for logging

        # KL loss with annealing weight
        f = (
            min(1.0, self.global_step / self.cfg_ae.loss.kl.patience)
            if self.cfg_ae.loss.kl.anneal
            else 1.0
        )
        self.log(
            "kl_weight",
            self.cfg_ae.loss.kl.weight * f,
            on_step=True,
            on_epoch=False,
            prog_bar=False,
            logger=True,
            sync_dist=True,
            add_dataloader_idx=False,
        )
        losses.update(
            self.compute_kl_penalty(
                mean=output_enc["mean"],
                log_scale=output_enc["log_scale"],
                mask=mask,
                w=self.cfg_ae.loss.kl.weight * f,
            )
        )

        # Structure loss
        losses.update(
            self.compute_struct_rec_loss(
                output_dec=output_dec,
                batch=batch,
                reduce_mode="sum",
                loss_ty=self.cfg_ae.loss.struct.type,
                weight=self.cfg_ae.loss.struct.weight,
            )
        )

        # Sequence loss
        losses.update(
            self.compute_seq_rec_loss(
                output_dec=output_dec,
                batch=batch,
                weight=self.cfg_ae.loss.seq.weight,
            )
        )

        # Log losses and training loss, losses with "_justlog" just for logging purposes
        self.log_losses(bs, losses, log_prefix)
        train_loss = sum([torch.mean(losses[k]) for k in losses if "_justlog" not in k])

        self.log(
            f"{log_prefix}/loss",
            train_loss,
            on_step=True,
            on_epoch=True,
            prog_bar=False,
            logger=True,
            batch_size=bs,
            sync_dist=True,
            add_dataloader_idx=False,
        )

        if not val_step:  # Don't log these for val step
            # Log Kl statistics
            componentwise_kl = self._per_component_kl(
                mean=output_enc["mean"],
                log_scale=output_enc["log_scale"],
                mask=mask,
            )  # [b, n, d]
            self.log_tensor_statistics(
                bs=bs,
                v=componentwise_kl,
                log_prefix=log_prefix + "_kl_latent",
                mask=mask,
                histogram_every_n=histogram_every_n,
            )
            self.log_tensor_statistics(
                bs=bs,
                v=(componentwise_kl > 0.1) * 1.0,
                log_prefix=log_prefix + "_kl_latent_active_thresh_0p1",
                mask=mask,
                do_hist=False,
                stats_to_log=["mean"],
            )
            # KL per aa type
            if per_aatype_kl and self.global_step % histogram_every_n == 2000:
                for i in range(20):
                    self.log_tensor_statistics(
                        bs=bs,
                        v=(componentwise_kl > 0.1) * 1.0,
                        log_prefix=log_prefix_stats
                        + "_kl_latent_active_thresh_0p1_per_aatype",
                        mask=mask * (batch["residue_type"] == i),
                        do_hist=False,
                        suffix=f"_aa_{i}",
                        stats_to_log=["mean"],
                    )

            self.log_train_loss_n_prog_bar(bs, train_loss)
            self.update_n_log_nsamples_processed(bs)
            self.log_nparams()

        if val_step:
            return train_loss, output_dec
        return train_loss

    def compute_struct_rec_loss(
        self,
        output_dec: Dict[str, torch.Tensor],
        batch: Dict[str, torch.Tensor],
        reduce_mode: Literal["sum", "mean"] = "sum",
        loss_ty: str = "l2",
        weight: float = 1.0,
    ) -> Dict[str, Float[Tensor, "b"]]:
        """
        Computes structural loss.

        Args:
            output_dec:
                {
                  "coors_nm": all atom coordinates, shape [b, n, 37, 3], in nm
                  "seq_logits": logits for the residue types, shape [b, n, 20]
                  "residue_mask": boolean [b, n]
                  "aatype_max": residue type by taking the most likely logit, shape [b, n], with integer values {0, ..., 19}
                  "atom_mask": boolean [b, n, 37, 3], atom37 mask corresponding to aatype_max
                }
            batch: data batch from dataloader
            reduce_mode: whether to apply mean over atoms or just sum, when reducing accross 37 atoms types
            loss_ty: whether to apply l1 loss, l2 loss, or both, l2cut (ignore atoms close to each other, to expand)

        Returns:
            Dictionary [str, tensor[b]] with multiple per batch element losses. If the key has "_justlog" then this loss
            will not be used to compute the total loss, but will just be logged.
        """

        def reduce_37(
            err: Float[torch.Tensor, "b n 37 3"],
            mask: Bool[torch.Tensor, "b n"],
            atom_mask: Bool[torch.Tensor, "b n 37"],
            mode: Literal["sum", "mean"] = "sum",
        ) -> Float[torch.Tensor, "b"]:
            nres = mask.sum(dim=-1)  # [b]
            nat = atom_mask.sum(dim=-1) * mask  # [b, n]
            err = torch.sum(err, dim=(-1, -2))  # [b, n]
            if mode == "mean":
                err = err / nat  # Take mean over existing atoms if mode == "mean"
            err = err.sum(dim=-1) / nres  # [b]
            return err

        coors_nm_pred = output_dec["coors_nm"]  # [b, n, 37, 3]
        coors_nm_true = batch["coords_nm"]  # [b, n, 37, 3]
        mask = output_dec["residue_mask"]  # [b, n] boolean
        atom_mask_true = batch["coord_mask"] * mask[..., None]  # [b, n, 37] boolean

        err = coors_nm_true - coors_nm_pred  # [b, n, 37, 3]
        err = err * mask[..., None, None] * atom_mask_true[..., None]  # [b, n, 37, 3]

        losses = {}

        # Compute RMSD in Å (without alignment)
        err_ang = nm_to_ang(err)  # [b, n, 37, 3]
        err_ang = torch.linalg.norm(err_ang, dim=-1) ** 2  # [b, n, 37]
        err_ang = err_ang * mask[..., None] * atom_mask_true  # [b, n, 37]
        nat = atom_mask_true.sum((-1, -2))  # [b]
        rmsd = torch.sqrt(torch.sum(err_ang, dim=(-1, -2)) / nat)  # [b]
        losses["rmsd_no_align_a37_ang_justlog"] = rmsd

        err_l1 = reduce_37(
            torch.abs(err), mask, atom_mask_true, mode=reduce_mode
        )  # L1 loss
        err_l2 = reduce_37(err**2, mask, atom_mask_true, mode=reduce_mode)  # L2 loss

        if loss_ty == "l1":
            losses["struct_l1"] = err_l1 * weight
            losses["struct_l2_justlog"] = err_l2 * weight
        elif loss_ty == "l2":
            losses["struct_l1_justlog"] = err_l1 * weight
            losses["struct_l2"] = err_l2 * weight
        elif loss_ty == "l12":
            losses["struct_l1"] = err_l1 * weight
            losses["struct_l2"] = err_l2 * weight
        else:
            raise IOError(f"Loss type {loss_ty} not recognized")
        return losses

    def compute_seq_rec_loss(
        self,
        output_dec: Dict[str, torch.Tensor],
        batch: Dict[str, torch.Tensor],
        weight: float = 1.0,
    ) -> Dict[str, Float[Tensor, "b"]]:
        """
        Computes cross entropy loss on predicted sequence.

        Args:
            output_dec:
                {
                  "coors_nm": all atom coordinates, shape [b, n, 37, 3], in nm
                  "seq_logits": logits for the residue types, shape [b, n, 20]
                  "residue_mask": boolean [b, n]
                  "aatype_max": residue type by taking the most likely logit, shape [b, n], with integer values {0, ..., 19}
                  "atom_mask": boolean [b, n, 37, 3], atom37 mask corresponding to aatype_max
                }
            batch: data batch from dataloader

        Returns:
            Dictionary [str, tensor[b]] with multiple per batch element losses. If the key has "_justlog" then this loss
            will not be used to compute the total loss, but will just be logged.
        """
        assert (
            "residue_type" in batch
        ), "`residue_type` not in batch, failed in compute_seq_rec_loss"
        mask = output_dec["residue_mask"]  # [b, n]
        nres = mask.sum(dim=-1)  # [b]
        logits_pred = output_dec["seq_logits"]  # [b, n, 20]
        target_aa = batch["residue_type"]  # [b, n]
        target_aa = (
            target_aa * mask
        )  # [b, n] gets rid of -1 for padding (issue with cross entropy loss below)

        assert logits_pred.shape[-1] == 20, "Wrong number of logits"

        # Compute cross entropy
        b, n = mask.shape[0], mask.shape[1]
        logits_pred_flat = logits_pred.view(b * n, 20)  # [b * n, 20]
        target_aa_flat = target_aa.view(b * n)  # [b * n]
        seq_loss_flat = torch.nn.functional.cross_entropy(
            input=logits_pred_flat,
            target=target_aa_flat,
            reduction="none",
        )  # [b * n]
        seq_loss = seq_loss_flat.view(b, n)  # [b, n]
        seq_loss = seq_loss * mask  # [b, n]
        seq_loss = torch.sum(seq_loss, dim=-1) / nres  # [b]

        # Compute seq recovery rate
        pred_aa = output_dec["aatype_max"]  # [b, n]
        seq_rec = pred_aa == target_aa  # [b, n]
        seq_rec = seq_rec * mask  # [b, n]
        seq_rec_rate = seq_rec.sum(dim=-1) / nres  # [b]
        return {
            "ce_seq": seq_loss * weight,
            "ce_seq_now": seq_loss,
            "seq_rec_rate_justlog": seq_rec_rate,
        }

    def _per_component_kl(
        self,
        mean: Float[Tensor, "b n d"],
        log_scale: Float[Tensor, "b n d"],
        mask: Bool[Tensor, "b n"],
    ) -> Float[Tensor, "b n d"]:
        """
        Computes KL penalty on the latent Gaussian distribution, per residue.

        Returns KL, per residue (masked) and latent dimension, shape [b n d].
        """
        scale = torch.exp(log_scale)  # [b, n, 3]
        kl_div = (scale**2 + mean**2 - 1.0 - 2.0 * log_scale) * 0.5  # [b, n, d]
        return kl_div * mask[..., None]  # [b, n, d]

    def compute_kl_penalty(
        self,
        mean: Float[Tensor, "b n d"],
        log_scale: Float[Tensor, "b n d"],
        mask: Bool[Tensor, "b n"],
        w: float,
    ) -> Dict[str, Float[Tensor, "b"]]:
        """
        Computes KL penalty on the latent Gaussian distribution.
        """
        nres = torch.sum(mask, dim=-1)  # [b]
        kl_div = self._per_component_kl(
            mean=mean,
            log_scale=log_scale,
            mask=mask,
        )  # [b, n, d]
        kl_div = torch.sum(kl_div, dim=(-1, -2))  # [b]
        kl_div = kl_div / nres  # [b]
        loss = {
            "kl_w": kl_div * w,
            "kl_now_justlog": kl_div,
        }
        return loss

    def log_tensor_statistics(
        self,
        bs: int,
        v: Float[torch.Tensor, "b n d"],
        log_prefix: str,
        mask: Bool[torch.Tensor, "b n"],
        do_stats: bool = True,
        do_hist: bool = True,
        suffix: str = "",
        histogram_every_n: int = 3000,
        stats_to_log: List[str] = ["max", "min", "mean", "median", "std"],
    ) -> None:
        """
        Logs tensor statistics: mean, median, max, min, std, ...

        If res_type is not None, then it also logs histograms per residue type
        """
        vals = v.clone()[mask]  # [num of unmasked residues, d]
        vals = torch.flatten(vals)  # 1D vector
        if do_stats:
            funs = {
                "max": torch.max,
                "min": torch.min,
                "mean": torch.mean,
                "median": torch.median,
                "std": torch.std,
            }
            for k in stats_to_log:
                self.log(
                    f"{log_prefix}/{k}{suffix}",
                    funs[k](vals),
                    on_step=True,
                    on_epoch=True,
                    prog_bar=False,
                    logger=True,
                    batch_size=bs,
                    sync_dist=True,
                    add_dataloader_idx=False,
                )

        if self.global_step % histogram_every_n == 0 and do_hist:
            self.log_histogram(id_log=f"{log_prefix}/histogram{suffix}", v=vals)

    def log_pca(
        self,
        v: Float[torch.Tensor, "b n d"],
        log_prefix: str,
        mask: Bool[torch.Tensor, "b n"],
        every_n: int = 3000,
    ) -> None:
        """
        Logs PCA components plot of latent variable z.
        """

        def _log_scatter(vals_x, vals_y, xlabel, ylabel, log_id):
            fig, ax = plt.subplots()
            ax.scatter(vals_x, vals_y)
            ax.set_xlabel(xlabel)
            ax.set_ylabel(ylabel)
            ax.set_title(f"{xlabel} - {ylabel} - Step: {self.trainer.global_step}")
            try:
                self.logger.experiment.log({log_id: wandb.Image(fig)})
            except:
                pass
            plt.close("all")

        if self.global_step % every_n != 0:
            return

        n_components = 4
        vals = (
            v.clone()[mask].cpu().detach().float().numpy()
        )  # [num of unmasked residues, d]
        vals_pca = PCA(n_components=n_components).fit_transform(
            vals
        )  # [num of unmasked residues, n_components]

        for i in range(n_components):
            for j in range(n_components):
                if j <= i:
                    continue
                _log_scatter(
                    vals_x=vals_pca[:, i],
                    vals_y=vals_pca[:, j],
                    xlabel=f"PC {i}",
                    ylabel=f"PC {j}",
                    log_id=f"{log_prefix}/{i}_{j}",
                )

    def log_pca_per_residue_type(
        self,
        v: Float[torch.Tensor, "b n d"],
        res_ty: Int[torch.Tensor, "b n"],
        log_prefix: str,
        mask: Bool[torch.Tensor, "b n"],
        every_n: int = 3000,
    ) -> None:
        """
        Logs PCA components plot of latent variable z. Not done.
        """
        return None


    def log_losses(
        self, bs: int, losses: Dict[str, Float[torch.Tensor, "b"]], log_prefix: str
    ) -> None:
        for k in losses:
            log_name = k[: -len("_justlog")] if k.endswith("_justlog") else k
            self.log(
                f"{log_prefix}/loss_{log_name}",
                torch.mean(losses[k]),
                on_step=True,
                on_epoch=True,
                prog_bar=False,
                logger=True,
                batch_size=bs,
                sync_dist=True,
                add_dataloader_idx=False,
            )

    def log_train_loss_n_prog_bar(self, b: int, train_loss: torch.Tensor) -> None:
        self.log(
            f"train_loss",
            train_loss,
            on_step=True,
            on_epoch=True,
            prog_bar=True,
            logger=True,
            batch_size=b,
            sync_dist=True,
            add_dataloader_idx=False,
        )

    def log_nparams(self):
        self.log(
            "scaling/nparams_enc",
            self.nparams_enc * 1.0,
            on_step=True,
            on_epoch=False,
            prog_bar=False,
            logger=True,
            batch_size=1,
            sync_dist=True,
        )  # constant line but ok, easy to compare # params

        self.log(
            "scaling/nparams_dec",
            self.nparams_dec * 1.0,
            on_step=True,
            on_epoch=False,
            prog_bar=False,
            logger=True,
            batch_size=1,
            sync_dist=True,
        )  # constant line but ok, easy to compare # params

    def update_n_log_nsamples_processed(self, b: int):
        self.nsamples_processed = self.nsamples_processed + b * self.trainer.world_size
        self.log(
            "scaling/nsamples_processed",
            self.nsamples_processed * 1.0,
            on_step=True,
            on_epoch=False,
            prog_bar=False,
            logger=True,
            batch_size=1,
            sync_dist=True,
        )

    def validation_step(self, batch: Dict, batch_idx: int):
        """
        Evalaute validation loss.

        Args:
            batch: batch from dataset (see last argument)
            batch_idx: batch index (unused)
        """
        with torch.no_grad():
            bs = batch["coords_nm"].shape[0]
            loss, rec_sample = self.training_step(batch, batch_idx=-1)
            self.validation_output.append(loss.item())
            for i in range(bs):
                self.validation_data_samples.append(
                    {
                        "coors_nm": batch["coords_nm"][i, ...],  # [n, 37, 3]
                        "aatype": batch["residue_type"][i, ...],  # [n]
                        "atom_mask": batch["coord_mask"][i, ...],  # [n, 37]
                        "mask": batch["mask_dict"]["coords"][i, :, 0, 0],  # [n]
                    }
                )
                self.validation_rec_samples.append(
                    {
                        "coors_nm": rec_sample["coors_nm"][i, ...],  # [n, 37, 3]
                        "aatype": rec_sample["aatype_max"][i, ...],  # [n]
                        "atom_mask": rec_sample["atom_mask"][i, ...],  # [n, 37]
                        "mask": rec_sample["residue_mask"][
                            i, ...
                        ],  # [n], should be the same as mask from orig sample
                    }
                )

    def on_validation_epoch_end(self):
        """
        Stores samples as PDBs and cleans validation results.
        """
        count = 0
        for sample_data, sample_rec in zip(
            self.validation_data_samples, self.validation_rec_samples
        ):
            count += 1
            if count > 10:
                break
            coors_data = (
                nm_to_ang(sample_data["coors_nm"]).float().detach().cpu().numpy()
            )  # [n, 37, 3]
            aatype_data = sample_data["aatype"].detach().cpu().numpy()  # [n]
            mask_data = sample_data["mask"].detach().cpu().numpy()  # [n]
            atom_mask_data = sample_data["atom_mask"].detach().cpu().numpy()  # [n, 37]

            coors_rec = nm_to_ang(sample_rec["coors_nm"]).float().detach().cpu().numpy()
            aatype_rec = sample_rec["aatype"].detach().cpu().numpy()
            mask_rec = sample_rec["mask"].detach().cpu().numpy()
            atom_mask_rec = sample_rec["atom_mask"].detach().cpu().numpy()

            f_base = f"epoch_{self.current_epoch}_id_{count}_rank_{self.global_rank}"

            # Save original sample
            fname = f_base + "_data_sample.pdb"
            fpath = os.path.join(self.val_path_tmp, fname)
            write_prot_to_pdb(
                prot_pos=coors_data
                * mask_data[..., None, None]
                * atom_mask_data[..., None],
                file_path=fpath,
                aatype=aatype_data * mask_data,
                overwrite=True,
                no_indexing=True,
            )

            # Save recovered sample
            fname = f_base + "_rec_sample.pdb"
            fpath = os.path.join(self.val_path_tmp, fname)
            write_prot_to_pdb(
                prot_pos=coors_rec
                * mask_rec[..., None, None]
                * atom_mask_rec[..., None],
                file_path=fpath,
                aatype=aatype_rec * mask_rec,
                overwrite=True,
                no_indexing=True,
            )

            # Save recovered sample with true sequence
            fname = f_base + "_rec_sample_w_true_seq.pdb"
            fpath = os.path.join(self.val_path_tmp, fname)
            write_prot_to_pdb(
                prot_pos=coors_rec
                * mask_data[..., None, None]
                * atom_mask_data[..., None],
                file_path=fpath,
                aatype=aatype_data * mask_data,
                overwrite=True,
                no_indexing=True,
            )

        self.validation_data_samples = []
        self.validation_rec_samples = []
        self.validation_output_data = []
        # Should log here?

    def predict_step(self, batch: Dict, batch_idx: int) -> Dict:
        """
        Makes predictions. Given a data batch, encodes, and returns decoded batch.

        Args:
            batch: data batch.

        Returns:
            Dict representing the decoded batch, with all info from the encoder output.
        """
        mask = batch["mask_dict"]["coords"][..., 0, 0]  # [b, n] boolean
        batch["mask"] = mask
        ca_coors_nm = batch["coords_nm"][..., 1, :]  # [b, n, 3]
        ca_coors_nm = ca_coors_nm * mask[..., None]  # [b, n, 3]

        output_enc = self.encoder(batch)
        # {
        #   "z_latent": latent_sample, shape [b, n, d]
        #   "mean": mean of latent (diag) Gaussian dist, shape [b, n, d]
        #   "log_scale": log standard deviation of latent (diag) Gaussian dist, shape [b, n, d]
        # }

        input_decoder = {
            "z_latent": output_enc["z_latent"],
            "ca_coors_nm": ca_coors_nm,
            "residue_mask": mask,
            "mask": mask,
        }
        output = self.decoder(input_decoder)
        # {
        #   "coors_nm": all atom coordinates, shape [b, n, 37, 3], in nm
        #   "seq_logits": logits for the residue types, shape [b, n, 20]
        #   "residue_mask": boolean [b, n]
        #   "aatype_max": residue type by taking the most likely logit, shape [b, n], with integer values {0, ..., 19}
        #   "atom_mask": boolean [b, n, 37, 3], atom37 mask corresponding to aatype_max
        # }

        output.update(output_enc)
        return (batch, output)