initial commit

backend/nodes/level.py

@@ -0,0 +1,150 @@
+"""
+Leveling nodes — background removal and zero correction.
+
+Gwyddion equivalents:
+  PlaneLevelField → gwy_data_field_fit_plane + gwy_data_field_plane_level
+  PolyLevelField  → gwy_data_field_fit_polynom (via level.c polylevel module)
+  FixZero         → fix_zero in level.c
+
+Plane-fit algorithm follows Gwyddion's level.h definition:
+  z_fit = pa + pbx * x + pby * y   (least-squares over all pixels)
+"""
+
+from __future__ import annotations
+import numpy as np
+from backend.node_registry import register_node
+from backend.data_types import DataField
+
+
+# ---------------------------------------------------------------------------
+# PlaneLevelField
+# ---------------------------------------------------------------------------
+
+@register_node(display_name="Plane Level")
+class PlaneLevelField:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "field": ("DATA_FIELD",),
+            }
+        }
+
+    RETURN_TYPES = ("DATA_FIELD",)
+    RETURN_NAMES = ("leveled",)
+    FUNCTION = "process"
+    CATEGORY = "level"
+    DESCRIPTION = (
+        "Fit and subtract a least-squares plane from the data. "
+        "Equivalent to gwy_data_field_fit_plane + gwy_data_field_plane_level."
+    )
+
+    def process(self, field: DataField) -> tuple:
+        data = field.data.copy()
+        yres, xres = data.shape
+
+        # Normalised coordinate grids in [0, 1]
+        x = np.linspace(0.0, 1.0, xres)
+        y = np.linspace(0.0, 1.0, yres)
+        xx, yy = np.meshgrid(x, y)
+
+        # Design matrix: [1, x, y]  shape (N, 3)
+        A = np.column_stack([
+            np.ones(xres * yres),
+            xx.ravel(),
+            yy.ravel(),
+        ])
+        z = data.ravel()
+
+        # Least-squares: solve A @ [pa, pbx, pby] = z
+        coeffs, _, _, _ = np.linalg.lstsq(A, z, rcond=None)
+        pa, pbx, pby = coeffs
+
+        plane = (pa + pbx * xx + pby * yy)
+        return (field.replace(data=data - plane),)
+
+
+# ---------------------------------------------------------------------------
+# PolyLevelField
+# ---------------------------------------------------------------------------
+
+@register_node(display_name="Polynomial Level")
+class PolyLevelField:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "field": ("DATA_FIELD",),
+                "degree_x": ("INT", {"default": 2, "min": 0, "max": 5, "step": 1}),
+                "degree_y": ("INT", {"default": 2, "min": 0, "max": 5, "step": 1}),
+            }
+        }
+
+    RETURN_TYPES = ("DATA_FIELD", "DATA_FIELD")
+    RETURN_NAMES = ("leveled", "background")
+    FUNCTION = "process"
+    CATEGORY = "level"
+    DESCRIPTION = (
+        "Fit and subtract a polynomial background of given degree in x and y. "
+        "Equivalent to gwy_data_field_fit_polynom."
+    )
+
+    def process(self, field: DataField, degree_x: int, degree_y: int) -> tuple:
+        data = field.data.copy()
+        yres, xres = data.shape
+
+        x = np.linspace(0.0, 1.0, xres)
+        y = np.linspace(0.0, 1.0, yres)
+        xx, yy = np.meshgrid(x, y)
+
+        # Build Vandermonde-style design matrix with all monomials x^i * y^j
+        cols = []
+        for i in range(degree_x + 1):
+            for j in range(degree_y + 1):
+                cols.append((xx ** i * yy ** j).ravel())
+        A = np.column_stack(cols)
+        z = data.ravel()
+
+        coeffs, _, _, _ = np.linalg.lstsq(A, z, rcond=None)
+
+        background = (A @ coeffs).reshape(yres, xres)
+        leveled = data - background
+
+        return (field.replace(data=leveled), field.replace(data=background))
+
+
+# ---------------------------------------------------------------------------
+# FixZero
+# ---------------------------------------------------------------------------
+
+@register_node(display_name="Fix Zero")
+class FixZero:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "field": ("DATA_FIELD",),
+                "method": (["min", "mean", "median"],),
+            }
+        }
+
+    RETURN_TYPES = ("DATA_FIELD",)
+    RETURN_NAMES = ("zeroed",)
+    FUNCTION = "process"
+    CATEGORY = "level"
+    DESCRIPTION = (
+        "Shift data so that the minimum (or mean/median) is zero. "
+        "Equivalent to fix_zero in Gwyddion's level.c."
+    )
+
+    def process(self, field: DataField, method: str) -> tuple:
+        data = field.data.copy()
+        if method == "min":
+            data -= data.min()
+        elif method == "mean":
+            data -= data.mean()
+        elif method == "median":
+            data -= np.median(data)
+        else:
+            raise ValueError(f"Unknown method: {method}")
+        return (field.replace(data=data),)