modeles

This module defines models and solvers for 4D-VarNet.

4D-VarNet is a framework for solving inverse problems in data assimilation using deep learning and PyTorch Lightning.

Classes:

Name	Description
Lit4dVarNet	A PyTorch Lightning module for training and testing 4D-VarNet models.
GradSolver	A gradient-based solver for optimization in 4D-VarNet.
ConvLstmGradModel	A convolutional LSTM model for gradient modulation.
BaseObsCost	A base class for observation cost computation.
BilinAEPriorCost	A prior cost model using bilinear autoencoders.

BaseObsCost

Bases: Module

A base class for computing observation cost.

Attributes:

Name	Type	Description
w	float	Weight for the observation cost.

Source code in ocean4dvarnet/models.py

class BaseObsCost(nn.Module):
    """
    A base class for computing observation cost.

    Attributes:
        w (float): Weight for the observation cost.
    """

    def __init__(self, w=1) -> None:
        """
        Initialize the BaseObsCost module.

        Args:
            w (float, optional): Weight for the observation cost. Defaults to 1.
        """
        super().__init__()
        self.w = w

    def forward(self, state, batch):
        """
        Compute the observation cost.

        Args:
            state (torch.Tensor): The current state tensor.
            batch (dict): The input batch containing data.

        Returns:
            torch.Tensor: The computed observation cost.
        """
        msk = batch.input.isfinite()
        return self.w * F.mse_loss(state[msk], batch.input.nan_to_num()[msk])

__init__(w=1)

Initialize the BaseObsCost module.

Parameters:

Name	Type	Description	Default
w	float	Weight for the observation cost. Defaults to 1.	1

Source code in ocean4dvarnet/models.py

def __init__(self, w=1) -> None:
    """
    Initialize the BaseObsCost module.

    Args:
        w (float, optional): Weight for the observation cost. Defaults to 1.
    """
    super().__init__()
    self.w = w

forward(state, batch)

Compute the observation cost.

Parameters:

Name	Type	Description	Default
state	Tensor	The current state tensor.	required
batch	dict	The input batch containing data.	required

Returns:

Type	Description
	torch.Tensor: The computed observation cost.

Source code in ocean4dvarnet/models.py

def forward(self, state, batch):
    """
    Compute the observation cost.

    Args:
        state (torch.Tensor): The current state tensor.
        batch (dict): The input batch containing data.

    Returns:
        torch.Tensor: The computed observation cost.
    """
    msk = batch.input.isfinite()
    return self.w * F.mse_loss(state[msk], batch.input.nan_to_num()[msk])

BilinAEPriorCost

Bases: Module

A prior cost model using bilinear autoencoders.

Attributes:

Name	Type	Description
bilin_quad	bool	Whether to use bilinear quadratic terms.
conv_in	Conv2d	Convolutional layer for input.
conv_hidden	Conv2d	Convolutional layer for hidden states.
bilin_1	Conv2d	Bilinear layer 1.
bilin_21	Conv2d	Bilinear layer 2 (part 1).
bilin_22	Conv2d	Bilinear layer 2 (part 2).
conv_out	Conv2d	Convolutional layer for output.
down	Module	Downsampling layer.
up	Module	Upsampling layer.

Source code in ocean4dvarnet/models.py

class BilinAEPriorCost(nn.Module):
    """
    A prior cost model using bilinear autoencoders.

    Attributes:
        bilin_quad (bool): Whether to use bilinear quadratic terms.
        conv_in (nn.Conv2d): Convolutional layer for input.
        conv_hidden (nn.Conv2d): Convolutional layer for hidden states.
        bilin_1 (nn.Conv2d): Bilinear layer 1.
        bilin_21 (nn.Conv2d): Bilinear layer 2 (part 1).
        bilin_22 (nn.Conv2d): Bilinear layer 2 (part 2).
        conv_out (nn.Conv2d): Convolutional layer for output.
        down (nn.Module): Downsampling layer.
        up (nn.Module): Upsampling layer.
    """

    def __init__(self, dim_in, dim_hidden, kernel_size=3, downsamp=None, bilin_quad=True):
        """
        Initialize the BilinAEPriorCost module.

        Args:
            dim_in (int): Number of input dimensions.
            dim_hidden (int): Number of hidden dimensions.
            kernel_size (int, optional): Kernel size for convolutions. Defaults to 3.
            downsamp (int, optional): Downsampling factor. Defaults to None.
            bilin_quad (bool, optional): Whether to use bilinear quadratic terms. Defaults to True.
        """
        super().__init__()
        self.bilin_quad = bilin_quad
        self.conv_in = nn.Conv2d(dim_in, dim_hidden, kernel_size=kernel_size, padding=kernel_size // 2)
        self.conv_hidden = nn.Conv2d(dim_hidden, dim_hidden, kernel_size=kernel_size, padding=kernel_size // 2)

        self.bilin_1 = nn.Conv2d(dim_hidden, dim_hidden, kernel_size=kernel_size, padding=kernel_size // 2)
        self.bilin_21 = nn.Conv2d(dim_hidden, dim_hidden, kernel_size=kernel_size, padding=kernel_size // 2)
        self.bilin_22 = nn.Conv2d(dim_hidden, dim_hidden, kernel_size=kernel_size, padding=kernel_size // 2)

        self.conv_out = nn.Conv2d(2 * dim_hidden, dim_in, kernel_size=kernel_size, padding=kernel_size // 2)

        self.down = nn.AvgPool2d(downsamp) if downsamp is not None else nn.Identity()
        self.up = nn.UpsamplingBilinear2d(scale_factor=downsamp) if downsamp is not None else nn.Identity()

    def forward_ae(self, x):
        """
        Perform the forward pass through the autoencoder.

        Args:
            x (torch.Tensor): Input tensor.

        Returns:
            torch.Tensor: Output tensor after passing through the autoencoder.
        """
        x = self.down(x)
        x = self.conv_in(x)
        x = self.conv_hidden(F.relu(x))

        nonlin = self.bilin_21(x) ** 2 if self.bilin_quad else (self.bilin_21(x) * self.bilin_22(x))
        x = self.conv_out(torch.cat([self.bilin_1(x), nonlin], dim=1))
        x = self.up(x)
        return x

    def forward(self, state):
        """
        Compute the prior cost using the autoencoder.

        Args:
            state (torch.Tensor): The current state tensor.

        Returns:
            torch.Tensor: The computed prior cost.
        """
        return F.mse_loss(state, self.forward_ae(state))

__init__(dim_in, dim_hidden, kernel_size=3, downsamp=None, bilin_quad=True)

Initialize the BilinAEPriorCost module.

Parameters:

Name	Type	Description	Default
dim_in	int	Number of input dimensions.	required
dim_hidden	int	Number of hidden dimensions.	required
kernel_size	int	Kernel size for convolutions. Defaults to 3.	3
downsamp	int	Downsampling factor. Defaults to None.	None
bilin_quad	bool	Whether to use bilinear quadratic terms. Defaults to True.	True

Source code in ocean4dvarnet/models.py

def __init__(self, dim_in, dim_hidden, kernel_size=3, downsamp=None, bilin_quad=True):
    """
    Initialize the BilinAEPriorCost module.

    Args:
        dim_in (int): Number of input dimensions.
        dim_hidden (int): Number of hidden dimensions.
        kernel_size (int, optional): Kernel size for convolutions. Defaults to 3.
        downsamp (int, optional): Downsampling factor. Defaults to None.
        bilin_quad (bool, optional): Whether to use bilinear quadratic terms. Defaults to True.
    """
    super().__init__()
    self.bilin_quad = bilin_quad
    self.conv_in = nn.Conv2d(dim_in, dim_hidden, kernel_size=kernel_size, padding=kernel_size // 2)
    self.conv_hidden = nn.Conv2d(dim_hidden, dim_hidden, kernel_size=kernel_size, padding=kernel_size // 2)

    self.bilin_1 = nn.Conv2d(dim_hidden, dim_hidden, kernel_size=kernel_size, padding=kernel_size // 2)
    self.bilin_21 = nn.Conv2d(dim_hidden, dim_hidden, kernel_size=kernel_size, padding=kernel_size // 2)
    self.bilin_22 = nn.Conv2d(dim_hidden, dim_hidden, kernel_size=kernel_size, padding=kernel_size // 2)

    self.conv_out = nn.Conv2d(2 * dim_hidden, dim_in, kernel_size=kernel_size, padding=kernel_size // 2)

    self.down = nn.AvgPool2d(downsamp) if downsamp is not None else nn.Identity()
    self.up = nn.UpsamplingBilinear2d(scale_factor=downsamp) if downsamp is not None else nn.Identity()

forward(state)

Compute the prior cost using the autoencoder.

Parameters:

Name	Type	Description	Default
state	Tensor	The current state tensor.	required

Returns:

Type	Description
	torch.Tensor: The computed prior cost.

Source code in ocean4dvarnet/models.py

def forward(self, state):
    """
    Compute the prior cost using the autoencoder.

    Args:
        state (torch.Tensor): The current state tensor.

    Returns:
        torch.Tensor: The computed prior cost.
    """
    return F.mse_loss(state, self.forward_ae(state))

forward_ae(x)

Perform the forward pass through the autoencoder.

Parameters:

Name	Type	Description	Default
x	Tensor	Input tensor.	required

Returns:

Type	Description
	torch.Tensor: Output tensor after passing through the autoencoder.

Source code in ocean4dvarnet/models.py

def forward_ae(self, x):
    """
    Perform the forward pass through the autoencoder.

    Args:
        x (torch.Tensor): Input tensor.

    Returns:
        torch.Tensor: Output tensor after passing through the autoencoder.
    """
    x = self.down(x)
    x = self.conv_in(x)
    x = self.conv_hidden(F.relu(x))

    nonlin = self.bilin_21(x) ** 2 if self.bilin_quad else (self.bilin_21(x) * self.bilin_22(x))
    x = self.conv_out(torch.cat([self.bilin_1(x), nonlin], dim=1))
    x = self.up(x)
    return x

ConvLstmGradModel

Bases: Module

A convolutional LSTM model for gradient modulation.

Attributes:

Name	Type	Description
dim_hidden	int	Number of hidden dimensions.
gates	Conv2d	Convolutional gates for LSTM.
conv_out	Conv2d	Output convolutional layer.
dropout	Dropout	Dropout layer.
down	Module	Downsampling layer.
up	Module	Upsampling layer.

Source code in ocean4dvarnet/models.py

class ConvLstmGradModel(nn.Module):
    """
    A convolutional LSTM model for gradient modulation.

    Attributes:
        dim_hidden (int): Number of hidden dimensions.
        gates (nn.Conv2d): Convolutional gates for LSTM.
        conv_out (nn.Conv2d): Output convolutional layer.
        dropout (nn.Dropout): Dropout layer.
        down (nn.Module): Downsampling layer.
        up (nn.Module): Upsampling layer.
    """

    def __init__(self, dim_in, dim_out, dim_hidden, kernel_size=3, dropout=0.1, downsamp=None):
        """
        Initialize the ConvLstmGradModel.

        Args:
            dim_in (int): Number of input dimensions.
            dim_hidden (int): Number of hidden dimensions.
            kernel_size (int, optional): Kernel size for convolutions. Defaults to 3.
            dropout (float, optional): Dropout rate. Defaults to 0.1.
            downsamp (int, optional): Downsampling factor. Defaults to None.
        """
        super().__init__()
        self.dim_hidden = dim_hidden
        self.gates = torch.nn.Conv2d(
            dim_in + dim_hidden,
            4 * dim_hidden,
            kernel_size=kernel_size,
            padding=kernel_size // 2,
        )

        self.conv_out = torch.nn.Conv2d(dim_hidden, dim_out, kernel_size=kernel_size, padding=kernel_size // 2)

        self.dropout = torch.nn.Dropout(dropout)
        self._state = []
        self.down = nn.AvgPool2d(downsamp) if downsamp is not None else nn.Identity()
        self.up = nn.UpsamplingBilinear2d(scale_factor=downsamp) if downsamp is not None else nn.Identity()

    def predict(self, x, timesteps=None, extra=[], hidden=None, cell=None):
        """
        Perform the forward pass of the LSTM.

        Args:
            x (torch.Tensor): Input tensor.

        Returns:
            torch.Tensor: Output tensor.
        """
        # if self._grad_norm is None:
        #     self._grad_norm = (x**2).mean().sqrt()
        # x = x / self._grad_norm
        # hidden, cell = self._state
        # x = self.dropout(x)
        x = self.down(x)
        gates = self.gates(torch.cat((x, hidden), 1))

        in_gate, remember_gate, out_gate, cell_gate = gates.chunk(4, 1)

        in_gate, remember_gate, out_gate = map(torch.sigmoid, [in_gate, remember_gate, out_gate])
        cell_gate = torch.tanh(cell_gate)

        cell = (remember_gate * cell) + (in_gate * cell_gate)
        hidden = out_gate * torch.tanh(cell)

        self._state = hidden, cell
        out = self.conv_out(hidden)
        out = self.up(out)
        return out

__init__(dim_in, dim_out, dim_hidden, kernel_size=3, dropout=0.1, downsamp=None)

Initialize the ConvLstmGradModel.

Parameters:

Name	Type	Description	Default
dim_in	int	Number of input dimensions.	required
dim_hidden	int	Number of hidden dimensions.	required
kernel_size	int	Kernel size for convolutions. Defaults to 3.	3
dropout	float	Dropout rate. Defaults to 0.1.	0.1
downsamp	int	Downsampling factor. Defaults to None.	None

Source code in ocean4dvarnet/models.py

def __init__(self, dim_in, dim_out, dim_hidden, kernel_size=3, dropout=0.1, downsamp=None):
    """
    Initialize the ConvLstmGradModel.

    Args:
        dim_in (int): Number of input dimensions.
        dim_hidden (int): Number of hidden dimensions.
        kernel_size (int, optional): Kernel size for convolutions. Defaults to 3.
        dropout (float, optional): Dropout rate. Defaults to 0.1.
        downsamp (int, optional): Downsampling factor. Defaults to None.
    """
    super().__init__()
    self.dim_hidden = dim_hidden
    self.gates = torch.nn.Conv2d(
        dim_in + dim_hidden,
        4 * dim_hidden,
        kernel_size=kernel_size,
        padding=kernel_size // 2,
    )

    self.conv_out = torch.nn.Conv2d(dim_hidden, dim_out, kernel_size=kernel_size, padding=kernel_size // 2)

    self.dropout = torch.nn.Dropout(dropout)
    self._state = []
    self.down = nn.AvgPool2d(downsamp) if downsamp is not None else nn.Identity()
    self.up = nn.UpsamplingBilinear2d(scale_factor=downsamp) if downsamp is not None else nn.Identity()

predict(x, timesteps=None, extra=[], hidden=None, cell=None)

Perform the forward pass of the LSTM.

Parameters:

Name	Type	Description	Default
x	Tensor	Input tensor.	required

Returns:

Type	Description
	torch.Tensor: Output tensor.

Source code in ocean4dvarnet/models.py

def predict(self, x, timesteps=None, extra=[], hidden=None, cell=None):
    """
    Perform the forward pass of the LSTM.

    Args:
        x (torch.Tensor): Input tensor.

    Returns:
        torch.Tensor: Output tensor.
    """
    # if self._grad_norm is None:
    #     self._grad_norm = (x**2).mean().sqrt()
    # x = x / self._grad_norm
    # hidden, cell = self._state
    # x = self.dropout(x)
    x = self.down(x)
    gates = self.gates(torch.cat((x, hidden), 1))

    in_gate, remember_gate, out_gate, cell_gate = gates.chunk(4, 1)

    in_gate, remember_gate, out_gate = map(torch.sigmoid, [in_gate, remember_gate, out_gate])
    cell_gate = torch.tanh(cell_gate)

    cell = (remember_gate * cell) + (in_gate * cell_gate)
    hidden = out_gate * torch.tanh(cell)

    self._state = hidden, cell
    out = self.conv_out(hidden)
    out = self.up(out)
    return out

GradModelWithCondition

Bases: Module

A generic conditional model for gradient modulation.

Attributes:

Name	Type	Description
grad_model		grad update model

Source code in ocean4dvarnet/models.py

class GradModelWithCondition(torch.nn.Module):
    """
    A generic conditional model for gradient modulation.

    Attributes:
        grad_model : grad update model
    """

    def __init__(self, grad_model=False, dropout=0.0, use_grad_norm=True):
        """
        Initialize the ConvLstmGradModel.

        Args:
            grad_model : grad update model
        """
        super().__init__()
        self.grad_model = grad_model
        self.dropout = torch.nn.Dropout(dropout)
        self.use_grad_norm = use_grad_norm

        if hasattr(self.grad_model, "dim_3d"):
            self.dim_3d = self.grad_model.dim_3d

        if hasattr(self.grad_model, "dims"):
            if self.grad_model.dims == 3:
                self.dim_3d = True

    def reset_state(self, inp):
        """
        Args:
            inp (torch.Tensor): Input tensor to determine state size.
        """
        # Initialize hidden and cell state for LSTM if use a ConvLstmGradModel, otherwise set to None
        if hasattr(self.grad_model, "dim_hidden"):
            size = [inp.shape[0], self.grad_model.dim_hidden, *inp.shape[-2:]]
            if hasattr(self.grad_model, "downsamp"):
                downsamp = self.grad_model.downsamp
            else:
                downsamp = None
            self.down = nn.AvgPool2d(downsamp) if downsamp is not None else nn.Identity()
            self._state = [
                self.down(torch.zeros(size, device=inp.device)),
                self.down(torch.zeros(size, device=inp.device)),
            ]
        else:
            self._state = None, None

        self._grad_norm = None

    def forward(self, x, timesteps=None, extra=[]):
        """
        Perform the forward pass of the LSTM.

        Args:
            x (torch.Tensor): Input tensor.

        Returns:
            torch.Tensor: Output tensor.
        """

        if self._grad_norm is None:
            if self.use_grad_norm:
                self._grad_norm = (x**2).mean().sqrt()
            else:
                self._grad_norm = 1.0

        # print('self._grad_norm in GradModelWithCondition:', self._grad_norm, flush=True)
        x = x / self._grad_norm
        hidden, cell = self._state
        x = self.dropout(x)
        out = self.grad_model.predict(x, timesteps=timesteps, extra=extra, hidden=hidden, cell=cell)

        return out

__init__(grad_model=False, dropout=0.0, use_grad_norm=True)

Initialize the ConvLstmGradModel.

Parameters:

Name	Type	Description	Default
grad_model		grad update model	required

Source code in ocean4dvarnet/models.py

def __init__(self, grad_model=False, dropout=0.0, use_grad_norm=True):
    """
    Initialize the ConvLstmGradModel.

    Args:
        grad_model : grad update model
    """
    super().__init__()
    self.grad_model = grad_model
    self.dropout = torch.nn.Dropout(dropout)
    self.use_grad_norm = use_grad_norm

    if hasattr(self.grad_model, "dim_3d"):
        self.dim_3d = self.grad_model.dim_3d

    if hasattr(self.grad_model, "dims"):
        if self.grad_model.dims == 3:
            self.dim_3d = True

forward(x, timesteps=None, extra=[])

Perform the forward pass of the LSTM.

Parameters:

Name	Type	Description	Default
x	Tensor	Input tensor.	required

Returns:

Type	Description
	torch.Tensor: Output tensor.

Source code in ocean4dvarnet/models.py

def forward(self, x, timesteps=None, extra=[]):
    """
    Perform the forward pass of the LSTM.

    Args:
        x (torch.Tensor): Input tensor.

    Returns:
        torch.Tensor: Output tensor.
    """

    if self._grad_norm is None:
        if self.use_grad_norm:
            self._grad_norm = (x**2).mean().sqrt()
        else:
            self._grad_norm = 1.0

    # print('self._grad_norm in GradModelWithCondition:', self._grad_norm, flush=True)
    x = x / self._grad_norm
    hidden, cell = self._state
    x = self.dropout(x)
    out = self.grad_model.predict(x, timesteps=timesteps, extra=extra, hidden=hidden, cell=cell)

    return out

reset_state(inp)

Parameters:

Name	Type	Description	Default
inp	Tensor	Input tensor to determine state size.	required

Source code in ocean4dvarnet/models.py

def reset_state(self, inp):
    """
    Args:
        inp (torch.Tensor): Input tensor to determine state size.
    """
    # Initialize hidden and cell state for LSTM if use a ConvLstmGradModel, otherwise set to None
    if hasattr(self.grad_model, "dim_hidden"):
        size = [inp.shape[0], self.grad_model.dim_hidden, *inp.shape[-2:]]
        if hasattr(self.grad_model, "downsamp"):
            downsamp = self.grad_model.downsamp
        else:
            downsamp = None
        self.down = nn.AvgPool2d(downsamp) if downsamp is not None else nn.Identity()
        self._state = [
            self.down(torch.zeros(size, device=inp.device)),
            self.down(torch.zeros(size, device=inp.device)),
        ]
    else:
        self._state = None, None

    self._grad_norm = None

GradSolver

Bases: Module

A gradient-based solver for optimization in unrolled architectures.

Attributes:

Name	Type	Description
prior_cost	Module	The prior cost function.
obs_cost	Module	The observation cost function.
grad_mod	Module	The gradient modulation model.
n_step	int	Number of optimization steps.
lr_grad	float	Learning rate for gradient updates.
lbd	float	Regularization parameter.

Source code in ocean4dvarnet/models.py

class GradSolver(nn.Module):
    """
    A gradient-based solver for optimization in unrolled architectures.

    Attributes:
        prior_cost (nn.Module, optional): The prior cost function.
        obs_cost (nn.Module, optional): The observation cost function.
        grad_mod (nn.Module): The gradient modulation model.
        n_step (int): Number of optimization steps.
        lr_grad (float): Learning rate for gradient updates.
        lbd (float): Regularization parameter.

    """

    def __init__(
        self,
        grad_mod,
        n_step,
        lr_grad=0.2,
        lbd=1.0,
        input_grad_update="state",
        std_init=0.1,
        prior_cost: Optional[nn.Module] = None,
        obs_cost: Optional[nn.Module] = None,
        **kwargs,
    ):
        """
        Initialize the GradSolver.

        Args:
            prior_cost (nn.Module): The prior cost function.
            obs_cost (nn.Module): The observation cost function.
            grad_mod (nn.Module): The gradient modulation model.
            n_step (int): Number of optimization steps.
            lr_grad (float, optional): Learning rate for gradient updates. Defaults to 0.2.
            lbd (float, optional): Regularization parameter. Defaults to 1.0.
            input_grad_update (str, optional): Quantities added for updating the input gradient. Defaults to "state".
            std_init (float, optional): Standard deviation for initializing the state. Defaults to 0.1.
        """
        super().__init__()
        self.prior_cost = prior_cost
        self.obs_cost = obs_cost
        self.grad_mod = grad_mod

        self.n_step = n_step
        self.lr_grad = lr_grad
        self.lbd = lbd

        self.input_grad_update = input_grad_update
        self.std_init = std_init

        self._grad_norm = None

    def init_state(self, batch, x_init=None):
        """
        Initialize the state for optimization.

        Args:
            batch (dict): Input batch containing data.
            x_init (torch.Tensor, optional): Initial state. Defaults to None.

        Returns:
            torch.Tensor: Initialized state.
        """
        if x_init is not None:
            return x_init.detach().requires_grad_(True)

        if self.std_init > 0:
            x0 = self.std_init * torch.randn_like(batch.input)
            return x0.detach().requires_grad_(True)
        else:
            return torch.zeros_like(batch.input).detach().requires_grad_(True)

    def init_h_state(self, batch, h_state=None):
        """
        Initialize the state for optimization.

        Args:
            batch (dict): Input batch containing data.
            x_init (torch.Tensor, optional): Initial state. Defaults to None.

        Returns:
            torch.Tensor: Initialized state.
        """
        if h_state is not None:
            self.h_state = h_state
        else:
            self.h_state = torch.zeros_like(batch.input).detach().requires_grad_(True)

    def format2D_3D(self, x):
        if hasattr(self.grad_mod, "dim_3d"):
            if self.grad_mod.dim_3d:
                x = x.unsqueeze(1)

        return x

    def solver_step(self, state, batch, step, alpha_step=1.0):
        """
        Perform a single optimization step.

        Args:
            state (torch.Tensor): Current state.
            batch (dict): Input batch containing data.
            step (int): Current optimization step between 0 and 1.

        Returns:
            torch.Tensor: Updated state.
        """

        if isinstance(step, float):
            t = torch.tensor([step], device=state.device).repeat(state.shape[0])
        else:
            t = step

        if "subgrad" in self.input_grad_update:
            gobs = (batch.input - state).nan_to_num()

            gprior = state - self.prior_cost.forward_ae(state)
            grad = torch.concatenate((self.format2D_3D(gobs), self.format2D_3D(gprior)), dim=1)

            if "state" in self.input_grad_update:
                grad = torch.concatenate((grad, self.format2D_3D(state)), dim=1)

            if "previous" in self.input_grad_update:
                grad = torch.concatenate((grad, self.format2D_3D(self.h_state)), dim=1)

        elif "gradsplit" in self.input_grad_update:
            prior_cost = self.prior_cost(state)
            obs_cost = self.obs_cost(state, batch)
            # Compute full gradient
            grad_prior = torch.autograd.grad(prior_cost, state, create_graph=True)[0]
            grad_obs = torch.autograd.grad(obs_cost, state, create_graph=True)[0]
            grad = torch.concatenate((grad_prior, grad_obs), dim=1)
            if "state" in self.input_grad_update:
                grad = grad / ((grad**2).mean().sqrt().detach())
                grad = torch.concatenate((grad, state), dim=1)

        elif "grad" in self.input_grad_update:
            var_cost = self.prior_cost(state) + self.lbd**2 * self.obs_cost(state, batch)
            grad = torch.autograd.grad(var_cost, state, create_graph=True)[0]

            if "state" in self.input_grad_update:
                grad = grad / ((grad**2).mean().sqrt().detach())
                grad = torch.concatenate((grad, self.format2D_3D(state)), dim=1)

            if "previous" in self.input_grad_update:
                grad = torch.concatenate((grad, self.format2D_3D(self.h_state)), dim=1)

        elif self.input_grad_update == "obs-only":
            grad = batch.input.nan_to_num()

        elif self.input_grad_update == "obs+state":
            grad = torch.concatenate(
                (self.format2D_3D(state), self.format2D_3D(batch.input.nan_to_num())),
                dim=1,
            )

        gmod = self.grad_mod(grad, timesteps=t, extra=None)
        if hasattr(self.grad_mod, "dim_3d"):
            if self.grad_mod.dim_3d:
                gmod = gmod.squeeze(1)

        state_update = alpha_step * gmod
        if ("grad" in self.input_grad_update) and (self.lr_grad > 0.0):
            state_update += self.lr_grad * (step + 1) / self.n_step * grad[:, : state.shape[1], :, :]

        self.h_state = state_update

        return state - state_update

    def forward(self, batch, x_init=None, h_state=None, phase="test"):
        """
        Perform the forward pass of the solver.

        Args:
            batch (dict): Input batch containing data.

        Returns:
            torch.Tensor: Final optimized state.
        """
        with torch.set_grad_enabled(True):
            state = self.init_state(batch, x_init=x_init)
            self.init_h_state(batch, h_state=h_state)
            self.grad_mod.reset_state(batch.input)

            if not self.training:
                if ("subgrad" in self.input_grad_update) or ("grad" not in self.input_grad_update):
                    state.requires_grad_(False)
                    self.h_state.requires_grad_(False)

            for step in range(self.n_step):
                alpha_step = 1.0 / self.n_step
                state = self.solver_step(state, batch, step=step / self.n_step, alpha_step=alpha_step)
                if (not self.training) and ("grad" in self.input_grad_update):
                    if "subgrad" in self.input_grad_update:
                        state = state.detach().requires_grad_(False)
                    else:
                        state = state.detach().requires_grad_(True)

        return state

__init__(grad_mod, n_step, lr_grad=0.2, lbd=1.0, input_grad_update='state', std_init=0.1, prior_cost=None, obs_cost=None, **kwargs)

Initialize the GradSolver.

Parameters:

Name	Type	Description	Default
prior_cost	Module	The prior cost function.	None
obs_cost	Module	The observation cost function.	None
grad_mod	Module	The gradient modulation model.	required
n_step	int	Number of optimization steps.	required
lr_grad	float	Learning rate for gradient updates. Defaults to 0.2.	0.2
lbd	float	Regularization parameter. Defaults to 1.0.	1.0
input_grad_update	str	Quantities added for updating the input gradient. Defaults to "state".	'state'
std_init	float	Standard deviation for initializing the state. Defaults to 0.1.	0.1

Source code in ocean4dvarnet/models.py

def __init__(
    self,
    grad_mod,
    n_step,
    lr_grad=0.2,
    lbd=1.0,
    input_grad_update="state",
    std_init=0.1,
    prior_cost: Optional[nn.Module] = None,
    obs_cost: Optional[nn.Module] = None,
    **kwargs,
):
    """
    Initialize the GradSolver.

    Args:
        prior_cost (nn.Module): The prior cost function.
        obs_cost (nn.Module): The observation cost function.
        grad_mod (nn.Module): The gradient modulation model.
        n_step (int): Number of optimization steps.
        lr_grad (float, optional): Learning rate for gradient updates. Defaults to 0.2.
        lbd (float, optional): Regularization parameter. Defaults to 1.0.
        input_grad_update (str, optional): Quantities added for updating the input gradient. Defaults to "state".
        std_init (float, optional): Standard deviation for initializing the state. Defaults to 0.1.
    """
    super().__init__()
    self.prior_cost = prior_cost
    self.obs_cost = obs_cost
    self.grad_mod = grad_mod

    self.n_step = n_step
    self.lr_grad = lr_grad
    self.lbd = lbd

    self.input_grad_update = input_grad_update
    self.std_init = std_init

    self._grad_norm = None

forward(batch, x_init=None, h_state=None, phase='test')

Perform the forward pass of the solver.

Parameters:

Name	Type	Description	Default
batch	dict	Input batch containing data.	required

Returns:

Type	Description
	torch.Tensor: Final optimized state.

Source code in ocean4dvarnet/models.py

def forward(self, batch, x_init=None, h_state=None, phase="test"):
    """
    Perform the forward pass of the solver.

    Args:
        batch (dict): Input batch containing data.

    Returns:
        torch.Tensor: Final optimized state.
    """
    with torch.set_grad_enabled(True):
        state = self.init_state(batch, x_init=x_init)
        self.init_h_state(batch, h_state=h_state)
        self.grad_mod.reset_state(batch.input)

        if not self.training:
            if ("subgrad" in self.input_grad_update) or ("grad" not in self.input_grad_update):
                state.requires_grad_(False)
                self.h_state.requires_grad_(False)

        for step in range(self.n_step):
            alpha_step = 1.0 / self.n_step
            state = self.solver_step(state, batch, step=step / self.n_step, alpha_step=alpha_step)
            if (not self.training) and ("grad" in self.input_grad_update):
                if "subgrad" in self.input_grad_update:
                    state = state.detach().requires_grad_(False)
                else:
                    state = state.detach().requires_grad_(True)

    return state

init_h_state(batch, h_state=None)

Initialize the state for optimization.

Parameters:

Name	Type	Description	Default
batch	dict	Input batch containing data.	required
x_init	Tensor	Initial state. Defaults to None.	required

Returns:

Type	Description
	torch.Tensor: Initialized state.

Source code in ocean4dvarnet/models.py

def init_h_state(self, batch, h_state=None):
    """
    Initialize the state for optimization.

    Args:
        batch (dict): Input batch containing data.
        x_init (torch.Tensor, optional): Initial state. Defaults to None.

    Returns:
        torch.Tensor: Initialized state.
    """
    if h_state is not None:
        self.h_state = h_state
    else:
        self.h_state = torch.zeros_like(batch.input).detach().requires_grad_(True)

init_state(batch, x_init=None)

Initialize the state for optimization.

Parameters:

Name	Type	Description	Default
batch	dict	Input batch containing data.	required
x_init	Tensor	Initial state. Defaults to None.	None

Returns:

Type	Description
	torch.Tensor: Initialized state.

Source code in ocean4dvarnet/models.py

def init_state(self, batch, x_init=None):
    """
    Initialize the state for optimization.

    Args:
        batch (dict): Input batch containing data.
        x_init (torch.Tensor, optional): Initial state. Defaults to None.

    Returns:
        torch.Tensor: Initialized state.
    """
    if x_init is not None:
        return x_init.detach().requires_grad_(True)

    if self.std_init > 0:
        x0 = self.std_init * torch.randn_like(batch.input)
        return x0.detach().requires_grad_(True)
    else:
        return torch.zeros_like(batch.input).detach().requires_grad_(True)

solver_step(state, batch, step, alpha_step=1.0)

Perform a single optimization step.

Parameters:

Name	Type	Description	Default
state	Tensor	Current state.	required
batch	dict	Input batch containing data.	required
step	int	Current optimization step between 0 and 1.	required

Returns:

Type	Description
	torch.Tensor: Updated state.

Source code in ocean4dvarnet/models.py

def solver_step(self, state, batch, step, alpha_step=1.0):
    """
    Perform a single optimization step.

    Args:
        state (torch.Tensor): Current state.
        batch (dict): Input batch containing data.
        step (int): Current optimization step between 0 and 1.

    Returns:
        torch.Tensor: Updated state.
    """

    if isinstance(step, float):
        t = torch.tensor([step], device=state.device).repeat(state.shape[0])
    else:
        t = step

    if "subgrad" in self.input_grad_update:
        gobs = (batch.input - state).nan_to_num()

        gprior = state - self.prior_cost.forward_ae(state)
        grad = torch.concatenate((self.format2D_3D(gobs), self.format2D_3D(gprior)), dim=1)

        if "state" in self.input_grad_update:
            grad = torch.concatenate((grad, self.format2D_3D(state)), dim=1)

        if "previous" in self.input_grad_update:
            grad = torch.concatenate((grad, self.format2D_3D(self.h_state)), dim=1)

    elif "gradsplit" in self.input_grad_update:
        prior_cost = self.prior_cost(state)
        obs_cost = self.obs_cost(state, batch)
        # Compute full gradient
        grad_prior = torch.autograd.grad(prior_cost, state, create_graph=True)[0]
        grad_obs = torch.autograd.grad(obs_cost, state, create_graph=True)[0]
        grad = torch.concatenate((grad_prior, grad_obs), dim=1)
        if "state" in self.input_grad_update:
            grad = grad / ((grad**2).mean().sqrt().detach())
            grad = torch.concatenate((grad, state), dim=1)

    elif "grad" in self.input_grad_update:
        var_cost = self.prior_cost(state) + self.lbd**2 * self.obs_cost(state, batch)
        grad = torch.autograd.grad(var_cost, state, create_graph=True)[0]

        if "state" in self.input_grad_update:
            grad = grad / ((grad**2).mean().sqrt().detach())
            grad = torch.concatenate((grad, self.format2D_3D(state)), dim=1)

        if "previous" in self.input_grad_update:
            grad = torch.concatenate((grad, self.format2D_3D(self.h_state)), dim=1)

    elif self.input_grad_update == "obs-only":
        grad = batch.input.nan_to_num()

    elif self.input_grad_update == "obs+state":
        grad = torch.concatenate(
            (self.format2D_3D(state), self.format2D_3D(batch.input.nan_to_num())),
            dim=1,
        )

    gmod = self.grad_mod(grad, timesteps=t, extra=None)
    if hasattr(self.grad_mod, "dim_3d"):
        if self.grad_mod.dim_3d:
            gmod = gmod.squeeze(1)

    state_update = alpha_step * gmod
    if ("grad" in self.input_grad_update) and (self.lr_grad > 0.0):
        state_update += self.lr_grad * (step + 1) / self.n_step * grad[:, : state.shape[1], :, :]

    self.h_state = state_update

    return state - state_update

Lit4dVarNet

Bases: LitModel

A PyTorch Lightning module for training and testing 4DVarNet models. See LitMod for further details.

Source code in ocean4dvarnet/models.py

class Lit4dVarNet(LitModel):
    """
    A PyTorch Lightning module for training and testing 4DVarNet models.
    See LitMod for further details.
    """

    def training_loss(self, batch, out, phase):
        """
        Compute the training loss of 4DVarNet.

        Args:
            batch (dict): Input batch.
            out (tensor): Reconstruction from given observations.
            phase (str): Phase ("train" or "val").

        Returns:
            tuple: Loss.
        """
        base_loss = super().training_loss(batch, out, phase)
        prior_cost = self.solver.prior_cost(self.solver.init_state(batch, out))

        return base_loss + 1.0 * prior_cost

training_loss(batch, out, phase)

Compute the training loss of 4DVarNet.

Parameters:

Name	Type	Description	Default
batch	dict	Input batch.	required
out	tensor	Reconstruction from given observations.	required
phase	str	Phase ("train" or "val").	required

Returns:

Name	Type	Description
tuple		Loss.

Source code in ocean4dvarnet/models.py

def training_loss(self, batch, out, phase):
    """
    Compute the training loss of 4DVarNet.

    Args:
        batch (dict): Input batch.
        out (tensor): Reconstruction from given observations.
        phase (str): Phase ("train" or "val").

    Returns:
        tuple: Loss.
    """
    base_loss = super().training_loss(batch, out, phase)
    prior_cost = self.solver.prior_cost(self.solver.init_state(batch, out))

    return base_loss + 1.0 * prior_cost

LitModel

Bases: LightningModule

A PyTorch Lightning module for training and testing models.

Attributes:

Name	Type	Description
solver	GradSolver	The solver used for optimization.
rec_weight	Tensor	Reconstruction weight for loss computation.
opt_fn	callable	Function to configure the optimizer.
test_metrics	dict	Dictionary of test metrics.
pre_metric_fn	callable	Preprocessing function for metrics.
norm_stats	tuple	Normalization statistics (mean, std).
persist_rw	bool	Whether to persist reconstruction weight as a buffer.

Source code in ocean4dvarnet/models.py

class LitModel(pl.LightningModule):
    """
    A PyTorch Lightning module for training and testing models.

    Attributes:
        solver (GradSolver): The solver used for optimization.
        rec_weight (torch.Tensor): Reconstruction weight for loss computation.
        opt_fn (callable): Function to configure the optimizer.
        test_metrics (dict): Dictionary of test metrics.
        pre_metric_fn (callable): Preprocessing function for metrics.
        norm_stats (tuple): Normalization statistics (mean, std).
        persist_rw (bool): Whether to persist reconstruction weight as a buffer.
    """

    def __init__(
        self, solver, rec_weight, opt_fn, test_metrics=None, pre_metric_fn=None, norm_stats=None, persist_rw=True
    ):
        """
        Initialize the Lit4dVarNet module.

        Args:
            solver (GradSolver): The solver used for optimization.
            rec_weight (numpy.ndarray): Reconstruction weight for loss computation.
            opt_fn (callable): Function to configure the optimizer.
            test_metrics (dict, optional): Dictionary of test metrics.
            pre_metric_fn (callable, optional): Preprocessing function for metrics.
            norm_stats (tuple, optional): Normalization statistics (mean, std).
            persist_rw (bool, optional): Whether to persist reconstruction weight as a buffer.
        """
        super().__init__()
        self.solver = solver
        self.register_buffer("rec_weight", torch.from_numpy(rec_weight), persistent=persist_rw)
        self.test_data = None
        self._norm_stats = norm_stats
        self.opt_fn = opt_fn
        self.metrics = test_metrics or {}
        self.pre_metric_fn = pre_metric_fn or (lambda x: x)

    @property
    def norm_stats(self):
        """
        Retrieve normalization statistics (mean, std).

        Returns:
            tuple: Normalization statistics (mean, std).
        """
        if self._norm_stats is not None:
            return self._norm_stats
        elif self.trainer.datamodule is not None:
            return self.trainer.datamodule.norm_stats()
        return (0.0, 1.0)

    @staticmethod
    def weighted_mse(err, weight):
        """
        Compute the weighted mean squared error.

        Args:
            err (torch.Tensor): Error tensor.
            weight (torch.Tensor): Weight tensor.

        Returns:
            torch.Tensor: Weighted MSE loss.
        """
        err_w = err * weight[None, ...]
        non_zeros = (torch.ones_like(err) * weight[None, ...]) == 0.0
        err_num = err.isfinite() & ~non_zeros
        if err_num.sum() == 0:
            return torch.scalar_tensor(1000.0, device=err_num.device).requires_grad_()
        loss = F.mse_loss(err_w[err_num], torch.zeros_like(err_w[err_num]))
        return loss

    def training_step(self, batch, batch_idx):
        """
        Perform a single training step.

        Args:
            batch (dict): Input batch.
            batch_idx (int): Batch index.

        Returns:
            torch.Tensor: Training loss.
        """
        return self.step(batch, "train")[0]

    def validation_step(self, batch, batch_idx):
        """
        Perform a single validation step.

        Args:
            batch (dict): Input batch.
            batch_idx (int): Batch index.

        Returns:
            torch.Tensor: Validation loss.
        """
        return self.step(batch, "val")[0]

    def forward(self, batch):
        """
        Forward pass through the solver.

        Args:
            batch (dict): Input batch.

        Returns:
            torch.Tensor: Solver output.
        """
        return self.solver(batch)

    def step(self, batch, phase=""):
        """
        Perform a single step for training or validation.

        Args:
            batch (dict): Input batch.
            phase (str, optional): Phase ("train" or "val").

        Returns:
            tuple: Loss and output tensor.
        """
        if self.training and batch.tgt.isfinite().float().mean() < 0.9:
            return None, None

        out = self(batch=batch)

        return self.training_loss(batch, out, phase), out

    def training_loss(self, batch, out, phase):
        """
        Compute the training loss to be backpropagated.

        Args:
            batch (dict): Input batch.
            out (tensor): Reconstruction from given observations.
            phase (str): Phase ("train" or "val").

        Returns:
            tuple: Loss.
        """
        loss = self.weighted_mse(out - batch.tgt, self.rec_weight)
        grad_loss = self.weighted_mse(kfilts.sobel(out) - kfilts.sobel(batch.tgt), self.rec_weight)

        with torch.no_grad():
            self.log(
                f"{phase}_mse", 10000 * loss * self.norm_stats[1] ** 2, prog_bar=True, on_step=False, on_epoch=True
            )
            self.log(f"{phase}_loss", loss, prog_bar=True, on_step=False, on_epoch=True)
            self.log(f"{phase}_gloss", grad_loss, prog_bar=True, on_step=False, on_epoch=True)

        return 50 * loss + 1000 * grad_loss

    def configure_optimizers(self):
        """
        Configure the optimizer.

        Returns:
            torch.optim.Optimizer: Optimizer instance.
        """
        return self.opt_fn(self)

    def test_step(self, batch, batch_idx):
        """
        Perform a single test step.

        Args:
            batch (dict): Input batch.
            batch_idx (int): Batch index.
        """
        if batch_idx == 0:
            self.test_data = []
        out = self(batch=batch)
        m, s = self.norm_stats

        self.test_data.append(
            torch.stack(
                [
                    batch.input.cpu() * s + m,
                    batch.tgt.cpu() * s + m,
                    out.squeeze(dim=-1).detach().cpu() * s + m,
                ],
                dim=1,
            )
        )

    @property
    def test_quantities(self):
        """
        Retrieve the names of test quantities.

        Returns:
            list: List of test quantity names.
        """
        return ["inp", "tgt", "out"]

    def on_test_epoch_end(self):
        """
        Perform actions at the end of the test epoch.

        This includes logging metrics and saving test data.
        """
        rec_da = self.trainer.test_dataloaders.dataset.reconstruct(self.test_data, self.rec_weight.cpu().numpy())

        if isinstance(rec_da, list):
            rec_da = rec_da[0]

        self.test_data = rec_da.assign_coords(dict(v0=self.test_quantities)).to_dataset(dim="v0")

        metric_data = self.test_data.pipe(self.pre_metric_fn)
        metrics = pd.Series({metric_n: metric_fn(metric_data) for metric_n, metric_fn in self.metrics.items()})

        print(metrics.to_frame(name="Metrics").to_markdown())
        if self.logger:
            self.test_data.to_netcdf(Path(self.logger.log_dir) / "test_data.nc")
            print(Path(self.trainer.log_dir) / "test_data.nc")
            self.logger.log_metrics(metrics.to_dict())

norm_stats property

Retrieve normalization statistics (mean, std).

Returns:

Name	Type	Description
tuple		Normalization statistics (mean, std).

test_quantities property

Retrieve the names of test quantities.

Returns:

Name	Type	Description
list		List of test quantity names.

__init__(solver, rec_weight, opt_fn, test_metrics=None, pre_metric_fn=None, norm_stats=None, persist_rw=True)

Initialize the Lit4dVarNet module.

Parameters:

Name	Type	Description	Default
solver	GradSolver	The solver used for optimization.	required
rec_weight	ndarray	Reconstruction weight for loss computation.	required
opt_fn	callable	Function to configure the optimizer.	required
test_metrics	dict	Dictionary of test metrics.	None
pre_metric_fn	callable	Preprocessing function for metrics.	None
norm_stats	tuple	Normalization statistics (mean, std).	None
persist_rw	bool	Whether to persist reconstruction weight as a buffer.	True

Source code in ocean4dvarnet/models.py

def __init__(
    self, solver, rec_weight, opt_fn, test_metrics=None, pre_metric_fn=None, norm_stats=None, persist_rw=True
):
    """
    Initialize the Lit4dVarNet module.

    Args:
        solver (GradSolver): The solver used for optimization.
        rec_weight (numpy.ndarray): Reconstruction weight for loss computation.
        opt_fn (callable): Function to configure the optimizer.
        test_metrics (dict, optional): Dictionary of test metrics.
        pre_metric_fn (callable, optional): Preprocessing function for metrics.
        norm_stats (tuple, optional): Normalization statistics (mean, std).
        persist_rw (bool, optional): Whether to persist reconstruction weight as a buffer.
    """
    super().__init__()
    self.solver = solver
    self.register_buffer("rec_weight", torch.from_numpy(rec_weight), persistent=persist_rw)
    self.test_data = None
    self._norm_stats = norm_stats
    self.opt_fn = opt_fn
    self.metrics = test_metrics or {}
    self.pre_metric_fn = pre_metric_fn or (lambda x: x)

configure_optimizers()

Configure the optimizer.

Returns:

Type	Description
	torch.optim.Optimizer: Optimizer instance.

Source code in ocean4dvarnet/models.py

def configure_optimizers(self):
    """
    Configure the optimizer.

    Returns:
        torch.optim.Optimizer: Optimizer instance.
    """
    return self.opt_fn(self)

forward(batch)

Forward pass through the solver.

Parameters:

Name	Type	Description	Default
batch	dict	Input batch.	required

Returns:

Type	Description
	torch.Tensor: Solver output.

Source code in ocean4dvarnet/models.py

def forward(self, batch):
    """
    Forward pass through the solver.

    Args:
        batch (dict): Input batch.

    Returns:
        torch.Tensor: Solver output.
    """
    return self.solver(batch)

on_test_epoch_end()

Perform actions at the end of the test epoch.

This includes logging metrics and saving test data.

Source code in ocean4dvarnet/models.py

def on_test_epoch_end(self):
    """
    Perform actions at the end of the test epoch.

    This includes logging metrics and saving test data.
    """
    rec_da = self.trainer.test_dataloaders.dataset.reconstruct(self.test_data, self.rec_weight.cpu().numpy())

    if isinstance(rec_da, list):
        rec_da = rec_da[0]

    self.test_data = rec_da.assign_coords(dict(v0=self.test_quantities)).to_dataset(dim="v0")

    metric_data = self.test_data.pipe(self.pre_metric_fn)
    metrics = pd.Series({metric_n: metric_fn(metric_data) for metric_n, metric_fn in self.metrics.items()})

    print(metrics.to_frame(name="Metrics").to_markdown())
    if self.logger:
        self.test_data.to_netcdf(Path(self.logger.log_dir) / "test_data.nc")
        print(Path(self.trainer.log_dir) / "test_data.nc")
        self.logger.log_metrics(metrics.to_dict())

step(batch, phase='')

Perform a single step for training or validation.

Parameters:

Name	Type	Description	Default
batch	dict	Input batch.	required
phase	str	Phase ("train" or "val").	''

Returns:

Name	Type	Description
tuple		Loss and output tensor.

Source code in ocean4dvarnet/models.py

def step(self, batch, phase=""):
    """
    Perform a single step for training or validation.

    Args:
        batch (dict): Input batch.
        phase (str, optional): Phase ("train" or "val").

    Returns:
        tuple: Loss and output tensor.
    """
    if self.training and batch.tgt.isfinite().float().mean() < 0.9:
        return None, None

    out = self(batch=batch)

    return self.training_loss(batch, out, phase), out

test_step(batch, batch_idx)

Perform a single test step.

Parameters:

Name	Type	Description	Default
batch	dict	Input batch.	required
batch_idx	int	Batch index.	required

Source code in ocean4dvarnet/models.py

def test_step(self, batch, batch_idx):
    """
    Perform a single test step.

    Args:
        batch (dict): Input batch.
        batch_idx (int): Batch index.
    """
    if batch_idx == 0:
        self.test_data = []
    out = self(batch=batch)
    m, s = self.norm_stats

    self.test_data.append(
        torch.stack(
            [
                batch.input.cpu() * s + m,
                batch.tgt.cpu() * s + m,
                out.squeeze(dim=-1).detach().cpu() * s + m,
            ],
            dim=1,
        )
    )

training_loss(batch, out, phase)

Compute the training loss to be backpropagated.

Parameters:

Name	Type	Description	Default
batch	dict	Input batch.	required
out	tensor	Reconstruction from given observations.	required
phase	str	Phase ("train" or "val").	required

Returns:

Name	Type	Description
tuple		Loss.

Source code in ocean4dvarnet/models.py

def training_loss(self, batch, out, phase):
    """
    Compute the training loss to be backpropagated.

    Args:
        batch (dict): Input batch.
        out (tensor): Reconstruction from given observations.
        phase (str): Phase ("train" or "val").

    Returns:
        tuple: Loss.
    """
    loss = self.weighted_mse(out - batch.tgt, self.rec_weight)
    grad_loss = self.weighted_mse(kfilts.sobel(out) - kfilts.sobel(batch.tgt), self.rec_weight)

    with torch.no_grad():
        self.log(
            f"{phase}_mse", 10000 * loss * self.norm_stats[1] ** 2, prog_bar=True, on_step=False, on_epoch=True
        )
        self.log(f"{phase}_loss", loss, prog_bar=True, on_step=False, on_epoch=True)
        self.log(f"{phase}_gloss", grad_loss, prog_bar=True, on_step=False, on_epoch=True)

    return 50 * loss + 1000 * grad_loss

training_step(batch, batch_idx)

Perform a single training step.

Parameters:

Name	Type	Description	Default
batch	dict	Input batch.	required
batch_idx	int	Batch index.	required

Returns:

Type	Description
	torch.Tensor: Training loss.

Source code in ocean4dvarnet/models.py

def training_step(self, batch, batch_idx):
    """
    Perform a single training step.

    Args:
        batch (dict): Input batch.
        batch_idx (int): Batch index.

    Returns:
        torch.Tensor: Training loss.
    """
    return self.step(batch, "train")[0]

validation_step(batch, batch_idx)

Perform a single validation step.

Parameters:

Name	Type	Description	Default
batch	dict	Input batch.	required
batch_idx	int	Batch index.	required

Returns:

Type	Description
	torch.Tensor: Validation loss.

Source code in ocean4dvarnet/models.py

def validation_step(self, batch, batch_idx):
    """
    Perform a single validation step.

    Args:
        batch (dict): Input batch.
        batch_idx (int): Batch index.

    Returns:
        torch.Tensor: Validation loss.
    """
    return self.step(batch, "val")[0]

weighted_mse(err, weight) staticmethod

Compute the weighted mean squared error.

Parameters:

Name	Type	Description	Default
err	Tensor	Error tensor.	required
weight	Tensor	Weight tensor.	required

Returns:

Type	Description
	torch.Tensor: Weighted MSE loss.

Source code in ocean4dvarnet/models.py

@staticmethod
def weighted_mse(err, weight):
    """
    Compute the weighted mean squared error.

    Args:
        err (torch.Tensor): Error tensor.
        weight (torch.Tensor): Weight tensor.

    Returns:
        torch.Tensor: Weighted MSE loss.
    """
    err_w = err * weight[None, ...]
    non_zeros = (torch.ones_like(err) * weight[None, ...]) == 0.0
    err_num = err.isfinite() & ~non_zeros
    if err_num.sum() == 0:
        return torch.scalar_tensor(1000.0, device=err_num.device).requires_grad_()
    loss = F.mse_loss(err_w[err_num], torch.zeros_like(err_w[err_num]))
    return loss