80 lines
2.5 KiB
Python
80 lines
2.5 KiB
Python
"""Log-polar PSDF — power spectral density in log-polar coordinates."""
|
||
|
||
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="Log-Polar PSDF")
|
||
class LogPolarPSDF:
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"field": ("DATA_FIELD",),
|
||
"n_phi": ("INT", {"default": 180, "min": 36, "max": 720, "step": 1}),
|
||
"n_r": ("INT", {"default": 100, "min": 20, "max": 500, "step": 1}),
|
||
}
|
||
}
|
||
|
||
OUTPUTS = (
|
||
('DATA_FIELD', 'psdf'),
|
||
)
|
||
FUNCTION = "process"
|
||
|
||
DESCRIPTION = (
|
||
"Compute the power spectral density function in log-polar coordinates. "
|
||
"The x-axis is the azimuthal angle (0–360°) and y-axis is log(frequency). "
|
||
"Better than Cartesian PSDF for anisotropy analysis. "
|
||
)
|
||
|
||
def process(self, field: DataField, n_phi: int, n_r: int) -> tuple:
|
||
data = np.asarray(field.data, dtype=np.float64)
|
||
yres, xres = data.shape
|
||
|
||
# Compute 2D power spectrum
|
||
fft = np.fft.fft2(data - data.mean())
|
||
power = np.abs(np.fft.fftshift(fft))**2
|
||
|
||
cy, cx = yres // 2, xres // 2
|
||
|
||
# Build log-polar grid
|
||
r_max = min(cx, cy)
|
||
log_r = np.linspace(0, np.log(r_max), n_r)
|
||
phi = np.linspace(0, 2 * np.pi, n_phi, endpoint=False)
|
||
|
||
result = np.zeros((n_r, n_phi))
|
||
|
||
for ir in range(n_r):
|
||
r = np.exp(log_r[ir])
|
||
for ip in range(n_phi):
|
||
fx = cx + r * np.cos(phi[ip])
|
||
fy = cy + r * np.sin(phi[ip])
|
||
# Bilinear interpolation
|
||
ix = int(fx)
|
||
iy = int(fy)
|
||
if 0 <= ix < xres - 1 and 0 <= iy < yres - 1:
|
||
dx = fx - ix
|
||
dy = fy - iy
|
||
val = (power[iy, ix] * (1 - dx) * (1 - dy) +
|
||
power[iy, ix + 1] * dx * (1 - dy) +
|
||
power[iy + 1, ix] * (1 - dx) * dy +
|
||
power[iy + 1, ix + 1] * dx * dy)
|
||
result[ir, ip] = val
|
||
|
||
# Log scale for display
|
||
result = np.log1p(result)
|
||
|
||
psdf_field = DataField(
|
||
data=result,
|
||
xreal=360.0,
|
||
yreal=float(np.log(r_max)),
|
||
si_unit_xy="deg",
|
||
si_unit_z="",
|
||
domain="frequency",
|
||
)
|
||
return (psdf_field,)
|