parallax aberration surface

%23%20Parallax%20aberration%20surface%20%E2%80%94%20port%20of%20the%20%22Parallax%20Aberration%20Surface%22%0A%23%20notebook%20by%20Georgios%20Varnavides%20(2025-02-01).%0A%23%0A%23%20Clean-room%20reimplementation%3A%20the%20original%20notebook%20drives%20py4DSTEM's%0A%23%20%60Parallax.guess_common_aberrations%60%20(GPL).%20This%20app%20instead%20builds%20the%20same%0A%23%20demonstration%20from%20first%20principles%20%E2%80%94%20the%20aberration%20ray-shift%20is%20derived%20from%0A%23%20the%20standard%20wave-aberration%20function%2C%20and%20the%20super-resolution%20reconstruction%0A%23%20uses%20a%20drizzle%20(bilinear-splat%20%2B%20hit-map%20normalization)%20adapted%20from%20ECLIPSE's%0A%23%20own%20scatterem%20(Apache-2.0)%2C%20not%20py4DSTEM.%20The%20tilt-corrected%20bright-field%20stack%0A%23%20is%20simulated%20in-app%20from%20the%20apoferritin%20projected%20potential%20in%20public%2F%2C%20so%20the%0A%23%20app%20is%20self-contained%20and%20runs%20in%20WASM.%0A%23%0A%23%20Pedagogy%3A%20each%20virtual%20bright-field%20image%20(one%20per%20bright-field%20detector%20pixel%20%2F%0A%23%20beam%20tilt)%20is%20displaced%20by%20the%20gradient%20of%20the%20aberration%20surface%20at%20that%20tilt.%0A%23%20Dial%20the%20aberration%20coefficients%20and%20scan%2Fdetector%20rotation%20to%20match%20the%20true%0A%23%20optics%3B%20when%20the%20predicted%20shifts%20match%2C%20drizzling%20the%20stack%20re-aligns%20every%0A%23%20image%20and%20the%20reconstruction%20sharpens%20back%20to%20the%20specimen.%0A%0Aimport%20marimo%0A%0A__generated_with%20%3D%20%220.23.13%22%0Aapp%20%3D%20marimo.App(width%3D%22medium%22)%0A%0A%0A%40app.cell%0Adef%20_(mo)%3A%0A%20%20%20%20%23%20controls%20%E2%80%94%20microscope%20geometry%20%2B%20plotting%0A%20%20%20%20rotation%20%3D%20mo.ui.slider(%0A%20%20%20%20%20%20%20%20start%3D-90%2C%20stop%3D90%2C%20step%3D1%2C%20value%3D-15%2C%0A%20%20%20%20%20%20%20%20label%3D%22rotation%20angle%20%5B%C2%B0%5D%22%2C%20show_value%3DTrue%2C%0A%20%20%20%20)%0A%20%20%20%20arrow_freq%20%3D%20mo.ui.slider(%0A%20%20%20%20%20%20%20%20start%3D1%2C%20stop%3D16%2C%20step%3D1%2C%20value%3D4%2C%0A%20%20%20%20%20%20%20%20label%3D%22plot%20arrow%20frequency%22%2C%20show_value%3DTrue%2C%0A%20%20%20%20)%0A%20%20%20%20upsampling%20%3D%20mo.ui.slider(%0A%20%20%20%20%20%20%20%20start%3D1%2C%20stop%3D4%2C%20step%3D1%2C%20value%3D2%2C%0A%20%20%20%20%20%20%20%20label%3D%22upsampling%20factor%22%2C%20show_value%3DTrue%2C%0A%20%20%20%20)%0A%20%20%20%20return%20arrow_freq%2C%20rotation%2C%20upsampling%0A%0A%0A%40app.cell%0Adef%20_(mo)%3A%0A%20%20%20%20%23%20controls%20%E2%80%94%20aberrations%20(defocus%20defaults%20to%20the%20true%201.5e4%20%C3%85%3B%20the%20rest%20to%200)%0A%20%20%20%20defocus%20%3D%20mo.ui.slider(%0A%20%20%20%20%20%20%20%20start%3D-2e4%2C%20stop%3D2e4%2C%20step%3D5e2%2C%20value%3D1.5e4%2C%0A%20%20%20%20%20%20%20%20label%3D%22defocus%20%5B%C3%85%5D%22%2C%20show_value%3DTrue%2C%0A%20%20%20%20)%0A%20%20%20%20spherical%20%3D%20mo.ui.slider(%0A%20%20%20%20%20%20%20%20start%3D-2e7%2C%20stop%3D2e7%2C%20step%3D5e5%2C%20value%3D0%2C%0A%20%20%20%20%20%20%20%20label%3D%22spherical%20aberration%20Cs%20%5B%C3%85%5D%22%2C%20show_value%3DTrue%2C%0A%20%20%20%20)%0A%20%20%20%20astigmatism%20%3D%20mo.ui.slider(%0A%20%20%20%20%20%20%20%20start%3D0%2C%20stop%3D1e4%2C%20step%3D5e2%2C%20value%3D0%2C%0A%20%20%20%20%20%20%20%20label%3D%22astigmatism%20%5B%C3%85%5D%22%2C%20show_value%3DTrue%2C%0A%20%20%20%20)%0A%20%20%20%20astigmatism_angle%20%3D%20mo.ui.slider(%0A%20%20%20%20%20%20%20%20start%3D-90%2C%20stop%3D90%2C%20step%3D1%2C%20value%3D45%2C%0A%20%20%20%20%20%20%20%20label%3D%22astigmatism%20angle%20%5B%C2%B0%5D%22%2C%20show_value%3DTrue%2C%0A%20%20%20%20)%0A%20%20%20%20coma%20%3D%20mo.ui.slider(%0A%20%20%20%20%20%20%20%20start%3D0%2C%20stop%3D1e7%2C%20step%3D5e3%2C%20value%3D0%2C%0A%20%20%20%20%20%20%20%20label%3D%22coma%20%5B%C3%85%5D%22%2C%20show_value%3DTrue%2C%0A%20%20%20%20)%0A%20%20%20%20coma_angle%20%3D%20mo.ui.slider(%0A%20%20%20%20%20%20%20%20start%3D-180%2C%20stop%3D180%2C%20step%3D1%2C%20value%3D90%2C%0A%20%20%20%20%20%20%20%20label%3D%22coma%20angle%20%5B%C2%B0%5D%22%2C%20show_value%3DTrue%2C%0A%20%20%20%20)%0A%20%20%20%20return%20astigmatism%2C%20astigmatism_angle%2C%20coma%2C%20coma_angle%2C%20defocus%2C%20spherical%0A%0A%0A%40app.cell%0Adef%20_(%0A%20%20%20%20arrow_freq%2C%0A%20%20%20%20astigmatism%2C%0A%20%20%20%20astigmatism_angle%2C%0A%20%20%20%20bf_shifts%2C%0A%20%20%20%20coma%2C%0A%20%20%20%20coma_angle%2C%0A%20%20%20%20defocus%2C%0A%20%20%20%20detector_mask%2C%0A%20%20%20%20drizzle%2C%0A%20%20%20%20fov%2C%0A%20%20%20%20kxy%2C%0A%20%20%20%20mo%2C%0A%20%20%20%20np%2C%0A%20%20%20%20plt%2C%0A%20%20%20%20rotation%2C%0A%20%20%20%20scan_sampling%2C%0A%20%20%20%20spherical%2C%0A%20%20%20%20stack%2C%0A%20%20%20%20upsampling%2C%0A)%3A%0A%20%20%20%20_ab%20%3D%20dict(%0A%20%20%20%20%20%20%20%20defocus%3Ddefocus.value%2C%0A%20%20%20%20%20%20%20%20Cs%3Dspherical.value%2C%0A%20%20%20%20%20%20%20%20astig%3Dastigmatism.value%2C%0A%20%20%20%20%20%20%20%20astig_angle%3Dnp.deg2rad(astigmatism_angle.value)%2C%0A%20%20%20%20%20%20%20%20coma%3Dcoma.value%2C%0A%20%20%20%20%20%20%20%20coma_angle%3Dnp.deg2rad(coma_angle.value)%2C%0A%20%20%20%20%20%20%20%20rot_deg%3Drotation.value%2C%0A%20%20%20%20)%0A%0A%20%20%20%20%23%20predicted%20per-tilt%20shifts%20%5Bscan%20px%5D%3B%20drizzle%20by%20their%20negative%20to%20re-align%0A%20%20%20%20shifts_px%20%3D%20bf_shifts(scan_sampling%2C%20**_ab)%0A%20%20%20%20recon%20%3D%20drizzle(stack%2C%20-shifts_px%2C%20upsampling.value)%0A%0A%20%20%20%20%23%20---%20figure%20(dark%20panels%20%2B%20white%20annotations%2C%20over%20the%20white%20app%20canvas)%20---%0A%20%20%20%20_bg%20%3D%20%22%230b0f19%22%0A%20%20%20%20plt.rcParams.update(%0A%20%20%20%20%20%20%20%20%7B%0A%20%20%20%20%20%20%20%20%20%20%20%20%22text.color%22%3A%20%22white%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%22axes.labelcolor%22%3A%20%22white%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%22axes.edgecolor%22%3A%20%22white%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%22xtick.color%22%3A%20%22white%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%22ytick.color%22%3A%20%22white%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%22xtick.labelcolor%22%3A%20%22white%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%22ytick.labelcolor%22%3A%20%22white%22%2C%0A%20%20%20%20%20%20%20%20%7D%0A%20%20%20%20)%0A%20%20%20%20fig%2C%20(ax_q%2C%20ax_im)%20%3D%20plt.subplots(1%2C%202%2C%20figsize%3D(9%2C%204.5)%2C%20dpi%3D100%2C%20layout%3D%22constrained%22)%0A%20%20%20%20fig.patch.set_facecolor(_bg)%0A%0A%20%20%20%20%23%20left%3A%20predicted%20bright-field%20shifts%20as%20a%20quiver%20over%20the%20detector%20disk%0A%20%20%20%20_ny%2C%20_nx%20%3D%20detector_mask.shape%0A%20%20%20%20_yy%2C%20_xx%20%3D%20np.meshgrid(np.arange(_ny)%2C%20np.arange(_nx)%2C%20indexing%3D%22ij%22)%0A%20%20%20%20_freq%20%3D%20np.logical_and(_xx%20%25%20arrow_freq.value%20%3D%3D%200%2C%20_yy%20%25%20arrow_freq.value%20%3D%3D%200)%0A%20%20%20%20_plot%20%3D%20_freq%5Bdetector_mask%5D%0A%20%20%20%20_kr_max%20%3D%201.2%20*%20np.abs(kxy).max()%0A%20%20%20%20ax_q.set_facecolor(_bg)%0A%20%20%20%20ax_q.quiver(%0A%20%20%20%20%20%20%20%20kxy%5B_plot%2C%201%5D%2C%20kxy%5B_plot%2C%200%5D%2C%0A%20%20%20%20%20%20%20%20shifts_px%5B_plot%2C%201%5D%2C%20shifts_px%5B_plot%2C%200%5D%2C%0A%20%20%20%20%20%20%20%20color%3D%22%23ff5555%22%2C%20angles%3D%22xy%22%2C%20scale_units%3D%22xy%22%2C%0A%20%20%20%20%20%20%20%20scale%3D1%20%2F%20(0.02%20*%20_kr_max%20%2B%201e-9)%2C%20width%3D0.006%2C%0A%20%20%20%20)%0A%20%20%20%20ax_q.set_xlim(-_kr_max%2C%20_kr_max)%0A%20%20%20%20ax_q.set_ylim(-_kr_max%2C%20_kr_max)%0A%20%20%20%20ax_q.set_aspect(%22equal%22)%0A%20%20%20%20ax_q.set_title(%22predicted%20BF%20shifts%22%2C%20fontsize%3D12)%0A%20%20%20%20ax_q.set_xlabel(r%22%24k_y%24%20%5B%C3%85%24%5E%7B-1%7D%24%5D%22%2C%20fontsize%3D9)%0A%20%20%20%20ax_q.set_ylabel(r%22%24k_x%24%20%5B%C3%85%24%5E%7B-1%7D%24%5D%22%2C%20fontsize%3D9)%0A%0A%20%20%20%20%23%20right%3A%20predicted%20aligned%20(drizzled)%20bright-field%20image%0A%20%20%20%20ax_im.imshow(recon%2C%20cmap%3D%22magma%22%2C%20extent%3D%5B0%2C%20fov%2C%20fov%2C%200%5D)%0A%20%20%20%20ax_im.set_title(%22predicted%20aligned%20BF%20image%22%2C%20fontsize%3D12)%0A%20%20%20%20ax_im.set_xlabel(%22y%20%5B%C3%85%5D%22%2C%20fontsize%3D9)%0A%20%20%20%20ax_im.set_ylabel(%22x%20%5B%C3%85%5D%22%2C%20fontsize%3D9)%0A%20%20%20%20ax_im.set_facecolor(_bg)%0A%0A%20%20%20%20%23%20---%20layout%20---%0A%20%20%20%20_geometry%20%3D%20mo.vstack(%0A%20%20%20%20%20%20%20%20%5Bmo.md(%22**Microscope%20geometry%20%26%20plotting**%22)%2C%20rotation%2C%20arrow_freq%2C%20upsampling%5D%2C%0A%20%20%20%20%20%20%20%20gap%3D0.4%2C%0A%20%20%20%20)%0A%20%20%20%20_aberrations%20%3D%20mo.vstack(%0A%20%20%20%20%20%20%20%20%5B%0A%20%20%20%20%20%20%20%20%20%20%20%20mo.md(%22**Aberrations**%22)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20defocus%2C%20spherical%2C%20astigmatism%2C%20astigmatism_angle%2C%20coma%2C%20coma_angle%2C%0A%20%20%20%20%20%20%20%20%5D%2C%0A%20%20%20%20%20%20%20%20gap%3D0.4%2C%0A%20%20%20%20)%0A%20%20%20%20_controls%20%3D%20mo.hstack(%5B_geometry%2C%20_aberrations%5D%2C%20justify%3D%22center%22%2C%20gap%3D2%2C%20wrap%3DTrue)%0A%20%20%20%20mo.vstack(%5B_controls%2C%20fig%5D%2C%20align%3D%22center%22%2C%20gap%3D0.75)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(np)%3A%0A%20%20%20%20%23%20Aberration%20ray-shift%20%2B%20drizzle%20(self-contained%20for%20WASM).%0A%20%20%20%20def%20energy2wavelength(energy)%3A%0A%20%20%20%20%20%20%20%20%22%22%22Relativistic%20electron%20wavelength%20in%20%C3%85.%22%22%22%0A%20%20%20%20%20%20%20%20hplanck%20%3D%206.62607e-34%0A%20%20%20%20%20%20%20%20c%20%3D%20299792458.0%0A%20%20%20%20%20%20%20%20me%20%3D%209.1093856e-31%0A%20%20%20%20%20%20%20%20e%20%3D%201.6021766208e-19%0A%20%20%20%20%20%20%20%20return%20(%0A%20%20%20%20%20%20%20%20%20%20%20%20hplanck%20*%20c%20%2F%20np.sqrt(energy%20*%20(2%20*%20me%20*%20c**2%20%2F%20e%20%2B%20energy))%20%2F%20e%20*%201.0e10%0A%20%20%20%20%20%20%20%20)%0A%0A%20%20%20%20def%20aberration_shift_field(q_grid%2C%20wavelength%2C%20recip_sampling%2C%20rot_deg%3D0.0%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20defocus%3D0.0%2C%20Cs%3D0.0%2C%20astig%3D0.0%2C%20astig_angle%3D0.0%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20coma%3D0.0%2C%20coma_angle%3D0.0)%3A%0A%20%20%20%20%20%20%20%20%22%22%22Real-space%20ray%20displacement%20(%C3%85)%20at%20each%20detector%20spatial%20frequency.%0A%0A%20%20%20%20%20%20%20%20The%20wave-aberration%20surface%20(radians)%20for%20the%20common%20low-order%20aberrations%0A%20%20%20%20%20%20%20%20is%20%CF%87(%CE%B1%2C%CF%86)%20%3D%202%CF%80%2F%CE%BB%20%C2%B7%20%5B%20%C2%BDC1%20%CE%B1%C2%B2%20%2B%20%C2%BDA1%20%CE%B1%C2%B2%20cos2(%CF%86%E2%88%92%CF%86_A1)%20%2B%20%E2%85%93B2%20%CE%B1%C2%B3%20cos(%CF%86%E2%88%92%CF%86_B2)%0A%20%20%20%20%20%20%20%20%2B%20%C2%BCC3%20%CE%B1%E2%81%B4%20%5D%2C%20with%20%CE%B1%20%3D%20%CE%BB%7Cq%7C%20the%20scattering%20angle%20and%20defocus%20C1%20%3D%20%E2%88%92defocus.%0A%20%20%20%20%20%20%20%20A%20ray%20at%20spatial%20frequency%20q%20lands%20displaced%20by%20%E2%88%87_q%20%CF%87%20%2F%202%CF%80%20%E2%80%94%20the%20parallax%0A%20%20%20%20%20%20%20%20shift%20this%20app%20visualizes.%20Evaluated%20numerically%20so%20the%20expression%20stays%0A%20%20%20%20%20%20%20%20transparent%20and%20any%20surface%20could%20be%20dropped%20in.%0A%20%20%20%20%20%20%20%20%22%22%22%0A%20%20%20%20%20%20%20%20QX%2C%20QY%20%3D%20q_grid%0A%20%20%20%20%20%20%20%20r%20%3D%20np.deg2rad(rot_deg)%0A%20%20%20%20%20%20%20%20%23%20passive%20rotation%20of%20the%20scan%2Fdetector%20coordinate%20systems%0A%20%20%20%20%20%20%20%20qx%20%3D%20QX%20*%20np.cos(-r)%20%2B%20QY%20*%20np.sin(-r)%0A%20%20%20%20%20%20%20%20qy%20%3D%20-QX%20*%20np.sin(-r)%20%2B%20QY%20*%20np.cos(-r)%0A%20%20%20%20%20%20%20%20alpha%20%3D%20np.hypot(qx%2C%20qy)%20*%20wavelength%0A%20%20%20%20%20%20%20%20phi%20%3D%20np.arctan2(qy%2C%20qx)%0A%20%20%20%20%20%20%20%20chi%20%3D%20(2%20*%20np.pi%20%2F%20wavelength)%20*%20(%0A%20%20%20%20%20%20%20%20%20%20%20%200.5%20*%20(-defocus)%20*%20alpha**2%0A%20%20%20%20%20%20%20%20%20%20%20%20%2B%200.5%20*%20astig%20*%20alpha**2%20*%20np.cos(2%20*%20(phi%20-%20astig_angle))%0A%20%20%20%20%20%20%20%20%20%20%20%20%2B%20(1%20%2F%203)%20*%20coma%20*%20alpha**3%20*%20np.cos(phi%20-%20coma_angle)%0A%20%20%20%20%20%20%20%20%20%20%20%20%2B%200.25%20*%20Cs%20*%20alpha**4%0A%20%20%20%20%20%20%20%20)%0A%20%20%20%20%20%20%20%20grad_x%2C%20grad_y%20%3D%20np.gradient(chi%2C%20recip_sampling%2C%20recip_sampling)%0A%20%20%20%20%20%20%20%20return%20grad_x%20%2F%20(2%20*%20np.pi)%2C%20grad_y%20%2F%20(2%20*%20np.pi)%20%20%23%20%C3%85%0A%0A%20%20%20%20def%20drizzle(images%2C%20shifts%2C%20upsample%2C%20kde_sigma%3D0.5%2C%20eps%3D1e-12)%3A%0A%20%20%20%20%20%20%20%20%22%22%22Drizzle%20(bilinear-splat%20%2B%20hit-map%20normalization)%20super-resolution.%0A%0A%20%20%20%20%20%20%20%20Adapted%20from%20ECLIPSE's%20scatterem%20(Apache-2.0).%20Each%20low-resolution%20frame's%0A%20%20%20%20%20%20%20%20pixels%20are%20splatted%20onto%20an%20%60upsample%60%C3%97%20finer%20grid%20at%20sub-pixel%20positions%0A%20%20%20%20%20%20%20%20%60(pixel%20%2B%20shift)%20%C2%B7%20upsample%60%3B%20accumulated%20intensity%20is%20normalized%20by%20the%0A%20%20%20%20%20%20%20%20accumulated%20bilinear%20weight%20(the%20hit%20map)%2C%20then%20lightly%20Gaussian-smoothed%0A%20%20%20%20%20%20%20%20(Nadaraya%E2%80%93Watson).%20Shifts%20are%20(dx%2C%20dy)%20in%20low-resolution%20pixels.%0A%20%20%20%20%20%20%20%20%22%22%22%0A%20%20%20%20%20%20%20%20def%20gaussian_blur(array%2C%20sigma)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%22%22%22Periodic%20Gaussian%20blur%20via%20FFT%20(numpy-only%2C%20matches%20the%20drizzle's%20wrap).%22%22%22%0A%20%20%20%20%20%20%20%20%20%20%20%20ny%2C%20nx%20%3D%20array.shape%0A%20%20%20%20%20%20%20%20%20%20%20%20gy%20%3D%20np.fft.fftfreq(ny)%5B%3A%2C%20None%5D%0A%20%20%20%20%20%20%20%20%20%20%20%20gx%20%3D%20np.fft.fftfreq(nx)%5BNone%2C%20%3A%5D%0A%20%20%20%20%20%20%20%20%20%20%20%20kernel%20%3D%20np.exp(-2%20*%20(np.pi%20*%20sigma)%20**%202%20*%20(gy**2%20%2B%20gx**2))%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20np.real(np.fft.ifft2(np.fft.fft2(array)%20*%20kernel))%0A%0A%20%20%20%20%20%20%20%20n%2C%20h%2C%20w%20%3D%20images.shape%0A%20%20%20%20%20%20%20%20hh%2C%20ww%20%3D%20h%20*%20upsample%2C%20w%20*%20upsample%0A%20%20%20%20%20%20%20%20ys%2C%20xs%20%3D%20np.meshgrid(%0A%20%20%20%20%20%20%20%20%20%20%20%20np.arange(h%2C%20dtype%3Dnp.float64)%2C%20np.arange(w%2C%20dtype%3Dnp.float64)%2C%20indexing%3D%22ij%22%0A%20%20%20%20%20%20%20%20)%0A%20%20%20%20%20%20%20%20ya%20%3D%20(ys%5BNone%5D%20%2B%20shifts%5B%3A%2C%200%2C%20None%2C%20None%5D)%20*%20upsample%0A%20%20%20%20%20%20%20%20xa%20%3D%20(xs%5BNone%5D%20%2B%20shifts%5B%3A%2C%201%2C%20None%2C%20None%5D)%20*%20upsample%0A%20%20%20%20%20%20%20%20y0%20%3D%20np.floor(ya).astype(int)%0A%20%20%20%20%20%20%20%20x0%20%3D%20np.floor(xa).astype(int)%0A%20%20%20%20%20%20%20%20fy%20%3D%20ya%20-%20y0%0A%20%20%20%20%20%20%20%20fx%20%3D%20xa%20-%20x0%0A%0A%20%20%20%20%20%20%20%20accum%20%3D%20np.zeros(hh%20*%20ww)%0A%20%20%20%20%20%20%20%20hits%20%3D%20np.zeros(hh%20*%20ww)%0A%20%20%20%20%20%20%20%20vals%20%3D%20images.reshape(n%2C%20h%20*%20w)%0A%20%20%20%20%20%20%20%20for%20oy%2C%20ox%20in%20((0%2C%200)%2C%20(1%2C%200)%2C%20(0%2C%201)%2C%20(1%2C%201))%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20weight%20%3D%20((1%20-%20fy)%20if%20oy%20%3D%3D%200%20else%20fy)%20*%20((1%20-%20fx)%20if%20ox%20%3D%3D%200%20else%20fx)%0A%20%20%20%20%20%20%20%20%20%20%20%20yc%20%3D%20(y0%20%2B%20oy)%20%25%20hh%0A%20%20%20%20%20%20%20%20%20%20%20%20xc%20%3D%20(x0%20%2B%20ox)%20%25%20ww%0A%20%20%20%20%20%20%20%20%20%20%20%20idx%20%3D%20(yc%20*%20ww%20%2B%20xc).reshape(n%2C%20h%20*%20w).ravel()%0A%20%20%20%20%20%20%20%20%20%20%20%20wv%20%3D%20weight.reshape(n%2C%20h%20*%20w).ravel()%0A%20%20%20%20%20%20%20%20%20%20%20%20accum%20%2B%3D%20np.bincount(idx%2C%20weights%3Dwv%20*%20vals.ravel()%2C%20minlength%3Dhh%20*%20ww)%0A%20%20%20%20%20%20%20%20%20%20%20%20hits%20%2B%3D%20np.bincount(idx%2C%20weights%3Dwv%2C%20minlength%3Dhh%20*%20ww)%0A%0A%20%20%20%20%20%20%20%20accum%20%3D%20accum.reshape(hh%2C%20ww)%0A%20%20%20%20%20%20%20%20hits%20%3D%20hits.reshape(hh%2C%20ww)%0A%20%20%20%20%20%20%20%20if%20kde_sigma%20%3E%200%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20sigma%20%3D%20kde_sigma%20*%20upsample%0A%20%20%20%20%20%20%20%20%20%20%20%20accum%20%3D%20gaussian_blur(accum%2C%20sigma)%0A%20%20%20%20%20%20%20%20%20%20%20%20hits%20%3D%20gaussian_blur(hits%2C%20sigma)%0A%20%20%20%20%20%20%20%20return%20accum%20%2F%20np.clip(hits%2C%20eps%2C%20None)%0A%0A%20%20%20%20return%20aberration_shift_field%2C%20drizzle%2C%20energy2wavelength%0A%0A%0A%40app.cell%0Adef%20_(aberration_shift_field%2C%20energy2wavelength%2C%20mo%2C%20np)%3A%0A%20%20%20%20%23%20---%20simulated%20tilt-corrected%20bright-field%20stack%20(apoferritin)%20---%0A%20%20%20%20%23%20Each%20bright-field%20detector%20pixel%20(a%20beam%20tilt)%20forms%20one%20virtual%20BF%20image%2C%0A%20%20%20%20%23%20here%20the%20specimen%20decimated%20to%20the%20scan%20grid%20with%20a%20sub-pixel%20offset%20equal%20to%0A%20%20%20%20%23%20the%20true%20parallax%20shift%20(defocus%201.5e4%20%C3%85%2C%20scan%2Fdetector%20rotation%20%E2%88%9215%C2%B0).%20The%0A%20%20%20%20%23%20sub-pixel%20diversity%20across%20tilts%20is%20what%20the%20drizzle%20later%20exploits.%0A%20%20%20%20energy%20%3D%20300e3%20%20%23%20eV%0A%20%20%20%20wavelength%20%3D%20energy2wavelength(energy)%20%20%23%20%C3%85%0A%0A%20%20%20%20detector_gpts%20%3D%2064%0A%20%20%20%20bf_radius_px%20%3D%2010%20%20%23%20bright-field%20disk%20radius%20on%20the%20detector%0A%20%20%20%20recip_sampling%20%3D%200.006%20%20%23%20%C3%85%E2%81%BB%C2%B9%20per%20detector%20pixel%0A%20%20%20%20downsample%20%3D%204%20%20%23%20object%20px%20per%20scan%20px%0A%20%20%20%20dose_scale%20%3D%20100.0%20%20%23%20electrons%20per%20scan%20pixel%20(SNR%20knob)%0A%0A%20%20%20%20true_defocus%20%3D%201.5e4%20%20%23%20%C3%85%0A%20%20%20%20true_rotation_deg%20%3D%20-15.0%0A%0A%20%20%20%20%23%20detector%20bright-field%20disk%20-%3E%20one%20virtual%20BF%20image%20per%20pixel%0A%20%20%20%20_c%20%3D%20detector_gpts%20%2F%2F%202%0A%20%20%20%20_yy%2C%20_xx%20%3D%20np.mgrid%5B0%3Adetector_gpts%2C%200%3Adetector_gpts%5D%0A%20%20%20%20detector_mask%20%3D%20np.hypot(_xx%20-%20_c%2C%20_yy%20-%20_c)%20%3C%3D%20bf_radius_px%0A%20%20%20%20_inds%20%3D%20np.argwhere(detector_mask)%0A%20%20%20%20kxy%20%3D%20(_inds%20-%20_inds.mean(0))%20*%20recip_sampling%20%20%23%20%C3%85%E2%81%BB%C2%B9%20coordinate%20of%20each%20tilt%0A%20%20%20%20num_bf%20%3D%20len(_inds)%0A%0A%20%20%20%20_q1%20%3D%20(np.arange(detector_gpts)%20-%20_c)%20*%20recip_sampling%0A%20%20%20%20q_grid%20%3D%20np.meshgrid(_q1%2C%20_q1%2C%20indexing%3D%22ij%22)%0A%0A%20%20%20%20def%20bf_shifts(scan_sampling%2C%20**aberrations)%3A%0A%20%20%20%20%20%20%20%20%22%22%22Per-tilt%20parallax%20shift%20%5Bscan%20px%5D%20for%20the%20bright-field%20disk%20pixels.%22%22%22%0A%20%20%20%20%20%20%20%20grad_x%2C%20grad_y%20%3D%20aberration_shift_field(%0A%20%20%20%20%20%20%20%20%20%20%20%20q_grid%2C%20wavelength%2C%20recip_sampling%2C%20**aberrations%0A%20%20%20%20%20%20%20%20)%0A%20%20%20%20%20%20%20%20return%20np.stack(%5Bgrad_x%5Bdetector_mask%5D%2C%20grad_y%5Bdetector_mask%5D%5D%2C%201)%20%2F%20scan_sampling%0A%0A%20%20%20%20%23%20object%3A%20apoferritin%20projected%20potential%20(normalized%20to%20%5B0%2C%201%5D)%0A%20%20%20%20def%20_load(fname)%3A%0A%20%20%20%20%20%20%20%20import%20os%0A%0A%20%20%20%20%20%20%20%20loc%20%3D%20str(mo.notebook_location())%0A%20%20%20%20%20%20%20%20if%20loc.startswith(%22http%3A%2F%2F%22)%20or%20loc.startswith(%22https%3A%2F%2F%22)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20import%20io%0A%20%20%20%20%20%20%20%20%20%20%20%20import%20urllib.request%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20base%20%3D%20loc.rstrip(%22%2F%22)%0A%20%20%20%20%20%20%20%20%20%20%20%20candidates%20%3D%20%5B%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20base%20%2B%20%22%2Fpublic%2F%22%20%2B%20fname%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20base.rsplit(%22%2F%22%2C%201)%5B0%5D%20%2B%20%22%2Fpublic%2F%22%20%2B%20fname%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%5D%0A%20%20%20%20%20%20%20%20%20%20%20%20last_err%20%3D%20None%0A%20%20%20%20%20%20%20%20%20%20%20%20for%20url%20in%20candidates%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20try%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20with%20urllib.request.urlopen(url)%20as%20resp%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20return%20np.load(io.BytesIO(resp.read()))%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20except%20Exception%20as%20exc%3A%20%20%23%20noqa%3A%20BLE001%20-%20try%20the%20next%20candidate%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20last_err%20%3D%20exc%0A%20%20%20%20%20%20%20%20%20%20%20%20raise%20last_err%0A%20%20%20%20%20%20%20%20for%20root%20in%20(loc%2C%20os.path.dirname(loc)%2C%20os.getcwd())%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20path%20%3D%20os.path.join(root%2C%20%22public%22%2C%20fname)%0A%20%20%20%20%20%20%20%20%20%20%20%20if%20os.path.exists(path)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20return%20np.load(path)%0A%20%20%20%20%20%20%20%20raise%20FileNotFoundError(f%22could%20not%20locate%20public%2F%7Bfname%7D%22)%0A%0A%20%20%20%20_obj%20%3D%20_load(%22apoF_projected-potential_192x192_4qprobe.npy%22).astype(np.float64)%0A%20%20%20%20_obj%20%3D%20(_obj%20-%20_obj.min())%20%2F%20np.ptp(_obj)%0A%20%20%20%20object_gpts%20%3D%20_obj.shape%5B0%5D%0A%20%20%20%20scan_gpts%20%3D%20object_gpts%20%2F%2F%20downsample%0A%20%20%20%20object_sampling%20%3D%201.0%20%20%23%20%C3%85%20per%20object%20px%0A%20%20%20%20scan_sampling%20%3D%20downsample%20*%20object_sampling%20%20%23%20%C3%85%20per%20scan%20px%0A%20%20%20%20fov%20%3D%20object_gpts%20*%20object_sampling%20%20%23%20%C3%85%0A%0A%20%20%20%20%23%20Each%20virtual%20BF%20image%20is%20the%20specimen%20decimated%20to%20the%20scan%20grid%20with%20a%0A%20%20%20%20%23%20sub-pixel%20offset%20equal%20to%20the%20true%20parallax%20shift%20%E2%80%94%20a%20periodic%20bilinear%0A%20%20%20%20%23%20gather%20at%20object%20coordinates%20(scan_idx%20%E2%88%92%20shift)%C2%B7downsample.%20Vectorized%20over%0A%20%20%20%20%23%20all%20tilts%20(no%20per-frame%20FFT)%20so%20the%20one-time%20simulation%20stays%20fast%20in%20WASM.%0A%20%20%20%20_s_true%20%3D%20bf_shifts(scan_sampling%2C%20defocus%3Dtrue_defocus%2C%20rot_deg%3Dtrue_rotation_deg)%0A%20%20%20%20_ii%2C%20_jj%20%3D%20np.meshgrid(np.arange(scan_gpts)%2C%20np.arange(scan_gpts)%2C%20indexing%3D%22ij%22)%0A%20%20%20%20_row%20%3D%20(_ii%5BNone%5D%20-%20_s_true%5B%3A%2C%200%2C%20None%2C%20None%5D)%20*%20downsample%0A%20%20%20%20_col%20%3D%20(_jj%5BNone%5D%20-%20_s_true%5B%3A%2C%201%2C%20None%2C%20None%5D)%20*%20downsample%0A%20%20%20%20_r0%20%3D%20np.floor(_row).astype(int)%0A%20%20%20%20_c0%20%3D%20np.floor(_col).astype(int)%0A%20%20%20%20_fr%20%3D%20_row%20-%20_r0%0A%20%20%20%20_fc%20%3D%20_col%20-%20_c0%0A%20%20%20%20_r0%20%25%3D%20object_gpts%0A%20%20%20%20_c0%20%25%3D%20object_gpts%0A%20%20%20%20_r1%20%3D%20(_r0%20%2B%201)%20%25%20object_gpts%0A%20%20%20%20_c1%20%3D%20(_c0%20%2B%201)%20%25%20object_gpts%0A%20%20%20%20_clean%20%3D%20(%0A%20%20%20%20%20%20%20%20_obj%5B_r0%2C%20_c0%5D%20*%20(1%20-%20_fr)%20*%20(1%20-%20_fc)%0A%20%20%20%20%20%20%20%20%2B%20_obj%5B_r0%2C%20_c1%5D%20*%20(1%20-%20_fr)%20*%20_fc%0A%20%20%20%20%20%20%20%20%2B%20_obj%5B_r1%2C%20_c0%5D%20*%20_fr%20*%20(1%20-%20_fc)%0A%20%20%20%20%20%20%20%20%2B%20_obj%5B_r1%2C%20_c1%5D%20*%20_fr%20*%20_fc%0A%20%20%20%20)%0A%20%20%20%20_rng%20%3D%20np.random.default_rng(0)%0A%20%20%20%20stack%20%3D%20(%0A%20%20%20%20%20%20%20%20_rng.poisson(np.clip(_clean%2C%200%2C%20None)%20*%20dose_scale)%20%2F%20dose_scale%0A%20%20%20%20).astype(np.float32)%0A%20%20%20%20return%20bf_shifts%2C%20detector_mask%2C%20fov%2C%20kxy%2C%20scan_sampling%2C%20stack%0A%0A%0A%40app.cell%0Adef%20_()%3A%0A%20%20%20%20%23%20Imports%0A%20%20%20%20import%20marimo%20as%20mo%0A%20%20%20%20import%20numpy%20as%20np%0A%20%20%20%20import%20matplotlib.pyplot%20as%20plt%0A%0A%20%20%20%20return%20mo%2C%20np%2C%20plt%0A%0A%0Aif%20__name__%20%3D%3D%20%22__main__%22%3A%0A%20%20%20%20app.run()%0A