51 lines
1.5 KiB
Python
51 lines
1.5 KiB
Python
from __future__ import annotations
|
|
import numpy as np
|
|
from backend.node_registry import register_node
|
|
from backend.data_types import DataField
|
|
|
|
|
|
@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}),
|
|
}
|
|
}
|
|
|
|
OUTPUTS = (
|
|
('DATA_FIELD', 'leveled'),
|
|
('DATA_FIELD', '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)
|
|
|
|
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))
|