Torch EnKF

I used to write a EnKF in pytorch. May be useful in someday.

# -*- coding: utf-8 -*-
"""
Created on Mon Oct 12 09:56:07 2023
@author: hanjingye
"""
import torch 
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
 
def enkf(priori_state, obs_state, obs_std, H=None, mask=None):
    """
    Batched Ensemble Kalman Filter (EnKF) using observation sampling for state estimation.
 
    Args:
    prior_states (torch.Tensor): The prior ensemble state matrix of size (batch_size, ensemble_size, state_size).
    obs_states (torch.Tensor): The observation state vector of size (batch_size, observation_size).
    obs_std (torch.Tensor): The standard deviation of observation states of size (batch_size, observation_size).
    H (torch.Tensor, optional): The projection matrix from state space to observation space. Size (observation_size, state_size).
    mask (torch.Tensor, optional): The update control matrix. Size (state_size, observation_size).
 
    Returns:
    torch.Tensor: The posterior ensemble state matrix.
    """
    # Ensure all tensors are on the same device as specified by global variable `device`
    prior_states = priori_state.to(device)
    obs_states = obs_state.to(device)
    obs_std = obs_std.to(device)
 
    batch_size, ensemble_size, state_size = prior_states.shape
    _, observation_size = obs_states.shape
 
    # Regularize observation standard deviation
    obs_std += 1e-5
    obs_std[torch.isnan(obs_std)] = 99999999.9
     
    # Default H matrix if not provided
    if H is None:
        H = torch.eye(observation_size, state_size).to(device)
 
    # Default mask matrix if not provided
    if mask is None:
        mask = torch.ones((state_size, observation_size)).to(device)
 
    # Generate observation samples for each ensemble member
    obs_samples = obs_states.unsqueeze(1).repeat(1, ensemble_size, 1) + torch.randn(batch_size, ensemble_size, observation_size).to(device) * obs_std.unsqueeze(1)
 
    # Replace missing observations (nan) in samples with the corresponding ensemble values
    nan_mask = torch.isnan(obs_samples)
    obs_samples[nan_mask] = prior_states[:, :, :observation_size][nan_mask]
 
    # Calculate ensemble mean and anomalies
    ensemble_mean = prior_states.mean(dim=1, keepdim=True)
    anomalies = prior_states - ensemble_mean
 
    # Compute Covariance Matrices
    CovStaSta = torch.matmul(anomalies.transpose(1, 2), anomalies) / (ensemble_size - 1)
    CovObsObs = torch.matmul(torch.matmul(H, CovStaSta), H.transpose(0, 1)) + torch.diag_embed(obs_std**2)
    P = torch.matmul(CovStaSta, H.transpose(0, 1))
 
    # Compute Kalman Gain
    K = torch.matmul(P, torch.linalg.inv(CovObsObs))
 
    # Apply mask to Kalman Gain
    K = K * mask
    innovation = obs_samples - torch.matmul(prior_states, H.transpose(0, 1))
    posterior_states = prior_states + torch.matmul(K, innovation.transpose(1, 2)).transpose(1, 2)
    return posterior_states

No target file, test, entry point, or integration requirement is named; start by reviewing the repository structure and deciding where this PyTorch EnKF would belong. Done would require an agreed integration point, defined expected behavior, and tests, none of which are specified in the issue.