170 lines
6.8 KiB
Python
170 lines
6.8 KiB
Python
import numpy as np
|
||
import pytest
|
||
from tests.node_tests._shared import make_field
|
||
from backend.data_types import DataField
|
||
|
||
|
||
def run_blind(field, n_pixels=17, threshold=0.0, method="partial", use_edges=False):
|
||
from backend.nodes.tip_blind_estimate import BlindTipEstimate
|
||
node = BlindTipEstimate()
|
||
tip, certainty = node.process(
|
||
field=field, n_pixels=n_pixels, threshold=threshold,
|
||
method=method, use_edges=use_edges,
|
||
)
|
||
return tip, certainty
|
||
|
||
|
||
# ── Output types and dimensions ──────────────────────────────────────────────
|
||
|
||
def test_outputs_are_correct_types():
|
||
field = make_field(shape=(32, 32), xreal=32e-9, yreal=32e-9)
|
||
tip, certainty = run_blind(field, n_pixels=9)
|
||
assert isinstance(tip, DataField)
|
||
assert isinstance(certainty, np.ndarray)
|
||
assert certainty.dtype == np.uint8
|
||
|
||
|
||
def test_tip_output_shape():
|
||
"""Tip must be (n_pixels, n_pixels) — odd, bumped if even."""
|
||
field = make_field(shape=(32, 32), xreal=32e-9, yreal=32e-9)
|
||
for n in (7, 11, 17):
|
||
tip, _ = run_blind(field, n_pixels=n)
|
||
assert tip.data.shape == (n, n), f"Expected ({n},{n}), got {tip.data.shape}"
|
||
|
||
|
||
def test_tip_n_pixels_even_bumped():
|
||
field = make_field(shape=(32, 32), xreal=32e-9, yreal=32e-9)
|
||
tip, _ = run_blind(field, n_pixels=16)
|
||
assert tip.data.shape[0] == 17
|
||
|
||
|
||
def test_certainty_output_matches_field_shape():
|
||
field = make_field(shape=(48, 64))
|
||
_, certainty = run_blind(field, n_pixels=9)
|
||
assert certainty.shape == field.data.shape
|
||
|
||
|
||
def test_certainty_is_binary():
|
||
"""Certainty mask values must all be 0 or 255."""
|
||
field = make_field(shape=(32, 32), xreal=32e-9, yreal=32e-9)
|
||
_, certainty = run_blind(field, n_pixels=9)
|
||
vals = np.unique(certainty)
|
||
for v in vals:
|
||
assert v in (0, 255), f"Non-binary certainty value: {v}"
|
||
|
||
|
||
# ── Tip conventions ───────────────────────────────────────────────────────────
|
||
|
||
def test_tip_min_is_zero():
|
||
field = make_field(shape=(32, 32), xreal=32e-9, yreal=32e-9)
|
||
tip, _ = run_blind(field, n_pixels=9)
|
||
assert tip.data.min() >= -1e-15
|
||
|
||
|
||
def test_tip_max_at_centre():
|
||
"""Apex (centre pixel) must be the maximum of the estimated tip."""
|
||
field = make_field(shape=(32, 32), xreal=32e-9, yreal=32e-9)
|
||
tip, _ = run_blind(field, n_pixels=9)
|
||
ci = tip.data.shape[0] // 2
|
||
assert tip.data[ci, ci] == pytest.approx(tip.data.max(), abs=1e-20)
|
||
|
||
|
||
def test_tip_units_inherited():
|
||
field = make_field(xreal=1e-6, yreal=1e-6)
|
||
field.si_unit_xy = "nm"
|
||
field.si_unit_z = "V"
|
||
tip, _ = run_blind(field, n_pixels=9)
|
||
assert tip.si_unit_xy == "nm"
|
||
assert tip.si_unit_z == "V"
|
||
|
||
|
||
# ── Flat field gives flat (zero) tip ─────────────────────────────────────────
|
||
|
||
def test_flat_field_gives_zero_tip():
|
||
"""A flat image has no features → blind estimation cannot refine the tip → stays flat."""
|
||
flat = make_field(data=np.full((32, 32), 3.14))
|
||
tip, _ = run_blind(flat, n_pixels=7)
|
||
# Tip should be all zeros (no refinement possible)
|
||
assert np.allclose(tip.data, 0.0, atol=1e-12)
|
||
|
||
|
||
# ── Single spike → estimated tip matches expected shape ──────────────────────
|
||
|
||
def test_spike_gives_sharp_tip():
|
||
"""
|
||
A single sharp spike dilated by a known paraboloid tip gives a broadened image.
|
||
Blind estimation on the broadened image should recover a tip that is ≤ the true tip
|
||
everywhere (blind estimation is an upper bound on the true tip shape).
|
||
"""
|
||
from scipy.ndimage import grey_dilation
|
||
from backend.nodes.tip_model import TipModel
|
||
|
||
pixel_size = 1e-9
|
||
n = 64
|
||
field_ref = make_field(shape=(n, n), xreal=n * pixel_size, yreal=n * pixel_size)
|
||
|
||
# Create true parabolic tip (33×33, radius=20nm)
|
||
true_tip_node = TipModel()
|
||
true_tip, = true_tip_node.process(
|
||
field=field_ref, shape="parabola", radius=20e-9,
|
||
half_angle=20.0, n_pixels=33,
|
||
)
|
||
|
||
# Spike surface
|
||
surface = np.zeros((n, n))
|
||
surface[n // 2, n // 2] = 1e-9 # 1 nm spike
|
||
|
||
# Dilated (measured) image
|
||
dil_struct = true_tip.data - true_tip.data.max()
|
||
measured_data = grey_dilation(surface, structure=dil_struct)
|
||
measured = make_field(data=measured_data, xreal=n * pixel_size, yreal=n * pixel_size)
|
||
|
||
# Blind estimation
|
||
est_tip, certainty = run_blind(measured, n_pixels=33, threshold=0.0, method="partial")
|
||
|
||
# The estimated tip must be ≤ the true tip everywhere (blind est. is upper bound)
|
||
# Both are min=0, max=apex. Compare normalised shapes.
|
||
true_norm = true_tip.data / true_tip.data.max() if true_tip.data.max() > 0 else true_tip.data
|
||
est_norm = est_tip.data / est_tip.data.max() if est_tip.data.max() > 0 else est_tip.data
|
||
# After normalisation: est ≤ true everywhere (blind is conservative)
|
||
assert np.all(est_norm <= true_norm + 1e-6), \
|
||
"Blind estimate exceeded true tip (not a valid upper bound)"
|
||
|
||
|
||
# ── Partial vs full ───────────────────────────────────────────────────────────
|
||
|
||
def test_full_method_runs():
|
||
"""Full estimation should run without error on a small image."""
|
||
field = make_field(shape=(24, 24), xreal=24e-9, yreal=24e-9)
|
||
tip, certainty = run_blind(field, n_pixels=7, method="full")
|
||
assert isinstance(tip, DataField)
|
||
assert isinstance(certainty, np.ndarray)
|
||
|
||
|
||
# ── Certainty increases with sharp features ───────────────────────────────────
|
||
|
||
def test_certainty_nonzero_for_sharp_image():
|
||
"""An image with distinct features should produce some certain pixels."""
|
||
from scipy.ndimage import grey_dilation
|
||
from backend.nodes.tip_model import TipModel
|
||
|
||
pixel_size = 1e-9
|
||
n = 64
|
||
field_ref = make_field(shape=(n, n), xreal=n * pixel_size, yreal=n * pixel_size)
|
||
true_tip_node = TipModel()
|
||
true_tip, = true_tip_node.process(
|
||
field=field_ref, shape="parabola", radius=20e-9,
|
||
half_angle=20.0, n_pixels=17,
|
||
)
|
||
|
||
# Create a sharp pyramid surface
|
||
cx, cy = n // 2, n // 2
|
||
ys, xs = np.mgrid[0:n, 0:n]
|
||
surface = np.maximum(0, 5e-9 - 0.3e-9 * np.maximum(np.abs(xs - cx), np.abs(ys - cy)))
|
||
dil_struct = true_tip.data - true_tip.data.max()
|
||
measured_data = grey_dilation(surface, structure=dil_struct)
|
||
measured = make_field(data=measured_data, xreal=n * pixel_size, yreal=n * pixel_size)
|
||
|
||
_, certainty = run_blind(measured, n_pixels=17, method="partial")
|
||
assert certainty.sum() > 0, "No certain pixels found for a sharp image"
|