clean up node naming
This commit is contained in:
50
backend/nodes/level_poly.py
Normal file
50
backend/nodes/level_poly.py
Normal file
@@ -0,0 +1,50 @@
|
||||
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))
|
||||
Reference in New Issue
Block a user