109 lines
3.8 KiB
Python
109 lines
3.8 KiB
Python
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="Tip Model")
|
||
class TipModel:
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"field": ("DATA_FIELD",),
|
||
"shape": (["parabola", "cone", "sphere"], {"default": "parabola"}),
|
||
"radius": ("FLOAT", {
|
||
"default": 10e-9, "min": 1e-10, "max": 1e-6, "step": 1e-10,
|
||
}),
|
||
"half_angle": ("FLOAT", {
|
||
"default": 20.0, "min": 1.0, "max": 89.0, "step": 0.5,
|
||
}),
|
||
"n_pixels": ("INT", {"default": 65, "min": 3, "max": 511, "step": 2}),
|
||
}
|
||
}
|
||
|
||
OUTPUTS = (
|
||
('DATA_FIELD', 'tip'),
|
||
)
|
||
FUNCTION = "process"
|
||
|
||
DESCRIPTION = (
|
||
"Generate a synthetic AFM tip model DATA_FIELD. "
|
||
"The input field sets the pixel size for the tip. "
|
||
"The apex (centre pixel) is the maximum value; edges are shifted to zero. "
|
||
"Shapes: parabola — paraboloid with apex radius R; "
|
||
"cone — sphere-capped cone (radius R, half_angle from tip axis in degrees); "
|
||
"sphere — ball-on-stick (sphere cap only). "
|
||
)
|
||
|
||
def process(
|
||
self,
|
||
field: DataField,
|
||
shape: str,
|
||
radius: float,
|
||
half_angle: float,
|
||
n_pixels: int,
|
||
) -> tuple:
|
||
pixel_size = (field.dx + field.dy) * 0.5
|
||
|
||
# Ensure odd size so the centre pixel is well-defined
|
||
n = n_pixels if n_pixels % 2 == 1 else n_pixels + 1
|
||
ci = n // 2
|
||
|
||
# Physical offsets from the centre pixel
|
||
offsets = (np.arange(n) - ci) * pixel_size
|
||
gx, gy = np.meshgrid(offsets, offsets)
|
||
r2 = gx**2 + gy**2
|
||
r = np.sqrt(r2)
|
||
|
||
if shape == "parabola":
|
||
# Gwyddion parabola(): a = 1/(2R), z0 = 2a·x_half²
|
||
# z[y,x] = z0 − a·r² → min at corners is exactly 0
|
||
a = 0.5 / radius
|
||
x_half = ci * pixel_size # half-width in physical units
|
||
z0 = 2.0 * a * x_half**2
|
||
data = z0 - a * r2
|
||
data -= data.min() # shift so min = 0
|
||
|
||
elif shape == "cone":
|
||
# Gwyddion cone():
|
||
# angle = half-angle from horizontal = π/2 − half_angle_from_axis
|
||
# z0 = R/sin(angle)
|
||
# r_cross² = (R·cos(angle))²
|
||
# inner (r² < r_cross²): z = sqrt(R² − r²) ← spherical cap
|
||
# outer: z = z0 − r/tan(angle)
|
||
angle = np.radians(90.0 - half_angle) # slope from horizontal
|
||
z0 = radius / np.sin(angle)
|
||
br2 = (radius * np.cos(angle))**2
|
||
ta = 1.0 / np.tan(angle)
|
||
inner = np.sqrt(np.maximum(0.0, radius**2 - r2))
|
||
outer = z0 - ta * r
|
||
data = np.where(r2 < br2, inner, outer)
|
||
data -= data.min()
|
||
|
||
elif shape == "sphere":
|
||
# Gwyddion sphere(): cone with near-vertical sides (angle ≈ 0 from horizontal)
|
||
angle = 1e-5 # radians — essentially vertical cone
|
||
z0 = radius / np.sin(angle)
|
||
br2 = (radius * np.cos(angle))**2
|
||
ta = 1.0 / np.tan(angle)
|
||
inner = np.sqrt(np.maximum(0.0, radius**2 - r2))
|
||
outer = z0 - ta * r
|
||
data = np.where(r2 < br2, inner, outer)
|
||
data -= data.min()
|
||
|
||
else:
|
||
raise ValueError(f"Unknown tip shape {shape!r}. Choose: parabola, cone, sphere.")
|
||
|
||
xreal = n * pixel_size
|
||
tip_field = DataField(
|
||
data=data,
|
||
xreal=xreal,
|
||
yreal=xreal,
|
||
si_unit_xy=field.si_unit_xy,
|
||
si_unit_z=field.si_unit_z,
|
||
)
|
||
return (tip_field,)
|