199 lines
6.1 KiB
Python
199 lines
6.1 KiB
Python
"""DWT anisotropy — quantify surface anisotropy using wavelet decomposition."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import numpy as np
|
|
|
|
from backend.data_types import DataField, RecordTable
|
|
from backend.node_registry import register_node
|
|
|
|
|
|
def _next_power_of_2(n: int) -> int:
|
|
"""Return the smallest power of 2 >= n."""
|
|
p = 1
|
|
while p < n:
|
|
p <<= 1
|
|
return p
|
|
|
|
|
|
def _pad_to_pow2(data: np.ndarray) -> np.ndarray:
|
|
"""Pad *data* to the next power-of-2 dimensions using edge values."""
|
|
rows, cols = data.shape
|
|
new_rows = _next_power_of_2(rows)
|
|
new_cols = _next_power_of_2(cols)
|
|
if new_rows == rows and new_cols == cols:
|
|
return data.copy()
|
|
out = np.zeros((new_rows, new_cols), dtype=np.float64)
|
|
out[:rows, :cols] = data
|
|
# edge-pad
|
|
if new_cols > cols:
|
|
out[:rows, cols:] = data[:, -1:]
|
|
if new_rows > rows:
|
|
out[rows:, :cols] = data[-1:, :]
|
|
if new_rows > rows and new_cols > cols:
|
|
out[rows:, cols:] = data[-1, -1]
|
|
return out
|
|
|
|
|
|
def _haar_decompose_2d(data: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
|
|
"""
|
|
One level of 2-D Haar wavelet decomposition.
|
|
|
|
Returns (LL, LH, HL, HH) each of shape (rows//2, cols//2).
|
|
LL = (a+b+c+d)/2 approximation
|
|
LH = (a+b-c-d)/2 horizontal detail (captures vertical features)
|
|
HL = (a-b+c-d)/2 vertical detail (captures horizontal features)
|
|
HH = (a-b-c+d)/2 diagonal detail
|
|
"""
|
|
rows, cols = data.shape
|
|
a = data[0::2, 0::2] # top-left
|
|
b = data[0::2, 1::2] # top-right
|
|
c = data[1::2, 0::2] # bottom-left
|
|
d = data[1::2, 1::2] # bottom-right
|
|
|
|
ll = (a + b + c + d) / 2.0
|
|
lh = (a + b - c - d) / 2.0
|
|
hl = (a - b + c - d) / 2.0
|
|
hh = (a - b - c + d) / 2.0
|
|
return ll, lh, hl, hh
|
|
|
|
|
|
def _compute_dwt_anisotropy(
|
|
data: np.ndarray,
|
|
n_levels: int,
|
|
) -> tuple[list[float], list[float], list[float], list[np.ndarray]]:
|
|
"""
|
|
Multi-level 2-D Haar decomposition with per-level energy ratios.
|
|
|
|
Returns
|
|
-------
|
|
x_energies : per-level sum(HL**2)
|
|
y_energies : per-level sum(LH**2)
|
|
ratios : per-level x_energy / y_energy
|
|
ratio_maps : per-level ratio arrays (at decomposition resolution)
|
|
"""
|
|
padded = _pad_to_pow2(np.asarray(data, dtype=np.float64))
|
|
current = padded
|
|
|
|
x_energies: list[float] = []
|
|
y_energies: list[float] = []
|
|
ratios: list[float] = []
|
|
ratio_maps: list[np.ndarray] = []
|
|
|
|
for _ in range(n_levels):
|
|
if current.shape[0] < 2 or current.shape[1] < 2:
|
|
break
|
|
ll, lh, hl, hh = _haar_decompose_2d(current)
|
|
|
|
x_energy = float(np.sum(hl ** 2))
|
|
y_energy = float(np.sum(lh ** 2))
|
|
ratio = x_energy / (y_energy + 1e-30)
|
|
|
|
x_energies.append(x_energy)
|
|
y_energies.append(y_energy)
|
|
ratios.append(ratio)
|
|
|
|
# Build a per-pixel ratio map at this level's resolution
|
|
hl_sq = hl ** 2
|
|
lh_sq = lh ** 2
|
|
level_map = hl_sq / (lh_sq + 1e-30)
|
|
ratio_maps.append(level_map)
|
|
|
|
current = ll
|
|
|
|
return x_energies, y_energies, ratios, ratio_maps
|
|
|
|
|
|
def _build_anisotropy_map(
|
|
ratio_maps: list[np.ndarray],
|
|
orig_rows: int,
|
|
orig_cols: int,
|
|
) -> np.ndarray:
|
|
"""
|
|
Combine per-level ratio maps into a single anisotropy map at
|
|
the original field resolution by upsampling and averaging.
|
|
"""
|
|
if not ratio_maps:
|
|
return np.ones((orig_rows, orig_cols), dtype=np.float64)
|
|
|
|
from scipy.ndimage import zoom
|
|
|
|
target = np.zeros((orig_rows, orig_cols), dtype=np.float64)
|
|
for level_map in ratio_maps:
|
|
zy = orig_rows / level_map.shape[0]
|
|
zx = orig_cols / level_map.shape[1]
|
|
upsampled = zoom(level_map, (zy, zx), order=1)
|
|
# zoom may produce shape off by one — trim or pad
|
|
upsampled = upsampled[:orig_rows, :orig_cols]
|
|
if upsampled.shape[0] < orig_rows or upsampled.shape[1] < orig_cols:
|
|
tmp = np.ones((orig_rows, orig_cols), dtype=np.float64)
|
|
tmp[: upsampled.shape[0], : upsampled.shape[1]] = upsampled
|
|
upsampled = tmp
|
|
target += upsampled
|
|
|
|
target /= len(ratio_maps)
|
|
return target
|
|
|
|
|
|
@register_node(display_name="DWT Anisotropy")
|
|
class DWTAnisotropy:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"field": ("DATA_FIELD",),
|
|
"n_levels": (
|
|
"INT",
|
|
{"default": 4, "min": 1, "max": 10},
|
|
),
|
|
"ratio_threshold": (
|
|
"FLOAT",
|
|
{"default": 0.2, "min": 0.001, "max": 10.0, "step": 0.01},
|
|
),
|
|
}
|
|
}
|
|
|
|
OUTPUTS = (
|
|
('DATA_FIELD', 'anisotropy_map'),
|
|
('RECORD_TABLE', 'statistics'),
|
|
)
|
|
FUNCTION = "process"
|
|
|
|
DESCRIPTION = (
|
|
"Quantify surface anisotropy using a multi-level 2-D Haar wavelet decomposition. "
|
|
"At each level, horizontal (HL) and vertical (LH) detail energies are compared to "
|
|
"produce an X/Y energy ratio. Ratio > 1 indicates more horizontal features; "
|
|
"ratio < 1 indicates more vertical features. Equivalent to Gwyddion's dwtanisotropy.c."
|
|
)
|
|
|
|
def process(
|
|
self,
|
|
field: DataField,
|
|
n_levels: int,
|
|
ratio_threshold: float,
|
|
) -> tuple:
|
|
data = np.asarray(field.data, dtype=np.float64)
|
|
orig_rows, orig_cols = data.shape
|
|
|
|
x_energies, y_energies, ratios, ratio_maps = _compute_dwt_anisotropy(
|
|
data, int(n_levels),
|
|
)
|
|
|
|
# Build per-pixel anisotropy map
|
|
aniso_map = _build_anisotropy_map(ratio_maps, orig_rows, orig_cols)
|
|
aniso_field = field.replace(data=aniso_map)
|
|
|
|
# Build statistics table
|
|
rows = []
|
|
for i, (xe, ye, r) in enumerate(zip(x_energies, y_energies, ratios)):
|
|
rows.append({
|
|
"level": i + 1,
|
|
"x_energy": float(xe),
|
|
"y_energy": float(ye),
|
|
"ratio": float(r),
|
|
"anisotropic": abs(r - 1.0) > float(ratio_threshold),
|
|
})
|
|
stats = RecordTable(rows)
|
|
|
|
return (aniso_field, stats)
|