Added custom and learnable filters support (#11)

Signed-off-by: Nicola VIGANÒ <nicola.vigano@cea.fr>

Author: Obi-Wan

src/autoden/losses.py

@@ -8,6 +8,7 @@
 from torch.nn.modules.loss import _Loss
 
 from autoden.transforms._wavelets import dwtn, swtn, wavelet_norm
+from autoden.transforms.custom_filters import CustomFilterDecomposition
 
 
 def _differentiate(inp: pt.Tensor, dim: int, position: str) -> pt.Tensor:
@@ -239,3 +240,52 @@ def forward(self, img: pt.Tensor) -> pt.Tensor:
             return loss_vals.sum()
         else:
             return loss_vals
+
+
+class LossCCF(LossRegularizer):
+    """Custom convolutional filter decomposition loss function."""
+
+    def __init__(
+        self,
+        lambda_val: float,
+        filters: pt.Tensor,
+        weights: pt.Tensor,
+        size_average=None,
+        reduce=None,
+        reduction: str = "mean",
+        min_approx: bool = False,
+    ) -> None:
+        super().__init__(size_average, reduce, reduction)
+        self.lambda_val = lambda_val
+        self.filters = filters
+        self.weights = weights
+        self.min_approx = min_approx
+
+        self.n_dims = filters.ndim - 2
+
+        if filters.shape[0] != weights.numel():
+            raise ValueError(
+                f"The number of convolution kernels ({filters.shape[0]}) does"
+                f" not match the number of weights ({weights.numel()})"
+            )
+
+        self.weights = weights.reshape([1, -1, *(1,) * self.n_dims])
+
+    def forward(self, img: pt.Tensor) -> pt.Tensor:
+        """Compute decomposition on current batch."""
+        _check_input_tensor(img, self.n_dims)
+        axes = list(range(-(self.n_dims + 1), 0))
+
+        decomp = CustomFilterDecomposition(kernels=self.filters[not self.min_approx :, ...], device=img.device)
+        weights = self.weights[:, not self.min_approx :, ...].to(img.device)
+
+        coeffs = decomp.analyze(img)
+
+        loss_vals: pt.Tensor = self.lambda_val * (weights * coeffs).abs().sum(1).sum(axes)
+
+        if self.reduction.lower() == "mean":
+            return loss_vals.mean()
+        elif self.reduction.lower() == "sum":
+            return loss_vals.sum()
+        else:
+            return loss_vals

src/autoden/transforms/__init__.py

@@ -0,0 +1,5 @@
+# -*- coding: utf-8 -*-
+"""Transforms sub-package."""
+
+__author__ = """Nicola Vigano"""
+__email__ = "nicola.vigano@cea.fr"

src/autoden/transforms/custom_filters.py

@@ -0,0 +1,157 @@
+"""
+Custom and learnable filters decompositions.
+"""
+
+import math
+from abc import ABC, abstractmethod
+from typing import Callable, Literal
+
+import torch as pt
+import torch.nn as nn
+import torch.nn.functional as F
+from numpy.typing import NDArray
+
+
+class ConvolutionalDecompositionBase(ABC, nn.Module):
+    """Base class for all decompositions."""
+
+    m: int
+    k: int
+    in_ch: int
+    n_dims: int
+    norm: Literal["backward", "forward", "ortho"] | None
+
+    _ndconvs_d: dict[int, Callable[..., pt.Tensor]] = {1: F.conv1d, 2: F.conv2d, 3: F.conv3d}
+    _ndconvs_t: dict[int, Callable[..., pt.Tensor]] = {1: F.conv_transpose1d, 2: F.conv_transpose2d, 3: F.conv_transpose3d}
+
+    def __init__(
+        self, k: int, n_dims: int, in_ch: int, m: int, norm: Literal["backward", "forward", "ortho"] | None = "backward"
+    ) -> None:
+        """Initialize the ConvolutionalDecomposition.
+
+        Parameters
+        ----------
+        k : int
+            Kernel size.
+        n_dims : int
+            Number of dimensions for the convolution.
+        in_ch : int
+            Number of input channels.
+        m : int
+            Number of output channels.
+        norm : Literal["backward", "forward", "ortho"] | None, optional
+            Normalization type. Defaults to "backward".
+        """
+        super().__init__()
+        self.k = k
+        self.n_dims = n_dims
+        self.in_ch = in_ch
+        self.m = m
+        self.norm = norm
+
+    @abstractmethod
+    def get_kernels(self) -> pt.Tensor:
+        """Return the kernels to be used for the convolutions.
+
+        Returns
+        -------
+        pt.Tensor
+            The kernels for the convolutions.
+        """
+
+    def analyze(self, x: pt.Tensor) -> pt.Tensor:
+        """Apply the analysis (forward) transform using the kernels.
+
+        Parameters
+        ----------
+        x : pt.Tensor
+            Input tensor of shape (B, in_ch, [D, H], W).
+
+        Returns
+        -------
+        pt.Tensor
+            Output tensor of shape (B, m, [D, H], W).
+        """
+        w = self.get_kernels()
+        c = self._ndconvs_d[self.n_dims](x, w, padding=self.k // 2)
+        if self.norm is not None:
+            if self.norm.lower() == "ortho":
+                c = c / math.sqrt(self.m) * math.sqrt(self.in_ch)
+            elif self.norm.lower() == "forward":
+                c = c / float(self.m) * float(self.in_ch)
+        return c
+
+    def synthesize(self, c: pt.Tensor) -> pt.Tensor:
+        """Apply the synthesis (inverse) transform using the kernels.
+
+        Parameters
+        ----------
+        c : pt.Tensor
+            Input tensor of shape (B, m, [D, H], W).
+
+        Returns
+        -------
+        pt.Tensor
+            Output tensor of shape (B, in_ch, [D, H], W).
+        """
+        w = self.get_kernels()
+        x = self._ndconvs_t[self.n_dims](c, w, padding=self.k // 2)
+        if self.norm is not None:
+            if self.norm.lower() == "ortho":
+                x = x / math.sqrt(self.m) * math.sqrt(self.in_ch)
+            elif self.norm.lower() == "backward":
+                x = x / float(self.m) * float(self.in_ch)
+        return x
+
+
+class CustomFilterDecomposition(ConvolutionalDecompositionBase):
+    """Decomposition using custom filters (kernels)."""
+
+    kernels: pt.Tensor
+
+    def __init__(
+        self,
+        kernels: pt.Tensor | NDArray,
+        device: str = "cuda" if pt.cuda.is_available() else "cpu",
+        norm: Literal["backward", "forward", "ortho"] | None = "backward",
+    ) -> None:
+        """Initialize the CustomFilterDecomposition.
+
+        Parameters
+        ----------
+        kernels : pt.Tensor | NDArray
+            The kernels to be used for the convolutions. Should have shape (m, in_ch, *((k,) * n_dims)).
+        device : str, optional
+            The device to use for the kernels. Defaults to "cuda" if available, otherwise "cpu".
+        norm : Literal["backward", "forward", "ortho"] | None, optional
+            Normalization type. Defaults to "backward".
+        """
+        m = kernels.shape[0]
+        in_ch = kernels.shape[1]
+        n_dims = kernels.ndim - 2
+        if n_dims < 1:
+            raise ValueError(f"Kernels should have shape (m, in_ch, *((k,) * n_dims)), but {kernels.shape} was passed")
+        k = kernels.shape[-1]
+        if any(s != k for s in kernels.shape[-n_dims:-1]):
+            raise ValueError(
+                f"Kernels should have the same size `k` in all directions, but {kernels.shape[-n_dims]} was passed."
+                f" Complete shape: {kernels.shape}"
+            )
+        super().__init__(k=k, in_ch=in_ch, n_dims=n_dims, m=m, norm=norm)
+
+        if not isinstance(kernels, pt.Tensor):
+            kernels = pt.tensor(kernels)
+        kernels = kernels.detach().to(device).clone()
+        self.register_buffer("kernels", kernels)
+
+        self.device = device
+
+    def get_kernels(self) -> pt.Tensor:
+        """Return the kernels to be used for the convolutions.
+
+        Returns
+        -------
+        pt.Tensor
+            The kernels for the convolutions.
+        """
+        return self.kernels