import torch
import torch.nn as nn
from .sparsity import scaling


def reg_loss(time_deriv_list, sparse_theta_list, coeff_vector_list):
    '''Loss function for the regularisation loss. Calculates loss for each term in list.'''
    loss = torch.stack([torch.mean((time_deriv - theta @ coeff_vector)**2) for time_deriv, theta, coeff_vector in zip(time_deriv_list, sparse_theta_list, coeff_vector_list)])
    return loss

def mse_loss(prediction, target):
    '''Loss functions for the MSE loss. Calculates loss for each term in list.'''
    loss = torch.mean((prediction - target)**2, dim=0)
    return loss

def l1_loss(coeff_vector_list, l1):
    '''Loss functions for the L1 loss on the coefficients. Calculates loss for each term in list.'''
    loss = torch.stack([torch.sum(torch.abs(coeff_vector)) for coeff_vector in coeff_vector_list])
    return l1 * loss
   
    



 
  

import torch
import torch.nn as nn
from .sparsity import scaling


def reg_loss(time_deriv_list, sparse_theta_list, coeff_vector_list):
    '''Loss function for the regularisation loss. Calculates loss for each term in list.'''
    loss = torch.stack([torch.mean((time_deriv - theta @ coeff_vector)**2) for time_deriv, theta, coeff_vector in zip(time_deriv_list, sparse_theta_list, coeff_vector_list)])
    return loss

def mse_loss(prediction, target):
    '''Loss functions for the MSE loss. Calculates loss for each term in list.'''
    loss = torch.mean((prediction - target)**2, dim=0)
    return loss

def l1_loss(coeff_vector_list, l1):
    '''Loss functions for the L1 loss on the coefficients. Calculates loss for each term in list.'''
    loss = torch.stack([torch.sum(torch.abs(coeff_vector)) for coeff_vector in coeff_vector_list])
    return l1 * loss
   
    



 
  

Loss function for the regularisation loss. Calculates loss for each term in list. Loss functions for the MSE loss. Calculates loss for each term in list. Loss functions for the L1 loss on the coefficients. Calculates loss for each term in list.