plmodules.py

from typing import Union, List
from omegaconf import DictConfig
import numpy as np
import rootutils
import itertools
import os.path
import wandb
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import pytorch_lightning as pl
from torchmetrics.functional.image import \
    structural_similarity_index_measure as ssim
from torchmetrics.functional.image import \
    peak_signal_noise_ratio as psnr
from models import Func_C2P, \
    Generator_D2O, Generator_O2D, \
    Discriminator_O, Discriminator_D, \
    Encoder_O, Encoder_D
from func import L_ID_SCHEDULE
base = rootutils.find_root(search_from=__file__, indicator=".project-root")



class BaseModule(pl.LightningModule):
    '''
    Base module for all GAN-based models
    accept common args, implement common methods
    '''
    def __init__(self, 
                G_D2O_config: DictConfig = None, 
                D_O_config: DictConfig = None, 
                Encoder_config: DictConfig = None,
                GAN_mode: str = 'WGAN-GP',  # [ 'vanilla' | 'hinge' | 'lsGAN' | 'WGAN' | 'WGAN-GP' ]
                loss_id: str = 'L2',  # [ 'L1' | 'L2' ]
                optimizer: str = "RMSprop",  # [ 'Adam' | 'RMSprop' ]
                lr: float = 1e-5,
                b1: float = 0.5, 
                b2: float = 0.999, 
                weight_clipping: bool = False,
                clip_value: float = 0.01,
                lambda_gp: float = 10,
                n_opt_G: int = 1,
                n_opt_D: int = 5, 
                log_images: bool = True,
                embed_lead: bool = False,  # lead is in dataset but not to model
                idx_of_lead: int = 0,  # position of `lead time` in aux_info
                *args, **kwargs
                ):
        super(BaseModule, self).__init__()
        self.save_hyperparameters()
        # basic check inside config here, currently none
        pass
    
    
    def on_train_start(self):
        if self.hparams.debug:
            return
        # save critical code
        if '/home/sw/Work4-GAN' in str(base):
            wandb.save(os.path.join(base, 'scripts/train/*.py'))
        elif '/home/pod/shared-nvme/Work1-GAN' in str(base):
            wandb.save(os.path.join(base, 'train/*.py'))
    
    
    def weight_clipping(self, model, clip_value):
        for p in model.parameters():
            p.data.clamp_(-clip_value, clip_value)
    
    
    def compute_gradient_penalty(self, D, condition, real_samples, fake_samples):
        '''
        ref1: https://github.com/nocotan/pytorch-lightning-gans/blob/master/models/wgan_gp.py
        ref2 to scale to patchGAN: https://github.com/p-hss/weather-gan/blob/main/src/model.py 
        
        conditions, real_samples, fake_samples: Tensor or Tuple[Tensor]
        e.g. `compute_gradient_penalty(self.D, info, (obs_circu, obs_tp), (D2O_circu, D2O_tp))`
                                             static conditions,    real inputs,        fake inputs
        e.g. `compute_gradient_penalty(self.D, None, real_tp, fake_tp)`
        '''
        item_samples = real_samples[0] if isinstance(real_samples, tuple) else real_samples
        alpha = torch.rand(item_samples.size(0), 1, 1, 1).type_as(item_samples)
        interpolates = []
        for real_items, fake_items in zip(real_samples, fake_samples):
            interpolates += [
                (alpha * real_items + ((1 - alpha) * fake_items)).requires_grad_(True)
            ]
        if not isinstance(real_samples, tuple) or len(real_samples) == 1:
            interpolates = [interpolates[0]]  # strange but necessary
        d_interpolates = D(condition, *interpolates)
        gradients = torch.autograd.grad(
            outputs=d_interpolates,
            inputs=interpolates,
            grad_outputs=torch.ones(d_interpolates.size()).type_as(item_samples),
            create_graph=True,
            retain_graph=True,
            only_inputs=True,
        )[0]
        gradients = gradients.flatten(start_dim=1)
        if self.hparams.GAN_mode == 'WGAN-GP':
            slopes = gradients.norm(2, dim=1)
            gradient_penalty = ((slopes - 1) ** 2).mean()
        else: 
            raise NotImplementedError
        return gradient_penalty
    
    
    def pixelwise_score(self, fake, real):
        if self.hparams.loss_id == 'L2':
            return F.mse_loss(fake, real, reduce=False)
        elif self.hparams.loss_id == 'L1':
            return F.l1_loss(fake, real, reduce=False)
        else:
            raise NotImplementedError
    
    
    def valid_rmse_score(self, y, y_hat, valid_grid_2d):
        '''
        RMSE score for valid grid only
        input (y / y_hat): [B, C, H, W]
        valid_grid_2d: [B, 1, H, W]
        '''
        y = (y * valid_grid_2d).flatten(start_dim=1)
        y_hat = (y_hat * valid_grid_2d).flatten(start_dim=1)
        return torch.sqrt(F.mse_loss(y, y_hat))
    
    
    def valid_pcc_score(self, y, y_hat, valid_grid_2d):
        '''
        Pattern Corr Coeff score for valid grid only
        input (y / y_hat): [B, C, H, W]
        valid_grid_2d: [B, 1, H, W]
        '''
        y = (y * valid_grid_2d).flatten(start_dim=1)
        y_hat = (y_hat * valid_grid_2d).flatten(start_dim=1)
        y_mean = torch.mean(y, dim=-1, keepdim=True)
        y_hat_mean = torch.mean(y_hat, dim=-1, keepdim=True)
        return torch.sum((y - y_mean) * (y_hat - y_hat_mean)) / \
            torch.sqrt(torch.sum((y - y_mean) ** 2) * torch.sum((y_hat - y_hat_mean) ** 2))
    
    
    def valid_crps_score(self, y, y_hat, valid_grid_2d):
        '''
        CRPS score for valid grid only
        
        CRPS(F, x) = E_F|X - x| - 1/2 * E_F|X - X'|
        y: [B, C, H, W]
        y_hat: [E, B, C, H, W] or [B, C, H, W]
        valid_grid_2d: [B, 1, H, W]
        
        NOTE: mind that y & y_hat can not be exchanged
        '''
        # [_, H, W] -> [_, spatial]
        y = (y * valid_grid_2d).flatten(start_dim=1)
        if y.dim() != y_hat.dim():
            valid_grid_2d = valid_grid_2d[None]
            y_hat = (y_hat * valid_grid_2d).flatten(start_dim=2)
        else:
            y_hat = (y_hat * valid_grid_2d).flatten(start_dim=1)
        
        if y.dim() == y_hat.dim(): # single value CRPS degenrated to MAE
            return F.l1_loss(y, y_hat)
        else:
            # squeeze channel dim [B, C, S] -> [B, S], [E, B, C, S] -> [E, B, S]
            y = y.squeeze(1) if y.dim() == 3 else y
            y_hat = y_hat.squeeze(2) if y_hat.dim() == 4 else y_hat
            # make y_hat & y same shape
            y_expand = y.unsqueeze(0).expand(y_hat.shape)
            CRPS = torch.abs(y_hat - y_expand).mean(axis=(0, 2)) - \
                0.5 * (torch.abs(y_hat.unsqueeze(1) - y_hat.unsqueeze(0)).mean(axis=(0, 1, 3)))
            # leave batch dim when averaging over ensemble and spatial, then ave
            return CRPS.mean()
    
    
    def log_precip_map(self, names, imgs_list):
        '''
        names: List[str]
        imgs_list: List[torch.Tensor] -> [B, C, H, W]
        should be same length
        '''
        sample = imgs_list[0]
        B = sample.shape[0]
        idx_sel = torch.randint(0, B, (4,))
        imgs = [wandb.Image(imgs[idx_sel],
                            mode='L',
                            caption=f"{self.current_epoch:03d} {name}",
                            ) for name, imgs in zip(names, imgs_list)]
        wandb.log({'images': imgs})
    
    
    def get_optimizers(self, params):
        if self.hparams.optimizer == "Adam":
            opt = optim.Adam(params, 
                            lr=self.hparams.lr, 
                            betas=(self.hparams.b1, self.hparams.b2)
                            )
        elif self.hparams.optimizer == "RMSprop":
            opt = optim.RMSprop(params, 
                                lr=self.hparams.lr
                                )
        else:
            raise NotImplementedError
        return opt
    
    
    def reparameterize(self, mu, logvar, ensemble=1):
        if ensemble == 1:
            std = torch.exp(0.5 * logvar)
            eps = torch.randn_like(std)
            return mu + eps * std
        elif ensemble > 1: # [B, L] -> [E, B, L] -> [E*B, L]
            _, L = mu.shape
            std = torch.exp(0.5 * logvar)\
                .unsqueeze(0).expand(ensemble, -1, -1).reshape(-1, L)
            mu = mu.unsqueeze(0).expand(ensemble, -1, -1).reshape(-1, L)
            eps = torch.randn_like(std)
            return mu + eps * std
    
    
    def kl_divergence(self, mu1, logvar1, mu2, logvar2):
        kl_loss = 0.5 * (logvar2 - logvar1 + torch.exp(logvar1) / torch.exp(logvar2) + (mu1 - mu2)**2 / torch.exp(logvar2) - 1)
        return torch.mean(kl_loss)
    
    
    def get_adversarial_loss(self, critic, is_real):
        if self.hparams.GAN_mode == 'vanilla':
            if is_real:
                return nn.BCEWithLogitsLoss()(critic, torch.ones_like(critic).type_as(critic))
            else:
                return nn.BCEWithLogitsLoss()(critic, torch.zeros_like(critic).type_as(critic))
        elif self.hparams.GAN_mode == 'hinge':
            if is_real:
                return torch.mean(F.relu(1 - critic))
            else:
                return torch.mean(F.relu(1 + critic))
        elif self.hparams.GAN_mode in ['lsGAN']: 
            if is_real:
                return F.mse_loss(critic, torch.ones_like(critic).type_as(critic))
            else:
                return F.mse_loss(critic, -torch.ones_like(critic).type_as(critic))
        elif self.hparams.GAN_mode in ['WGAN', 'WGAN-GP']:
            if is_real:
                return -torch.mean(critic)
            else:
                return torch.mean(critic)
        else:
            raise NotImplementedError
    
    
    def get_coarse_bool(self, valid_grid_2d):
        valid_grid_2d = valid_grid_2d.float()
        return F.interpolate(valid_grid_2d, scale_factor=0.1,
                            mode="bilinear", align_corners=False)\
                            .bool()
    
    
    def drop_lead_in_aux(self, info):
        '''
        drop lead time in aux info
        '''
        S1, S2, valid_grid_2d, coord, aux = info
        mask = torch.ones(aux.size(), dtype=torch.bool)
        mask[:, self.hparams.idx_of_lead] = False
        aux = aux[mask].view(-1, aux.size(1) - 1)
        return S1, S2, valid_grid_2d, coord, aux



class PerfectProgModule(BaseModule):
    '''
    circu -> tp
    dont need lead weighting schedule
    '''
    def __init__(self, 
                Func_C2P_config: DictConfig = None,
                scheduler: str = None, # [ None | 'ReduceLROnPlateau' ]
                *args, **kwargs
                ):
        super(PerfectProgModule, self).__init__(
            *args, **kwargs
        )
        self.save_hyperparameters()
        print('haprams:', self.hparams)
        
        self.PPmodel = Func_C2P(**Func_C2P_config)
    
    
    # train mode
    def forward(self, info, circu):
        if not self.hparams.embed_lead:
            info = self.drop_lead_in_aux(info)
        # obs circu is finer than dyn grid
        # thus avgpool to //4 here
        circu = F.avg_pool2d(circu, 4)
        pred_tp = self.PPmodel(info, circu)
        _, _, valid_grid_2d, _, _ = info
        return pred_tp * valid_grid_2d
    
    
    @torch.no_grad()
    def inference(self, info, circu):
        if not self.hparams.embed_lead:
            info = self.drop_lead_in_aux(info)
        pred_tp = self.PPmodel(info, circu)
        _, _, valid_grid_2d, _, _ = info
        return pred_tp * valid_grid_2d
    
    
    def common_step(self, batch):
        info, circu, obs_tp = batch
        _, _, valid_grid_2d, _, _ = info
        obs_tp = obs_tp * valid_grid_2d
        
        circu = F.avg_pool2d(circu, 4) # during training
        
        pred_tp = self.PPmodel(info, circu)
        pred_tp = pred_tp * valid_grid_2d
        
        return pred_tp, obs_tp
    
    
    def common_log(self, stage, pred_tp, obs_tp, valid_grid_2d):
        self.log(f'{stage}_tp_ssim', ssim(pred_tp, obs_tp))
        self.log(f'{stage}_tp_psnr', psnr(pred_tp, obs_tp))
        self.log(f'{stage}_tp_rmse', self.valid_rmse_score(pred_tp, obs_tp, valid_grid_2d))
        self.log(f'{stage}_tp_pcc', self.valid_pcc_score(pred_tp, obs_tp, valid_grid_2d))
        self.log(f'{stage}_tp_crps', self.valid_crps_score(pred_tp, obs_tp, valid_grid_2d))
    
    
    def training_step(self, batch, batch_idx):
        pred_tp, obs_tp = self.common_step(batch)
        info, _, _ = batch
        _, _, valid_grid_2d, _, _ = info
        self.common_log('train', pred_tp, obs_tp, valid_grid_2d)
        
        tp_id_loss = self.pixelwise_score(pred_tp, obs_tp).mean()
        return tp_id_loss
    
    
    def validation_step(self, batch, batch_idx):
        pred_tp, obs_tp = self.common_step(batch)
        
        info, _, _ = batch
        _, _, valid_grid_2d, _, _ = info
        self.common_log('val', pred_tp, obs_tp, valid_grid_2d)
        
        if batch_idx == 0 and self.hparams.log_images:
            self.log_precip_map(['OBS tp', 'PRED tp'], [obs_tp, pred_tp])
    
    
    def test_step(self, batch, batch_idx):
        pred_tp, obs_tp = self.common_step(batch)
        
        info, _, _ = batch
        _, _, valid_grid_2d, _, _ = info
        self.common_log('test', pred_tp, obs_tp, valid_grid_2d)
    
    
    def configure_optimizers(self):
        trainable_params = self.PPmodel.parameters()
        if self.hparams.scheduler == 'ReduceLROnPlateau':
            opt = self.get_optimizers(trainable_params)
            sch = torch.optim.lr_scheduler.ReduceLROnPlateau(
                opt, 'min', patience=2, threshold=1e-3,
                threshold_mode='abs', min_lr=1e-5, verbose=True)
            return {
                'optimizer': opt,
                'lr_scheduler': sch,
                'monitor': 'val_crps'
            }   
        else:
            opt = self.get_optimizers(trainable_params)
            return opt



class DetermModule(BaseModule):
    '''
    model state (circu, tp) -> obs state (circu, tp)
    pairwise training with deterministic model
    
    constrain only tp here
    '''
    def __init__(self, 
                G_D2O_class: str = "Func_C2P", # [ "Func_C2P" | "Generator_D2O" ]
                G_D2O_config: DictConfig = None,
                scheduler: str = None, # [ None | 'ReduceLROnPlateau' ]
                lambda_id: Union[float, int, str] = 1,  # some fixed value or 'function'
                lead_thres: int = 3,  # control lead weighting function
                init_lambda_id: int = 1,  # control lead weighting function
                *args, **kwargs
                ):
        super(DetermModule, self).__init__(
            G_D2O_config=G_D2O_config, *args, **kwargs
        )
        self.save_hyperparameters()
        print('haprams:', self.hparams)
        
        G_D2O_config.Net_config.nz = 0  # deterministic
        if self.hparams.G_D2O_class == "Func_C2P":
            assert G_D2O_config.Net_class == "C2P_Net"
        elif self.hparams.G_D2O_class == "Generator_D2O":
            assert G_D2O_config.Net_class == "D2O_Net"
        
        G_class = globals()[G_D2O_class]
        self.Generator_D2O = G_class(**G_D2O_config)
    
    
    def forward(self, info, dyn_state):
        if not self.hparams.embed_lead:
            info = self.drop_lead_in_aux(info)
        _, _, valid_grid_2d, _, _ = info
        if self.hparams.G_D2O_class == "Func_C2P":
            D2O_tp = self.Generator_D2O(info, dyn_state) # natural deterministic
            return D2O_tp * valid_grid_2d
        elif self.hparams.G_D2O_class == "Generator_D2O":
            _, _, valid_grid_2d, _, _ = info
            D2O_circu, D2O_tp = self.Generator_D2O(info, dyn_state, None) # no z means deterministic
            return D2O_circu, D2O_tp * valid_grid_2d
    
    
    @torch.no_grad()
    def inference(self, info, dyn_state):
        return self.forward(info, dyn_state)
    
    
    def common_step(self, batch):
        info, dyn_state, obs_tp, obs_circu = batch
        _, _, valid_grid_2d, _, _ = info
        obs_tp = obs_tp * valid_grid_2d
        
        if self.hparams.G_D2O_class == "Func_C2P":
            D2O_tp = self.Generator_D2O(info, dyn_state)
            D2O_tp = D2O_tp * valid_grid_2d
            return D2O_tp, obs_tp
        elif self.hparams.G_D2O_class == "Generator_D2O":
            D2O_circu, D2O_tp = self.Generator_D2O(info, dyn_state, None) # no z means deterministic
            D2O_tp = D2O_tp * valid_grid_2d
            return D2O_circu, D2O_tp, obs_circu, obs_tp
    
    
    def common_log(self, stage, D2O_tp, obs_tp, valid_grid_2d):
        # we dont need to check circu in THIS module
        self.log(f'{stage}_tp_ssim', ssim(D2O_tp, obs_tp))
        self.log(f'{stage}_tp_psnr', psnr(D2O_tp, obs_tp))
        self.log(f'{stage}_tp_rmse', self.valid_rmse_score(D2O_tp, obs_tp, valid_grid_2d))
        self.log(f'{stage}_tp_pcc', self.valid_pcc_score(D2O_tp, obs_tp, valid_grid_2d))
        self.log(f'{stage}_tp_crps', self.valid_crps_score(D2O_tp, obs_tp, valid_grid_2d))
    
    
    def training_step(self, batch, batch_idx):
        if self.hparams.G_D2O_class == "Func_C2P":
            D2O_tp, obs_tp = self.common_step(batch)
        elif self.hparams.G_D2O_class == "Generator_D2O":
            D2O_circu, D2O_tp, obs_circu, obs_tp = self.common_step(batch)
        info, _, _, _ = batch
        _, _, valid_grid_2d, _, _ = info
        self.common_log('train', D2O_tp, obs_tp, valid_grid_2d)
        
        # ID LOSS
        tp_id_loss = self.pixelwise_score(D2O_tp, obs_tp)
        if self.hparams.lambda_id == 'function':
            info, _, _, _ = batch
            _, _, _, _, aux = info
            lead = aux[:, self.hparams.idx_of_lead]
            tp_id_loss = tp_id_loss.mean(axis=(1, 2, 3))
            tp_id_loss = L_ID_SCHEDULE(lead, thres=self.hparams.lead_thres, init_lambda=self.hparams.init_lambda_id) * tp_id_loss
        else: 
            tp_id_loss = tp_id_loss * self.hparams.lambda_id
        
        g_loss = tp_id_loss.mean()
        return g_loss
    
    
    def validation_step(self, batch, batch_idx):
        if self.hparams.G_D2O_class == "Func_C2P":
            D2O_tp, obs_tp = self.common_step(batch)
        elif self.hparams.G_D2O_class == "Generator_D2O":
            D2O_circu, D2O_tp, obs_circu, obs_tp = self.common_step(batch)
        info, _, _, _ = batch
        _, _, valid_grid_2d, _, _ = info
        self.common_log('val', D2O_tp, obs_tp, valid_grid_2d)
        
        _, dyn_state, _, _ = batch
        dyn_tp = dyn_state[:, -1, :, :].unsqueeze(1)
        dyn_tp = F.interpolate(dyn_tp, scale_factor=10)
        dyn_tp = dyn_tp * valid_grid_2d
        if batch_idx == 0 and self.hparams.log_images:
            self.log_precip_map(['DYN tp', 'OBS tp', 'PRED tp'], [dyn_tp, obs_tp, D2O_tp])
    
    
    def test_step(self, batch, batch_idx):
        if self.hparams.G_D2O_class == "Func_C2P":
            D2O_tp, obs_tp = self.common_step(batch)
        elif self.hparams.G_D2O_class == "Generator_D2O":
            D2O_circu, D2O_tp, obs_circu, obs_tp = self.common_step(batch)
        info, _, _, _ = batch
        _, _, valid_grid_2d, _, _ = info
        self.common_log('test', D2O_tp, obs_tp, valid_grid_2d)
    
    
    def configure_optimizers(self):
        trainable_params = self.Generator_D2O.parameters()
        if self.hparams.scheduler == 'ReduceLROnPlateau':
            opt_G = self.get_optimizers(trainable_params)
            sch_G = torch.optim.lr_scheduler.ReduceLROnPlateau(
                opt_G, 'min', patience=2, threshold=1e-3,
                threshold_mode='abs', min_lr=1e-5, verbose=True)
            return {
                'optimizer': opt_G,
                'lr_scheduler': sch_G,
                'monitor': 'val_tp_crps'
            }   
        else:
            opt_G = self.get_optimizers(trainable_params)
            return opt_G



class AdversarialModule_v1(BaseModule):
    '''
    model state (circu, tp) -> obs state (circu, tp)
    pairwise training with adversarial model
    adv_loss, id_loss (like pix2pix), inv_loss (constrain circu)
    
    v1: simply add adv_loss and etc
    '''
    def __init__(self,
                G_D2O_config: DictConfig = None,
                D_O_config: DictConfig = None,
                nz: int = 8,
                validation_ensemble: int = 1,
                lambda_id: Union[float, int, str] = 1,  # some fixed value or 'function'
                lambda_id_circu_proportion: Union[float, int] = 0.1, # share schedule w/ lambda_id of tp
                lead_thres: int = 3,  # control lead weighting function
                init_lambda_id: int = 10,  # control lead weighting function
                lambda_adv: Union[float, int] = 1,
                lambda_inv: Union[float, int] = 0.1, # constrain circu only
                delta_inv: float = 0,
                *args, **kwargs
                ):
        super(AdversarialModule_v1, self).__init__(
            G_D2O_config=G_D2O_config, D_O_config=D_O_config, *args, **kwargs
        )
        self.save_hyperparameters()
        print('haprams:', self.hparams)
        
        G_D2O_config.nz = nz  # itertively pass to D2O_Net
        self.G = Generator_D2O(**G_D2O_config)
        self.D = Discriminator_O(**D_O_config)
    
    
    def forward(self, info, dyn_state):
        B = dyn_state.shape[0]; L = self.hparams.nz
        z = torch.randn(B, L, 1, 1).type_as(dyn_state) if L > 0 else None
        return self.forward_Generator_D2O(self.G, info, dyn_state, z)
    
    
    @torch.no_grad()
    def inference(self, info, dyn_state):
        return self.forward(info, dyn_state)
    
    
    def forward_Generator_D2O(self, model, info, dyn_state, z):
        if not self.hparams.embed_lead:
            info = self.drop_lead_in_aux(info)
        D2O_circu, D2O_tp = model(info, dyn_state, z)
        _, _, valid_grid_2d, _, _ = info
        return D2O_circu, D2O_tp * valid_grid_2d
    
    
    def forward_Discriminator_O(self, model, info, obs_circu, obs_tp):
        if not self.hparams.embed_lead:
            info = self.drop_lead_in_aux(info)
        _, _, valid_grid_2d, _, _ = info
        obs_tp = obs_tp * valid_grid_2d
        return model(info, obs_circu, obs_tp)
    
    
    def common_log(self, stage, pred_tp, obs_tp, pred_circu, obs_circu, valid_grid_2d):
        self.log(f'{stage}_tp_ssim', ssim(pred_tp, obs_tp))
        self.log(f'{stage}_tp_psnr', psnr(pred_tp, obs_tp))
        self.log(f'{stage}_tp_rmse', self.valid_rmse_score(pred_tp, obs_tp, valid_grid_2d))
        self.log(f'{stage}_tp_pcc', self.valid_pcc_score(pred_tp, obs_tp, valid_grid_2d))
        self.log(f'{stage}_tp_crps', self.valid_crps_score(pred_tp, obs_tp, valid_grid_2d))
        # circu not on valid grid
        self.log(f'{stage}_circu_rmse', torch.sqrt(F.mse_loss(pred_circu, obs_circu)))
    
    
    def training_step(self,  batch, batch_idx, optimizer_idx):
        info, dyn_state, obs_tp, obs_circu = batch
        
        B = dyn_state.shape[0]; L = self.hparams.nz
        z = torch.randn(B, L, 1, 1).type_as(dyn_state) if L > 0 else None
        
        _, _, valid_grid_2d, _, aux = info
        obs_tp = obs_tp * valid_grid_2d
        dyn_circu = dyn_state[:, :-1, :, :]
        
        # train generator
        if optimizer_idx == 0:
            D2O_circu, D2O_tp = self.forward_Generator_D2O(self.G, info, dyn_state, z)
            
            # --- ADV LOSS ---
            fake_critic = self.forward_Discriminator_O(self.D, info, D2O_circu, D2O_tp)
            adv_loss = self.get_adversarial_loss(fake_critic, True)
            g_loss_adv = adv_loss * self.hparams.lambda_adv
            
            # --- ID LOSS ---
            # constrain tp & circu
            id_loss = self.pixelwise_score(D2O_tp, obs_tp)
            id_loss_circu = self.pixelwise_score(D2O_circu, obs_circu)
            if self.hparams.lambda_id == 'function':
                lead = aux[:, self.hparams.idx_of_lead]
                id_loss = id_loss.mean(axis=(1, 2, 3))
                id_loss = L_ID_SCHEDULE(lead, thres=self.hparams.lead_thres, init_lambda=self.hparams.init_lambda_id) * id_loss
                id_loss_circu = id_loss_circu.mean(axis=(1, 2, 3))
                id_loss_circu = L_ID_SCHEDULE(lead, thres=self.hparams.lead_thres, init_lambda=self.hparams.init_lambda_id) * id_loss_circu
            else:
                id_loss = id_loss * self.hparams.lambda_id
                id_loss_circu = id_loss_circu * self.hparams.lambda_id
            g_loss_id = id_loss.mean() + id_loss_circu.mean() * self.hparams.lambda_id_circu_proportion
            
            # --- INV LOSS ---
            # constrain pooled circu alone
            D2O_circu_pool = F.avg_pool2d(D2O_circu, 4)
            inv_loss = self.pixelwise_score(D2O_circu_pool, dyn_circu).mean()
            inv_loss = F.relu(inv_loss - self.hparams.delta_inv)
            g_loss_inv = inv_loss * self.hparams.lambda_inv
            
            g_loss = g_loss_adv + g_loss_id + g_loss_inv
            self.log('train_adv_loss', g_loss_adv)
            self.log('train_id_loss', g_loss_id)
            self.log('train_inv_loss', g_loss_inv)
            self.log('train_g_loss', g_loss, prog_bar=True)
            
            self.common_log('train', D2O_tp, obs_tp, D2O_circu, obs_circu, valid_grid_2d)
            
            return g_loss
        
        # train discriminator
        if optimizer_idx == 1:
            D2O_circu, D2O_tp = self.forward_Generator_D2O(self.G, info, dyn_state, z)
            
            real_critic = self.forward_Discriminator_O(self.D, info, obs_circu, obs_tp)
            fake_critic = self.forward_Discriminator_O(self.D, info, D2O_circu.detach(), D2O_tp.detach())
            real_loss = self.get_adversarial_loss(real_critic, True)
            fake_loss = self.get_adversarial_loss(fake_critic, False)
            
            d_loss = real_loss + fake_loss
            self.log('train_real_critic', torch.mean(real_critic))
            self.log('train_fake_critic', torch.mean(fake_critic))
            self.log('train_w_dist', torch.mean(real_critic) - torch.mean(fake_critic))
            self.log('train_d_loss', d_loss, prog_bar=True)
            
            # constraint on D 
            if self.hparams.weight_clipping or self.hparams.GAN_mode == 'WGAN':
                self.weight_clipping(self.D, self.hparams.clip_value)
            elif self.hparams.GAN_mode == 'WGAN-GP':
                if not self.hparams.embed_lead:
                    info = self.drop_lead_in_aux(info)
                gradient_penalty = self.compute_gradient_penalty(self.D, info, (obs_circu, obs_tp), (D2O_circu, D2O_tp))
                d_loss += self.hparams.lambda_gp * gradient_penalty
                self.log('loss_gp', gradient_penalty)
            
            return d_loss
    
    
    def validation_step(self, batch, batch_idx):
        info, dyn_state, obs_tp, obs_circu = batch
        _, _, valid_grid_2d, _, aux = info
        obs_tp = obs_tp * valid_grid_2d
        
        B = dyn_state.shape[0]; L = self.hparams.nz; E = self.hparams.validation_ensemble
        if E == 1:
            D2O_circu, D2O_tp = self.forward(info, dyn_state)
            self.common_log('val', D2O_tp, obs_tp, D2O_circu, obs_circu, valid_grid_2d)
            
            D2O_circu_pool = F.avg_pool2d(D2O_circu, 4)
            self.log('val_inv_rmse', torch.sqrt(F.mse_loss(D2O_circu_pool, dyn_state[:, :-1, :, :])))
        elif E > 1:
            # for each sample, generate E samples
            # z = torch.randn(E, B, L).type_as(dyn_state).reshape(-1, L)
            dyn_state = dyn_state.unsqueeze(0).expand(E, -1, -1, -1, -1).reshape(-1, *dyn_state.shape[1:])
            info = [item.unsqueeze(0).expand(E, *([-1]*len(item.shape))).reshape(-1, *item.shape[1:]) for item in info]
            D2O_circu, D2O_tp = self.forward(info, dyn_state)
            D2O_circu_pool = F.avg_pool2d(D2O_circu, 4)
            
            # prob skill for ensemble
            obs_tp_expand = obs_tp.unsqueeze(0).expand(E, -1, -1, -1, -1).reshape(-1, *obs_tp.shape[1:])
            obs_circu_expand = obs_circu.unsqueeze(0).expand(E, -1, -1, -1, -1).reshape(-1, *obs_circu.shape[1:])
            valid_grid_2d_expand = valid_grid_2d.unsqueeze(0).expand(E, -1, -1, -1, -1).reshape(-1, *valid_grid_2d.shape[1:])
            self.common_log('val', D2O_tp, obs_tp_expand, D2O_circu, obs_circu_expand, valid_grid_2d_expand)
            
            D2O_tp = D2O_tp.reshape(E, B, *D2O_tp.shape[1:])
            self.log('val_ensemble_crps', self.valid_crps_score(y_hat=D2O_tp, y=obs_tp, valid_grid_2d=valid_grid_2d))
            
            # determ skill for ensemble mean
            D2O_tp_mean = D2O_tp.mean(dim=0)
            D2O_circu_mean = D2O_circu.mean(dim=0)
            self.common_log('val_mean', D2O_tp_mean, obs_tp, D2O_circu_mean, obs_circu, valid_grid_2d)
            
            if len(dyn_state.size()) == 4:
                self.log('val_inv_rmse', torch.sqrt(F.mse_loss(D2O_circu_pool, dyn_state[:, :-1, :, :])))
            elif len(dyn_state.size()) == 5:
                self.log('val_inv_rmse', torch.sqrt(F.mse_loss(D2O_circu_pool, dyn_state[:, :, :-1, :, :])))
            
            # retrieve one ensemble member for subsequent image logging 
            D2O_tp = D2O_tp[0]
            dyn_state = dyn_state.reshape(E, B, *dyn_state.shape[1:])
            dyn_state = dyn_state[0]
        
        if batch_idx == 0 and self.hparams.log_images:
            dyn_tp = dyn_state[:, -1, :, :].unsqueeze(1)
            dyn_tp = F.interpolate(dyn_tp, scale_factor=10)
            dyn_tp = dyn_tp * valid_grid_2d
            self.log_precip_map(['DYN tp', 'OBS tp', 'PRED tp'], [dyn_tp, obs_tp, D2O_tp])
    
    
    def configure_optimizers(self):
        opt_G = self.get_optimizers(self.G.parameters())
        opt_D = self.get_optimizers(self.D.parameters())
        return (
            {'optimizer': opt_G, 'frequency': self.hparams.n_opt_G},
            {'optimizer': opt_D, 'frequency': self.hparams.n_opt_D}
        )



class AdversarialModule_v3(AdversarialModule_v2):
    '''
    use cVAE-GAN to constrain circu
    w/o O2D model
    '''
    def __init__(self,
                E_D_config: DictConfig = None,
                E_O_config: DictConfig = None,
                lambda_kl: Union[float, int] = 0.01, # kl loss between Latent_D & Latent_O
                *args, **kwargs
                ):
        super(AdversarialModule_v3, self).__init__(
            *args, **kwargs
        )
        self.save_hyperparameters()
        E_D_config.latent_dim = self.hparams.nz
        E_O_config.latent_dim = self.hparams.nz
        
        self.E_D = Encoder_D(**E_D_config)
        self.E_O = Encoder_O(**E_O_config)
    
    
    def forward(self, info, dyn_state):
        mu_D, logvar_D = self.forward_Encoder_D(self.E_D, info, dyn_state)
        z_D = self.reparameterize(mu_D, logvar_D)
        return self.forward_Generator_D2O(self.G, info, dyn_state, z_D)

    
    def forward_Encoder_O(self, model, info, obs_circu, obs_tp):
        if not self.hparams.embed_lead:
            info = self.drop_lead_in_aux(info)
        _, _, valid_grid_2d, _, _ = info
        obs_tp = obs_tp * valid_grid_2d
        return model(info, obs_circu, obs_tp)
    
    
    def forward_Encoder_D(self, model, info, dyn_state):
        if not self.hparams.embed_lead:
            info = self.drop_lead_in_aux(info)
        _, _, valid_grid_2d, _, _ = info
        coarse_valid_grid_2d = self.get_coarse_bool(valid_grid_2d)
        dyn_circu = dyn_state[:, :-1, :, :]
        dyn_tp = dyn_state[:, -1, :, :].unsqueeze(1)
        dyn_tp = dyn_tp * coarse_valid_grid_2d
        dyn_state = torch.cat([dyn_circu, dyn_tp], dim=1)
        return model(info, dyn_state)
    
    
    def configure_optimizers(self):
        opt_GE = self.get_optimizers(itertools.chain(self.G.parameters(), self.E_D.parameters(), self.E_O.parameters()))
        opt_D = self.get_optimizers(self.D.parameters())
        return (
            {'optimizer': opt_GE, 'frequency': self.hparams.n_opt_G},
            {'optimizer': opt_D, 'frequency': self.hparams.n_opt_D}
        )
    
    
    def training_step(self,  batch, batch_idx, optimizer_idx):
        info, dyn_state, obs_tp, obs_circu = batch
        _, _, valid_grid_2d, _, aux = info
        obs_tp = obs_tp * valid_grid_2d
        
        B = dyn_state.shape[0]; L = self.hparams.nz; E = self.hparams.validation_ensemble
        
        # latents: [E*B, L]
        mu_D, logvar_D = self.forward_Encoder_D(self.E_D, info, dyn_state)
        mu_O, logvar_O = self.forward_Encoder_O(self.E_O, info, obs_circu, obs_tp)
        # z_D = self.reparameterize(mu_D, logvar_D, ensemble=E)
        z_O = self.reparameterize(mu_O, logvar_O, ensemble=E)
        
        # inputs: [E*B, ...]
        dyn_state_expand = dyn_state.unsqueeze(0).expand(E, -1, -1, -1, -1).reshape(-1, *dyn_state.shape[1:])
        info_expand = [item.unsqueeze(0).expand(E, *([-1]*len(item.shape))).reshape(-1, *item.shape[1:]) for item in info]
        
        # prepare refs: [E*B, ...] for fields of obs space
        obs_tp_expand = obs_tp.unsqueeze(0).expand(E, -1, -1, -1, -1).reshape(-1, *obs_tp.shape[1:])
        obs_circu_expand = obs_circu.unsqueeze(0).expand(E, -1, -1, -1, -1).reshape(-1, *obs_circu.shape[1:])
        valid_grid_2d_expand = valid_grid_2d.unsqueeze(0).expand(E, -1, -1, -1, -1).reshape(-1, *valid_grid_2d.shape[1:])
        
        # forward
        D2O_circu_expand, D2O_tp_expand = self.forward_Generator_D2O(self.G, info_expand, dyn_state_expand, z_O)  # [E*B, ...]
        
        # prepare refs: [B, ...] w/ valid_grid for fields of dyn space
        dyn_circu = dyn_state[:, :-1, :, :]
        dyn_tp = dyn_state[:, -1, :, :].unsqueeze(1)
        coarse_valid_grid_2d = self.get_coarse_bool(valid_grid_2d)
        dyn_tp = dyn_tp * coarse_valid_grid_2d
        
        # prepare ensemble mean for fields of dyn space: [B, ...]
        D2O_tp_em = D2O_tp_expand.reshape(E, B, *D2O_tp_expand.shape[1:]).mean(dim=0)
        D2O_circu_em = D2O_circu_expand.reshape(E, B, *D2O_circu_expand.shape[1:]).mean(dim=0)
        
        # train generator & encoder
        if optimizer_idx == 0:
            # --- ADV LOSS ---
            # -- for every member --
            fake_critic_O = self.forward_Discriminator_O(self.D, info_expand, D2O_circu_expand, D2O_tp_expand)
            adv_loss = self.get_adversarial_loss(fake_critic_O, True)
            g_loss_adv = adv_loss * self.hparams.lambda_adv
            
            # --- ID LOSS ---
            # -- for ensemble mean -- 
            # constrain D2O & O2D
            id_loss = self.pixelwise_score(D2O_tp_em, obs_tp).mean(axis=(1, 2, 3))
            id_loss_circu = self.pixelwise_score(D2O_circu_em, obs_circu).mean(axis=(1, 2, 3))
            if self.hparams.lambda_id == 'function':
                lead = aux[:, self.hparams.idx_of_lead]
                id_loss = L_ID_SCHEDULE(lead, thres=self.hparams.lead_thres, init_lambda=self.hparams.init_lambda_id) * id_loss
                id_loss_circu = L_ID_SCHEDULE(lead, thres=self.hparams.lead_thres, init_lambda=self.hparams.init_lambda_id) * id_loss_circu
            else:
                id_loss = id_loss * self.hparams.lambda_id
                id_loss_circu = id_loss_circu * self.hparams.lambda_id
            g_loss_id = id_loss.mean() + id_loss_circu.mean() * self.hparams.lambda_id_circu_proportion
            
            # --- INV LOSS ---
            # -- for every member --
            # constrain circu alone, with relax limit
            D2O_circu_expand_pool = F.avg_pool2d(D2O_circu_expand, 4)
            dyn_circu_expand = dyn_state_expand[:, :-1]
            inv_loss_D2O = self.pixelwise_score(D2O_circu_expand_pool, dyn_circu_expand).mean()
            inv_loss_D2O = F.relu(inv_loss_D2O - self.hparams.delta_inv)
            g_loss_inv = inv_loss_D2O * self.hparams.lambda_inv
            
            # --- KL LOSS ---
            # -- for every member --
            kl_loss = self.kl_divergence(mu_D, logvar_D, mu_O, logvar_O)
            g_loss_kl = kl_loss * self.hparams.lambda_kl

            g_loss = g_loss_adv + g_loss_id + g_loss_inv +  g_loss_kl
            self.log('train_adv_loss', g_loss_adv)
            self.log('train_id_loss', g_loss_id)
            self.log('train_inv_loss', g_loss_inv)
            self.log('train_kl_loss', g_loss_kl)
            
            self.common_log('train', D2O_tp_expand, obs_tp_expand, D2O_circu_expand, obs_circu_expand, valid_grid_2d_expand)
            
            return g_loss
        
        # train discriminator O
        # -- fields on O space are probabilistic --
        # feed every member to D_o
        if optimizer_idx == 1:
            real_critic_O = self.forward_Discriminator_O(self.D, info, obs_circu, obs_tp)  
            fake_critic_O = self.forward_Discriminator_O(self.D, info_expand, D2O_circu_expand.detach(), D2O_tp_expand.detach())

            real_loss = self.get_adversarial_loss(real_critic_O, True)
            fake_loss = self.get_adversarial_loss(fake_critic_O, False)
            
            d_loss = real_loss + fake_loss
            
            self.log('train_real_critic', torch.mean(real_critic_O))
            self.log('train_fake_critic', torch.mean(fake_critic_O))
            self.log('train_w_dist', torch.mean(real_critic_O) - torch.mean(fake_critic_O))
            self.log('train_d_loss', d_loss, prog_bar=True)

            # constraint on D_o
            if self.hparams.weight_clipping or self.hparams.GAN_mode == 'WGAN':
                self.weight_clipping(self.D, self.hparams.clip_value)
            elif self.hparams.GAN_mode == 'WGAN-GP':
                if not self.hparams.embed_lead:
                    info_expand = self.drop_lead_in_aux(info_expand)
                gradient_penalty = self.compute_gradient_penalty(self.D, info_expand, (obs_circu_expand, obs_tp_expand), (D2O_circu_expand, D2O_tp_expand))
                d_loss += self.hparams.lambda_gp * gradient_penalty
                self.log('loss_gp_O', gradient_penalty)

            return d_loss
    
    
    def validation_step(self, batch, batch_idx):
        info, dyn_state, obs_tp, obs_circu = batch
        _, _, valid_grid_2d, _, _ = info
        obs_tp = obs_tp * valid_grid_2d
        
        # --- log sample for [ single ensemble member ] ---
        mu_D, logvar_D = self.forward_Encoder_D(self.E_D, info, dyn_state)
        z_D = self.reparameterize(mu_D, logvar_D)
        mu_O, logvar_O = self.forward_Encoder_O(self.E_O, info, obs_circu, obs_tp)
        z_O = self.reparameterize(mu_O, logvar_O)
        
        # Dyn -> Obs
        D2O_circu, D2O_tp = self.forward_Generator_D2O(self.G, info, dyn_state, z_O)
        _, inference_tp = self.forward_Generator_D2O(self.G, info, dyn_state, z_D)
        
        if batch_idx == 0 and self.hparams.log_images:
            dyn_tp = dyn_state[:, -1, :, :].unsqueeze(1)
            dyn_tp = F.interpolate(dyn_tp, scale_factor=10)
            dyn_tp = dyn_tp * valid_grid_2d
            self.log_precip_map(
                ['dyn tp', 'd2o tp', 'inference tp', 'obs tp'], 
                [dyn_tp, D2O_tp, inference_tp, obs_tp])

        # --- log score for [ ensemble members ] ---
        B = dyn_state.shape[0]; L = self.hparams.nz; E = self.hparams.validation_ensemble
        if E == 1:
            self.common_log('val', D2O_tp, obs_tp, D2O_circu, obs_circu, valid_grid_2d)
            
            D2O_circu_pool = F.avg_pool2d(D2O_circu, 4)
            self.log('val_inv_rmse', torch.sqrt(F.mse_loss(D2O_circu_pool, dyn_state[:, :-1, :, :])))
        elif E > 1:
            z_D = self.reparameterize(mu_D, logvar_D, ensemble=E) # [E*B, L]
            dyn_state_expand = dyn_state.unsqueeze(0).expand(E, -1, -1, -1, -1).reshape(-1, *dyn_state.shape[1:])
            info_expand = [item.unsqueeze(0).expand(E, *([-1]*len(item.shape))).reshape(-1, *item.shape[1:]) for item in info]
            D2O_circu_expand, D2O_tp_expand = self.forward_Generator_D2O(self.G, info_expand, dyn_state_expand, z_D)
            
            # prob skill for ensemble
            obs_tp_expand = obs_tp.unsqueeze(0).expand(E, -1, -1, -1, -1).reshape(-1, *obs_tp.shape[1:])
            obs_circu_expand = obs_circu.unsqueeze(0).expand(E, -1, -1, -1, -1).reshape(-1, *obs_circu.shape[1:])
            valid_grid_2d_expand = valid_grid_2d.unsqueeze(0).expand(E, -1, -1, -1, -1).reshape(-1, *valid_grid_2d.shape[1:])
            self.common_log('val', D2O_tp_expand, obs_tp_expand, D2O_circu_expand, obs_circu_expand, valid_grid_2d_expand)
            
            D2O_tp_ensemble = D2O_tp_expand.reshape(E, B, *inference_tp.shape[1:])
            D2O_circu_ensemble = D2O_circu_expand.reshape(E, B, *D2O_circu.shape[1:])
            self.log('val_ensemble_crps', self.valid_crps_score(y_hat=D2O_tp_ensemble, y=obs_tp, valid_grid_2d=valid_grid_2d))
            
            # determ skill for ensemble mean
            D2O_tp_em = D2O_tp_ensemble.mean(dim=0)
            D2O_circu_em = D2O_circu_ensemble.mean(dim=0)
            self.common_log('val_mean', D2O_tp_em, obs_tp, D2O_circu_em, obs_circu, valid_grid_2d)
            
            # inv skill for every member
            D2O_circu_pool_expand = F.avg_pool2d(D2O_circu_expand, 4)
            if len(dyn_state.size()) == 4:
                self.log('val_inv_rmse', torch.sqrt(F.mse_loss(D2O_circu_pool_expand, dyn_state_expand[:, :-1, :, :])))
            elif len(dyn_state.size()) == 5:
                self.log('val_inv_rmse', torch.sqrt(F.mse_loss(D2O_circu_pool_expand, dyn_state_expand[:, :, :-1, :, :])))



class CycleGanModule(BaseModule):
    '''
    w/ 2 generators and 2 discriminators
    ADV loss + ID loss + Cycle loss + INV loss
    ID loss constrain circu & tp
    INV loss constrain circu only, with fixed relax limit
    D2O & O2D models are deterministic
    '''
    def __init__(self,
                G_D2O_config: DictConfig = None,
                D_O_config: DictConfig = None,
                G_O2D_config: DictConfig = None,
                D_D_config: DictConfig = None,
                lambda_id: Union[float, int, str] = 1,  # some fixed value or 'function'
                lambda_id_circu_proportion: Union[float, int] = 1, # share schedule w/ lambda_id of tp
                lead_thres: int = 3,  # control lead weighting function
                init_lambda_id: int = 10,  # control lead weighting function
                lambda_adv: Union[float, int] = 1,
                lambda_cycle: Union[float, int] = 10,
                lambda_inv: Union[float, int] = 0.1, # constrain circu only
                delta_inv: float = 0.01,
                n_opt_Dd: int = None, 
                n_opt_Do: int = None, 
                lambda_gp_d: float = None,
                lambda_gp_o: float = None,
                *args, **kwargs
                ):
        super(CycleGanModule, self).__init__(
            G_D2O_config=G_D2O_config, D_O_config=D_O_config, 
            G_O2D_config=G_O2D_config, D_D_config=D_D_config,
            *args, **kwargs
        )
        self.save_hyperparameters()
        print('haprams:', self.hparams)
        
        if n_opt_Dd is None:
            self.hparams.n_opt_Dd = self.hparams.n_opt_D
        if n_opt_Do is None:
            self.hparams.n_opt_Do = self.hparams.n_opt_G
        if lambda_gp_d is None:
            self.hparams.lambda_gp_d = self.hparams.lambda_gp
        if lambda_gp_o is None:
            self.hparams.lambda_gp_o = self.hparams.lambda_gp
        
        G_D2O_config.nz = 0  # XXX: currently make everything deterministic
        self.G_D2O = Generator_D2O(**G_D2O_config)
        self.D_O = Discriminator_O(**D_O_config)
        self.G_O2D = Generator_O2D(**G_O2D_config)
        self.D_D = Discriminator_D(**D_D_config)
    
    
    def forward(self, info, dyn_state):
        # because deterministic, no need for z
        return self.forward_Generator_D2O(self.G_D2O, info, dyn_state, None)
    
    
    @torch.no_grad()
    def inference(self, info, dyn_state):
        return self.forward(info, dyn_state)
    
    
    def forward_Generator_D2O(self, model, info, dyn_state, z):
        if not self.hparams.embed_lead:
            info = self.drop_lead_in_aux(info)
        D2O_circu, D2O_tp = model(info, dyn_state, z)
        _, _, valid_grid_2d, _, _ = info
        return D2O_circu, D2O_tp * valid_grid_2d
    
    
    def forward_Discriminator_O(self, model, info, obs_circu, obs_tp):
        if not self.hparams.embed_lead:
            info = self.drop_lead_in_aux(info)
        _, _, valid_grid_2d, _, _ = info
        obs_tp = obs_tp * valid_grid_2d
        return model(info, obs_circu, obs_tp)
    
    
    def forward_Generator_O2D(self, model, info, obs_circu, obs_tp):
        if not self.hparams.embed_lead:
            info = self.drop_lead_in_aux(info)
        _, _, valid_grid_2d, _, _ = info
        coarse_valid_grid_2d = self.get_coarse_bool(valid_grid_2d)
        dyn_state = model(info, obs_circu, obs_tp)
        dyn_circu = dyn_state[:, :-1, :, :]
        dyn_tp = dyn_state[:, -1, :, :].unsqueeze(1)
        dyn_tp = dyn_tp * coarse_valid_grid_2d
        return torch.cat([dyn_circu, dyn_tp], dim=1)
    
    
    def forward_Discriminator_D(self, model, info, dyn_state):
        if not self.hparams.embed_lead:
            info = self.drop_lead_in_aux(info)
        _, _, valid_grid_2d, _, _ = info
        coarse_valid_grid_2d = self.get_coarse_bool(valid_grid_2d)
        dyn_circu = dyn_state[:, :-1, :, :]
        dyn_tp = dyn_state[:, -1, :, :].unsqueeze(1)
        dyn_tp = dyn_tp * coarse_valid_grid_2d
        dyn_state = torch.cat([dyn_circu, dyn_tp], dim=1)
        return model(info, dyn_state)
    
    
    def common_log(self, stage, pred_tp, obs_tp, pred_circu, obs_circu, valid_grid_2d):
        self.log(f'{stage}_tp_ssim', ssim(pred_tp, obs_tp))
        self.log(f'{stage}_tp_psnr', psnr(pred_tp, obs_tp))
        self.log(f'{stage}_tp_rmse', self.valid_rmse_score(pred_tp, obs_tp, valid_grid_2d))
        self.log(f'{stage}_tp_pcc', self.valid_pcc_score(pred_tp, obs_tp, valid_grid_2d))
        self.log(f'{stage}_tp_crps', self.valid_crps_score(pred_tp, obs_tp, valid_grid_2d))
        # circu not on valid grid
        self.log(f'{stage}_circu_rmse', torch.sqrt(F.mse_loss(pred_circu, obs_circu)))
    
    
    def training_step(self,  batch, batch_idx, optimizer_idx):
        info, dyn_state, obs_tp, obs_circu = batch
        _, _, valid_grid_2d, _, aux = info
        obs_tp = obs_tp * valid_grid_2d
        z = None
        
        D2O_circu, D2O_tp = self.forward_Generator_D2O(self.G_D2O, info, dyn_state, z)
        O2D_state = self.forward_Generator_O2D(self.G_O2D, info, obs_circu, obs_tp)
        
        O2D_circu = O2D_state[:, :-1, :, :]
        O2D_tp = O2D_state[:, -1, :, :].unsqueeze(1)
        dyn_circu = dyn_state[:, :-1, :, :]
        dyn_tp = dyn_state[:, -1, :, :].unsqueeze(1)
        coarse_valid_grid_2d = self.get_coarse_bool(valid_grid_2d)
        dyn_tp = dyn_tp * coarse_valid_grid_2d
        
        # train generator against D_d
        if optimizer_idx == 0:
            # --- ADV LOSS ---
            fake_critic_D = self.forward_Discriminator_D(self.D_D, info, O2D_state)
            adv_loss = self.get_adversarial_loss(fake_critic_D, True)
            g_loss_adv = adv_loss * self.hparams.lambda_adv

        # train generator against D_o
        if optimizer_idx == 2:
            # --- ADV LOSS ---
            fake_critic_O = self.forward_Discriminator_O(self.D_O, info, D2O_circu, D2O_tp)
            adv_loss = self.get_adversarial_loss(fake_critic_O, True)
            g_loss_adv = adv_loss * self.hparams.lambda_adv
        
        # train generator
        if optimizer_idx == 0 or optimizer_idx == 2:
            # --- ID LOSS ---
            # constrain tp & circu
            # constrain D2O & O2D
            id_loss = self.pixelwise_score(D2O_tp, obs_tp).mean(axis=(1, 2, 3))\
                        + self.pixelwise_score(O2D_tp, dyn_tp).mean(axis=(1, 2, 3))
            id_loss_circu = self.pixelwise_score(D2O_circu, obs_circu).mean(axis=(1, 2, 3))\
                        + self.pixelwise_score(O2D_circu, dyn_circu).mean(axis=(1, 2, 3))
            if self.hparams.lambda_id == 'function':
                lead = aux[:, self.hparams.idx_of_lead]
                id_loss = L_ID_SCHEDULE(lead, thres=self.hparams.lead_thres, init_lambda=self.hparams.init_lambda_id) * id_loss
                id_loss_circu = L_ID_SCHEDULE(lead, thres=self.hparams.lead_thres, init_lambda=self.hparams.init_lambda_id) * id_loss_circu
            else:
                id_loss = id_loss * self.hparams.lambda_id
                id_loss_circu = id_loss_circu * self.hparams.lambda_id
            g_loss_id = id_loss.mean() + id_loss_circu.mean() * self.hparams.lambda_id_circu_proportion
            
            # --- CYLCE LOSS ---
            dod_state = self.forward_Generator_O2D(self.G_O2D, info, D2O_circu, D2O_tp)
            odo_circu, odo_tp = self.forward_Generator_D2O(self.G_D2O, info, O2D_state, z)
            cycle_loss = self.pixelwise_score(odo_tp, obs_tp).mean() + self.pixelwise_score(odo_circu, obs_circu).mean() + self.pixelwise_score(dod_state, dyn_state).mean()
            g_loss_cycle = cycle_loss * self.hparams.lambda_cycle
            
            # --- INV LOSS ---
            # constrain circu alone, with relax limit
            obs_circu_pool = F.avg_pool2d(obs_circu, 4)
            D2O_circu_pool = F.avg_pool2d(D2O_circu, 4)
            inv_loss_D2O = self.pixelwise_score(D2O_circu_pool, dyn_circu).mean()
            inv_loss_D2O = F.relu(inv_loss_D2O - self.hparams.delta_inv)
            inv_loss_O2D = self.pixelwise_score(O2D_state[:,:-1], obs_circu_pool).mean()  # drop tp
            inv_loss_O2D = F.relu(inv_loss_O2D - self.hparams.delta_inv)
            g_loss_inv = (inv_loss_D2O + inv_loss_O2D) * self.hparams.lambda_inv
            
            g_loss = g_loss_adv + g_loss_id + g_loss_inv + g_loss_cycle
            self.log('train_adv_loss', g_loss_adv)
            self.log('train_id_loss', g_loss_id)
            self.log('train_inv_loss', g_loss_inv)
            self.log('train_cycle_loss', g_loss_cycle)
            
            self.common_log('train', D2O_tp, obs_tp, D2O_circu, obs_circu, valid_grid_2d)
            
            return g_loss
        
        # train discriminator D
        if optimizer_idx == 1:
            real_critic_D = self.forward_Discriminator_D(self.D_D, info, dyn_state)
            fake_critic_D = self.forward_Discriminator_D(self.D_D, info, O2D_state.detach())
            
            real_loss = self.get_adversarial_loss(real_critic_D, True)
            fake_loss = self.get_adversarial_loss(fake_critic_D, False)
            
            d_d_loss = real_loss + fake_loss

            self.log('train_real_critic_D', torch.mean(real_critic_D))
            self.log('train_fake_critic_D', torch.mean(fake_critic_D))
            self.log('train_w_dist_D', torch.mean(real_critic_D) - torch.mean(fake_critic_D))
            
            # constraint on D_d
            if self.hparams.weight_clipping or self.hparams.GAN_mode == 'WGAN':
                self.weight_clipping(self.D_D, self.hparams.clip_value)
            elif self.hparams.GAN_mode == 'WGAN-GP':
                if not self.hparams.embed_lead:
                    info = self.drop_lead_in_aux(info)
                gradient_penalty_D = self.compute_gradient_penalty(self.D_D, info, (dyn_state), (O2D_state))
                d_d_loss += self.hparams.lambda_gp_d * gradient_penalty_D
                self.log('loss_gp_D', gradient_penalty_D)
            
            return d_d_loss
        
        # train discriminator O
        if optimizer_idx == 3:
            real_critic_O = self.forward_Discriminator_O(self.D_O, info, obs_circu, obs_tp)
            fake_critic_O = self.forward_Discriminator_O(self.D_O, info, D2O_circu.detach(), D2O_tp.detach())

            real_loss = self.get_adversarial_loss(real_critic_O, True)
            fake_loss = self.get_adversarial_loss(fake_critic_O, False)
            
            d_o_loss = real_loss + fake_loss
            
            self.log('train_real_critic_O', torch.mean(real_critic_O))
            self.log('train_fake_critic_O', torch.mean(fake_critic_O))
            self.log('train_w_dist_O', torch.mean(real_critic_O) - torch.mean(fake_critic_O))

            # constraint on D_o
            if self.hparams.weight_clipping or self.hparams.GAN_mode == 'WGAN':
                self.weight_clipping(self.D_O, self.hparams.clip_value)
            elif self.hparams.GAN_mode == 'WGAN-GP':
                if not self.hparams.embed_lead:
                    info = self.drop_lead_in_aux(info)
                gradient_penalty_O = self.compute_gradient_penalty(self.D_O, info, (obs_circu, obs_tp), (D2O_circu, D2O_tp))
                d_o_loss += self.hparams.lambda_gp_o * gradient_penalty_O
                self.log('loss_gp_O', gradient_penalty_O)

            return d_o_loss
    
    
    def validation_step(self, batch, batch_idx):
        info, dyn_state, obs_tp, obs_circu = batch
        _, _, valid_grid_2d, _, _ = info
        obs_tp = obs_tp * valid_grid_2d
        
        # Dyn -> Obs -> Dyn
        D2O_circu, D2O_tp = self.forward(info, dyn_state) # D2O
        dod_state = self.forward_Generator_O2D(self.G_O2D, info, D2O_circu, D2O_tp)
        
        # Obs -> Dyn -> Obs
        O2D_state = self.forward_Generator_O2D(self.G_O2D, info, obs_circu, obs_tp)
        _, odo_tp = self.forward_Generator_D2O(self.G_D2O, info, O2D_state, None)
        
        # log only inference scores
        self.common_log('val', D2O_tp, obs_tp, D2O_circu, obs_circu, valid_grid_2d)
        
        D2O_circu_pool = F.avg_pool2d(D2O_circu, 4)
        self.log('val_inv_rmse', torch.sqrt(F.mse_loss(D2O_circu_pool, dyn_state[:, :-1, :, :])))
        
        if batch_idx == 0 and self.hparams.log_images:
            dyn_tp = dyn_state[:, -1, :, :].unsqueeze(1)
            dyn_tp = F.interpolate(dyn_tp, scale_factor=10)
            dyn_tp = dyn_tp * valid_grid_2d
            dod_tp = dod_state[:, -1, :, :].unsqueeze(1)
            dod_tp = F.interpolate(dod_tp, scale_factor=10)
            dod_tp = dod_tp * valid_grid_2d
            odo_tp = odo_tp * valid_grid_2d
            O2D_tp = O2D_state[:, -1, :, :].unsqueeze(1)
            O2D_tp = F.interpolate(O2D_tp, scale_factor=10)
            O2D_tp = O2D_tp * valid_grid_2d
            self.log_precip_map(
                ['dyn tp', 'd2o tp', 'dod tp', 'obs tp', 'o2d tp', 'odo tp'], 
                [dyn_tp, D2O_tp, dod_tp, obs_tp, O2D_tp, odo_tp])
    
    
    def configure_optimizers(self):
        opt_Gd = self.get_optimizers(self.G_O2D.parameters())
        opt_Go = self.get_optimizers(self.G_D2O.parameters())
        opt_Dd = self.get_optimizers(self.D_D.parameters())
        opt_Do = self.get_optimizers(self.D_O.parameters())
        return (
            {'optimizer': opt_Gd, 'frequency': self.hparams.n_opt_G},
            {'optimizer': opt_Dd, 'frequency': self.hparams.n_opt_Dd},
            {'optimizer': opt_Go, 'frequency': self.hparams.n_opt_G},
            {'optimizer': opt_Do, 'frequency': self.hparams.n_opt_Do}
        )



class CvaeGanModule(BaseModule):
    '''
    cycleGAN module + uncertainty
    D2O is probabilistic, O2D is deterministic
    E_D & E_O both generate latent input to G_D2O only, not G_O2D
    
    id loss, adv loss, cycle loss, inv loss
    + kl loss between Latent_D & Latent_O

    '''
    def __init__(self,
                G_D2O_config: DictConfig = None,
                D_O_config: DictConfig = None,
                G_O2D_config: DictConfig = None,
                D_D_config: DictConfig = None,
                E_D_config: DictConfig = None,
                E_O_config: DictConfig = None,
                nz: int = 8,
                validation_ensemble: int = 1,
                n_opt_GE: int = 1, # frequency of optimizer step for (G & E) together
                lambda_id: Union[float, int, str] = 1,  # some fixed value or 'function'
                lambda_id_circu_proportion: Union[float, int] = 0.1, # share schedule w/ lambda_id of tp
                lead_thres: int = 3,  # control lead weighting function
                init_lambda_id: int = 10,  # control lead weighting function
                lambda_adv: Union[float, int] = 1,
                lambda_cycle: Union[float, int] = 10,
                lambda_inv: Union[float, int] = 0.1, # constrain circu only
                delta_inv: float = 0.1, # relax limit for inv loss
                lambda_kl: Union[float, int] = 0.01, # kl loss between Latent_D & Latent_O
                n_opt_Dd: int = None, 
                n_opt_Do: int = None, 
                lambda_gp_d: float = None,
                lambda_gp_o: float = None,
                *args, **kwargs
                ):
        super(CvaeGanModule, self).__init__(
            G_D2O_config=G_D2O_config, D_O_config=D_O_config, 
            G_O2D_config=G_O2D_config, D_D_config=D_D_config,
            *args, **kwargs
        )
        self.save_hyperparameters()
        print('haprams:', self.hparams)
        
        if n_opt_Dd is None:
            self.hparams.n_opt_Dd = self.hparams.n_opt_D
        if n_opt_Do is None:
            self.hparams.n_opt_Do = self.hparams.n_opt_G
        if lambda_gp_d is None:
            self.hparams.lambda_gp_d = self.hparams.lambda_gp
        if lambda_gp_o is None:
            self.hparams.lambda_gp_o = self.hparams.lambda_gp
        
        assert nz > 0
        # itertively pass to each Net
        G_D2O_config.nz = nz  
        E_D_config.latent_dim = nz
        E_O_config.latent_dim = nz
        
        self.G_D2O = Generator_D2O(**G_D2O_config)
        self.D_O = Discriminator_O(**D_O_config)
        self.G_O2D = Generator_O2D(**G_O2D_config)
        self.D_D = Discriminator_D(**D_D_config)
        self.E_D = Encoder_D(**E_D_config)
        self.E_O = Encoder_O(**E_O_config)
    
    
    def forward(self, info, dyn_state):
        mu_D, logvar_D = self.forward_Encoder_D(self.E_D, info, dyn_state)
        z_D = self.reparameterize(mu_D, logvar_D)
        return self.forward_Generator_D2O(self.G_D2O, info, dyn_state, z_D)
    
    
    @torch.no_grad()
    def inference(self, info, dyn_state):
        return self.forward(info, dyn_state)
    
    
    def forward_Generator_D2O(self, model, info, dyn_state, z):
        if not self.hparams.embed_lead:
            info = self.drop_lead_in_aux(info)
        D2O_circu, D2O_tp = model(info, dyn_state, z)
        _, _, valid_grid_2d, _, _ = info
        return D2O_circu, D2O_tp * valid_grid_2d
    
    
    def forward_Discriminator_O(self, model, info, obs_circu, obs_tp):
        if not self.hparams.embed_lead:
            info = self.drop_lead_in_aux(info)
        _, _, valid_grid_2d, _, _ = info
        obs_tp = obs_tp * valid_grid_2d
        return model(info, obs_circu, obs_tp)
    
    
    def forward_Encoder_O(self, model, info, obs_circu, obs_tp):
        if not self.hparams.embed_lead:
            info = self.drop_lead_in_aux(info)
        _, _, valid_grid_2d, _, _ = info
        obs_tp = obs_tp * valid_grid_2d
        return model(info, obs_circu, obs_tp)
    
    
    def forward_Generator_O2D(self, model, info, obs_circu, obs_tp):
        if not self.hparams.embed_lead:
            info = self.drop_lead_in_aux(info)
        _, _, valid_grid_2d, _, _ = info
        coarse_valid_grid_2d = self.get_coarse_bool(valid_grid_2d)
        dyn_state = model(info, obs_circu, obs_tp)
        dyn_circu = dyn_state[:, :-1, :, :]
        dyn_tp = dyn_state[:, -1, :, :].unsqueeze(1)
        dyn_tp = dyn_tp * coarse_valid_grid_2d
        return torch.cat([dyn_circu, dyn_tp], dim=1)
    
    
    def forward_Discriminator_D(self, model, info, dyn_state):
        if not self.hparams.embed_lead:
            info = self.drop_lead_in_aux(info)
        _, _, valid_grid_2d, _, _ = info
        coarse_valid_grid_2d = self.get_coarse_bool(valid_grid_2d)
        dyn_circu = dyn_state[:, :-1, :, :]
        dyn_tp = dyn_state[:, -1, :, :].unsqueeze(1)
        dyn_tp = dyn_tp * coarse_valid_grid_2d
        dyn_state = torch.cat([dyn_circu, dyn_tp], dim=1)
        return model(info, dyn_state)
    
    
    def forward_Encoder_D(self, model, info, dyn_state):
        if not self.hparams.embed_lead:
            info = self.drop_lead_in_aux(info)
        _, _, valid_grid_2d, _, _ = info
        coarse_valid_grid_2d = self.get_coarse_bool(valid_grid_2d)
        dyn_circu = dyn_state[:, :-1, :, :]
        dyn_tp = dyn_state[:, -1, :, :].unsqueeze(1)
        dyn_tp = dyn_tp * coarse_valid_grid_2d
        dyn_state = torch.cat([dyn_circu, dyn_tp], dim=1)
        return model(info, dyn_state)
    
    
    def common_log(self, stage, pred_tp, obs_tp, pred_circu, obs_circu, valid_grid_2d):
        self.log(f'{stage}_tp_ssim', ssim(pred_tp, obs_tp))
        self.log(f'{stage}_tp_psnr', psnr(pred_tp, obs_tp))
        self.log(f'{stage}_tp_rmse', self.valid_rmse_score(pred_tp, obs_tp, valid_grid_2d))
        self.log(f'{stage}_tp_pcc', self.valid_pcc_score(pred_tp, obs_tp, valid_grid_2d))
        self.log(f'{stage}_tp_crps', self.valid_crps_score(pred_tp, obs_tp, valid_grid_2d))
        # circu not on valid grid
        self.log(f'{stage}_circu_rmse', torch.sqrt(F.mse_loss(pred_circu, obs_circu)))
    
    
    def training_step(self,  batch, batch_idx, optimizer_idx):
        info, dyn_state, obs_tp, obs_circu = batch
        
        mu_D, logvar_D = self.forward_Encoder_D(self.E_D, info, dyn_state)
        # z_D = self.reparameterize(mu_D, logvar_D)
        mu_O, logvar_O = self.forward_Encoder_O(self.E_O, info, obs_circu, obs_tp)
        z_O = self.reparameterize(mu_O, logvar_O)
        
        D2O_circu, D2O_tp = self.forward_Generator_D2O(self.G_D2O, info, dyn_state, z_O)  # training stage: encode obs latent
        O2D_state = self.forward_Generator_O2D(self.G_O2D, info, obs_circu, obs_tp)
        
        _, _, valid_grid_2d, _, aux = info
        obs_tp = obs_tp * valid_grid_2d
        O2D_circu = O2D_state[:, :-1, :, :]
        O2D_tp = O2D_state[:, -1, :, :].unsqueeze(1)
        dyn_circu = dyn_state[:, :-1, :, :]
        dyn_tp = dyn_state[:, -1, :, :].unsqueeze(1)
        coarse_valid_grid_2d = self.get_coarse_bool(valid_grid_2d)
        dyn_tp = dyn_tp * coarse_valid_grid_2d
        
        # train generator against D_d
        if optimizer_idx == 0:
            # --- ADV LOSS ---
            fake_critic_D = self.forward_Discriminator_D(self.D_D, info, O2D_state)
            adv_loss = self.get_adversarial_loss(fake_critic_D, True)
            g_loss_adv = adv_loss * self.hparams.lambda_adv

        # train generator against D_o
        if optimizer_idx == 2:
            # --- ADV LOSS ---
            fake_critic_O = self.forward_Discriminator_O(self.D_O, info, D2O_circu, D2O_tp)
            adv_loss = self.get_adversarial_loss(fake_critic_O, True)
            g_loss_adv = adv_loss * self.hparams.lambda_adv
        
        # train generator $ encoder
        if optimizer_idx == 0 or optimizer_idx == 2:
            # --- ID LOSS ---
            # constrain D2O & O2D
            id_loss = self.pixelwise_score(D2O_tp, obs_tp).mean(axis=(1, 2, 3))\
                        + self.pixelwise_score(O2D_tp, dyn_tp).mean(axis=(1, 2, 3))
            id_loss_circu = self.pixelwise_score(D2O_circu, obs_circu).mean(axis=(1, 2, 3))\
                        + self.pixelwise_score(O2D_circu, dyn_circu).mean(axis=(1, 2, 3))
            if self.hparams.lambda_id == 'function':
                lead = aux[:, self.hparams.idx_of_lead]
                id_loss = L_ID_SCHEDULE(lead, thres=self.hparams.lead_thres, init_lambda=self.hparams.init_lambda_id) * id_loss
                id_loss_circu = L_ID_SCHEDULE(lead, thres=self.hparams.lead_thres, init_lambda=self.hparams.init_lambda_id) * id_loss_circu
            else:
                id_loss = id_loss * self.hparams.lambda_id
                id_loss_circu = id_loss_circu * self.hparams.lambda_id
            g_loss_id = id_loss.mean() + id_loss_circu.mean() * self.hparams.lambda_id_circu_proportion
            
            # --- CYLCE LOSS ---
            dod_state = self.forward_Generator_O2D(self.G_O2D, info, D2O_circu, D2O_tp)
            odo_circu, odo_tp = self.forward_Generator_D2O(self.G_D2O, info, O2D_state, z_O)
            cycle_loss = self.pixelwise_score(odo_tp, obs_tp).mean() + self.pixelwise_score(odo_circu, obs_circu).mean() + self.pixelwise_score(dod_state, dyn_state).mean()
            # TODO: try pooled cycle loss
            g_loss_cycle = cycle_loss * self.hparams.lambda_cycle
            
            # --- INV LOSS ---
            # constrain circu alone, with relax limit
            obs_circu_pool = F.avg_pool2d(obs_circu, 4)
            D2O_circu_pool = F.avg_pool2d(D2O_circu, 4)
            inv_loss_D2O = self.pixelwise_score(D2O_circu_pool, dyn_circu).mean()
            inv_loss_D2O = F.relu(inv_loss_D2O - self.hparams.delta_inv)
            inv_loss_O2D = self.pixelwise_score(O2D_state[:, :-1], obs_circu_pool).mean()  # drop tp
            inv_loss_O2D = F.relu(inv_loss_O2D - self.hparams.delta_inv)
            g_loss_inv = (inv_loss_D2O + inv_loss_O2D) * self.hparams.lambda_inv
            
            # --- KL LOSS ---
            kl_loss = self.kl_divergence(mu_D, logvar_D, mu_O, logvar_O)
            g_loss_kl = kl_loss * self.hparams.lambda_kl

            g_loss = g_loss_adv + g_loss_id + g_loss_inv + g_loss_cycle + g_loss_kl
            self.log('train_adv_loss', g_loss_adv)
            self.log('train_id_loss', g_loss_id)
            self.log('train_inv_loss', g_loss_inv)
            self.log('train_cycle_loss', g_loss_cycle)
            self.log('train_kl_loss', g_loss_kl)
            
            self.common_log('train', D2O_tp, obs_tp, D2O_circu, obs_circu, valid_grid_2d)
            
            return g_loss
        
        # train discriminator D
        if optimizer_idx == 1:
            real_critic_D = self.forward_Discriminator_D(self.D_D, info, dyn_state)
            fake_critic_D = self.forward_Discriminator_D(self.D_D, info, O2D_state.detach())
            
            real_loss = self.get_adversarial_loss(real_critic_D, True)
            fake_loss = self.get_adversarial_loss(fake_critic_D, False)
            
            d_d_loss = real_loss + fake_loss

            self.log('train_real_critic_D', torch.mean(real_critic_D))
            self.log('train_fake_critic_D', torch.mean(fake_critic_D))
            self.log('train_w_dist_D', torch.mean(real_critic_D) - torch.mean(fake_critic_D))
            
            # constraint on D_d
            if self.hparams.weight_clipping or self.hparams.GAN_mode == 'WGAN':
                self.weight_clipping(self.D_D, self.hparams.clip_value)
            elif self.hparams.GAN_mode == 'WGAN-GP':
                if not self.hparams.embed_lead:
                    info = self.drop_lead_in_aux(info)
                gradient_penalty_D = self.compute_gradient_penalty(self.D_D, info, (dyn_state), (O2D_state))
                d_d_loss += self.hparams.lambda_gp_d * gradient_penalty_D
                self.log('loss_gp_D', gradient_penalty_D)
            
            return d_d_loss
        
        # train discriminator O
        if optimizer_idx == 3:
            real_critic_O = self.forward_Discriminator_O(self.D_O, info, obs_circu, obs_tp)
            fake_critic_O = self.forward_Discriminator_O(self.D_O, info, D2O_circu.detach(), D2O_tp.detach())

            real_loss = self.get_adversarial_loss(real_critic_O, True)
            fake_loss = self.get_adversarial_loss(fake_critic_O, False)
            
            d_o_loss = real_loss + fake_loss
            
            self.log('train_real_critic_O', torch.mean(real_critic_O))
            self.log('train_fake_critic_O', torch.mean(fake_critic_O))
            self.log('train_w_dist_O', torch.mean(real_critic_O) - torch.mean(fake_critic_O))

            # constraint on D_o
            if self.hparams.weight_clipping or self.hparams.GAN_mode == 'WGAN':
                self.weight_clipping(self.D_O, self.hparams.clip_value)
            elif self.hparams.GAN_mode == 'WGAN-GP':
                if not self.hparams.embed_lead:
                    info = self.drop_lead_in_aux(info)
                gradient_penalty_O = self.compute_gradient_penalty(self.D_O, info, (obs_circu, obs_tp), (D2O_circu, D2O_tp))
                d_o_loss += self.hparams.lambda_gp_o * gradient_penalty_O
                self.log('loss_gp_O', gradient_penalty_O)

            return d_o_loss
    
    
    def validation_step(self, batch, batch_idx):
        info, dyn_state, obs_tp, obs_circu = batch
        _, _, valid_grid_2d, _, _ = info
        obs_tp = obs_tp * valid_grid_2d
        
        # --- log sample for [ single ensemble member ] ---
        mu_D, logvar_D = self.forward_Encoder_D(self.E_D, info, dyn_state)
        z_D = self.reparameterize(mu_D, logvar_D)
        mu_O, logvar_O = self.forward_Encoder_O(self.E_O, info, obs_circu, obs_tp)
        z_O = self.reparameterize(mu_O, logvar_O)
        
        # Dyn -> Obs -> Dyn
        D2O_circu, D2O_tp = self.forward_Generator_D2O(self.G_D2O, info, dyn_state, z_O)
        dod_state = self.forward_Generator_O2D(self.G_O2D, info, D2O_circu, D2O_tp)
        _, inference_tp = self.forward_Generator_D2O(self.G_D2O, info, dyn_state, z_D)
        
        # Obs -> Dyn -> Obs
        O2D_state = self.forward_Generator_O2D(self.G_O2D, info, obs_circu, obs_tp)
        _, odo_tp = self.forward_Generator_D2O(self.G_D2O, info, O2D_state, z_O)
        
        if batch_idx == 0 and self.hparams.log_images:
            dyn_tp = dyn_state[:, -1, :, :].unsqueeze(1)
            dyn_tp = F.interpolate(dyn_tp, scale_factor=10)
            dyn_tp = dyn_tp * valid_grid_2d
            dod_tp = dod_state[:, -1, :, :].unsqueeze(1)
            dod_tp = F.interpolate(dod_tp, scale_factor=10)
            dod_tp = dod_tp * valid_grid_2d
            odo_tp = odo_tp * valid_grid_2d
            O2D_tp = O2D_state[:, -1, :, :].unsqueeze(1)
            O2D_tp = F.interpolate(O2D_tp, scale_factor=10)
            O2D_tp = O2D_tp * valid_grid_2d
            self.log_precip_map(
                ['dyn tp', 'd2o tp', 'inference tp', 'dod tp', 'obs tp', 'o2d tp', 'odo tp'], 
                [dyn_tp, D2O_tp, inference_tp, dod_tp, obs_tp, O2D_tp, odo_tp])

        # --- log score for [ ensemble members ] ---
        B = dyn_state.shape[0]; L = self.hparams.nz; E = self.hparams.validation_ensemble
        if E == 1:
            self.common_log('val', D2O_tp, obs_tp, D2O_circu, obs_circu, valid_grid_2d)
            
            D2O_circu_pool = F.avg_pool2d(D2O_circu, 4)
            self.log('val_inv_rmse', torch.sqrt(F.mse_loss(D2O_circu_pool, dyn_state[:, :-1, :, :])))
        elif E > 1:
            z_D = self.reparameterize(mu_D, logvar_D, ensemble=E) # [E*B, L]
            dyn_state_expand = dyn_state.unsqueeze(0).expand(E, -1, -1, -1, -1).reshape(-1, *dyn_state.shape[1:])
            info_expand = [item.unsqueeze(0).expand(E, *([-1]*len(item.shape))).reshape(-1, *item.shape[1:]) for item in info]
            D2O_circu_expand, D2O_tp_expand = self.forward_Generator_D2O(self.G_D2O, info_expand, dyn_state_expand, z_D)
            
            # prob skill for ensemble
            obs_tp_expand = obs_tp.unsqueeze(0).expand(E, -1, -1, -1, -1).reshape(-1, *obs_tp.shape[1:])
            obs_circu_expand = obs_circu.unsqueeze(0).expand(E, -1, -1, -1, -1).reshape(-1, *obs_circu.shape[1:])
            valid_grid_2d_expand = valid_grid_2d.unsqueeze(0).expand(E, -1, -1, -1, -1).reshape(-1, *valid_grid_2d.shape[1:])
            self.common_log('val', D2O_tp_expand, obs_tp_expand, D2O_circu_expand, obs_circu_expand, valid_grid_2d_expand)
            
            D2O_tp_ensemble = D2O_tp_expand.reshape(E, B, *inference_tp.shape[1:])
            D2O_circu_ensemble = D2O_circu_expand.reshape(E, B, *D2O_circu.shape[1:])
            self.log('val_ensemble_crps', self.valid_crps_score(y_hat=D2O_tp_ensemble, y=obs_tp, valid_grid_2d=valid_grid_2d))
            
            # determ skill for ensemble mean
            D2O_tp_em = D2O_tp_ensemble.mean(dim=0)
            D2O_circu_em = D2O_circu_ensemble.mean(dim=0)
            self.common_log('val_mean', D2O_tp_em, obs_tp, D2O_circu_em, obs_circu, valid_grid_2d)
            
            # inv skill for every member
            D2O_circu_pool_expand = F.avg_pool2d(D2O_circu_expand, 4)
            if len(dyn_state.size()) == 4:
                self.log('val_inv_rmse', torch.sqrt(F.mse_loss(D2O_circu_pool_expand, dyn_state_expand[:, :-1, :, :])))
            elif len(dyn_state.size()) == 5:
                self.log('val_inv_rmse', torch.sqrt(F.mse_loss(D2O_circu_pool_expand, dyn_state_expand[:, :, :-1, :, :])))
    
    
    def configure_optimizers(self):
        opt_GdE = self.get_optimizers(itertools.chain(self.G_O2D.parameters(), self.E_O.parameters(), self.E_D.parameters()))
        opt_GoE = self.get_optimizers(itertools.chain(self.G_D2O.parameters(), self.E_D.parameters(), self.E_O.parameters()))
        opt_Dd = self.get_optimizers(self.D_D.parameters())
        opt_Do = self.get_optimizers(self.D_O.parameters())
        return (
            {'optimizer': opt_GdE, 'frequency': self.hparams.n_opt_GE},
            {'optimizer': opt_Dd, 'frequency': self.hparams.n_opt_Dd},
            {'optimizer': opt_GoE, 'frequency': self.hparams.n_opt_GE},
            {'optimizer': opt_Do, 'frequency': self.hparams.n_opt_Do}
        )



class CvaeGanModule_v2(CvaeGanModule):
    '''
    every train step, generate an ensemble for each sample
    id loss & cycle loss constrained on ensemble mean
    adv loss & inv loss constrained on every member
    '''
    def __init__(self, *args, **kwargs):
        super(CvaeGanModule_v2, self).__init__(
            *args, **kwargs
        )
        assert self.hparams.validation_ensemble > 1
    
    
    def training_step(self,  batch, batch_idx, optimizer_idx):
        info, dyn_state, obs_tp, obs_circu = batch
        _, _, valid_grid_2d, _, aux = info
        obs_tp = obs_tp * valid_grid_2d
        
        B = dyn_state.shape[0]; L = self.hparams.nz; E = self.hparams.validation_ensemble
        
        # latents: [E*B, L]
        mu_D, logvar_D = self.forward_Encoder_D(self.E_D, info, dyn_state)
        mu_O, logvar_O = self.forward_Encoder_O(self.E_O, info, obs_circu, obs_tp)
        #z_D = self.reparameterize(mu_D, logvar_D, ensemble=E)
        z_O = self.reparameterize(mu_O, logvar_O, ensemble=E)
        
        # inputs: [E*B, ...]
        dyn_state_expand = dyn_state.unsqueeze(0).expand(E, -1, -1, -1, -1).reshape(-1, *dyn_state.shape[1:])
        info_expand = [item.unsqueeze(0).expand(E, *([-1]*len(item.shape))).reshape(-1, *item.shape[1:]) for item in info]
        
        # prepare refs: [E*B, ...] for fields of obs space
        obs_tp_expand = obs_tp.unsqueeze(0).expand(E, -1, -1, -1, -1).reshape(-1, *obs_tp.shape[1:])
        obs_circu_expand = obs_circu.unsqueeze(0).expand(E, -1, -1, -1, -1).reshape(-1, *obs_circu.shape[1:])
        valid_grid_2d_expand = valid_grid_2d.unsqueeze(0).expand(E, -1, -1, -1, -1).reshape(-1, *valid_grid_2d.shape[1:])
        
        # forward
        D2O_circu_expand, D2O_tp_expand = self.forward_Generator_D2O(self.G_D2O, info_expand, dyn_state_expand, z_O)  # [E*B, ...]
        O2D_state = self.forward_Generator_O2D(self.G_O2D, info, obs_circu, obs_tp)  # [B, ...]
        
        # prepare refs: [B, ...] w/ valid_grid for fields of dyn space
        O2D_circu = O2D_state[:, :-1, :, :]
        O2D_tp = O2D_state[:, -1, :, :].unsqueeze(1)
        dyn_circu = dyn_state[:, :-1, :, :]
        dyn_tp = dyn_state[:, -1, :, :].unsqueeze(1)
        coarse_valid_grid_2d = self.get_coarse_bool(valid_grid_2d)
        dyn_tp = dyn_tp * coarse_valid_grid_2d
        
        # prepare ensemble mean for fields of dyn space: [B, ...]
        D2O_tp_em = D2O_tp_expand.reshape(E, B, *D2O_tp_expand.shape[1:]).mean(dim=0)
        D2O_circu_em = D2O_circu_expand.reshape(E, B, *D2O_circu_expand.shape[1:]).mean(dim=0)
        
        # train generator against D_d for every member
        if optimizer_idx == 0:
            # --- ADV LOSS ---
            # -- for every member --
            fake_critic_D = self.forward_Discriminator_D(self.D_D, info, O2D_state)
            adv_loss = self.get_adversarial_loss(fake_critic_D, True)
            g_loss_adv = adv_loss * self.hparams.lambda_adv
        
        # train generator against D_o
        if optimizer_idx == 2:
            # --- ADV LOSS ---
            # -- for every member --
            fake_critic_O = self.forward_Discriminator_O(self.D_O, info_expand, D2O_circu_expand, D2O_tp_expand)
            adv_loss = self.get_adversarial_loss(fake_critic_O, True)
            g_loss_adv = adv_loss * self.hparams.lambda_adv

        # train generator $ encoder
        if optimizer_idx == 0 or optimizer_idx == 2:
            # --- ID LOSS ---
            # -- for ensemble mean -- 
            # constrain D2O & O2D
            id_loss = self.pixelwise_score(D2O_tp_em, obs_tp).mean(axis=(1, 2, 3))\
                        + self.pixelwise_score(O2D_tp, dyn_tp).mean(axis=(1, 2, 3))
            id_loss_circu = self.pixelwise_score(D2O_circu_em, obs_circu).mean(axis=(1, 2, 3))\
                        + self.pixelwise_score(O2D_circu, dyn_circu).mean(axis=(1, 2, 3))
            if self.hparams.lambda_id == 'function':
                lead = aux[:, self.hparams.idx_of_lead]
                id_loss = L_ID_SCHEDULE(lead, thres=self.hparams.lead_thres, init_lambda=self.hparams.init_lambda_id) * id_loss
                id_loss_circu = L_ID_SCHEDULE(lead, thres=self.hparams.lead_thres, init_lambda=self.hparams.init_lambda_id) * id_loss_circu
            else:
                id_loss = id_loss * self.hparams.lambda_id
                id_loss_circu = id_loss_circu * self.hparams.lambda_id
            g_loss_id = id_loss.mean() + id_loss_circu.mean() * self.hparams.lambda_id_circu_proportion

            # --- CYLCE LOSS --- 
            # -- for ensemble mean --
            dod_state_expand = self.forward_Generator_O2D(self.G_O2D, info_expand, D2O_circu_expand, D2O_tp_expand)
            dod_state_em = dod_state_expand.reshape(E, B, *dod_state_expand.shape[1:]).mean(dim=0)
            
            O2D_state_expand = O2D_state.unsqueeze(0).expand(E, -1, -1, -1, -1).reshape(-1, *O2D_state.shape[1:])
            odo_circu_expand, odo_tp_expand = self.forward_Generator_D2O(self.G_D2O, info_expand, O2D_state_expand, z_O)
            odo_circu_em = odo_circu_expand.reshape(E, B, *odo_circu_expand.shape[1:]).mean(dim=0)
            odo_tp_em = odo_tp_expand.reshape(E, B, *odo_tp_expand.shape[1:]).mean(dim=0)
            
            cycle_loss = self.pixelwise_score(odo_tp_em, obs_tp).mean() + self.pixelwise_score(odo_circu_em, obs_circu).mean() + self.pixelwise_score(dod_state_em, dyn_state).mean()
            g_loss_cycle = cycle_loss * self.hparams.lambda_cycle
            
            # --- INV LOSS ---
            # -- for every member --
            # constrain circu alone, with relax limit
            D2O_circu_expand_pool = F.avg_pool2d(D2O_circu_expand, 4)
            obs_circu_pool = F.avg_pool2d(obs_circu, 4)
            dyn_circu_expand = dyn_state_expand[:, :-1]
            inv_loss_D2O = self.pixelwise_score(D2O_circu_expand_pool, dyn_circu_expand).mean()
            inv_loss_D2O = F.relu(inv_loss_D2O - self.hparams.delta_inv)
            inv_loss_O2D = self.pixelwise_score(O2D_state[:, :-1], obs_circu_pool).mean()
            inv_loss_O2D = F.relu(inv_loss_O2D - self.hparams.delta_inv)
            g_loss_inv = (inv_loss_D2O + inv_loss_O2D) * self.hparams.lambda_inv
            
            # --- KL LOSS ---
            # -- for every member --
            kl_loss = self.kl_divergence(mu_D, logvar_D, mu_O, logvar_O)
            g_loss_kl = kl_loss * self.hparams.lambda_kl

            g_loss = g_loss_adv + g_loss_id + g_loss_inv + g_loss_cycle + g_loss_kl
            self.log('train_adv_loss', g_loss_adv)
            self.log('train_id_loss', g_loss_id)
            self.log('train_inv_loss', g_loss_inv)
            self.log('train_cycle_loss', g_loss_cycle)
            self.log('train_kl_loss', g_loss_kl)
            
            self.common_log('train', D2O_tp_expand, obs_tp_expand, D2O_circu_expand, obs_circu_expand, valid_grid_2d_expand)
            
            return g_loss

        # train discriminator D
        # -- fields on D space are deterministic --
        if optimizer_idx == 1:
            real_critic_D = self.forward_Discriminator_D(self.D_D, info, dyn_state)
            fake_critic_D = self.forward_Discriminator_D(self.D_D, info, O2D_state.detach())
            
            real_loss = self.get_adversarial_loss(real_critic_D, True)
            fake_loss = self.get_adversarial_loss(fake_critic_D, False)
            
            d_d_loss = real_loss + fake_loss

            self.log('train_real_critic_D', torch.mean(real_critic_D))
            self.log('train_fake_critic_D', torch.mean(fake_critic_D))
            self.log('train_w_dist_D', torch.mean(real_critic_D) - torch.mean(fake_critic_D))
            
            # constraint on D_d
            if self.hparams.weight_clipping or self.hparams.GAN_mode == 'WGAN':
                self.weight_clipping(self.D_D, self.hparams.clip_value)
            elif self.hparams.GAN_mode == 'WGAN-GP':
                if not self.hparams.embed_lead:
                    info = self.drop_lead_in_aux(info)
                gradient_penalty_D = self.compute_gradient_penalty(self.D_D, info, (dyn_state), (O2D_state))
                d_d_loss += self.hparams.lambda_gp_d * gradient_penalty_D
                self.log('loss_gp_D', gradient_penalty_D)
            
            return d_d_loss
        
        # train discriminator O
        # -- fields on O space are probabilistic --
        # feed every member to D_o
        if optimizer_idx == 3:
            real_critic_O = self.forward_Discriminator_O(self.D_O, info, obs_circu, obs_tp)  # XXX: should not need expand ?
            fake_critic_O = self.forward_Discriminator_O(self.D_O, info_expand, D2O_circu_expand.detach(), D2O_tp_expand.detach())

            real_loss = self.get_adversarial_loss(real_critic_O, True)
            fake_loss = self.get_adversarial_loss(fake_critic_O, False)
            
            d_o_loss = real_loss + fake_loss
            
            self.log('train_real_critic_O', torch.mean(real_critic_O))
            self.log('train_fake_critic_O', torch.mean(fake_critic_O))
            self.log('train_w_dist_O', torch.mean(real_critic_O) - torch.mean(fake_critic_O))

            # constraint on D_o
            if self.hparams.weight_clipping or self.hparams.GAN_mode == 'WGAN':
                self.weight_clipping(self.D_O, self.hparams.clip_value)
            elif self.hparams.GAN_mode == 'WGAN-GP':
                if not self.hparams.embed_lead:
                    info_expand = self.drop_lead_in_aux(info_expand)
                gradient_penalty_O = self.compute_gradient_penalty(self.D_O, info_expand, (obs_circu_expand, obs_tp_expand), (D2O_circu_expand, D2O_tp_expand))
                d_o_loss += self.hparams.lambda_gp_o * gradient_penalty_O
                self.log('loss_gp_O', gradient_penalty_O)

            return d_o_loss



if __name__ == '__main__':
    '''
    e.g. ` python plmodules.py cfg=cfg.PP.yaml `
    '''
    from omegaconf import OmegaConf
    from torch.utils.data import DataLoader
    from pytorch_lightning import Trainer
    from dataset import ModelCorrectionDataset
    
    def print_contents(var):
        if isinstance(var, torch.Tensor):
            print(var.shape)
        elif isinstance(var, tuple) or isinstance(var, list):
            for v in var:
                print_contents(v)

    
    '''
    load config
    '''
    config_cli = OmegaConf.from_cli() 
    if config_cli.get('cfg'):
        config = OmegaConf.load(config_cli.cfg)
    else:
        config = OmegaConf.load('cfg/base.yaml')
    # update layer-2 nested config
    for key, nested_config in config.items():
        if key in config_cli:
            nested_config.update(config_cli[key])
    print(OmegaConf.to_yaml(config))
    
    '''
    load data
    '''
    trainset = ModelCorrectionDataset(
        year_sel = [2010], mth_sel = [10],
        **config.dataset_config
    )
    valset = ModelCorrectionDataset(
        year_sel = [2010], mth_sel = [11],
        **config.dataset_config
    )
    
    trainloader = DataLoader(trainset, shuffle=True, **config.dataloader_config)
    valloader = DataLoader(valset, **config.dataloader_config)
    
    batch = next(iter(trainloader))
    print('load data: ...')
    print_contents(batch)
    
    '''
    feed module
    '''
    print("\n--- single batch forward test... ---")
    module_class = globals()[config.training_config.module_class]
    print(f'initiate module {module_class} for testing...')
    module = module_class( # plz switch to the module that you want to test
        debug=config.training_config.debug,
        log_images=False, 
        **config.plmodule_config
    )
    
    outputs = module.forward(*batch[:2])
    print('after inference: ...')
    print_contents(outputs)
    
    '''
    train
    '''
    print("\n--- training test... --- ")
    trainer = Trainer(
        max_epochs=10,
        accelerator='gpu', devices=[0],
        fast_dev_run=True,
    )
    
    trainer.fit(module, trainloader, valloader)