data

This module provides data handling utilities for 4D-VarNet models.

It includes classes and functions for creating datasets, augmenting data, managing data loading pipelines, and reconstructing data from patches. These utilities are designed to work seamlessly with PyTorch and xarray, enabling efficient data preprocessing and loading for machine learning tasks.

Classes:

Name	Description
- XrDataset	A PyTorch Dataset for extracting patches from xarray.DataArray objects.
- XrConcatDataset	A concatenation of multiple XrDatasets.
- AugmentedDataset	A dataset wrapper for applying data augmentation.
- BaseDataModule	A PyTorch Lightning DataModule for managing datasets and data loaders.
- ConcatDataModule	A DataModule for combining datasets from multiple domains.
- RandValDataModule	A DataModule for random splitting of training data into training and validation sets.

Raises:

Type	Description
-IncompleteScanConfiguration	Raised when the scan configuration does not cover the entire domain.
-DangerousDimOrdering	Raised when the dimension ordering of the input data is incorrect.

Key Features

• Patch extraction: Efficiently extract patches from large xarray.DataArray objects for training.
• Data augmentation: Support for augmenting datasets with noise and transformations.
• Reconstruction: Reconstruct the original data from extracted patches.
• Seamless integration: Designed to work with PyTorch Lightning for streamlined training pipelines.

AugmentedDataset

Bases: Dataset

A dataset that applies data augmentation to an input dataset.

Attributes:

Name	Type	Description
inp_ds	Dataset	The input dataset.
aug_factor	int	The number of augmented copies to generate.
aug_only	bool	Whether to include only augmented data.
noise_sigma	float	Standard deviation of noise to add to augmented data.

Source code in ocean4dvarnet/data.py

class AugmentedDataset(torch.utils.data.Dataset):
    """
    A dataset that applies data augmentation to an input dataset.

    Attributes:
        inp_ds (torch.utils.data.Dataset): The input dataset.
        aug_factor (int): The number of augmented copies to generate.
        aug_only (bool): Whether to include only augmented data.
        noise_sigma (float): Standard deviation of noise to add to augmented data.
    """

    def __init__(self, inp_ds, aug_factor, aug_only=False, noise_sigma=None):
        """
        Initialize the AugmentedDataset.

        Args:
            inp_ds (torch.utils.data.Dataset): The input dataset.
            aug_factor (int): The number of augmented copies to generate.
            aug_only (bool, optional): Whether to include only augmented data.
            noise_sigma (float, optional): Standard deviation of noise to add to augmented data.
        """
        self.aug_factor = aug_factor
        self.aug_only = aug_only
        self.inp_ds = inp_ds
        self.perm = np.random.permutation(len(self.inp_ds))
        self.noise_sigma = noise_sigma

    def __len__(self):
        """
        Return the total number of items in the dataset.

        Returns:
            int: Total number of items.
        """
        return len(self.inp_ds) * (1 + self.aug_factor - int(self.aug_only))

    def __getitem__(self, idx):
        """
        Get an item from the dataset.

        Args:
            idx (int): Index of the item.

        Returns:
            TrainingItem: The requested item.
        """
        if self.aug_only:
            idx = idx + len(self.inp_ds)

        if idx < len(self.inp_ds):
            return self.inp_ds[idx]

        tgt_idx = idx % len(self.inp_ds)
        perm_idx = tgt_idx
        for _ in range(idx // len(self.inp_ds)):
            perm_idx = self.perm[perm_idx]

        item = self.inp_ds[tgt_idx]
        perm_item = self.inp_ds[perm_idx]

        noise = np.zeros_like(item.input, dtype=np.float32)
        if self.noise_sigma is not None:
            noise = np.random.randn(*item.input.shape).astype(np.float32) * self.noise_sigma

        return item._replace(input=noise + np.where(np.isfinite(perm_item.input),
                             item.tgt, np.full_like(item.tgt, np.nan)))

__getitem__(idx)

Get an item from the dataset.

Parameters:

Name	Type	Description	Default
idx	int	Index of the item.	required

Returns:

Name	Type	Description
TrainingItem		The requested item.

Source code in ocean4dvarnet/data.py

def __getitem__(self, idx):
    """
    Get an item from the dataset.

    Args:
        idx (int): Index of the item.

    Returns:
        TrainingItem: The requested item.
    """
    if self.aug_only:
        idx = idx + len(self.inp_ds)

    if idx < len(self.inp_ds):
        return self.inp_ds[idx]

    tgt_idx = idx % len(self.inp_ds)
    perm_idx = tgt_idx
    for _ in range(idx // len(self.inp_ds)):
        perm_idx = self.perm[perm_idx]

    item = self.inp_ds[tgt_idx]
    perm_item = self.inp_ds[perm_idx]

    noise = np.zeros_like(item.input, dtype=np.float32)
    if self.noise_sigma is not None:
        noise = np.random.randn(*item.input.shape).astype(np.float32) * self.noise_sigma

    return item._replace(input=noise + np.where(np.isfinite(perm_item.input),
                         item.tgt, np.full_like(item.tgt, np.nan)))

__init__(inp_ds, aug_factor, aug_only=False, noise_sigma=None)

Initialize the AugmentedDataset.

Parameters:

Name	Type	Description	Default
inp_ds	Dataset	The input dataset.	required
aug_factor	int	The number of augmented copies to generate.	required
aug_only	bool	Whether to include only augmented data.	False
noise_sigma	float	Standard deviation of noise to add to augmented data.	None

Source code in ocean4dvarnet/data.py

def __init__(self, inp_ds, aug_factor, aug_only=False, noise_sigma=None):
    """
    Initialize the AugmentedDataset.

    Args:
        inp_ds (torch.utils.data.Dataset): The input dataset.
        aug_factor (int): The number of augmented copies to generate.
        aug_only (bool, optional): Whether to include only augmented data.
        noise_sigma (float, optional): Standard deviation of noise to add to augmented data.
    """
    self.aug_factor = aug_factor
    self.aug_only = aug_only
    self.inp_ds = inp_ds
    self.perm = np.random.permutation(len(self.inp_ds))
    self.noise_sigma = noise_sigma

__len__()

Return the total number of items in the dataset.

Returns:

Name	Type	Description
int		Total number of items.

Source code in ocean4dvarnet/data.py

def __len__(self):
    """
    Return the total number of items in the dataset.

    Returns:
        int: Total number of items.
    """
    return len(self.inp_ds) * (1 + self.aug_factor - int(self.aug_only))

BaseDataModule

Bases: LightningDataModule

A base data module for managing datasets and data loaders in PyTorch Lightning.

Attributes:

Name	Type	Description
input_da	DataArray	The input data array.
domains	dict	Dictionary of domain splits (train, val, test).
xrds_kw	dict	Keyword arguments for XrDataset.
dl_kw	dict	Keyword arguments for DataLoader.
aug_kw	dict	Keyword arguments for AugmentedDataset.
norm_stats	tuple	Normalization statistics (mean, std).

Source code in ocean4dvarnet/data.py

class BaseDataModule(pl.LightningDataModule):
    """
    A base data module for managing datasets and data loaders in PyTorch Lightning.

    Attributes:
        input_da (xarray.DataArray): The input data array.
        domains (dict): Dictionary of domain splits (train, val, test).
        xrds_kw (dict): Keyword arguments for XrDataset.
        dl_kw (dict): Keyword arguments for DataLoader.
        aug_kw (dict): Keyword arguments for AugmentedDataset.
        norm_stats (tuple): Normalization statistics (mean, std).
    """

    def __init__(self, input_da, domains, xrds_kw, dl_kw, aug_kw=None, norm_stats=None, **kwargs):
        """
        Initialize the BaseDataModule.

        Args:
            input_da (xarray.DataArray): The input data array.
            domains (dict): Dictionary of domain splits (train, val, test).
            xrds_kw (dict): Keyword arguments for XrDataset.
            dl_kw (dict): Keyword arguments for DataLoader.
            aug_kw (dict, optional): Keyword arguments for AugmentedDataset.
            norm_stats (tuple, optional): Normalization statistics (mean, std).
        """
        super().__init__()
        self.input_da = input_da
        self.domains = domains
        self.xrds_kw = xrds_kw
        self.dl_kw = dl_kw
        self.aug_kw = aug_kw if aug_kw is not None else {}
        self._norm_stats = norm_stats

        self.train_ds = None
        self.val_ds = None
        self.test_ds = None
        self._post_fn = None

    def norm_stats(self):
        """
        Compute or retrieve normalization statistics (mean, std).

        Returns:
            tuple: Normalization statistics (mean, std).
        """
        if self._norm_stats is None:
            self._norm_stats = self.train_mean_std()
            logger.info(f"Normalisation parameters: {self._norm_stats}")
        return self._norm_stats

    def train_mean_std(self, variable='tgt'):
        """
        Compute the mean and standard deviation of the training data.

        Args:
            variable (str, optional): Variable to compute statistics for.

        Returns:
            tuple: Mean and standard deviation.
        """
        train_data = self.input_da.sel(self.xrds_kw.get('domain_limits', {})).sel(self.domains['train'])
        return train_data.sel(variable=variable).pipe(lambda da: (da.mean().values.item(), da.std().values.item()))

    def post_fn(self):
        """
        Create a post-processing function for normalizing data.

        Returns:
            callable: Post-processing function.
        """
        m, s = self.norm_stats()
        def normalize(item): return (item - m) / s
        return ft.partial(ft.reduce, lambda i, f: f(i), [
            TrainingItem._make,
            lambda item: item._replace(tgt=normalize(item.tgt)),
            lambda item: item._replace(input=normalize(item.input)),
        ])

    def setup(self, stage='test'):
        """
        Set up the datasets for training, validation, and testing.

        Args:
            stage (str, optional): Stage of the setup ('train', 'val', 'test').
        """
        train_data = self.input_da.sel(self.domains['train'])
        post_fn = self.post_fn()
        self.train_ds = XrDataset(
            train_data, **self.xrds_kw, postpro_fn=post_fn,
        )
        if self.aug_kw:
            self.train_ds = AugmentedDataset(self.train_ds, **self.aug_kw)

        self.val_ds = XrDataset(
            self.input_da.sel(self.domains['val']), **self.xrds_kw, postpro_fn=post_fn,
        )
        self.test_ds = XrDataset(
            self.input_da.sel(self.domains['test']), **self.xrds_kw, postpro_fn=post_fn,
        )

    def train_dataloader(self):
        """
        Create a DataLoader for the training dataset.

        Returns:
            DataLoader: Training DataLoader.
        """
        return torch.utils.data.DataLoader(self.train_ds, shuffle=True, **self.dl_kw)

    def val_dataloader(self):
        """
        Create a DataLoader for the validation dataset.

        Returns:
            DataLoader: Validation DataLoader.
        """
        return torch.utils.data.DataLoader(self.val_ds, shuffle=False, **self.dl_kw)

    def test_dataloader(self):
        """
        Create a DataLoader for the testing dataset.

        Returns:
            DataLoader: Testing DataLoader.
        """
        return torch.utils.data.DataLoader(self.test_ds, shuffle=False, **self.dl_kw)

__init__(input_da, domains, xrds_kw, dl_kw, aug_kw=None, norm_stats=None, **kwargs)

Initialize the BaseDataModule.

Parameters:

Name	Type	Description	Default
input_da	DataArray	The input data array.	required
domains	dict	Dictionary of domain splits (train, val, test).	required
xrds_kw	dict	Keyword arguments for XrDataset.	required
dl_kw	dict	Keyword arguments for DataLoader.	required
aug_kw	dict	Keyword arguments for AugmentedDataset.	None
norm_stats	tuple	Normalization statistics (mean, std).	None

Source code in ocean4dvarnet/data.py

def __init__(self, input_da, domains, xrds_kw, dl_kw, aug_kw=None, norm_stats=None, **kwargs):
    """
    Initialize the BaseDataModule.

    Args:
        input_da (xarray.DataArray): The input data array.
        domains (dict): Dictionary of domain splits (train, val, test).
        xrds_kw (dict): Keyword arguments for XrDataset.
        dl_kw (dict): Keyword arguments for DataLoader.
        aug_kw (dict, optional): Keyword arguments for AugmentedDataset.
        norm_stats (tuple, optional): Normalization statistics (mean, std).
    """
    super().__init__()
    self.input_da = input_da
    self.domains = domains
    self.xrds_kw = xrds_kw
    self.dl_kw = dl_kw
    self.aug_kw = aug_kw if aug_kw is not None else {}
    self._norm_stats = norm_stats

    self.train_ds = None
    self.val_ds = None
    self.test_ds = None
    self._post_fn = None

norm_stats()

Compute or retrieve normalization statistics (mean, std).

Returns:

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

Source code in ocean4dvarnet/data.py

def norm_stats(self):
    """
    Compute or retrieve normalization statistics (mean, std).

    Returns:
        tuple: Normalization statistics (mean, std).
    """
    if self._norm_stats is None:
        self._norm_stats = self.train_mean_std()
        logger.info(f"Normalisation parameters: {self._norm_stats}")
    return self._norm_stats

post_fn()

Create a post-processing function for normalizing data.

Returns:

Name	Type	Description
callable		Post-processing function.

Source code in ocean4dvarnet/data.py

def post_fn(self):
    """
    Create a post-processing function for normalizing data.

    Returns:
        callable: Post-processing function.
    """
    m, s = self.norm_stats()
    def normalize(item): return (item - m) / s
    return ft.partial(ft.reduce, lambda i, f: f(i), [
        TrainingItem._make,
        lambda item: item._replace(tgt=normalize(item.tgt)),
        lambda item: item._replace(input=normalize(item.input)),
    ])

setup(stage='test')

Set up the datasets for training, validation, and testing.

Parameters:

Name	Type	Description	Default
stage	str	Stage of the setup ('train', 'val', 'test').	'test'

Source code in ocean4dvarnet/data.py

def setup(self, stage='test'):
    """
    Set up the datasets for training, validation, and testing.

    Args:
        stage (str, optional): Stage of the setup ('train', 'val', 'test').
    """
    train_data = self.input_da.sel(self.domains['train'])
    post_fn = self.post_fn()
    self.train_ds = XrDataset(
        train_data, **self.xrds_kw, postpro_fn=post_fn,
    )
    if self.aug_kw:
        self.train_ds = AugmentedDataset(self.train_ds, **self.aug_kw)

    self.val_ds = XrDataset(
        self.input_da.sel(self.domains['val']), **self.xrds_kw, postpro_fn=post_fn,
    )
    self.test_ds = XrDataset(
        self.input_da.sel(self.domains['test']), **self.xrds_kw, postpro_fn=post_fn,
    )

test_dataloader()

Create a DataLoader for the testing dataset.

Returns:

Name	Type	Description
DataLoader		Testing DataLoader.

Source code in ocean4dvarnet/data.py

def test_dataloader(self):
    """
    Create a DataLoader for the testing dataset.

    Returns:
        DataLoader: Testing DataLoader.
    """
    return torch.utils.data.DataLoader(self.test_ds, shuffle=False, **self.dl_kw)

train_dataloader()

Create a DataLoader for the training dataset.

Returns:

Name	Type	Description
DataLoader		Training DataLoader.

Source code in ocean4dvarnet/data.py

def train_dataloader(self):
    """
    Create a DataLoader for the training dataset.

    Returns:
        DataLoader: Training DataLoader.
    """
    return torch.utils.data.DataLoader(self.train_ds, shuffle=True, **self.dl_kw)

train_mean_std(variable='tgt')

Compute the mean and standard deviation of the training data.

Parameters:

Name	Type	Description	Default
variable	str	Variable to compute statistics for.	'tgt'

Returns:

Name	Type	Description
tuple		Mean and standard deviation.

Source code in ocean4dvarnet/data.py

def train_mean_std(self, variable='tgt'):
    """
    Compute the mean and standard deviation of the training data.

    Args:
        variable (str, optional): Variable to compute statistics for.

    Returns:
        tuple: Mean and standard deviation.
    """
    train_data = self.input_da.sel(self.xrds_kw.get('domain_limits', {})).sel(self.domains['train'])
    return train_data.sel(variable=variable).pipe(lambda da: (da.mean().values.item(), da.std().values.item()))

val_dataloader()

Create a DataLoader for the validation dataset.

Returns:

Name	Type	Description
DataLoader		Validation DataLoader.

Source code in ocean4dvarnet/data.py

def val_dataloader(self):
    """
    Create a DataLoader for the validation dataset.

    Returns:
        DataLoader: Validation DataLoader.
    """
    return torch.utils.data.DataLoader(self.val_ds, shuffle=False, **self.dl_kw)

ConcatDataModule

Bases: BaseDataModule

A data module for concatenating datasets from multiple domains.

Source code in ocean4dvarnet/data.py

class ConcatDataModule(BaseDataModule):
    """A data module for concatenating datasets from multiple domains."""

    def train_mean_std(self):
        """
        Compute the mean and standard deviation of the training data across domains.

        Returns:
            tuple: Mean and standard deviation.
        """
        sum, count = 0, 0
        train_data = self.input_da.sel(self.xrds_kw.get('domain_limits', {}))
        for domain in self.domains['train']:
            _sum, _count = train_data.sel(domain).sel(variable='tgt').pipe(
                lambda da: (da.sum(), da.pipe(np.isfinite).sum())
            )
            sum += _sum
            count += _count

        mean = sum / count
        sum = 0
        for domain in self.domains['train']:
            _sum = train_data.sel(domain).sel(variable='tgt').pipe(lambda da: da - mean).pipe(np.square).sum()
            sum += _sum
        std = (sum / count)**0.5
        return mean.values.item(), std.values.item()

    def setup(self, stage='test'):
        """
        Set up the datasets for training, validation, and testing.

        Args:
            stage (str, optional): Stage of the setup ('train', 'val', 'test').
        """
        post_fn = self.post_fn()
        self.train_ds = XrConcatDataset([
            XrDataset(self.input_da.sel(domain), **self.xrds_kw, postpro_fn=post_fn,)
            for domain in self.domains['train']
        ])
        if self.aug_factor >= 1:
            self.train_ds = AugmentedDataset(self.train_ds, **self.aug_kw)

        self.val_ds = XrConcatDataset([
            XrDataset(self.input_da.sel(domain), **self.xrds_kw, postpro_fn=post_fn,)
            for domain in self.domains['val']
        ])
        self.test_ds = XrConcatDataset([
            XrDataset(self.input_da.sel(domain), **self.xrds_kw, postpro_fn=post_fn,)
            for domain in self.domains['test']
        ])

setup(stage='test')

Set up the datasets for training, validation, and testing.

Parameters:

Name	Type	Description	Default
stage	str	Stage of the setup ('train', 'val', 'test').	'test'

Source code in ocean4dvarnet/data.py

def setup(self, stage='test'):
    """
    Set up the datasets for training, validation, and testing.

    Args:
        stage (str, optional): Stage of the setup ('train', 'val', 'test').
    """
    post_fn = self.post_fn()
    self.train_ds = XrConcatDataset([
        XrDataset(self.input_da.sel(domain), **self.xrds_kw, postpro_fn=post_fn,)
        for domain in self.domains['train']
    ])
    if self.aug_factor >= 1:
        self.train_ds = AugmentedDataset(self.train_ds, **self.aug_kw)

    self.val_ds = XrConcatDataset([
        XrDataset(self.input_da.sel(domain), **self.xrds_kw, postpro_fn=post_fn,)
        for domain in self.domains['val']
    ])
    self.test_ds = XrConcatDataset([
        XrDataset(self.input_da.sel(domain), **self.xrds_kw, postpro_fn=post_fn,)
        for domain in self.domains['test']
    ])

train_mean_std()

Compute the mean and standard deviation of the training data across domains.

Returns:

Name	Type	Description
tuple		Mean and standard deviation.

Source code in ocean4dvarnet/data.py

def train_mean_std(self):
    """
    Compute the mean and standard deviation of the training data across domains.

    Returns:
        tuple: Mean and standard deviation.
    """
    sum, count = 0, 0
    train_data = self.input_da.sel(self.xrds_kw.get('domain_limits', {}))
    for domain in self.domains['train']:
        _sum, _count = train_data.sel(domain).sel(variable='tgt').pipe(
            lambda da: (da.sum(), da.pipe(np.isfinite).sum())
        )
        sum += _sum
        count += _count

    mean = sum / count
    sum = 0
    for domain in self.domains['train']:
        _sum = train_data.sel(domain).sel(variable='tgt').pipe(lambda da: da - mean).pipe(np.square).sum()
        sum += _sum
    std = (sum / count)**0.5
    return mean.values.item(), std.values.item()

DangerousDimOrdering

Bases: Exception

Exception raised when the dimension ordering of the input data is incorrect.

Source code in ocean4dvarnet/data.py

class DangerousDimOrdering(Exception):
    """Exception raised when the dimension ordering of the input data is incorrect."""
    pass

IncompleteScanConfiguration

Bases: Exception

Exception raised when the scan configuration does not cover the entire domain.

Source code in ocean4dvarnet/data.py

class IncompleteScanConfiguration(Exception):
    """Exception raised when the scan configuration does not cover the entire domain."""
    pass

LazyDataModule

Bases: BaseDataModule

A data module loading datasets in lazy mode ("on the fly").

Attributes: see BaseDataModule.

Source code in ocean4dvarnet/data.py

class LazyDataModule(BaseDataModule):
    """
    A data module loading datasets in lazy mode ("on the fly").

    Attributes: see BaseDataModule.
    """
    def __init__(self, *args, **kwargs):
        """
        See BaseDataModule.__init__.

        Differences are:
        - `input_da` must contain a xr.Dataset-valued dictionary;
        - `norm_stats` must be indicated, otherwise a TypeError will be
            raised.
        """
        super().__init__(*args, **kwargs)

        if not isinstance(self.input_da, dict):
            raise TypeError(
                "Argument `input_da` is expected to be a `dict`, "
                + f"got `{type(self.input_da)}`."
            )

        if self._norm_stats is None:
            raise TypeError(
                "Normalisation parameters (argument `norm_stats`) must "
                + "be provided in lazy loading"
            )

    def setup(self, stage='test'):
        """
        Set up the datasets for training, validation, and testing.

        Args:
            stage (str, optional): Stage of the setup ('train', 'val', 'test').
        """
        self.train_ds = LazyXrDataset(
            {k: v.sel(self.domains['train']) for (k, v) in self.input_da.items()},
            **self.xrds_kw["train"], postpro_fn=self.post_fn('train'),
        )
        if self.aug_kw:
            self.train_ds = AugmentedDataset(self.train_ds, **self.aug_kw)

        self.val_ds = LazyXrDataset(
            {k: v.sel(self.domains['val']) for (k, v) in self.input_da.items()},
            **self.xrds_kw["val"], postpro_fn=self.post_fn('val'),
        )
        self.test_ds = LazyXrDataset(
            {k: v.sel(self.domains['test']) for (k, v) in self.input_da.items()},
            **self.xrds_kw["test"], postpro_fn=self.post_fn('test'),
        )

    def norm_stats(self, phase=None):
        """
        Compute or retrieve normalization statistics (mean, std).

        Returns:
            tuple: Normalization statistics (mean, std).
        """
        if self._norm_stats is None:
            self._norm_stats = self.train_mean_std()
            logger.info(f"Normalisation parameters: {self._norm_stats}")
        return self._norm_stats[phase]

    def post_fn(self, phase=None):
        """
        Create a post-processing function for normalizing data depending
        on the dataset (training, validation or test dataset).

        Returns:
            callable: Post-processing function.
        """
        m, s = self.norm_stats(phase)
        def normalize(item): return (item - m) / s
        return ft.partial(ft.reduce, lambda i, f: f(i), [
            TrainingItem._make,
            lambda item: item._replace(tgt=normalize(item.tgt)),
            lambda item: item._replace(input=normalize(item.input)),
        ])

__init__(*args, **kwargs)

See BaseDataModule.init.

Differences are: - input_da must contain a xr.Dataset-valued dictionary; - norm_stats must be indicated, otherwise a TypeError will be raised.

Source code in ocean4dvarnet/data.py

def __init__(self, *args, **kwargs):
    """
    See BaseDataModule.__init__.

    Differences are:
    - `input_da` must contain a xr.Dataset-valued dictionary;
    - `norm_stats` must be indicated, otherwise a TypeError will be
        raised.
    """
    super().__init__(*args, **kwargs)

    if not isinstance(self.input_da, dict):
        raise TypeError(
            "Argument `input_da` is expected to be a `dict`, "
            + f"got `{type(self.input_da)}`."
        )

    if self._norm_stats is None:
        raise TypeError(
            "Normalisation parameters (argument `norm_stats`) must "
            + "be provided in lazy loading"
        )

norm_stats(phase=None)

Compute or retrieve normalization statistics (mean, std).

Returns:

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

Source code in ocean4dvarnet/data.py

def norm_stats(self, phase=None):
    """
    Compute or retrieve normalization statistics (mean, std).

    Returns:
        tuple: Normalization statistics (mean, std).
    """
    if self._norm_stats is None:
        self._norm_stats = self.train_mean_std()
        logger.info(f"Normalisation parameters: {self._norm_stats}")
    return self._norm_stats[phase]

post_fn(phase=None)

Create a post-processing function for normalizing data depending on the dataset (training, validation or test dataset).

Returns:

Name	Type	Description
callable		Post-processing function.

Source code in ocean4dvarnet/data.py

def post_fn(self, phase=None):
    """
    Create a post-processing function for normalizing data depending
    on the dataset (training, validation or test dataset).

    Returns:
        callable: Post-processing function.
    """
    m, s = self.norm_stats(phase)
    def normalize(item): return (item - m) / s
    return ft.partial(ft.reduce, lambda i, f: f(i), [
        TrainingItem._make,
        lambda item: item._replace(tgt=normalize(item.tgt)),
        lambda item: item._replace(input=normalize(item.input)),
    ])

setup(stage='test')

Set up the datasets for training, validation, and testing.

Parameters:

Name	Type	Description	Default
stage	str	Stage of the setup ('train', 'val', 'test').	'test'

Source code in ocean4dvarnet/data.py

def setup(self, stage='test'):
    """
    Set up the datasets for training, validation, and testing.

    Args:
        stage (str, optional): Stage of the setup ('train', 'val', 'test').
    """
    self.train_ds = LazyXrDataset(
        {k: v.sel(self.domains['train']) for (k, v) in self.input_da.items()},
        **self.xrds_kw["train"], postpro_fn=self.post_fn('train'),
    )
    if self.aug_kw:
        self.train_ds = AugmentedDataset(self.train_ds, **self.aug_kw)

    self.val_ds = LazyXrDataset(
        {k: v.sel(self.domains['val']) for (k, v) in self.input_da.items()},
        **self.xrds_kw["val"], postpro_fn=self.post_fn('val'),
    )
    self.test_ds = LazyXrDataset(
        {k: v.sel(self.domains['test']) for (k, v) in self.input_da.items()},
        **self.xrds_kw["test"], postpro_fn=self.post_fn('test'),
    )

LazyXrDataset

Bases: XrDataset

A PyTorch Dataset loading data in lazy mode ("on the fly"). If the sampled patch is smaller than the indicated patch dimensions (e. g. a patch sampled at the edge of the domain), the user can complete with nan or periodic repetition in order to return a patch of the specified dimension.

Attributes: see XrDataset.

Source code in ocean4dvarnet/data.py

class LazyXrDataset(XrDataset):
    """
    A PyTorch Dataset loading data in lazy mode ("on the fly").
    If the sampled patch is smaller than the indicated patch dimensions
    (e. g. a patch sampled at the edge of the domain), the user can
    complete with nan or periodic repetition in order to return a patch
    of the specified dimension.

    Attributes: see XrDataset.
    """

    def __init__(
        self, das, patch_dims, domain_limits=None, strides=None,
        postpro_fn=None, **kwargs,
    ):
        """
        Initialize the LazyXrDataset.

        Args:
            das (dict): dictionary containing xr.DataArray to be used.
            patch_dims (dict): da dimension and sizes of patches to extract.
            domain_limits (dict, optional): da dimension slices of domain, to
                Limits for selecting a subset of the domain. for patch
                extractions
            strides (dict, optional): dims to strides size for patch extraction.
                (default to one)
            postpro_fn (callable, optional): A function for post-processing
                extracted patches.
            edges (dict): if the theoretical coverage of a patch exceeds
                the domain's limits, this parameter can fill the "blank"
                area of the patch by repeating the other side of the
                domain ("periodic") or fill the gap with `nan` values
                ("fill_nan")).
                Example: `edges=dict(lat="fill_nan", lon="periodic")`.
        """
        self.return_coords = False
        self.postpro_fn = postpro_fn
        self.da = {k: v.sel(**(domain_limits)) for (k, v) in das.items()}
        self._check_dims_and_coords()
        self.patch_dims = patch_dims
        self.strides = strides or {}
        ref = next(iter(self.da))
        da_dims = dict(zip(self.da[ref].dims, self.da[ref].shape))
        self.ds_size = {
            dim: max((da_dims[dim] - patch_dims[dim]) // self.strides.get(dim, 1) + 1, 0)
            for dim in patch_dims
        }

        # If an edge-behaviour is specified for a dimension, increment
        # by one self.ds_size along this dimension
        self.edges = kwargs.get('edges', dict())
        for dim, behaviour in self.edges.items():
            if behaviour not in ('periodic', 'fill_nan'):
                raise ValueError(
                    f"edges[{dim}] must be either 'periodic' or 'fill_nan', "
                    + f"got {behaviour}"
                )

            if (da_dims[dim] - patch_dims[dim]) % strides[dim] != 0:
                self.ds_size[dim] += 1

    def __getitem__(self, item):
        """
        Get a specific patch by index.

        Args:
            item (int): Index of the patch.

        Returns:
            Patch data or coordinates, depending on the mode.
        """
        sl = {}
        _zip = zip(
            self.ds_size.keys(),
            np.unravel_index(item, tuple(self.ds_size.values())),
        )

        for dim, idx in _zip:
            sl[dim] = slice(
                self.strides.get(dim, 1) * idx,
                self.strides.get(dim, 1) * idx + self.patch_dims[dim]
            )

        ref = next(iter(self.da))
        sliced_domain = self.da[ref].isel(**sl)

        if self.return_coords:
            item = sliced_domain
            return item.coords.to_dataset()[list(self.patch_dims)]

        das = {k: self.da[k].isel(**sl) for k in self.da}

        # Handling edge behaviour
        for dim, behaviour in self.edges.items():
            if len(sliced_domain) >= self.patch_dims[dim]:
                continue
            offset = self.patch_dims[dim] - len(sliced_domain[dim])

            if behaviour == 'periodic':
                sl[dim] = slice(
                    sl[dim].start - offset, sl[dim].stop - offset,
                )

                for var in das:
                    das[var] = (
                        self.da[var]
                        .isel(time=sl['time'])  # TODO How to make it generic?
                        .roll({dim: -offset})
                        .assign_coords({dim: lambda x: x[dim] - offset})
                        .sortby(dim)
                        .isel({k: v for k, v in sl.items() if k != 'time'})
                    )
            elif behaviour == 'fill_nan':
                for var in das:
                    das[var] = das[var].pad(
                        pad_width=dict({dim: (0, offset)}),
                        mode='constant',
                        constant_values=np.nan,
                    )

        item = (
            xr.Dataset(
                data_vars=das,
                coords=next(iter(das.values())).coords,
            )
            .to_dataarray()
            .sortby('variable')
            .data
            .astype(np.float32)
        )

        if self.postpro_fn is not None:
            return self.postpro_fn(item)
        return item

    def _check_dims_and_coords(self):
        """
        Check that `self.da`'s xr.DataArrays all share the same dims and
        coords.
        """
        ref = next(iter(self.da))
        ref_val = self.da[ref]

        for k, v in self.da.items():
            if ref == k:
                continue

            if ref_val.dims != v.dims:
                raise ValueError(
                    "All provided xr.DataArray must share the same dimensions "
                    + f"({ref} != {k})"
                )

            if not ref_val.coords.equals(v.coords):
                raise ValueError(
                    "All provided xr.DataArray must share the same coordinates "
                    + f"({ref} != {k})"
                )

    def _get_da_sample(self):
        return self.da[next(iter(self.da))]

__getitem__(item)

Get a specific patch by index.

Parameters:

Name	Type	Description	Default
item	int	Index of the patch.	required

Returns:

Type	Description
	Patch data or coordinates, depending on the mode.

Source code in ocean4dvarnet/data.py

def __getitem__(self, item):
    """
    Get a specific patch by index.

    Args:
        item (int): Index of the patch.

    Returns:
        Patch data or coordinates, depending on the mode.
    """
    sl = {}
    _zip = zip(
        self.ds_size.keys(),
        np.unravel_index(item, tuple(self.ds_size.values())),
    )

    for dim, idx in _zip:
        sl[dim] = slice(
            self.strides.get(dim, 1) * idx,
            self.strides.get(dim, 1) * idx + self.patch_dims[dim]
        )

    ref = next(iter(self.da))
    sliced_domain = self.da[ref].isel(**sl)

    if self.return_coords:
        item = sliced_domain
        return item.coords.to_dataset()[list(self.patch_dims)]

    das = {k: self.da[k].isel(**sl) for k in self.da}

    # Handling edge behaviour
    for dim, behaviour in self.edges.items():
        if len(sliced_domain) >= self.patch_dims[dim]:
            continue
        offset = self.patch_dims[dim] - len(sliced_domain[dim])

        if behaviour == 'periodic':
            sl[dim] = slice(
                sl[dim].start - offset, sl[dim].stop - offset,
            )

            for var in das:
                das[var] = (
                    self.da[var]
                    .isel(time=sl['time'])  # TODO How to make it generic?
                    .roll({dim: -offset})
                    .assign_coords({dim: lambda x: x[dim] - offset})
                    .sortby(dim)
                    .isel({k: v for k, v in sl.items() if k != 'time'})
                )
        elif behaviour == 'fill_nan':
            for var in das:
                das[var] = das[var].pad(
                    pad_width=dict({dim: (0, offset)}),
                    mode='constant',
                    constant_values=np.nan,
                )

    item = (
        xr.Dataset(
            data_vars=das,
            coords=next(iter(das.values())).coords,
        )
        .to_dataarray()
        .sortby('variable')
        .data
        .astype(np.float32)
    )

    if self.postpro_fn is not None:
        return self.postpro_fn(item)
    return item

__init__(das, patch_dims, domain_limits=None, strides=None, postpro_fn=None, **kwargs)

Initialize the LazyXrDataset.

Parameters:

Name	Type	Description	Default
das	dict	dictionary containing xr.DataArray to be used.	required
patch_dims	dict	da dimension and sizes of patches to extract.	required
domain_limits	dict	da dimension slices of domain, to Limits for selecting a subset of the domain. for patch extractions	None
strides	dict	dims to strides size for patch extraction. (default to one)	None
postpro_fn	callable	A function for post-processing extracted patches.	None
edges	dict	if the theoretical coverage of a patch exceeds the domain's limits, this parameter can fill the "blank" area of the patch by repeating the other side of the domain ("periodic") or fill the gap with nan values ("fill_nan")). Example: edges=dict(lat="fill_nan", lon="periodic").	required

Source code in ocean4dvarnet/data.py

def __init__(
    self, das, patch_dims, domain_limits=None, strides=None,
    postpro_fn=None, **kwargs,
):
    """
    Initialize the LazyXrDataset.

    Args:
        das (dict): dictionary containing xr.DataArray to be used.
        patch_dims (dict): da dimension and sizes of patches to extract.
        domain_limits (dict, optional): da dimension slices of domain, to
            Limits for selecting a subset of the domain. for patch
            extractions
        strides (dict, optional): dims to strides size for patch extraction.
            (default to one)
        postpro_fn (callable, optional): A function for post-processing
            extracted patches.
        edges (dict): if the theoretical coverage of a patch exceeds
            the domain's limits, this parameter can fill the "blank"
            area of the patch by repeating the other side of the
            domain ("periodic") or fill the gap with `nan` values
            ("fill_nan")).
            Example: `edges=dict(lat="fill_nan", lon="periodic")`.
    """
    self.return_coords = False
    self.postpro_fn = postpro_fn
    self.da = {k: v.sel(**(domain_limits)) for (k, v) in das.items()}
    self._check_dims_and_coords()
    self.patch_dims = patch_dims
    self.strides = strides or {}
    ref = next(iter(self.da))
    da_dims = dict(zip(self.da[ref].dims, self.da[ref].shape))
    self.ds_size = {
        dim: max((da_dims[dim] - patch_dims[dim]) // self.strides.get(dim, 1) + 1, 0)
        for dim in patch_dims
    }

    # If an edge-behaviour is specified for a dimension, increment
    # by one self.ds_size along this dimension
    self.edges = kwargs.get('edges', dict())
    for dim, behaviour in self.edges.items():
        if behaviour not in ('periodic', 'fill_nan'):
            raise ValueError(
                f"edges[{dim}] must be either 'periodic' or 'fill_nan', "
                + f"got {behaviour}"
            )

        if (da_dims[dim] - patch_dims[dim]) % strides[dim] != 0:
            self.ds_size[dim] += 1

NoisyLazyDataModule

Bases: LazyDataModule

A data module that adds noise to input data if specified.

The noise added to the training data is a standardised gaussian noise scaled by a provided noise level. The noise added to the validation data is uniformly drawn from the interval [-n, n] where n is the noise level.

see LazyDataModule.

Name	Type	Description
noise	float	Noise level to apply.

Source code in ocean4dvarnet/data.py

class NoisyLazyDataModule(LazyDataModule):
    """
    A data module that adds noise to input data if specified.

    The noise added to the training data is a standardised gaussian noise
    scaled by a provided noise level.
    The noise added to the validation data is uniformly drawn from the
    interval [-n, n] where n is the noise level.

    Attributes: see LazyDataModule.
        noise (float): Noise level to apply.
    """
    def __init__(self, *args, **kwargs):
        """
        Initialize the NoisyLazyDataModule.

        Args:
            input_da (xarray.DataArray): The input data array.
            domains (dict): Dictionary of domain splits (train, val, test).
            xrds_kw (dict): Keyword arguments for XrDataset.
            dl_kw (dict): Keyword arguments for DataLoader.
            aug_kw (dict, optional): Keyword arguments for AugmentedDataset.
            norm_stats (tuple, optional): Normalization statistics (mean, std).
            noise (float, optional): Noise level to be added to the input.
        """
        super().__init__(*args, **kwargs)
        self._rng = np.random.default_rng()
        self.noise = kwargs.get('noise')  # in meters

        if self.noise:
            logging.info(
                f"Adding noise level of {self.noise} m to input data"
            )
        else:
            logging.warning(
                "You are using NoisyLazyDataModule and yet, you did not "
                + "provide a noise level!"
            )

    def post_fn(self, phase=None):
        m, s = self.norm_stats(phase)

        def add_noise(x):
            nl = self.noise

            if not nl:
                return x  # identity if no noise

            if phase == 'train':
                scale = self._rng.uniform(0., nl)
                noise = scale * self._rng.normal(0., 1., x.shape)
            elif phase == 'val':
                noise = self._rng.uniform(-nl, nl, x.shape)
            else:
                noise = 0.

            return x + noise.astype(np.float32)

        return ft.partial(ft.reduce, lambda i, f: f(i), [
            TrainingItem._make,
            lambda item: item._replace(tgt=(item.tgt - m) / s),
            lambda item: item._replace(input=(add_noise(item.input) - m) / s),
        ])

__init__(*args, **kwargs)

Initialize the NoisyLazyDataModule.

Parameters:

Name	Type	Description	Default
input_da	DataArray	The input data array.	required
domains	dict	Dictionary of domain splits (train, val, test).	required
xrds_kw	dict	Keyword arguments for XrDataset.	required
dl_kw	dict	Keyword arguments for DataLoader.	required
aug_kw	dict	Keyword arguments for AugmentedDataset.	required
norm_stats	tuple	Normalization statistics (mean, std).	required
noise	float	Noise level to be added to the input.	required

Source code in ocean4dvarnet/data.py

def __init__(self, *args, **kwargs):
    """
    Initialize the NoisyLazyDataModule.

    Args:
        input_da (xarray.DataArray): The input data array.
        domains (dict): Dictionary of domain splits (train, val, test).
        xrds_kw (dict): Keyword arguments for XrDataset.
        dl_kw (dict): Keyword arguments for DataLoader.
        aug_kw (dict, optional): Keyword arguments for AugmentedDataset.
        norm_stats (tuple, optional): Normalization statistics (mean, std).
        noise (float, optional): Noise level to be added to the input.
    """
    super().__init__(*args, **kwargs)
    self._rng = np.random.default_rng()
    self.noise = kwargs.get('noise')  # in meters

    if self.noise:
        logging.info(
            f"Adding noise level of {self.noise} m to input data"
        )
    else:
        logging.warning(
            "You are using NoisyLazyDataModule and yet, you did not "
            + "provide a noise level!"
        )

RandValDataModule

Bases: BaseDataModule

A data module that randomly splits the training data into training and validation sets.

Attributes:

Name	Type	Description
val_prop	float	Proportion of data to use for validation.

Source code in ocean4dvarnet/data.py

class RandValDataModule(BaseDataModule):
    """
    A data module that randomly splits the training data into training and validation sets.

    Attributes:
        val_prop (float): Proportion of data to use for validation.
    """

    def __init__(self, val_prop, *args, **kwargs):
        """
        Initialize the RandValDataModule.

        Args:
            val_prop (float): Proportion of data to use for validation.
        """
        super().__init__(*args, **kwargs)
        self.val_prop = val_prop

    def setup(self, stage='test'):
        """
        Set up the datasets for training, validation, and testing.

        Args:
            stage (str, optional): Stage of the setup ('train', 'val', 'test').
        """
        post_fn = self.post_fn()
        train_ds = XrDataset(self.input_da.sel(self.domains['train']), **self.xrds_kw, postpro_fn=post_fn,)
        n_val = int(self.val_prop * len(train_ds))
        n_train = len(train_ds) - n_val
        self.train_ds, self.val_ds = torch.utils.data.random_split(train_ds, [n_train, n_val])

        if self.aug_factor > 1:
            self.train_ds = AugmentedDataset(self.train_ds, **self.aug_kw)

        self.test_ds = XrDataset(self.input_da.sel(self.domains['test']), **self.xrds_kw, postpro_fn=post_fn,)

__init__(val_prop, *args, **kwargs)

Initialize the RandValDataModule.

Parameters:

Name	Type	Description	Default
val_prop	float	Proportion of data to use for validation.	required

Source code in ocean4dvarnet/data.py

def __init__(self, val_prop, *args, **kwargs):
    """
    Initialize the RandValDataModule.

    Args:
        val_prop (float): Proportion of data to use for validation.
    """
    super().__init__(*args, **kwargs)
    self.val_prop = val_prop

setup(stage='test')

Set up the datasets for training, validation, and testing.

Parameters:

Name	Type	Description	Default
stage	str	Stage of the setup ('train', 'val', 'test').	'test'

Source code in ocean4dvarnet/data.py

def setup(self, stage='test'):
    """
    Set up the datasets for training, validation, and testing.

    Args:
        stage (str, optional): Stage of the setup ('train', 'val', 'test').
    """
    post_fn = self.post_fn()
    train_ds = XrDataset(self.input_da.sel(self.domains['train']), **self.xrds_kw, postpro_fn=post_fn,)
    n_val = int(self.val_prop * len(train_ds))
    n_train = len(train_ds) - n_val
    self.train_ds, self.val_ds = torch.utils.data.random_split(train_ds, [n_train, n_val])

    if self.aug_factor > 1:
        self.train_ds = AugmentedDataset(self.train_ds, **self.aug_kw)

    self.test_ds = XrDataset(self.input_da.sel(self.domains['test']), **self.xrds_kw, postpro_fn=post_fn,)

XrConcatDataset

Bases: ConcatDataset

A concatenation of multiple XrDatasets.

This class allows combining multiple datasets into one for training or evaluation.

Source code in ocean4dvarnet/data.py

class XrConcatDataset(torch.utils.data.ConcatDataset):
    """
    A concatenation of multiple XrDatasets.

    This class allows combining multiple datasets into one for training or evaluation.
    """

    def reconstruct(self, batches, weight=None):
        """
        Reconstruct the original data arrays from batches.

        Args:
            batches (list): List of batches.
            weight (np.ndarray, optional): Weighting for overlapping patches.

        Returns:
            list: List of reconstructed xarray.DataArray objects.
        """
        items_iter = itertools.chain(*batches)
        rec_das = []
        for ds in self.datasets:
            ds_items = list(itertools.islice(items_iter, len(ds)))
            rec_das.append(ds.reconstruct_from_items(ds_items, weight))

        return rec_das

reconstruct(batches, weight=None)

Reconstruct the original data arrays from batches.

Parameters:

Name	Type	Description	Default
batches	list	List of batches.	required
weight	ndarray	Weighting for overlapping patches.	None

Returns:

Name	Type	Description
list		List of reconstructed xarray.DataArray objects.

Source code in ocean4dvarnet/data.py

def reconstruct(self, batches, weight=None):
    """
    Reconstruct the original data arrays from batches.

    Args:
        batches (list): List of batches.
        weight (np.ndarray, optional): Weighting for overlapping patches.

    Returns:
        list: List of reconstructed xarray.DataArray objects.
    """
    items_iter = itertools.chain(*batches)
    rec_das = []
    for ds in self.datasets:
        ds_items = list(itertools.islice(items_iter, len(ds)))
        rec_das.append(ds.reconstruct_from_items(ds_items, weight))

    return rec_das

XrDataset

Bases: Dataset

A PyTorch Dataset based on an xarray.DataArray with on-the-fly slicing.

This class allows efficient extraction of patches from an xarray.DataArray for training machine learning models.

Usage

If you want to be able to reconstruct the input, the input xr.DataArray should: - Have coordinates. - Have the last dims correspond to the patch dims in the same order. - Have, for each dim of patch_dim, (size(dim) - patch_dim(dim)) divisible by stride(dim).

The batches passed to self.reconstruct should: - Have the last dims correspond to the patch dims in the same order.

Attributes:

Name	Type	Description
da	DataArray	The input data array.
patch_dims	dict	Dimensions and sizes of patches to extract.
domain_limits	dict	Limits for selecting a subset of the domain.
strides	dict	Strides for patch extraction.
check_full_scan	bool	Whether to check if the entire domain is scanned.
check_dim_order	bool	Whether to check the dimension ordering.
postpro_fn	callable	A function for post-processing extracted patches.

Source code in ocean4dvarnet/data.py

class XrDataset(torch.utils.data.Dataset):
    """
    A PyTorch Dataset based on an xarray.DataArray with on-the-fly slicing.

    This class allows efficient extraction of patches from an xarray.DataArray
    for training machine learning models.

    Usage:
        If you want to be able to reconstruct the input, the input xr.DataArray should:
        - Have coordinates.
        - Have the last dims correspond to the patch dims in the same order.
        - Have, for each dim of patch_dim, (size(dim) - patch_dim(dim)) divisible by stride(dim).

        The batches passed to self.reconstruct should:
        - Have the last dims correspond to the patch dims in the same order.


    Attributes:
        da (xarray.DataArray): The input data array.
        patch_dims (dict): Dimensions and sizes of patches to extract.
        domain_limits (dict): Limits for selecting a subset of the domain.
        strides (dict): Strides for patch extraction.
        check_full_scan (bool): Whether to check if the entire domain is scanned.
        check_dim_order (bool): Whether to check the dimension ordering.
        postpro_fn (callable): A function for post-processing extracted patches.

    """

    def __init__(
            self, da, patch_dims, domain_limits=None, strides=None,
            check_full_scan=False, check_dim_order=False,
            postpro_fn=None
    ):
        """
        Initialize the XrDataset.

        Args:
            da (xarray.DataArray): Input data, with patch dims at the end in the dim orders
            patch_dims (dict):  da dimension and sizes of patches to extract.
            domain_limits (dict, optional): da dimension slices of domain, to Limits for selecting
                                            a subset of the domain. for patch extractions
            strides (dict, optional): dims to strides size for patch extraction.(default to one)
            check_full_scan (bool, optional): if True raise an error if the whole domain is not scanned by the patch.
            check_dim_order (bool, optional): Whether to check the dimension ordering.
            postpro_fn (callable, optional): A function for post-processing extracted patches.
        """
        super().__init__()
        self.return_coords = False
        self.postpro_fn = postpro_fn
        self.da = da.sel(**(domain_limits or {}))
        self.patch_dims = patch_dims
        self.strides = strides or {}
        da_dims = dict(zip(self.da.dims, self.da.shape))
        self.ds_size = {
            dim: max((da_dims[dim] - patch_dims[dim]) // self.strides.get(dim, 1) + 1, 0)
            for dim in patch_dims
        }

        if check_full_scan:
            for dim in patch_dims:
                if (da_dims[dim] - self.patch_dims[dim]) % self.strides.get(dim, 1) != 0:
                    raise IncompleteScanConfiguration(
                        f"""
                        Incomplete scan in dimension dim {dim}:
                        dataarray shape on this dim {da_dims[dim]}
                        patch_size along this dim {self.patch_dims[dim]}
                        stride along this dim {self.strides.get(dim, 1)}
                        [shape - patch_size] should be divisible by stride
                        """
                    )

        if check_dim_order:
            for dim in patch_dims:
                if not '#'.join(da.dims).endswith('#'.join(list(patch_dims))):
                    raise DangerousDimOrdering(
                        f"""
                        input dataarray's dims should end with patch_dims
                        dataarray's dim {da.dims}:
                        patch_dims {list(patch_dims)}
                        """
                    )

    def __len__(self):
        """
        Return the total number of patches in the dataset.

        Returns:
            int: Number of patches.
        """
        size = 1
        for v in self.ds_size.values():
            size *= v
        return size

    def __iter__(self):
        """
        Iterate over the dataset.

        Yields:
            Patch data for each index.
        """
        for i in range(len(self)):
            yield self[i]

    def get_coords(self):
        """
        Get the coordinates of all patches in the dataset.

        Returns:
            list: List of coordinates for each patch.
        """
        self.return_coords = True
        coords = []
        try:
            for i in range(len(self)):
                 coords.append(self[i])
        finally:
            self.return_coords = False
            return coords

    def __getitem__(self, item):
        """
        Get a specific patch by index.

        Args:
            item (int): Index of the patch.

        Returns:
            Patch data or coordinates, depending on the mode.
        """
        sl = {
            dim: slice(self.strides.get(dim, 1) * idx,
                       self.strides.get(dim, 1) * idx + self.patch_dims[dim])
            for dim, idx in zip(self.ds_size.keys(),
                                np.unravel_index(item, tuple(self.ds_size.values())))
        }
        item = self.da.isel(**sl)

        if self.return_coords:
            return item.coords.to_dataset()[list(self.patch_dims)]

        item = item.data.astype(np.float32)
        if self.postpro_fn is not None:
            return self.postpro_fn(item)
        return item

    def reconstruct(self, batches, weight=None):
        """
        Reconstruct the original data array from patches.

        Takes as input a list of np.ndarray of dimensions (b, *, *patch_dims).

        Args:
            batches (list): List of patches (torch tensor) corresponding to batches without shuffle.
            weight (np.ndarray, optional): Tensor of size patch_dims corresponding to the weight of a prediction
                depending on the position on the patch (default to ones everywhere). Overlapping patches will
                be averaged with weighting.

        Returns:
            xarray.DataArray: Reconstructed data array. A stitched xarray.DataArray with the coords of patch_dims.
        """
        items = list(itertools.chain(*batches))
        return self.reconstruct_from_items(items, weight)

    def reconstruct_from_items(self, items, weight=None):
        """
        Reconstruct the original data array from individual items.

        Args:
            items (list): List of individual patches.
            weight (np.ndarray, optional): Weighting for overlapping patches.

        Returns:
            xarray.DataArray: Reconstructed data array.
        """
        if weight is None:
            weight = np.ones(list(self.patch_dims.values()))
        w = xr.DataArray(weight, dims=list(self.patch_dims.keys()))

        coords = self.get_coords()

        new_dims = [f'v{i}' for i in range(len(items[0].shape) - len(coords[0].dims))]
        dims = new_dims + list(coords[0].dims)

        das = [xr.DataArray(it.numpy(), dims=dims, coords=co.coords)
               for it, co in zip(items, coords)]

        da_shape = dict(zip(coords[0].dims, self._get_da_sample().shape[-len(coords[0].dims):]))
        new_shape = dict(zip(new_dims, items[0].shape[:len(new_dims)]))

        rec_da = xr.DataArray(
            np.zeros([*new_shape.values(), *da_shape.values()]),
            dims=dims,
            coords={d: self._get_da_sample()[d] for d in self.patch_dims}
        )
        count_da = xr.zeros_like(rec_da)

        for da in das:
            rec_da.loc[da.coords] = rec_da.sel(da.coords) + da * w
            count_da.loc[da.coords] = count_da.sel(da.coords) + w

        return rec_da / count_da

    def _get_da_sample(self):
        return self.da

__getitem__(item)

Get a specific patch by index.

Parameters:

Name	Type	Description	Default
item	int	Index of the patch.	required

Returns:

Type	Description
	Patch data or coordinates, depending on the mode.

Source code in ocean4dvarnet/data.py

def __getitem__(self, item):
    """
    Get a specific patch by index.

    Args:
        item (int): Index of the patch.

    Returns:
        Patch data or coordinates, depending on the mode.
    """
    sl = {
        dim: slice(self.strides.get(dim, 1) * idx,
                   self.strides.get(dim, 1) * idx + self.patch_dims[dim])
        for dim, idx in zip(self.ds_size.keys(),
                            np.unravel_index(item, tuple(self.ds_size.values())))
    }
    item = self.da.isel(**sl)

    if self.return_coords:
        return item.coords.to_dataset()[list(self.patch_dims)]

    item = item.data.astype(np.float32)
    if self.postpro_fn is not None:
        return self.postpro_fn(item)
    return item

__init__(da, patch_dims, domain_limits=None, strides=None, check_full_scan=False, check_dim_order=False, postpro_fn=None)

Initialize the XrDataset.

Parameters:

Name	Type	Description	Default
da	DataArray	Input data, with patch dims at the end in the dim orders	required
patch_dims	dict	da dimension and sizes of patches to extract.	required
domain_limits	dict	da dimension slices of domain, to Limits for selecting a subset of the domain. for patch extractions	None
strides	dict	dims to strides size for patch extraction.(default to one)	None
check_full_scan	bool	if True raise an error if the whole domain is not scanned by the patch.	False
check_dim_order	bool	Whether to check the dimension ordering.	False
postpro_fn	callable	A function for post-processing extracted patches.	None

Source code in ocean4dvarnet/data.py

def __init__(
        self, da, patch_dims, domain_limits=None, strides=None,
        check_full_scan=False, check_dim_order=False,
        postpro_fn=None
):
    """
    Initialize the XrDataset.

    Args:
        da (xarray.DataArray): Input data, with patch dims at the end in the dim orders
        patch_dims (dict):  da dimension and sizes of patches to extract.
        domain_limits (dict, optional): da dimension slices of domain, to Limits for selecting
                                        a subset of the domain. for patch extractions
        strides (dict, optional): dims to strides size for patch extraction.(default to one)
        check_full_scan (bool, optional): if True raise an error if the whole domain is not scanned by the patch.
        check_dim_order (bool, optional): Whether to check the dimension ordering.
        postpro_fn (callable, optional): A function for post-processing extracted patches.
    """
    super().__init__()
    self.return_coords = False
    self.postpro_fn = postpro_fn
    self.da = da.sel(**(domain_limits or {}))
    self.patch_dims = patch_dims
    self.strides = strides or {}
    da_dims = dict(zip(self.da.dims, self.da.shape))
    self.ds_size = {
        dim: max((da_dims[dim] - patch_dims[dim]) // self.strides.get(dim, 1) + 1, 0)
        for dim in patch_dims
    }

    if check_full_scan:
        for dim in patch_dims:
            if (da_dims[dim] - self.patch_dims[dim]) % self.strides.get(dim, 1) != 0:
                raise IncompleteScanConfiguration(
                    f"""
                    Incomplete scan in dimension dim {dim}:
                    dataarray shape on this dim {da_dims[dim]}
                    patch_size along this dim {self.patch_dims[dim]}
                    stride along this dim {self.strides.get(dim, 1)}
                    [shape - patch_size] should be divisible by stride
                    """
                )

    if check_dim_order:
        for dim in patch_dims:
            if not '#'.join(da.dims).endswith('#'.join(list(patch_dims))):
                raise DangerousDimOrdering(
                    f"""
                    input dataarray's dims should end with patch_dims
                    dataarray's dim {da.dims}:
                    patch_dims {list(patch_dims)}
                    """
                )

__iter__()

Iterate over the dataset.

Yields:

Type	Description
	Patch data for each index.

Source code in ocean4dvarnet/data.py

def __iter__(self):
    """
    Iterate over the dataset.

    Yields:
        Patch data for each index.
    """
    for i in range(len(self)):
        yield self[i]

__len__()

Return the total number of patches in the dataset.

Returns:

Name	Type	Description
int		Number of patches.

Source code in ocean4dvarnet/data.py

def __len__(self):
    """
    Return the total number of patches in the dataset.

    Returns:
        int: Number of patches.
    """
    size = 1
    for v in self.ds_size.values():
        size *= v
    return size

get_coords()

Get the coordinates of all patches in the dataset.

Returns:

Name	Type	Description
list		List of coordinates for each patch.

Source code in ocean4dvarnet/data.py

def get_coords(self):
    """
    Get the coordinates of all patches in the dataset.

    Returns:
        list: List of coordinates for each patch.
    """
    self.return_coords = True
    coords = []
    try:
        for i in range(len(self)):
             coords.append(self[i])
    finally:
        self.return_coords = False
        return coords

reconstruct(batches, weight=None)

Reconstruct the original data array from patches.

Takes as input a list of np.ndarray of dimensions (b, , patch_dims).

Parameters:

Name	Type	Description	Default
batches	list	List of patches (torch tensor) corresponding to batches without shuffle.	required
weight	ndarray	Tensor of size patch_dims corresponding to the weight of a prediction depending on the position on the patch (default to ones everywhere). Overlapping patches will be averaged with weighting.	None

Returns:

Type	Description
	xarray.DataArray: Reconstructed data array. A stitched xarray.DataArray with the coords of patch_dims.

Source code in ocean4dvarnet/data.py

def reconstruct(self, batches, weight=None):
    """
    Reconstruct the original data array from patches.

    Takes as input a list of np.ndarray of dimensions (b, *, *patch_dims).

    Args:
        batches (list): List of patches (torch tensor) corresponding to batches without shuffle.
        weight (np.ndarray, optional): Tensor of size patch_dims corresponding to the weight of a prediction
            depending on the position on the patch (default to ones everywhere). Overlapping patches will
            be averaged with weighting.

    Returns:
        xarray.DataArray: Reconstructed data array. A stitched xarray.DataArray with the coords of patch_dims.
    """
    items = list(itertools.chain(*batches))
    return self.reconstruct_from_items(items, weight)

reconstruct_from_items(items, weight=None)

Reconstruct the original data array from individual items.

Parameters:

Name	Type	Description	Default
items	list	List of individual patches.	required
weight	ndarray	Weighting for overlapping patches.	None

Returns:

Type	Description
	xarray.DataArray: Reconstructed data array.

Source code in ocean4dvarnet/data.py

def reconstruct_from_items(self, items, weight=None):
    """
    Reconstruct the original data array from individual items.

    Args:
        items (list): List of individual patches.
        weight (np.ndarray, optional): Weighting for overlapping patches.

    Returns:
        xarray.DataArray: Reconstructed data array.
    """
    if weight is None:
        weight = np.ones(list(self.patch_dims.values()))
    w = xr.DataArray(weight, dims=list(self.patch_dims.keys()))

    coords = self.get_coords()

    new_dims = [f'v{i}' for i in range(len(items[0].shape) - len(coords[0].dims))]
    dims = new_dims + list(coords[0].dims)

    das = [xr.DataArray(it.numpy(), dims=dims, coords=co.coords)
           for it, co in zip(items, coords)]

    da_shape = dict(zip(coords[0].dims, self._get_da_sample().shape[-len(coords[0].dims):]))
    new_shape = dict(zip(new_dims, items[0].shape[:len(new_dims)]))

    rec_da = xr.DataArray(
        np.zeros([*new_shape.values(), *da_shape.values()]),
        dims=dims,
        coords={d: self._get_da_sample()[d] for d in self.patch_dims}
    )
    count_da = xr.zeros_like(rec_da)

    for da in das:
        rec_da.loc[da.coords] = rec_da.sel(da.coords) + da * w
        count_da.loc[da.coords] = count_da.sel(da.coords) + w

    return rec_da / count_da