Skip to content

elkins-lab/synth-afm

Repository files navigation

🧬 synth-afm: Differentiable HS-AFM Simulation

PyPI version Python versions Tests codecov License: MIT JAX Code style: black

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.


🧪 For Structural Biologists

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.

🤖 For Machine Learning Geeks

synth-afm treats the entire AFM scanning process as a differentiable operator $\mathcal{H}: \mathbb{R}^{N \times 3} \rightarrow \mathbb{R}^{H \times W}$:

  • 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.

🚀 Key Features

  • 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.scan for constant-memory simulation of long trajectories.
  • Flexible Tip Geometries: Supports spherical tip-shape dilation.
  • Integration: Reads PDB/mmCIF files via biotite and integrates with synth-pdb and synth-dynamics.

📦 Installation

pip install synth-afm

📖 Tutorials

Get started immediately with our interactive Jupyter notebooks:

🛠 Quick Start

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)

🧪 Scientific Validation

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).

🔗 Related Projects

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

📜 License

Distributed under the MIT License. See LICENSE for more information.

✍️ Citation

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}
}

About

Differentiable High-Speed AFM (HS-AFM) simulation for protein structures. JAX-powered toolkit for generating synthetic movies with realistic tip-dilation and scanning lag.

Topics

Resources

Stars

1 star

Watchers

0 watching

Forks

Packages

 
 
 

Contributors