41 lines
1.5 KiB
Python
41 lines
1.5 KiB
Python
import numpy as np
|
|
from tests.node_tests._shared import make_field
|
|
|
|
|
|
def test_plane_level():
|
|
from backend.nodes.level_plane import PlaneLevelField
|
|
node = PlaneLevelField()
|
|
|
|
N = 64
|
|
y, x = np.mgrid[0:N, 0:N] / N
|
|
signal = np.sin(2 * np.pi * 5 * x)
|
|
data = 100 * x + 50 * y + signal
|
|
field = make_field(data=data)
|
|
|
|
result, = node.process(field)
|
|
assert result.data.shape == field.data.shape
|
|
assert abs(result.data.mean()) < 1e-10
|
|
corr = np.corrcoef(result.data.ravel(), signal.ravel())[0, 1]
|
|
assert corr > 0.98
|
|
|
|
yy_px, xx_px = np.mgrid[0:N, 0:N]
|
|
|
|
def fit_pixel_plane(data_in, region):
|
|
A = np.column_stack([np.ones(int(np.count_nonzero(region))), xx_px[region].astype(np.float64), yy_px[region].astype(np.float64)])
|
|
coeffs, _, _, _ = np.linalg.lstsq(A, data_in[region].ravel().astype(np.float64), rcond=None)
|
|
return float(coeffs[0]), float(coeffs[1]), float(coeffs[2])
|
|
|
|
mask = np.zeros((N, N), dtype=np.uint8)
|
|
mask[20:44, 22:46] = 255
|
|
feature = np.zeros((N, N), dtype=np.float64)
|
|
feature[mask > 0] = 35.0
|
|
masked_field = make_field(data=100 * x + 50 * y + feature)
|
|
|
|
unmasked, = node.process(masked_field)
|
|
masked, = node.process(masked_field, masking="exclude", mask=mask)
|
|
|
|
outside = mask == 0
|
|
_, unmasked_bx, unmasked_by = fit_pixel_plane(unmasked.data, outside)
|
|
_, masked_bx, masked_by = fit_pixel_plane(masked.data, outside)
|
|
assert np.hypot(masked_bx, masked_by) < np.hypot(unmasked_bx, unmasked_by) * 1e-3
|