Project: 044_On_the_Identification_of_Temporal_Causal_Represent | Function: InstantaneousProcess.loss_function (iclr)

system: You are an expert Python programmer implementing algorithms from scientific papers.
Complete the masked function based on the algorithm description.
Follow Python best practices and maintain consistent code style.
Return **only** the complete Python function code.

user: Complete this Python function:

**DESCRIPTION:** Implements the comprehensive ELBO loss function for the synthetic instantaneous causal process learning experiments. This function combines reconstruction loss, KL divergence constraints, and weighted sparsity regularization with flexible sparsity parameters. It enforces both historical and instantaneous sparse dependencies through L1 penalties on Jacobian matrices.

**KEY TERMS:** {
  "recon_loss": "reconstruction loss ℓrecon",
  "kld_normal": "past KL divergence ℓKL_past",
  "kld_future": "future KL divergence ℓKL_future",
  "sparsity_loss": "gradient-based sparsity penalty ℓsparsity",
  "w_hist": "historical sparsity weights",
  "w_inst": "instantaneous sparsity weights"
}

**CONTEXT CODE:**         self.net = BetaVAE_MLP(input_dim=input_dim, 
                                z_dim=z_dim, 
                                hidden_dim=hidden_dim)

        self.transition_prior_fix = NPInstantaneousTransitionPrior(lags=lag, 
                                                      latent_size=z_dim, 
                                                      num_layers=3, 
                                                      hidden_dim=hidden_dim)
        self.transition_prior_change = NPChangeInstantaneousTransitionPrior(lags=lag, 
                                                        latent_size=z_dim,
                                                        embedding_dim=embedding_dim,
                                                        num_layers=3, 
                                                        hidden_dim=hidden_dim)
        
                                                            
        # base distribution for calculation of log prob under the model
        self.register_buffer('base_dist_mean', torch.zeros(self.z_dim))
        self.register_buffer('base_dist_var', torch.eye(self.z_dim))

    @property
    def base_dist(self):
        # Noise density function
        return D.MultivariateNormal(self.base_dist_mean, self.base_dist_var)
    
    def reparameterize(self, mean, logvar, random_sampling=True):
        if random_sampling:
            eps = torch.randn_like(logvar)
            std = torch.exp(0.5*logvar)
            z = mean + eps*std
            return z
        else:
            return mean

    def reconstruction_loss(self, x, x_recon, distribution):
        batch_size = x.size(0)
        assert batch_size != 0

        if distribution == 'bernoulli':
            recon_loss = F.binary_cross_entropy_with_logits(
                x_recon, x, size_average=False).div(batch_size)

        elif distribution == 'gaussian':
            recon_loss = F.mse_loss(x_recon, x, size_average=False).div(batch_size)

        elif distribution == 'sigmoid_gaussian':
            x_recon = F.sigmoid(x_recon)
            recon_loss = F.mse_loss(x_recon, x, size_average=False).div(batch_size)

        return recon_loss

    def loss_function(self, x, x_recon, mus, logvars, zs, embeddings):
        '''
        VAE ELBO loss: recon_loss + kld_loss (past: N(0,1), future: N(0,1) after flow) + sparsity_loss
        '''
        raise NotImplementedError("Function implementation has been masked for testing")
    
    def forward(self, batch):
        # Prepare data
        if self.z_dim_change>0:
            x, y, c = batch['xt'], batch['yt'], batch['ct']
            c = torch.squeeze(c).to(torch.int64)
            embeddings = self.embed_func(c)
        else:
            x, y = batch['xt'], batch['yt']
            embeddings = None
        batch_size, length, _ = x.shape
        x_flat = x.view(-1, self.input_dim)
        

        # Inference & Reshape to time-series format
        x_recon, mus, logvars, zs = self.net(x_flat)
        x_recon = x_recon.view(batch_size, length, self.input_dim)
        mus = mus.reshape(batch_size, length, self.z_dim)
        logvars  = logvars.reshape(batch_size, length, self.z_dim)
        zs = zs.reshape(batch_size, length, self.z_dim)

        return x, x_recon, mus, logvars, zs, embeddings

    def training_step(self, batch, batch_idx):
        x, x_recon, mus, logvars, zs, embeddings = self.forward(batch)
        sparsity_loss, recon_loss, kld_normal, kld_future = self.loss_function(x, x_recon, mus, logvars, zs, embeddings)

        # VAE training
        loss = recon_loss + self.beta * kld_normal + self.gamma * kld_future + self.theta * sparsity_loss
        self.log("train_elbo_loss", loss)
        self.log("train_recon_loss", recon_loss)
        self.log("train_kld_normal", self.beta * kld_normal)
        self.log("train_kld_future", self.gamma * kld_future)
        self.log("train_sparsity_loss", self.theta * sparsity_loss)

        return loss
    
    def validation_step(self, batch, batch_idx):
        x, x_recon, mus, logvars, zs, embeddings = self.forward(batch)
        sparsity_loss, recon_loss, kld_normal, kld_future = self.loss_function(x, x_recon, mus, logvars, zs, embeddings)

        # Compute Mean Correlation Coefficient (MCC)
        zt_recon = mus.view(-1, self.z_dim).T.detach().cpu().numpy()
        zt_true = batch["yt"].view(-1, self.z_dim).T.detach().cpu().numpy()
        mcc = compute_mcc(zt_recon, zt_true, self.correlation)

        # Draw sparsity matrix
        # import pdb
        # pdb.set_trace()
        # softmax through dim 1 on numpy
        # hist_jac_soft = torch.softmax(torch.abs(self.hist_jac), dim=1)
        # hist_jac_fig = plot_sparsity_matrix(hist_jac_soft, "Sparsity Matrix (hist_jac_fig)")
        # self.logger.experiment.log({'alpha_softmax': wandb.Image(hist_jac_fig)})

        trans_show_dir = torch.abs(self.trans_show)
        # trans_show_fig = plot_sparsity_matrix(trans_show_dir, "Sparsity Matrix (hist_jac_fig)")
        trans_show_fig = plot_sparsity_matrix(trans_show_dir, "Transition Matrix")
        self.logger.experiment.log({'trans': wandb.Image(trans_show_fig)})

        inst_show_dir = torch.abs(self.inst_show)
        # inst_show_fig = plot_sparsity_matrix(inst_show_dir, "Sparsity Matrix (inst_jac_fig)")
        inst_show_fig = plot_sparsity_matrix(inst_show_dir, "Instantaneous Matrix")
        self.logger.experiment.log({'inst': wandb.Image(inst_show_fig)})

        # if self.z_dim_fix>0:
        #     fig_fix = plot_sparsity_matrix(self.alphas_fix().detach().cpu(), "Sparsity Matrix (Fix)")
        #     self.logger.experiment.log({'alpha': wandb.Image(fig_fix)})
        # if self.z_dim_change>0:
        #     fig_change = plot_sparsity_matrix(self.alphas_change().detach().cpu(), "Sparsity Matrix (Change)")
        #     self.logger.experiment.log({'alpha': wandb.Image(fig_change)})

        # Validation loss
        loss = recon_loss + self.beta * kld_normal + self.gamma * kld_future + self.theta * sparsity_loss
        self.log("val_mcc", mcc) 
        self.log("val_elbo_loss", loss)
        self.log("val_recon_loss", recon_loss)
        self.log("val_kld_normal", self.beta * kld_normal)
        self.log("val_kld_future", self.gamma * kld_future)
        self.log("val_sparsity_loss", self.theta * sparsity_loss)

        return loss

    def configure_optimizers(self):
        opt_v = torch.optim.AdamW(filter(lambda p: p.requires_grad, self.parameters()), lr=self.lr, betas=(0.9, 0.999), weight_decay=0.0001)
        return [opt_v], []


Complete the function implementation replacing "raise NotImplementedError".