[modules] Add layer which computes Curl of 3D vector field

Author: aschuh-hf

src/deepali/modules/__init__.py

@@ -12,6 +12,7 @@
 from .basic import Reshape
 from .basic import View
 
+from .flow import Curl
 from .flow import ExpFlow
 
 from .image import BlurImage
@@ -37,6 +38,7 @@
 __all__ = (
     "AlignImage",
     "BlurImage",
+    "Curl",
     "DeviceProperty",
     "ExpFlow",
     "FilterImage",

src/deepali/modules/flow.py

@@ -3,15 +3,48 @@
 from __future__ import annotations
 
 from copy import copy as shallow_copy
-from typing import Optional
+from typing import Optional, Union
 
 from torch import Tensor
 from torch.nn import Module
 
 from deepali.core import ALIGN_CORNERS
+from deepali.core import Array, Scalar, ScalarOrTuple
 from deepali.core import functional as U
 
 
+class Curl(Module):
+    r"""Layer which calculates the curl of a vector field."""
+
+    def __init__(
+        self,
+        mode: Optional[str] = None,
+        sigma: Optional[float] = None,
+        spacing: Optional[Union[Scalar, Array]] = None,
+        stride: Optional[ScalarOrTuple[int]] = None,
+    ) -> None:
+        super().__init__()
+        self.mode = mode
+        self.sigma = sigma
+        self.spacing = spacing
+        self.stride = stride
+
+    def forward(self, x: Tensor) -> Tensor:
+        return U.curl(x, mode=self.mode, sigma=self.sigma, spacing=self.spacing, stride=self.stride)
+
+    def extra_repr(self) -> str:
+        args = []
+        if self.mode is not None:
+            args.append(f"mode={self.mode!r}")
+        if self.sigma is not None:
+            args.append(f"sigma={self.sigma!r}")
+        if self.spacing is not None:
+            args.append(f"spacing={self.spacing!r}")
+        if self.stride is not None:
+            args.append(f"stride={self.stride!r}")
+        return ", ".join(args)
+
+
 class ExpFlow(Module):
     r"""Layer that computes exponential map of flow field."""