fourier-transformed 2D rectangle map

Xeadriel · Nov 2, 2022 · 1efece8 · 1efece8
1 parent 6b6fc2c
commit 1efece8
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 0 deletions.
diff --git a/firelang/map/__init__.py b/firelang/map/__init__.py
@@ -1 +1,2 @@
 from .rect import *
+from .fourier import *
diff --git a/firelang/map/fourier.py b/firelang/map/fourier.py
@@ -0,0 +1,17 @@
+import torch
+from .rect import (
+    SmoothedRectMap,
+)
+
+__all__ = [
+    'SmoothedRectFourier2DMap',
+]
+
+class SmoothedRectFourier2DMap(SmoothedRectMap):
+
+    @property
+    def grid(self):
+        g = SmoothedRectMap.get_grid(self)
+        assert g.ndim == 2
+
+        g = torch.fft.fft2(g).real