ColabBTR

import torch
from scipy.spatial.transform import Rotation
from colabbtr.morphology import (
    load_pdb_ca, define_tip, afmize_supersampled, ierosion, differentiable_btr,
)

# 1. Load protein structure (CA atoms) from PDB file
xyz, radii = load_pdb_ca("data/3A5I.pdb")  # xyz: (N, 3) in nm, radii: (N,) in nm
xyz = xyz - xyz.mean(dim=0, keepdim=True)   # center the molecule

# 2. Define a conical probe tip (15x15 pixels, 1.0 nm/pixel, half-cone angle 0.3 rad)
tip = define_tip(torch.zeros(15, 15), resolution_x=1.0, resolution_y=1.0,
                 probeRadius=2.0, probeAngle=0.3)

# 3. Compute molecular surfaces and simulate AFM images (20 frames)
#
# afmize_supersampled builds the surface and the tip on a 5x finer grid, dilates
# there, and samples the coarse pixel centers last. Doing it the other way round
# -- surfing() then idilation() on the coarse grid -- point-samples the surface
# before the tip ever sees it, and every CA bead here (radius 0.225-0.340 nm) is
# smaller than half the 1 nm pitch, so about a third of the molecular pixels are
# dropped outright and the peaks come out too low.
config = {
    "min_x": -15.0, "max_x": 15.0,
    "min_y": -15.0, "max_y": 15.0,
    "resolution_x": 1.0, "resolution_y": 1.0,
}
nframe = 20
surfaces, images = [], []
for _ in range(nframe):
    R = torch.tensor(Rotation.random().as_matrix(), dtype=torch.float32)
    xyz_rot = xyz @ R.T
    image, surface = afmize_supersampled(   # both (H, W) on the coarse grid
        xyz_rot, radii, config,
        probe_radius=2.0, probe_angle=0.3, tip_size=15, factor=5,
    )
    surfaces.append(surface)
    images.append(image)
surfaces = torch.stack(surfaces)  # (nframe, H, W)
images = torch.stack(images)      # (nframe, H, W)

# 4. Reconstruct tip shape from AFM images via BTR (200 epochs)
tip_est, loss = differentiable_btr(
    images, tip_size=(15, 15), nepoch=200, lr=0.1, weight_decay=0.001,
)

# 5. Deconvolve: remove tip artifacts to recover the sample surfaces
surfaces_recovered = torch.stack([ierosion(images[i], tip_est) for i in range(nframe)])

# 6. Compare tip cross-sections (ground truth vs reconstructed)
import matplotlib.pyplot as plt

center = tip.shape[0] // 2
pixels = range(tip.shape[1])
plt.plot(pixels, tip[center, :].numpy(), label="Ground truth")
plt.plot(pixels, tip_est[center, :].numpy(), label="BTR estimate")
plt.xlabel("Pixel")
plt.ylabel("Height (nm)")
plt.legend()
plt.title("Tip cross-section (center row)")
plt.savefig("tip_cross_section.png", dpi=150, bbox_inches="tight")
plt.show()

import libasd
import numpy as np

# Scanning range (nm) = pixel count × resolution (nm/pixel)
x_range = images.shape[2]  # 30 pixels × 1.0 nm/pixel = 30 nm
y_range = images.shape[1]  # 30 pixels × 1.0 nm/pixel = 30 nm

libasd.write_asd("images.asd", images.numpy(), x_scanning_range=x_range, y_scanning_range=y_range)
libasd.write_asd("surfaces.asd", surfaces.numpy(), x_scanning_range=x_range, y_scanning_range=y_range)
libasd.write_asd("surfaces_recovered.asd", surfaces_recovered.numpy(), x_scanning_range=x_range, y_scanning_range=y_range)

# Tip shape (single frame)
tip_range = tip.shape[1]  # 15 pixels × 1.0 nm/pixel = 15 nm
libasd.write_asd("tip.asd", tip.numpy(), x_scanning_range=tip_range, y_scanning_range=tip_range)
libasd.write_asd("tip_est.asd", tip_est.numpy(), x_scanning_range=tip_range, y_scanning_range=tip_range)