Skip to content

synth_afm.simulator

The simulator module provides the high-level AFMSimulator class, which coordinates the height-mapping kernels, scanning lag, noise, and substrate tilt into a single, easy-to-use interface.

AFMSimulator

High-level simulator for generating HS-AFM images and movies.

Source code in synth_afm/simulator.py 
  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
 51
 52
 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
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
class AFMSimulator:
    """
    High-level simulator for generating HS-AFM images and movies.
    """

    def __init__(
        self,
        pixel_size: float = 1.0,
        grid_size: Tuple[int, int] = (64, 64),
        tip_radius: float = 20.0,
        smoothness: float = 0.5,
        noise_level: float = 0.0,
        substrate_tilt: Tuple[float, float] = (0.0, 0.0),
    ):
        """
        Initialize the AFM simulator.

        Args:
            pixel_size: Size of a single pixel in Angstroms.
            grid_size: (Height, Width) of the resulting image.
            tip_radius: Radius of the AFM tip in Angstroms.
            smoothness: Smoothness of the Log-Sum-Exp max approximation.
            noise_level: Standard deviation of Gaussian noise added to heights.
            substrate_tilt: (tilt_x, tilt_y) slope of the substrate.
        """
        self.pixel_size = pixel_size
        self.grid_size = grid_size
        self.tip_radius = tip_radius
        self.smoothness = smoothness
        self.noise_level = noise_level
        self.substrate_tilt = substrate_tilt

        # Pre-compute grid
        h, w = grid_size
        x = jnp.arange(h) * pixel_size
        y = jnp.arange(w) * pixel_size
        self.grid_x, self.grid_y = jnp.meshgrid(x, y, indexing="ij")

        # Pre-compute tilt background
        self.tilt_bg = self.substrate_tilt[0] * self.grid_x + self.substrate_tilt[1] * self.grid_y

    def scan(
        self,
        positions: jax.Array,
        radii: Optional[jax.Array] = None,
        use_tip_dilation: bool = True,
        key: Optional[jax.Array] = None,
    ) -> jax.Array:
        """
        Generates a height map for the given atomic positions.

        Args:
            positions: (N, 3) atomic coordinates.
            radii: (N,) atomic radii. If None, uses default 1.7A (Carbon).
            use_tip_dilation: If True, uses the realistic tip-dilation kernel.
            key: PRNG key for noise generation.
        """
        if radii is None:
            radii = jnp.full((positions.shape[0],), 1.7)

        if use_tip_dilation:
            img = tip_sample_height_map(
                positions,
                radii,
                self.grid_x,
                self.grid_y,
                self.tip_radius,
                self.smoothness,
            )
        else:
            img = simple_height_map_kernel(
                positions, radii, self.grid_x, self.grid_y, self.smoothness
            )

        # Apply tilt
        img = img + self.tilt_bg

        # Apply noise
        if self.noise_level > 0.0:
            if key is None:
                key = jax.random.PRNGKey(0)
            img = img + jax.random.normal(key, img.shape) * self.noise_level

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

    def scan_movie(
        self,
        trajectory: jax.Array,
        radii: Optional[jax.Array] = None,
        frames_per_second: float = 10.0,
        trajectory_dt: Optional[float] = None,
        use_tip_dilation: bool = True,
        key: Optional[jax.Array] = None,
    ) -> jax.Array:
        """
        Generates a movie of HS-AFM images from a coordinate trajectory.
        Simulates scanning lag by sampling different time-steps for each line.
        Uses jax.lax.scan for memory-efficient iteration over frames.

        Args:
            trajectory: (T, N, 3) atomic coordinates over time.
            radii: (N,) atomic radii.
            frames_per_second: Nominal scan rate (FPS).
            trajectory_dt: Time between trajectory steps (seconds).
                If None, assumes 1.0 / FPS.
            use_tip_dilation: Whether to use the realistic tip model.
            key: PRNG key for noise generation.

        Returns:
            (T, H, W) array of height maps.
        """
        t_steps, n_atoms, _ = trajectory.shape
        h, w = self.grid_size

        if radii is None:
            radii = jnp.full((n_atoms,), 1.7)

        # Time per line in seconds. Scan is along y-axis (columns).
        dt_line = 1.0 / (frames_per_second * w)

        # If trajectory_dt is not provided, we assume one trajectory step per frame duration
        if trajectory_dt is None:
            trajectory_dt = 1.0 / frames_per_second

        def _scan_one_frame(frame_idx: jax.Array, frame_key: Optional[jax.Array]) -> jax.Array:
            t_start_idx = frame_idx

            def _scan_one_column(y_idx: jax.Array) -> jax.Array:
                # Calculate time offset in units of trajectory steps
                t_offset_exact = (y_idx * dt_line) / trajectory_dt
                t_offset = jnp.round(t_offset_exact).astype(jnp.int32)
                t_curr = jnp.minimum(t_start_idx + t_offset, t_steps - 1)

                # Extract column grid
                col_x = jax.lax.dynamic_slice(self.grid_x, (0, y_idx), (h, 1))
                col_y = jax.lax.dynamic_slice(self.grid_y, (0, y_idx), (h, 1))

                if use_tip_dilation:
                    res = tip_sample_height_map(
                        trajectory[t_curr], radii, col_x, col_y, self.tip_radius, self.smoothness
                    )
                else:
                    res = simple_height_map_kernel(
                        trajectory[t_curr], radii, col_x, col_y, self.smoothness
                    )
                return res[:, 0]

            img = jax.vmap(_scan_one_column)(jnp.arange(w)).T

            # Apply tilt
            img = img + self.tilt_bg

            # Apply noise
            if self.noise_level > 0.0 and frame_key is not None:
                img = img + jax.random.normal(frame_key, img.shape) * self.noise_level

            return jnp.maximum(img, 0.0)

        if key is None and self.noise_level > 0.0:
            key = jax.random.PRNGKey(0)

        if key is not None:
            keys = jax.random.split(key, t_steps)
        else:
            keys = jnp.zeros((t_steps, 2), dtype=jnp.uint32)  # Dummy keys

        def _body_fn(carry: None, x: Tuple[jax.Array, jax.Array]) -> Tuple[None, jax.Array]:
            idx, k = x
            # Use key if it's not a dummy (meaning key was provided or defaulted)
            actual_key = k if key is not None else None
            return carry, _scan_one_frame(idx, actual_key)

        _, movie = jax.lax.scan(_body_fn, None, (jnp.arange(t_steps), keys))

        return cast(jax.Array, movie)

    def compute_force_map(
        self,
        positions: jax.Array,
        radii: jax.Array,
        z_heights: jax.Array,
        k_cantilever: float = 0.1,
    ) -> jax.Array:
        """
        Experimental: Computes a force map (deflection) based on tip-sample
        repulsion.

        This uses a simplified Hookean elastic model where the tip acts as a
        linear spring. The force is proportional to the compression depth
        (sample_height - z_heights) multiplied by the cantilever stiffness
        (k_cantilever). Returns force in arbitrary units (or pN if heights
        are in nm and k is in N/m, depending on provided constants).
        """
        sample_height = self.scan(positions, radii)
        compression = jnp.maximum(sample_height - z_heights, 0.0)
        return cast(jax.Array, k_cantilever * compression)

Methods:

scan(positions, radii=None, use_tip_dilation=True, key=None)

Generates a height map for the given atomic positions.

Parameters:

Name Type Description Default
positions Array

(N, 3) atomic coordinates.

 required
radii Optional[Array]

(N,) atomic radii. If None, uses default 1.7A (Carbon).

 None
use_tip_dilation bool

If True, uses the realistic tip-dilation kernel.

 True
key Optional[Array]

PRNG key for noise generation.

 None
Source code in synth_afm/simulator.py 
48
49
50
51
52
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
def scan(
    self,
    positions: jax.Array,
    radii: Optional[jax.Array] = None,
    use_tip_dilation: bool = True,
    key: Optional[jax.Array] = None,
) -> jax.Array:
    """
    Generates a height map for the given atomic positions.

    Args:
        positions: (N, 3) atomic coordinates.
        radii: (N,) atomic radii. If None, uses default 1.7A (Carbon).
        use_tip_dilation: If True, uses the realistic tip-dilation kernel.
        key: PRNG key for noise generation.
    """
    if radii is None:
        radii = jnp.full((positions.shape[0],), 1.7)

    if use_tip_dilation:
        img = tip_sample_height_map(
            positions,
            radii,
            self.grid_x,
            self.grid_y,
            self.tip_radius,
            self.smoothness,
        )
    else:
        img = simple_height_map_kernel(
            positions, radii, self.grid_x, self.grid_y, self.smoothness
        )

    # Apply tilt
    img = img + self.tilt_bg

    # Apply noise
    if self.noise_level > 0.0:
        if key is None:
            key = jax.random.PRNGKey(0)
        img = img + jax.random.normal(key, img.shape) * self.noise_level

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

scan_movie(trajectory, radii=None, frames_per_second=10.0, trajectory_dt=None, use_tip_dilation=True, key=None)

Generates a movie of HS-AFM images from a coordinate trajectory. Simulates scanning lag by sampling different time-steps for each line. Uses jax.lax.scan for memory-efficient iteration over frames.

Parameters:

Name Type Description Default
trajectory Array

(T, N, 3) atomic coordinates over time.

 required
radii Optional[Array]

(N,) atomic radii.

 None
frames_per_second float

Nominal scan rate (FPS).

 10.0
trajectory_dt Optional[float]

Time between trajectory steps (seconds). If None, assumes 1.0 / FPS.

 None
use_tip_dilation bool

Whether to use the realistic tip model.

 True
key Optional[Array]

PRNG key for noise generation.

 None

Returns:

Type Description
Array

(T, H, W) array of height maps.
Source code in synth_afm/simulator.py 
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
def scan_movie(
    self,
    trajectory: jax.Array,
    radii: Optional[jax.Array] = None,
    frames_per_second: float = 10.0,
    trajectory_dt: Optional[float] = None,
    use_tip_dilation: bool = True,
    key: Optional[jax.Array] = None,
) -> jax.Array:
    """
    Generates a movie of HS-AFM images from a coordinate trajectory.
    Simulates scanning lag by sampling different time-steps for each line.
    Uses jax.lax.scan for memory-efficient iteration over frames.

    Args:
        trajectory: (T, N, 3) atomic coordinates over time.
        radii: (N,) atomic radii.
        frames_per_second: Nominal scan rate (FPS).
        trajectory_dt: Time between trajectory steps (seconds).
            If None, assumes 1.0 / FPS.
        use_tip_dilation: Whether to use the realistic tip model.
        key: PRNG key for noise generation.

    Returns:
        (T, H, W) array of height maps.
    """
    t_steps, n_atoms, _ = trajectory.shape
    h, w = self.grid_size

    if radii is None:
        radii = jnp.full((n_atoms,), 1.7)

    # Time per line in seconds. Scan is along y-axis (columns).
    dt_line = 1.0 / (frames_per_second * w)

    # If trajectory_dt is not provided, we assume one trajectory step per frame duration
    if trajectory_dt is None:
        trajectory_dt = 1.0 / frames_per_second

    def _scan_one_frame(frame_idx: jax.Array, frame_key: Optional[jax.Array]) -> jax.Array:
        t_start_idx = frame_idx

        def _scan_one_column(y_idx: jax.Array) -> jax.Array:
            # Calculate time offset in units of trajectory steps
            t_offset_exact = (y_idx * dt_line) / trajectory_dt
            t_offset = jnp.round(t_offset_exact).astype(jnp.int32)
            t_curr = jnp.minimum(t_start_idx + t_offset, t_steps - 1)

            # Extract column grid
            col_x = jax.lax.dynamic_slice(self.grid_x, (0, y_idx), (h, 1))
            col_y = jax.lax.dynamic_slice(self.grid_y, (0, y_idx), (h, 1))

            if use_tip_dilation:
                res = tip_sample_height_map(
                    trajectory[t_curr], radii, col_x, col_y, self.tip_radius, self.smoothness
                )
            else:
                res = simple_height_map_kernel(
                    trajectory[t_curr], radii, col_x, col_y, self.smoothness
                )
            return res[:, 0]

        img = jax.vmap(_scan_one_column)(jnp.arange(w)).T

        # Apply tilt
        img = img + self.tilt_bg

        # Apply noise
        if self.noise_level > 0.0 and frame_key is not None:
            img = img + jax.random.normal(frame_key, img.shape) * self.noise_level

        return jnp.maximum(img, 0.0)

    if key is None and self.noise_level > 0.0:
        key = jax.random.PRNGKey(0)

    if key is not None:
        keys = jax.random.split(key, t_steps)
    else:
        keys = jnp.zeros((t_steps, 2), dtype=jnp.uint32)  # Dummy keys

    def _body_fn(carry: None, x: Tuple[jax.Array, jax.Array]) -> Tuple[None, jax.Array]:
        idx, k = x
        # Use key if it's not a dummy (meaning key was provided or defaulted)
        actual_key = k if key is not None else None
        return carry, _scan_one_frame(idx, actual_key)

    _, movie = jax.lax.scan(_body_fn, None, (jnp.arange(t_steps), keys))

    return cast(jax.Array, movie)

compute_force_map(positions, radii, z_heights, k_cantilever=0.1)

Experimental: Computes a force map (deflection) based on tip-sample repulsion.

This uses a simplified Hookean elastic model where the tip acts as a linear spring. The force is proportional to the compression depth (sample_height - z_heights) multiplied by the cantilever stiffness (k_cantilever). Returns force in arbitrary units (or pN if heights are in nm and k is in N/m, depending on provided constants).

Source code in synth_afm/simulator.py 
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
def compute_force_map(
    self,
    positions: jax.Array,
    radii: jax.Array,
    z_heights: jax.Array,
    k_cantilever: float = 0.1,
) -> jax.Array:
    """
    Experimental: Computes a force map (deflection) based on tip-sample
    repulsion.

    This uses a simplified Hookean elastic model where the tip acts as a
    linear spring. The force is proportional to the compression depth
    (sample_height - z_heights) multiplied by the cantilever stiffness
    (k_cantilever). Returns force in arbitrary units (or pN if heights
    are in nm and k is in N/m, depending on provided constants).
    """
    sample_height = self.scan(positions, radii)
    compression = jnp.maximum(sample_height - z_heights, 0.0)
    return cast(jax.Array, k_cantilever * compression)