64 lines
2.2 KiB
Python
64 lines
2.2 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="Gradient")
|
|
class Gradient:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"field": ("DATA_FIELD",),
|
|
"component": (["magnitude", "x", "y", "azimuth"],),
|
|
}
|
|
}
|
|
|
|
OUTPUTS = (
|
|
('DATA_FIELD', 'gradient'),
|
|
)
|
|
FUNCTION = "process"
|
|
|
|
DESCRIPTION = (
|
|
"Compute the spatial gradient using a Sobel operator. "
|
|
"'x'/'y' give the physical gradient components (z_unit/xy_unit); "
|
|
"'magnitude' gives sqrt(gx²+gy²); "
|
|
"'azimuth' gives the local slope direction in radians via atan2(gy, gx). "
|
|
"Equivalent to gwy_data_field_filter_sobel in Gwyddion (gradient.c)."
|
|
)
|
|
|
|
def process(self, field: DataField, component: str) -> tuple:
|
|
from scipy.ndimage import sobel
|
|
|
|
data = field.data
|
|
# Sobel kernel sums to ±8 over 2-pixel span; divide by 8·dx to get z/xy slope.
|
|
gx = sobel(data, axis=1) / (8.0 * field.dx)
|
|
gy = sobel(data, axis=0) / (8.0 * field.dy)
|
|
|
|
if component == "magnitude":
|
|
result = np.hypot(gx, gy)
|
|
z = str(field.si_unit_z or "").strip()
|
|
xy = str(field.si_unit_xy or "").strip()
|
|
out_unit_z = f"{z}/{xy}" if z and xy else (z or xy)
|
|
elif component == "x":
|
|
result = gx
|
|
z = str(field.si_unit_z or "").strip()
|
|
xy = str(field.si_unit_xy or "").strip()
|
|
out_unit_z = f"{z}/{xy}" if z and xy else (z or xy)
|
|
elif component == "y":
|
|
result = gy
|
|
z = str(field.si_unit_z or "").strip()
|
|
xy = str(field.si_unit_xy or "").strip()
|
|
out_unit_z = f"{z}/{xy}" if z and xy else (z or xy)
|
|
elif component == "azimuth":
|
|
# Azimuth: local slope direction, radians, matches Gwyddion's filter_azimuth
|
|
result = np.arctan2(gy, gx)
|
|
out_unit_z = "rad"
|
|
else:
|
|
raise ValueError(f"Unknown gradient component: {component!r}")
|
|
|
|
return (field.replace(data=result, si_unit_z=out_unit_z),)
|