Skip to content

synth_afm.kernels

The kernels module contains the low-level, fully differentiable JAX functions that implement the AFM imaging physics. These can be used directly for advanced workflows or custom simulators.

simple_height_map_kernel(positions, radii, grid_x, grid_y, smoothness=0.1)

A differentiable forward kernel to compute a molecular height map. Uses a Log-Sum-Exp approximation to smoothly combine atomic heights. Note: The background substrate is assumed to be at z=0.0, and the final height map is clipped to be non-negative.

Parameters:

Name Type Description Default
positions Array

(N, 3) array of atomic coordinates.

 required
radii Array

(N,) array of atomic van der Waals radii.

 required
grid_x Array

(H, W) grid of X coordinates.

 required
grid_y Array

(H, W) grid of Y coordinates.

 required
smoothness float

Controls the sharpness of the surface (lower is sharper).

 0.1

Returns:

Type Description
Array

(H, W) height map array.
Source code in synth_afm/kernels.py 
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
def simple_height_map_kernel(
    positions: jax.Array,
    radii: jax.Array,
    grid_x: jax.Array,
    grid_y: jax.Array,
    smoothness: float = 0.1,
) -> jax.Array:
    """
    A differentiable forward kernel to compute a molecular height map.
    Uses a Log-Sum-Exp approximation to smoothly combine atomic heights.
    Note: The background substrate is assumed to be at z=0.0, and the final
    height map is clipped to be non-negative.

    Args:
        positions: (N, 3) array of atomic coordinates.
        radii: (N,) array of atomic van der Waals radii.
        grid_x: (H, W) grid of X coordinates.
        grid_y: (H, W) grid of Y coordinates.
        smoothness: Controls the sharpness of the surface (lower is sharper).

    Returns:
        (H, W) height map array.
    """
    # 1. Compute horizontal distance from each atom to each grid point
    atom_x = positions[:, 0, None, None]
    atom_y = positions[:, 1, None, None]
    atom_z = positions[:, 2, None, None]

    dx = grid_x - atom_x
    dy = grid_y - atom_y
    r_sq = dx**2 + dy**2

    # 2. Compute the theoretical height of the surface above each atom
    vdw_sq = radii[:, None, None] ** 2
    heights = atom_z + jnp.sqrt(jnp.maximum(vdw_sq - r_sq, 0.0))

    # Mask out areas where the atom doesn't contribute (r > r_vdw)
    mask = (r_sq < vdw_sq).astype(jnp.float32)
    masked_heights = (heights - 1000.0) * (1.0 - mask) + heights * mask

    # 3. Combine heights using Log-Sum-Exp for differentiability
    lse_height = smoothness * jax.nn.logsumexp(masked_heights / smoothness, axis=0)

    # Floor the height map at zero (the substrate surface)
    return cast(jax.Array, jnp.maximum(lse_height, 0.0))

tip_sample_height_map(positions, radii, grid_x, grid_y, tip_radius=20.0, smoothness=0.5)

Computes a height map where the tip radius is explicitly modeled. This simulates 'tip dilation' where the image is the dilation of the sample surface by the tip shape.

Note: The background substrate is assumed to be at z=0.0, and the final height map is clipped to be non-negative.

Technically, the height at grid point (x, y) is: H(x, y) = max_{atom i} (z_i + sqrt((r_i + r_tip)^2 - d_xy_i^2)) - r_tip

Source code in synth_afm/kernels.py 
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
def tip_sample_height_map(
    positions: jax.Array,
    radii: jax.Array,
    grid_x: jax.Array,
    grid_y: jax.Array,
    tip_radius: float = 20.0,
    smoothness: float = 0.5,
) -> jax.Array:
    """
    Computes a height map where the tip radius is explicitly modeled.
    This simulates 'tip dilation' where the image is the dilation of the
    sample surface by the tip shape.

    Note: The background substrate is assumed to be at z=0.0, and the final
    height map is clipped to be non-negative.

    Technically, the height at grid point (x, y) is:
        H(x, y) = max_{atom i} (z_i + sqrt((r_i + r_tip)^2 - d_xy_i^2)) - r_tip
    """
    atom_x = positions[:, 0, None, None]
    atom_y = positions[:, 1, None, None]
    atom_z = positions[:, 2, None, None]

    dx = grid_x - atom_x
    dy = grid_y - atom_y
    r_sq = dx**2 + dy**2

    # Effective radius is the sum of atom radius and tip radius
    eff_radius = radii[:, None, None] + tip_radius
    eff_radius_sq = eff_radius**2

    # Theoretical height of the center of the tip sphere
    tip_center_z = atom_z + jnp.sqrt(jnp.maximum(eff_radius_sq - r_sq, 0.0))

    # Mask out areas where the tip doesn't touch the atom
    mask = (r_sq < eff_radius_sq).astype(jnp.float32)
    masked_heights = (tip_center_z - 1000.0) * (1.0 - mask) + tip_center_z * mask

    # Smooth max across atoms
    max_tip_z = smoothness * jax.nn.logsumexp(masked_heights / smoothness, axis=0)

    # The actual height measured is the tip center minus the tip radius
    measured_height = max_tip_z - tip_radius

    return cast(jax.Array, jnp.maximum(measured_height, 0.0))