151 lines
4.5 KiB
Python
151 lines
4.5 KiB
Python
"""
|
|
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"
|
|
|
|
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"
|
|
|
|
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"
|
|
|
|
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),)
|