combine fft filter into a single node, fix tests
This commit is contained in:
@@ -4,8 +4,7 @@ from backend.nodes import (
|
||||
colormap,
|
||||
crop_resize,
|
||||
fft_2d_inverse,
|
||||
filter_fft_1d,
|
||||
filter_fft_2d,
|
||||
filter_fft,
|
||||
filter_gaussian,
|
||||
filter_median,
|
||||
flip,
|
||||
|
||||
89
backend/nodes/filter_fft.py
Normal file
89
backend/nodes/filter_fft.py
Normal file
@@ -0,0 +1,89 @@
|
||||
from __future__ import annotations
|
||||
import numpy as np
|
||||
from backend.node_registry import register_node
|
||||
from backend.data_types import DataField, LineData
|
||||
from backend.nodes.helpers import _cached_1d_transfer, _cached_2d_transfer
|
||||
|
||||
|
||||
@register_node(display_name="FFT Filter")
|
||||
class FFTFilter:
|
||||
"""Frequency-domain filtering of a line profile or 2-D data field.
|
||||
|
||||
Accepts either a LINE or DATA_FIELD and returns a filtered output of the
|
||||
same type. Uses a Butterworth transfer function with configurable order
|
||||
for a smooth roll-off. Equivalent to Gwyddion fft_filter_1d / fft_filter_2d.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"input": ("LINE", {
|
||||
"label": "input",
|
||||
"accepted_types": ["DATA_FIELD"],
|
||||
}),
|
||||
"filter_type": (["lowpass", "highpass", "bandpass", "notch"],),
|
||||
"cutoff": ("FLOAT", {
|
||||
"default": 0.1, "min": 0.001, "max": 1.0, "step": 0.001,
|
||||
}),
|
||||
"cutoff_high": ("FLOAT", {
|
||||
"default": 0.4, "min": 0.001, "max": 1.0, "step": 0.001,
|
||||
}),
|
||||
"order": ("INT", {"default": 2, "min": 1, "max": 10, "step": 1}),
|
||||
}
|
||||
}
|
||||
|
||||
OUTPUTS = (
|
||||
('LINE', 'filtered', {"accepted_types": ["DATA_FIELD"]}),
|
||||
)
|
||||
FUNCTION = "process"
|
||||
|
||||
DESCRIPTION = (
|
||||
"Frequency-domain filtering of a line profile or 2-D data field. "
|
||||
"Connect a LINE for 1-D filtering or a DATA_FIELD for 2-D filtering — "
|
||||
"the output mirrors the input type. "
|
||||
"Supports lowpass, highpass, bandpass, and notch (band-reject) modes "
|
||||
"with a Butterworth roll-off. Cutoffs are fractions of the Nyquist frequency."
|
||||
)
|
||||
|
||||
def process(self, input, filter_type: str, cutoff: float,
|
||||
cutoff_high: float, order: int) -> tuple:
|
||||
if isinstance(input, DataField):
|
||||
return self._process_field(input, filter_type, float(cutoff), float(cutoff_high), int(order))
|
||||
return self._process_line(input, filter_type, float(cutoff), float(cutoff_high), int(order))
|
||||
|
||||
def _process_line(self, line, filter_type: str, cutoff: float,
|
||||
cutoff_high: float, order: int) -> tuple:
|
||||
z = np.asarray(line, dtype=np.float64).ravel()
|
||||
n = len(z)
|
||||
|
||||
Z = np.fft.rfft(z)
|
||||
H = _cached_1d_transfer(n, filter_type, cutoff, cutoff_high, order)
|
||||
Z *= H
|
||||
filtered = np.fft.irfft(Z, n=n)
|
||||
|
||||
if isinstance(line, LineData):
|
||||
return (
|
||||
LineData(
|
||||
data=filtered,
|
||||
x_axis=line.x_axis.copy() if line.x_axis is not None else None,
|
||||
x_unit=line.x_unit,
|
||||
y_unit=line.y_unit,
|
||||
),
|
||||
)
|
||||
return (filtered,)
|
||||
|
||||
def _process_field(self, field: DataField, filter_type: str, cutoff: float,
|
||||
cutoff_high: float, order: int) -> tuple:
|
||||
data = field.data
|
||||
yres, xres = data.shape
|
||||
|
||||
mean_val = float(data.mean())
|
||||
centered = data - mean_val
|
||||
|
||||
spectrum = np.fft.rfft2(centered)
|
||||
transfer = _cached_2d_transfer(yres, xres, filter_type, cutoff, cutoff_high, order)
|
||||
result = np.fft.irfft2(spectrum * transfer, s=(yres, xres))
|
||||
result += mean_val
|
||||
|
||||
return (field.replace(data=result),)
|
||||
@@ -1,63 +0,0 @@
|
||||
from __future__ import annotations
|
||||
import numpy as np
|
||||
from backend.node_registry import register_node
|
||||
from backend.data_types import LineData
|
||||
from backend.nodes.helpers import _cached_1d_transfer
|
||||
|
||||
|
||||
@register_node(display_name="FFT Filter 1D")
|
||||
class FFTFilter1D:
|
||||
"""Bandpass / lowpass / highpass / notch filtering of 1-D line profiles.
|
||||
|
||||
Equivalent to Gwyddion's fft_filter_1d module. Uses a Butterworth
|
||||
transfer function with configurable order for a smooth roll-off.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"line": ("LINE",),
|
||||
"filter_type": (["lowpass", "highpass", "bandpass", "notch"],),
|
||||
"cutoff": ("FLOAT", {
|
||||
"default": 0.1, "min": 0.001, "max": 1.0, "step": 0.001,
|
||||
}),
|
||||
"cutoff_high": ("FLOAT", {
|
||||
"default": 0.4, "min": 0.001, "max": 1.0, "step": 0.001,
|
||||
}),
|
||||
"order": ("INT", {"default": 2, "min": 1, "max": 10, "step": 1}),
|
||||
}
|
||||
}
|
||||
|
||||
OUTPUTS = (
|
||||
('LINE', 'filtered'),
|
||||
)
|
||||
FUNCTION = "process"
|
||||
|
||||
DESCRIPTION = (
|
||||
"Frequency-domain filtering of a 1-D line profile. "
|
||||
"Supports lowpass, highpass, bandpass, and notch (band-reject) modes "
|
||||
"with a Butterworth roll-off. Cutoffs are fractions of the Nyquist frequency. "
|
||||
"Equivalent to Gwyddion fft_filter_1d."
|
||||
)
|
||||
|
||||
def process(self, line, filter_type: str, cutoff: float,
|
||||
cutoff_high: float, order: int) -> tuple:
|
||||
z = np.asarray(line, dtype=np.float64).ravel()
|
||||
n = len(z)
|
||||
|
||||
Z = np.fft.rfft(z)
|
||||
H = _cached_1d_transfer(n, filter_type, float(cutoff), float(cutoff_high), int(order))
|
||||
Z *= H
|
||||
filtered = np.fft.irfft(Z, n=n)
|
||||
|
||||
if isinstance(line, LineData):
|
||||
return (
|
||||
LineData(
|
||||
data=filtered,
|
||||
x_axis=line.x_axis.copy() if line.x_axis is not None else None,
|
||||
x_unit=line.x_unit,
|
||||
y_unit=line.y_unit,
|
||||
),
|
||||
)
|
||||
return (filtered,)
|
||||
@@ -1,63 +0,0 @@
|
||||
from __future__ import annotations
|
||||
import numpy as np
|
||||
from backend.node_registry import register_node
|
||||
from backend.data_types import DataField
|
||||
from backend.nodes.helpers import _cached_2d_transfer
|
||||
|
||||
|
||||
@register_node(display_name="FFT Filter 2D")
|
||||
class FFTFilter2D:
|
||||
"""Frequency-domain filtering of 2-D data fields (images).
|
||||
|
||||
Equivalent to Gwyddion's fft_filter_2d module. Applies a radial
|
||||
Butterworth transfer function in the frequency domain to remove or
|
||||
isolate periodic features.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"field": ("DATA_FIELD",),
|
||||
"filter_type": (["lowpass", "highpass", "bandpass", "notch"],),
|
||||
"cutoff": ("FLOAT", {
|
||||
"default": 0.1, "min": 0.001, "max": 1.0, "step": 0.001,
|
||||
}),
|
||||
"cutoff_high": ("FLOAT", {
|
||||
"default": 0.4, "min": 0.001, "max": 1.0, "step": 0.001,
|
||||
}),
|
||||
"order": ("INT", {"default": 2, "min": 1, "max": 10, "step": 1}),
|
||||
}
|
||||
}
|
||||
|
||||
OUTPUTS = (
|
||||
('DATA_FIELD', 'filtered'),
|
||||
)
|
||||
FUNCTION = "process"
|
||||
|
||||
DESCRIPTION = (
|
||||
"Frequency-domain filtering of a 2-D data field. "
|
||||
"Supports lowpass, highpass, bandpass, and notch (band-reject) modes "
|
||||
"with a radial Butterworth roll-off. Cutoffs are fractions of the "
|
||||
"Nyquist frequency. Use lowpass to smooth, highpass to sharpen, or "
|
||||
"bandpass/notch to isolate or remove periodic noise. "
|
||||
"Equivalent to Gwyddion fft_filter_2d."
|
||||
)
|
||||
|
||||
def process(self, field: DataField, filter_type: str, cutoff: float,
|
||||
cutoff_high: float, order: int) -> tuple:
|
||||
data = field.data
|
||||
yres, xres = data.shape
|
||||
|
||||
mean_val = float(data.mean())
|
||||
centered = data - mean_val
|
||||
|
||||
spectrum = np.fft.rfft2(centered)
|
||||
transfer = _cached_2d_transfer(
|
||||
yres, xres, filter_type,
|
||||
float(cutoff), float(cutoff_high), int(order),
|
||||
)
|
||||
result = np.fft.irfft2(spectrum * transfer, s=(yres, xres))
|
||||
result += mean_val
|
||||
|
||||
return (field.replace(data=result),)
|
||||
Reference in New Issue
Block a user