Class PlotsCallback

astroclip/callbacks.py:57–103

class PlotsCallback(Callback):
    # TODO: Update with latest code
    def __init__(self) -> None:
        super().__init__()

    def plot_spectrum(self, batch, output):
        sample_id = 3

        bs = len(batch["spectrum"])
        sp_rec = batch["target"][:, 1:, 2:99].reshape(bs, -1)[sample_id]
        in_rec = batch["input"][:, 1:, 2:99].reshape(bs, -1)[sample_id]
        out_rec = output[:, 1:, 2:99].reshape(bs, -1)[sample_id]

        # plot the moving average of the spectrum
        win = 20

        sp_rec = [sp_rec[i : i + win].mean().item() for i in range(0, len(sp_rec), win)]
        in_rec = [in_rec[i : i + win].mean().item() for i in range(0, len(in_rec), win)]
        out_rec = [
            out_rec[i : i + win].mean().item() for i in range(0, len(out_rec), win)
        ]

        fig = plt.figure()
        plt.plot(sp_rec, label="original")
        plt.plot(in_rec, label="dropped")
        plt.plot(out_rec, label="reconstructed")
        plt.legend()
        return fig

    def on_validation_batch_start(
        self,
        trainer: Trainer,
        pl_module: LightningModule,
        batch: Any,
        batch_idx: int,
        dataloader_idx: int = 0,
    ) -> None:
        if batch_idx == 0:
            output = trainer.model(batch["input"])
            fig = self.plot_spectrum(batch, output)
            for logger in trainer.loggers:
                # Check WandbLogger and Enabled
                if issubclass(type(logger), WandbLogger) and not issubclass(
                    type(logger.experiment), wandb.sdk.lib.disabled.RunDisabled
                ):
                    logger: WandbLogger = logger
                    logger.experiment.log({f"plot/{pl_module.current_epoch:03d}": fig})