synth-afm is a JAX-powered toolkit for generating synthetic High-Speed Atomic Force Microscopy (HS-AFM) images and movies from atomistic protein structures.
Built with the differentiable biophysics philosophy, every step—from coordinate rotation to tip-collision height mapping—is end-to-end differentiable.
HS-AFM provides a unique look at "proteins at work," but interpreting noisy movies is challenging. synth-afm helps you bridge the resolution gap:
- Realistic Tip Physics: Uses a spherical-tip dilation model to account for the broadening effect of the AFM probe.
- Atomic Rigor: Automatically assigns van der Waals radii based on element (Bondi, 1964) for accurate topography.
- Temporal Distortion: Models Scanning Lag, simulating how protein dynamics during a scan cause the "shear" artifacts seen in real HS-AFM movies.
- Force Maps: Go beyond height-maps with experimental support for tip-sample repulsion (deflection) modeling.
synth-afm treats the entire AFM scanning process as a differentiable operator
-
End-to-End Differentiable: Built entirely in JAX, allowing you to flow gradients from an experimental AFM image
$\mathbf{I}_{exp}$ back to atomic coordinates$\mathbf{X}$ . -
Flexible Fitting: Enable gradient-based optimization of molecular structures using experimental AFM data as a loss term:
$\mathcal{L} = |\mathcal{H}(\mathbf{X}) - \mathbf{I}_{exp}|^2$ . - Synthetic Benchmarking: Generate large-scale, ground-truth datasets of "corrupted" AFM movies (with lag, noise, and dilation) to train denoising or state-detection models.
- Differentiable Height Mapping: Efficient Log-Sum-Exp collision detection for sub-nanometer topography.
- Physical Realism: Simulate cantilever noise and substrate tilt (linear gradients) to match experimental conditions.
- Scanning Lag Simulation: Models the line-by-line temporal delay inherent in pixel-by-pixel acquisition.
- Memory Efficiency: Uses
jax.lax.scanfor constant-memory simulation of long trajectories. - Flexible Tip Geometries: Supports spherical tip-shape dilation.
- Integration: Reads PDB/mmCIF files via
biotiteand integrates withsynth-pdbandsynth-dynamics.
pip install synth-afmGet started immediately with our interactive Jupyter notebooks:
- Quick Start: Differentiable HS-AFM Simulation: Learn how to generate height maps with tip dilation and scanning lag.
import jax.numpy as jnp
from synth_afm.simulator import AFMSimulator
from synth_afm.io import load_coords_and_radii
# 1. Load your structure (N, 3) and radii (N,)
coords, radii = load_coords_and_radii("molecule.pdb")
# 2. Initialize simulator (1A pixel size, 2nm tip radius, 0.5A noise, slight tilt)
sim = AFMSimulator(
pixel_size=1.0,
tip_radius=20.0,
noise_level=0.5,
substrate_tilt=(0.01, 0.0)
)
# 3. Generate height map (Differentiable!)
height_map = sim.scan(coords, radii)The height-mapping kernels are validated against the standard Villarrubia algorithm and verified to preserve atomic heights within 0.01 Å precision. The temporal lag simulation correctly reproduces the stroboscopic shearing effects documented in high-speed biological AFM (Ando et al., 2011).
synth-afm is part of a broader ecosystem for synthetic biophysics data generation:
| Project | Purpose |
|---|---|
| synth-pdb | Foundation: Realistic protein structure generation and PDB/mmCIF handling |
| synth-nmr | NMR observables (NOE, RDC, chemical shifts, J-couplings, relaxation) |
| synth-saxs | SAXS profile simulation via Debye formula |
| synth-cryo-em | Cryo-EM density map generation with CTF/noise modeling |
| synth-dynamics | ANM/Langevin dynamics for conformational ensembles |
| diff-biophys | Differentiable JAX implementations of all biophysics kernels |
Distributed under the MIT License. See LICENSE for more information.
If you use synth-afm in your research, please cite:
@software{synth_afm,
author = {Elkins, George},
title = {synth-afm: Differentiable HS-AFM Simulation},
year = {2026},
url = {https://github.com/elkins-lab/synth-afm}
}