ptychographic coherence

import%20marimo%0A%0A__generated_with%20%3D%20%220.14.12%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%0A%20%20%20%20max_semiangle%20%3D%20mo.ui.dropdown(options%3D%7Br%221%20alpha%22%3A1%2C%20%222%20alpha%22%3A2%2C%20%223%20alpha%22%3A3%2C%20%224%20alpha%22%3A4%2C%20%225%20alpha%22%3A5%2C%20%226%20alpha%22%3A6%7D%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%20value%3D%221%20alpha%22%2C%20%23%20initial%20value%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%20label%3D%22Max.%20Semiangle%22)%0A%20%20%20%20energy%20%3D%20mo.ui.dropdown(options%3D%7B%22300%20keV%22%3A300e3%2C%20%22200%20keV%22%3A200e3%2C%20%2280%20keV%22%3A80e3%2C%20%2260%20keV%22%3A60e3%2C%20%2230%20keV%22%3A30e3%7D%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%20value%3D%22300%20keV%22%2C%20%23%20initial%20value%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%20label%3D%22Electron%20energy%22)%0A%20%20%20%20detector%20%3D%20mo.ui.dropdown(options%3D%7B%22256%22%3A256%2C%20%22192%22%3A192%2C%20%22128%22%3A128%2C%20%2296%22%3A96%2C%20%2248%22%3A48%2C%20%2216%22%3A16%7D%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%20value%3D%2248%22%2C%20%23%20initial%20value%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%20label%3D%22Detector%20shape%22)%0A%0A%20%20%20%20deltaE%20%3D%20mo.ui.slider(%0A%20%20%20%20%20%20%20%20start%3D0%2C%20stop%3D1.0%2C%20step%3D0.1%2C%20label%3D%22energy%20spread%20%5BeV%5D%22%2C%20show_value%3DTrue%2C%20value%3D0.4%0A%20%20%20%20)%0A%0A%20%20%20%20Cc%20%3D%20mo.ui.slider(%0A%20%20%20%20%20%20%20%20start%3D0%2C%20stop%3D2%2C%20step%3D0.05%2C%20label%3D%22Cc%20%5Bmm%5D%22%2C%20show_value%3DTrue%2C%20value%3D1.2%0A%20%20%20%20)%0A%0A%0A%20%20%20%20source_size%20%3D%20mo.ui.slider(%0A%20%20%20%20%20%20%20%20start%3D0%2C%20stop%3D1.2%2C%20step%3D0.05%2C%20label%3D%22source%20size%20%5B%C3%85%5D%22%2C%20show_value%3DTrue%2C%20value%3D0.0%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%3D200%2C%20step%3D5%2C%20label%3D%22astigmatism%20%5B%C3%85%5D%22%2C%20show_value%3DTrue%0A%20%20%20%20)%0A%0A%20%20%20%20astigmatism_angle_slider%20%3D%20mo.ui.slider(%0A%20%20%20%20%20%20%20%20start%3D-90%2C%20stop%3D90%2C%20step%3D1%2C%20label%3D%22astigmatism%20angle%20%5B%C2%B0%5D%22%2C%20show_value%3DTrue%0A%20%20%20%20)%0A%0A%20%20%20%20convergence_angle%20%3D%20mo.ui.slider(%0A%20%20%20%20%20%20%20%20start%3D5%2C%0A%20%20%20%20%20%20%20%20stop%3D80%2C%0A%20%20%20%20%20%20%20%20value%3D20%2C%0A%20%20%20%20%20%20%20%20step%3D0.5%2C%0A%20%20%20%20%20%20%20%20label%3D%22conv.%20semiangle%20%5Bmrad%5D%22%2C%0A%20%20%20%20%20%20%20%20show_value%3DTrue%2C%0A%20%20%20%20)%0A%0A%0A%20%20%20%20%23%20mo.md(%0A%20%20%20%20%23%20%20%20%20%20f%22%22%22%0A%0A%20%20%20%20%23%20%20%20%20%20%7Bconvergence_angle.center()%7D%0A%20%20%20%20%23%20%20%20%20%20%7Bdefocus.center()%7D%0A%20%20%20%20%23%20%20%20%20%20%7Bastigmatism.center()%7D%0A%20%20%20%20%23%20%20%20%20%20%7Bastigmatism_angle.center()%7D%0A%20%20%20%20%23%20%20%20%20%20%7Bmo.as_html(fig).center()%7D%0A%20%20%20%20%23%20%20%20%20%20%7Bmo.as_html(fig_metrics).center()%7D%0A%0A%20%20%20%20%23%20%22%22%22%0A%20%20%20%20%23%20).center()%0A%20%20%20%20return%20(%0A%20%20%20%20%20%20%20%20Cc%2C%0A%20%20%20%20%20%20%20%20astigmatism%2C%0A%20%20%20%20%20%20%20%20astigmatism_angle_slider%2C%0A%20%20%20%20%20%20%20%20convergence_angle%2C%0A%20%20%20%20%20%20%20%20deltaE%2C%0A%20%20%20%20%20%20%20%20detector%2C%0A%20%20%20%20%20%20%20%20energy%2C%0A%20%20%20%20%20%20%20%20max_semiangle%2C%0A%20%20%20%20%20%20%20%20source_size%2C%0A%20%20%20%20)%0A%0A%0A%40app.cell%0Adef%20_(%0A%20%20%20%20Cc%2C%0A%20%20%20%20convergence_angle%2C%0A%20%20%20%20deltaE%2C%0A%20%20%20%20detector%2C%0A%20%20%20%20energy%2C%0A%20%20%20%20energy2wavelength%2C%0A%20%20%20%20max_semiangle%2C%0A%20%20%20%20mo%2C%0A%20%20%20%20np%2C%0A)%3A%0A%20%20%20%20lam%20%3D%20energy2wavelength(energy.value)%0A%20%20%20%20max_detector_alpha%20%3D%20convergence_angle.value%20*%20max_semiangle.value%20*%201e-3%0A%20%20%20%20k_max%20%3D%20max_detector_alpha%20%2F%20lam%20%0A%20%20%20%20sampling%20%3D%201%20%2F%202%20%2F%20k_max%20%0A%20%20%20%20dk%20%3D%201%20%2F%20(detector.value%20*%20sampling)%0A%20%20%20%20FOV%20%3D%201%20%2F%20dk%20%0A%20%20%20%20max_defocus%20%3D%20(FOV%2F2%20)%20%2F%20np.tan(convergence_angle.value%20*%201e-3)%0A%20%20%20%20%23%20defocus.start%20%3D%20-%20max_defocus%0A%20%20%20%20%23%20defocus.stop%20%3D%20max_defocus%0A%20%20%20%20focal_spread%20%3D%20Cc.value%20*%201e-3%20*%20deltaE.value%20%2F%20energy.value%20*%201e10%0A%20%20%20%20print(f'focal%20spread%20%7Bfocal_spread%7D')%0A%20%20%20%20defocus%20%3D%20mo.ui.slider(%0A%20%20%20%20%20%20%20%20start%3D-%20max_defocus%2C%20stop%3Dmax_defocus%2C%20step%3D5%2C%20label%3D%22defocus%20%5B%C3%85%5D%22%2C%20show_value%3DTrue%2C%20value%20%3D%201%0A%20%20%20%20)%0A%20%20%20%20scan_step%20%3D%20mo.ui.slider(%0A%20%20%20%20%20%20%20%20start%3D0.05%2C%20stop%3DFOV%2F2%2C%20step%3D0.05%2C%20label%3D%22scan%20step%20%5B%C3%85%5D%22%2C%20show_value%3DTrue%2C%20value%20%3D%200.5%0A%20%20%20%20)%0A%20%20%20%20return%20(%0A%20%20%20%20%20%20%20%20defocus%2C%0A%20%20%20%20%20%20%20%20dk%2C%0A%20%20%20%20%20%20%20%20focal_spread%2C%0A%20%20%20%20%20%20%20%20lam%2C%0A%20%20%20%20%20%20%20%20max_detector_alpha%2C%0A%20%20%20%20%20%20%20%20sampling%2C%0A%20%20%20%20%20%20%20%20scan_step%2C%0A%20%20%20%20)%0A%0A%0A%40app.cell%0Adef%20_(%0A%20%20%20%20Cc%2C%0A%20%20%20%20ComplexProbe%2C%0A%20%20%20%20add_scalebar%2C%0A%20%20%20%20astigmatism%2C%0A%20%20%20%20astigmatism_angle_slider%2C%0A%20%20%20%20convergence_angle%2C%0A%20%20%20%20defocus%2C%0A%20%20%20%20deltaE%2C%0A%20%20%20%20detector%2C%0A%20%20%20%20dk%2C%0A%20%20%20%20energy%2C%0A%20%20%20%20focal_spread%2C%0A%20%20%20%20lam%2C%0A%20%20%20%20max_detector_alpha%2C%0A%20%20%20%20mcf_modes%2C%0A%20%20%20%20mo%2C%0A%20%20%20%20np%2C%0A%20%20%20%20plt%2C%0A%20%20%20%20sampling%2C%0A%20%20%20%20scan_step%2C%0A%20%20%20%20show_complex%2C%0A%20%20%20%20source_size%2C%0A)%3A%0A%0A%20%20%20%20probe_radius%20%3D%20np.abs(defocus.value)%20*%20np.tan(convergence_angle.value%20*%201e-3)%0A%20%20%20%20areal_oversampling%20%3D%20np.pi%20*%20probe_radius**2%2Fscan_step.value**2%0A%20%20%20%20dalpha%20%3D%20dk%20*%20lam%20%20%0A%0A%20%20%20%20ronchi_mag%20%3D%20convergence_angle.value%20*%201e-3%20%2F%20(%20dalpha%20%20*%20probe_radius%20)%0A%20%20%20%20fund_samp%20%3D%201%20%2F%20(dk%20*%20scan_step.value)%0A%20%20%20%20probe_sampling%20%3D%201%20%2F%20(2%20*%20probe_radius%20*%20dk)%0A%0A%0A%20%20%20%20probe%20%3D%20ComplexProbe(%0A%20%20%20%20%20%20%20%20energy1%3Denergy.value%2C%0A%20%20%20%20%20%20%20%20gpts%3D(detector.value%2C%20detector.value)%2C%0A%20%20%20%20%20%20%20%20sampling%3D(sampling%2C%20sampling)%2C%0A%20%20%20%20%20%20%20%20semiangle_cutoff%3Dconvergence_angle.value%2C%0A%20%20%20%20%20%20%20%20defocus%3Ddefocus.value%2C%0A%20%20%20%20%20%20%20%20astigmatism%3Dastigmatism.value%2C%0A%20%20%20%20%20%20%20%20astigmatism_angle%3Dnp.deg2rad(astigmatism_angle_slider.value)%2C%0A%0A%20%20%20%20).build()%20%0A%20%20%20%20print(f'probe_radius%20%20%20%20%20%20%3A%20%7Bprobe_radius%7D')%0A%20%20%20%20print(f'probe_sampling%20%20%20%20%3A%20%7Bprobe_sampling%7D')%0A%20%20%20%20print(f'ronchi_mag%20%20%20%20%20%20%20%20%3A%20%7Bronchi_mag%7D')%0A%20%20%20%20print(f'areal_oversamp%20%20%20%20%3A%20%7Bareal_oversampling%7D')%0A%20%20%20%20print(f'fund_samp%20%20%20%20%20%20%20%20%20%3A%20%7Bfund_samp%7D')%0A%20%20%20%20print()%0A%20%20%20%20print(f'lam%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%3A%20%7Blam%7D')%20%0A%20%20%20%20print(f'det%20shape%20%20%20%20%20%20%20%20%20%3A%20%7Bdetector.value%7D')%0A%20%20%20%20print(f'max_detector_alpha%3A%20%7Bmax_detector_alpha*1e3%7D')%0A%20%20%20%20print(f'dk%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%3A%20%7Bdk%7D')%0A%20%20%20%20print(f'dk%20probe%20%20%20%20%20%20%20%20%20%20%3A%20%7Bprobe.reciprocal_space_sampling%7D')%0A%20%20%20%20print(f'sampling%20%20%20%20%20%20%20%20%20%20%3A%20%7Bsampling%7D')%0A%0A%20%20%20%20print(f'max_detector_alpha%3A%20%7Bmax_detector_alpha%7D')%0A%20%20%20%20print(f'dalpha%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%5Bmrad%5D%3A%20%7Bdalpha*1e3%7D')%0A%20%20%20%20print(f'probe.angular_sampling%5B0%5D%20%5Bmrad%5D%3A%20%7Bprobe.angular_sampling%7D')%0A%20%20%20%20probe_real%20%3D%20np.fft.fftshift(probe._array)%0A%20%20%20%20probe_fourier%20%3D%20np.fft.fftshift(probe._array_fourier)%0A%0A%20%20%20%20fig%2C%20ax_fourier%20%3D%20plt.subplots(1%2C%201%2C%20figsize%3D(4%2C%204))%0A%0A%20%20%20%20ax_fourier%20%3D%20show_complex(%0A%20%20%20%20%20%20%20%20probe_fourier%2C%0A%20%20%20%20%20%20%20%20figax%3D(fig%2C%20ax_fourier)%2C%0A%20%20%20%20%20%20%20%20ticks%3DFalse%2C%0A%20%20%20%20%20%20%20%20vmin%3D0%2C%0A%20%20%20%20%20%20%20%20vmax%3D1%2C%0A%20%20%20%20)%0A%0A%20%20%20%20ax_fourier.set_title(%22reciprocal-space%20wave%22)%0A%0A%0A%0A%0A%20%20%20%20rectangular_offset%20%3D%20(0.0%2C0.0)%0A%20%20%20%20evs%20%3D%20np.arange(8)%0A%20%20%20%20mcf_vals%2C%20mcf_modes_res%20%3D%20mcf_modes(probe%2C%20focal_spread%2C%20source_size.value%2C%20rectangular_offset%2C%20eigenvectors%3Devs)%0A%20%20%20%20mcf_modes_res2%20%3D%20mcf_modes_res%20%23*%20np.fft.fftshift(probe_fourier)%0A%20%20%20%20mcf_modes_res2%20%3D%20np.fft.fftshift(mcf_modes_res2%2C%20(-2%2C-1))%0A%20%20%20%20%23%20intens%20%3D%20np.linalg.norm(mcf_modes_res2%2C%20axis%3D(-2%2C-1))**2%0A%20%20%20%20mcf_vals%20%3D%20np.abs(mcf_vals)%0A%20%20%20%20mcf_percent%20%3D%20mcf_vals%20%2F%20mcf_vals.sum()%0A%0A%20%20%20%20width%20%3D%201%0A%20%20%20%20labels%20%3D%20%5Bf'Mode%20%7Bi%7D'%20for%20i%20in%20range(len(mcf_modes_res))%5D%0A%20%20%20%20fig_bar%2C%20ax_bar%20%3D%20plt.subplots(1%2C1%2C%20figsize%3D(3%2C%203))%0A%20%20%20%20ax_bar.bar(np.arange(len(mcf_modes_res))%2C%20mcf_percent%2C%20width%2C%20label%3Dlabels)%0A%20%20%20%20ax_bar.set_title(%22Relative%20Power%20of%20MCF%20Eigenvectors%22)%0A%20%20%20%20fig_modes%2C%20ax_modes%20%3D%20plt.subplots(1%2C%20len(evs)%2C%20figsize%3D(len(evs)*4%2C4))%0A%20%20%20%20for%20i%2C%20(v%2C%20m)%20in%20enumerate(zip(mcf_vals%2C%20mcf_modes_res2))%3A%20%20%20%20%0A%20%20%20%20%20%20%20%20%23%20ax_modes%5Bi%5D%20%3D%20show_complex(%0A%20%20%20%20%20%20%20%20%23%20%20%20%20%20mcf_modes_res2%5Bi%5D%2C%0A%20%20%20%20%20%20%20%20%23%20%20%20%20%20figax%3D(fig_modes%2C%20ax_modes%5Bi%5D)%2C%0A%20%20%20%20%20%20%20%20%23%20%20%20%20%20ticks%3DFalse%2C%0A%20%20%20%20%20%20%20%20%23%20%20%20%20%20vmin%3D0%2C%0A%20%20%20%20%20%20%20%20%23%20%20%20%20%20vmax%3D1%2C%0A%20%20%20%20%20%20%20%20%23%20)%0A%20%20%20%20%20%20%20%20ax_modes%5Bi%5D.imshow(mcf_modes_res2%5Bi%5D)%0A%20%20%20%20%20%20%20%20ax_modes%5Bi%5D.set_title(f'MCF%20Eigenvector%20%7Bi%7D%20Power%20%7Bmcf_percent%5Bi%5D*100%3A2.1f%7D')%0A%0A%20%20%20%20print(f'mcf_modes.shape%20%7Bmcf_modes_res2.shape%7D')%0A%20%20%20%20print(f'mcf_vals%20%20%20%20%20%20%20%20%7Bmcf_vals%7D')%0A%0A%20%20%20%20add_scalebar(ax_fourier%2C%20probe.angular_sampling%5B0%5D%2C%20%22mrad%22)%0A%20%20%20%20text_nyquist%20%3D%20mo.md(%0A%20%20%20%20%20%20%20%20f%22%22%22%0A%20%20%20%20%20%20%20%20Nyquist%20res.%20%20%20%20%20%20%3A%20%7B2*sampling%3A2.2f%7D%C3%85%20%0A%20%20%20%20%22%22%22%0A%20%20%20%20)%0A%20%20%20%20text1%20%3D%20mo.md(%0A%20%20%20%20%20%20%20%20f%22%22%22%0A%20%20%20%20%20%20%20%20Probe%20Radius%20%20%20%20%20%20%3A%20%7Bprobe_radius%3A2.2f%7D%C3%85%20%0A%20%20%20%20%22%22%22%0A%20%20%20%20)%0A%20%20%20%20text_focal%20%3D%20mo.md(%0A%20%20%20%20%20%20%20%20f%22%22%22%0A%20%20%20%20%20%20%20%20focal%20spread%20%20%20%20%20%20%3A%20%7Bfocal_spread%3A2.2f%7D%C3%85%20%0A%20%20%20%20%22%22%22%0A%20%20%20%20)%0A%0A%20%20%20%20text5%20%3D%20mo.md(%0A%20%20%20%20%20%20%20%20f%22%22%22%20%0A%20%20%20%20%20%20%20%20References%3A%20%0A%0A%20%20%20%20%20%20%20%20-%20Li%20et%20al%20(2023)%20%5B10.1017%2FS1431927622012387%5D(https%3A%2F%2Fdoi.org%2F10.1017%2FS1431927622012387)%20%0A%20%20%20%20%22%22%22%0A%20%20%20%20).center()%0A%0A%0A%20%20%20%20vertical%20%3D%20mo.vstack(%0A%20%20%20%20%20%20%20%20%5Benergy%2Cdetector%2Cconvergence_angle%2Cdefocus%2CdeltaE%2CCc%2Csource_size%2C%20text_nyquist%2Ctext1%2Ctext_focal%5D%2C%0A%20%20%20%20%20%20%20%20align%3D%22end%22%2C%0A%20%20%20%20%20%20%20%20gap%3D0%0A%20%20%20%20)%0A%20%20%20%20horizontal%20%3D%20mo.hstack(%0A%20%20%20%20%20%20%20%20%5B%20vertical%2C%20fig%2C%20fig_bar%5D%2C%0A%20%20%20%20%20%20%20%20align%3D%22center%22%2C%0A%20%20%20%20%20%20%20%20justify%3D%22center%22%2C%0A%20%20%20%20%20%20%20%20gap%3D0%2C%0A%20%20%20%20%20%20%20%20wrap%3DFalse%2C%0A%20%20%20%20)%0A%20%20%20%20horizontal2%20%3D%20mo.hstack(%0A%20%20%20%20%20%20%20%20%5B%20text5%5D%2C%0A%20%20%20%20%20%20%20%20align%3D%22center%22%2C%0A%20%20%20%20%20%20%20%20justify%3D%22center%22%2C%0A%20%20%20%20%20%20%20%20gap%3D0%2C%0A%20%20%20%20%20%20%20%20wrap%3DFalse%2C%0A%20%20%20%20)%0A%20%20%20%20vertical2%20%3D%20mo.vstack(%0A%20%20%20%20%20%20%20%20%5Bhorizontal%2C%20fig_modes%2C%20text5%5D%2C%0A%20%20%20%20%20%20%20%20align%3D%22start%22%2C%0A%20%20%20%20%20%20%20%20gap%3D0%0A%20%20%20%20)%0A%0A%20%20%20%20vertical2%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(fft_interpolation_masks%2C%20np)%3A%0A%0A%20%20%20%20from%20scipy.sparse.linalg%20import%20eigsh%0A%0A%0A%20%20%20%20def%20spatial_frequencies(%0A%20%20%20%20%20%20%20%20gpts%3A%20tuple%5Bint%2C%20...%5D%2C%0A%20%20%20%20%20%20%20%20sampling%3A%20tuple%5Bfloat%2C%20...%5D%2C%0A%20%20%20%20%20%20%20%20return_grid%3A%20bool%20%3D%20False%2C%0A%20%20%20%20%20%20%20%20xp%3A%20np.ndarray%20%7C%20str%20%7C%20None%20%3D%20np%2C%0A%20%20%20%20)%3A%0A%20%20%20%20%20%20%20%20%22%22%22%0A%20%20%20%20%20%20%20%20Return%20the%20spatial%20frequencies%20of%20a%20grid.%0A%0A%20%20%20%20%20%20%20%20Parameters%0A%20%20%20%20%20%20%20%20----------%0A%20%20%20%20%20%20%20%20gpts%20%3A%20tuple%20of%20int%0A%20%20%20%20%20%20%20%20%20%20%20%20Number%20of%20grid%20points.%0A%20%20%20%20%20%20%20%20sampling%20%3A%20tuple%20of%20float%0A%20%20%20%20%20%20%20%20%20%20%20%20Sampling%20of%20the%20grid%20%5B%C3%85%5D.%0A%20%20%20%20%20%20%20%20return_grid%20%3A%20bool%0A%20%20%20%20%20%20%20%20%20%20%20%20If%20True%2C%20return%20the%20grid%20as%20a%20single%20meshgrid%20array.%0A%20%20%20%20%20%20%20%20xp%20%3A%20module%0A%20%20%20%20%20%20%20%20%20%20%20%20Array%20module%20to%20use%2C%20options%20are%20numpy%20or%20cupy.%20Default%20is%20numpy.%0A%0A%20%20%20%20%20%20%20%20Returns%0A%20%20%20%20%20%20%20%20-------%0A%20%20%20%20%20%20%20%20spatial_frequencies%20%3A%20tuple%20of%20np.ndarray%0A%20%20%20%20%20%20%20%20%20%20%20%20Tuple%20of%20spatial%20frequencies%20in%20each%20dimension.%0A%20%20%20%20%20%20%20%20spatial_frequencies_grid%20%3A%20np.ndarray%0A%20%20%20%20%20%20%20%20%20%20%20%20If%20return_grid%20is%20True%2C%20the%20spatial%20frequencies%20as%20a%20single%20meshgrid%20array.%0A%20%20%20%20%20%20%20%20%22%22%22%0A%20%20%20%20%20%20%20%20dtype%20%3D%20np.float32%0A%0A%0A%0A%20%20%20%20%20%20%20%20out%20%3D%20tuple(xp.fft.fftfreq(n%2C%20d).astype(dtype)%20for%20n%2C%20d%20in%20zip(gpts%2C%20sampling))%0A%0A%20%20%20%20%20%20%20%20if%20return_grid%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20xp.meshgrid(*out%2C%20indexing%3D%22ij%22)%0A%20%20%20%20%20%20%20%20else%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20out%0A%20%20%20%20def%20fft_crop(array%3A%20np.ndarray%2C%20new_shape%3A%20tuple%5Bint%2C%20...%5D%2C%20normalize%3A%20bool%20%3D%20False)%3A%0A%20%20%20%20%20%20%20%20%22%22%22%0A%20%20%20%20%20%20%20%20Crop%20an%20array.%20It%20is%20assumed%20that%20the%20array%20is%20centered%20in%20Fourier%20space%2C%20this%20is%0A%20%20%20%20%20%20%20%20used%20for%20real-space%20interpolation.%0A%0A%20%20%20%20%20%20%20%20Parameters%0A%20%20%20%20%20%20%20%20----------%0A%20%20%20%20%20%20%20%20array%20%3A%20np.ndarray%0A%20%20%20%20%20%20%20%20%20%20%20%20Array%20to%20crop.%0A%20%20%20%20%20%20%20%20new_shape%20%3A%20tuple%20of%20int%0A%20%20%20%20%20%20%20%20%20%20%20%20New%20shape%20of%20the%20array.%20If%20the%20new%20shape%20is%20smaller%20than%20the%20input%20array%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20each%20preceding%20dimension%20is%20treated%20as%20a%20batch%20dimension.%0A%20%20%20%20%20%20%20%20normalize%20%3A%20bool%2C%20optional%0A%20%20%20%20%20%20%20%20%20%20%20%20If%20True%2C%20renormalize%20the%20array%20to%20conserve%20the%20total%20amplitude.%0A%0A%20%20%20%20%20%20%20%20Returns%0A%20%20%20%20%20%20%20%20-------%0A%20%20%20%20%20%20%20%20np.ndarray%0A%20%20%20%20%20%20%20%20%20%20%20%20Cropped%20array.%0A%20%20%20%20%20%20%20%20%22%22%22%0A%0A%0A%20%20%20%20%20%20%20%20if%20len(new_shape)%20%3C%20len(array.shape)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20new_shape%20%3D%20array.shape%5B%3A%20-len(new_shape)%5D%20%2B%20new_shape%0A%0A%20%20%20%20%20%20%20%20mask_in%2C%20mask_out%20%3D%20fft_interpolation_masks(array.shape%2C%20new_shape)%0A%0A%20%20%20%20%20%20%20%20new_array%20%3D%20np.zeros(new_shape%2C%20dtype%3Darray.dtype)%0A%0A%20%20%20%20%20%20%20%20new_array%5Bmask_out%5D%20%3D%20array%5Bmask_in%5D%0A%0A%20%20%20%20%20%20%20%20if%20normalize%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20new_array%20%3D%20new_array%20*%20np.prod(new_array.shape)%20%2F%20np.prod(array.shape)%0A%0A%20%20%20%20%20%20%20%20return%20new_array%0A%0A%20%20%20%20def%20_evaluate_flat_cropped_mcf(complex_probe%2C%20focal_spread%2C%20source_size%2C%20rectangular_offset)%20-%3E%20np.ndarray%3A%0A%0A%20%20%20%20%20%20%20%20kx%2C%20ky%20%3D%20spatial_frequencies(gpts%3Dcomplex_probe._gpts%2C%20sampling%3Dcomplex_probe._sampling%2C%20xp%3Dnp)%0A%0A%20%20%20%20%20%20%20%20k2%20%3D%20kx%5B%3A%2C%20None%5D%20**%202%20%2B%20ky%5BNone%5D%20**%202%0A%20%20%20%20%20%20%20%20kx%2C%20ky%20%3D%20np.meshgrid(kx%2C%20ky%2C%20indexing%3D%22ij%22)%0A%0A%20%20%20%20%20%20%20%20%23%20A%20%3D%20k2%20%3C%20(complex_probe._semiangle_cutoff%20%2F%20complex_probe._wavelength%20%2F%201e3)%20**%202%0A%20%20%20%20%20%20%20%20alpha%2C%20phi%20%3D%20complex_probe.get_scattering_angles()%0A%20%20%20%20%20%20%20%20A%20%3D%20complex_probe.evaluate_aperture(alpha%2C%20phi)%0A%0A%20%20%20%20%20%20%20%20A%2C%20kx%2C%20ky%2C%20k2%20%3D%20(arr.ravel().astype(np.float32)%20for%20arr%20in%20(A%2C%20kx%2C%20ky%2C%20k2))%0A%0A%20%20%20%20%20%20%20%20A%20%3D%20np.multiply.outer(A%2C%20A)%0A%20%20%20%20%20%20%20%20kx%20%3D%20np.subtract.outer(kx%2C%20kx)%0A%20%20%20%20%20%20%20%20ky%20%3D%20np.subtract.outer(ky%2C%20ky)%0A%20%20%20%20%20%20%20%20k2%20%3D%20np.subtract.outer(k2%2C%20k2)%0A%0A%20%20%20%20%20%20%20%20E%20%3D%20A%0A%20%20%20%20%20%20%20%20if%20focal_spread%20%3E%200.0%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20E%20*%3D%20np.exp(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20-((0.5%20*%20np.pi%20*%20complex_probe._wavelength%20*%20focal_spread)%20**%202)%20*%20k2**2%0A%20%20%20%20%20%20%20%20%20%20%20%20)%0A%0A%20%20%20%20%20%20%20%20if%20source_size%20%3E%200.0%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20E%20*%3D%20np.exp(-((np.pi%20*%20source_size)%20**%202)%20*%20(kx**2%20%2B%20ky**2))%0A%0A%20%20%20%20%20%20%20%20if%20rectangular_offset%20!%3D%20(0.0%2C%200.0)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20E%20*%3D%20np.sinc(kx%20*%20rectangular_offset%5B0%5D)%20*%20np.sinc(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20ky%20*%20rectangular_offset%5B1%5D%0A%20%20%20%20%20%20%20%20%20%20%20%20)%0A%0A%20%20%20%20%20%20%20%20return%20E%0A%0A%20%20%20%20def%20_cropped_shape(extent%2C%20semiangle_cutoff%2C%20wavelength)%3A%0A%20%20%20%20%20%20%20%20fourier_space_sampling%20%3D%201%20%2F%20extent%5B0%5D%2C%201%20%2F%20extent%5B1%5D%0A%20%20%20%20%20%20%20%20return%20(%0A%20%20%20%20%20%20%20%20%20%20%20%20int(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20np.ceil(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%202%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20*%20semiangle_cutoff%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%2F%20(fourier_space_sampling%5B0%5D%20*%20wavelength%20*%201e3)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20)%0A%20%20%20%20%20%20%20%20%20%20%20%20)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20int(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20np.ceil(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%202%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20*%20semiangle_cutoff%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%2F%20(fourier_space_sampling%5B1%5D%20*%20wavelength%20*%201e3)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20)%0A%20%20%20%20%20%20%20%20%20%20%20%20)%2C%0A%20%20%20%20%20%20%20%20)%0A%0A%20%20%20%20def%20mcf_modes(waves%2C%20focal_spread%2C%20source_size%2C%20rectangular_offset%2C%20eigenvectors%20%3D%20np.arange(8)%2C%20return_correlation%3A%20bool%20%3D%20False)%3A%0A%20%20%20%20%20%20%20%20%22%22%22%0A%20%20%20%20%20%20%20%20Evaluate%20the%20diagonal%20mixed%20coherence%20function%20for%20given%20wave%20functions.%0A%0A%20%20%20%20%20%20%20%20Parameters%0A%20%20%20%20%20%20%20%20----------%0A%20%20%20%20%20%20%20%20waves%20%3A%20Waves%0A%20%20%20%20%20%20%20%20%20%20%20%20Wave%20functions%20to%20which%20the%20diagonal%20mixed%20coherence%20function%20is%20applied.%0A%20%20%20%20%20%20%20%20return_correlation%20%3A%20bool%0A%20%20%20%20%20%20%20%20%20%20%20%20Return%20correlation%20coefficients%20(default%20is%20False).%0A%0A%20%20%20%20%20%20%20%20Returns%0A%20%20%20%20%20%20%20%20-------%0A%20%20%20%20%20%20%20%20mcf%20%3A%20np.ndarray%0A%20%20%20%20%20%20%20%20%20%20%20%20Array%20representing%20the%20diagonal%20mixed%20coherence%20function.%0A%20%20%20%20%20%20%20%20%22%22%22%0A%0A%20%20%20%20%20%20%20%20E%20%3D%20_evaluate_flat_cropped_mcf(waves%2C%20focal_spread%2C%20source_size%2C%20rectangular_offset)%0A%0A%20%20%20%20%20%20%20%20if%20max(eigenvectors)%20%2B%201%20%3E%3D%20E.shape%5B0%5D%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20raise%20RuntimeError()%0A%0A%20%20%20%20%20%20%20%20values%2C%20vectors%20%3D%20eigsh(E%2C%20k%3Dmax(eigenvectors)%20%2B%201)%0A%20%20%20%20%20%20%20%20order%20%3D%20np.argsort(-values)%0A%0A%20%20%20%20%20%20%20%20selected%20%3D%20order%5Bnp.array(eigenvectors)%5D%0A%20%20%20%20%20%20%20%20vectors%20%3D%20vectors%5B%3A%2C%20selected%5D.T.reshape(%0A%20%20%20%20%20%20%20%20%20%20%20%20(len(selected)%2C)%0A%20%20%20%20%20%20%20%20%20%20%20%20%2B%20waves._gpts%0A%20%20%20%20%20%20%20%20)%0A%20%20%20%20%20%20%20%20values%20%3D%20values%5Bselected%5D%0A%0A%20%20%20%20%20%20%20%20%23%20vectors%20%3D%20fft_crop(vectors%2C%20waves.gpts)%0A%0A%20%20%20%20%20%20%20%20%23%20TODO%3A%20implement%20returing%20correlation%20coefficients%0A%20%20%20%20%20%20%20%20%23%20R%20%3D%20np.corrcoef(E.ravel()%2C%20S.ravel())%0A%0A%20%20%20%20%20%20%20%20if%20return_correlation%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20raise%20NotImplementedError%0A%0A%20%20%20%20%20%20%20%20vectors%20%3D%20np.array(vectors)%0A%20%20%20%20%20%20%20%20values%20%3D%20np.array(values)%0A%20%20%20%20%20%20%20%20v%20%3D%20np.sqrt(np.abs(values))%0A%20%20%20%20%20%20%20%20return%20values%2C%20v%5B%3A%2C%20None%2C%20None%5D%20*%20vectors%0A%20%20%20%20return%20(mcf_modes%2C)%0A%0A%0A%40app.cell%0Adef%20_(np)%3A%0A%20%20%20%20%23%20Complex%20Probes%20Utilities%0A%20%20%20%20def%20energy2wavelength(energy1)%3A%0A%20%20%20%20%20%20%20%20%22%22%22%20%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%0A%20%20%20%20%20%20%20%20return%20(%0A%20%20%20%20%20%20%20%20%20%20%20%20hplanck%0A%20%20%20%20%20%20%20%20%20%20%20%20*%20c%0A%20%20%20%20%20%20%20%20%20%20%20%20%2F%20np.sqrt(energy1%20*%20(2%20*%20me%20*%20c**2%20%2F%20e%20%2B%20energy1))%0A%20%20%20%20%20%20%20%20%20%20%20%20%2F%20e%0A%20%20%20%20%20%20%20%20%20%20%20%20*%201.0e10%0A%20%20%20%20%20%20%20%20)%0A%0A%0A%20%20%20%20class%20ComplexProbe%3A%0A%20%20%20%20%20%20%20%20%22%22%22%20%22%22%22%0A%0A%20%20%20%20%20%20%20%20%23%20fmt%3A%20off%0A%20%20%20%20%20%20%20%20_polar_symbols%20%3D%20(%0A%20%20%20%20%20%20%20%20%20%20%20%20%22C10%22%2C%20%22C12%22%2C%20%22phi12%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%22C21%22%2C%20%22phi21%22%2C%20%22C23%22%2C%20%22phi23%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%22C30%22%2C%20%22C32%22%2C%20%22phi32%22%2C%20%22C34%22%2C%20%22phi34%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%22C41%22%2C%20%22phi41%22%2C%20%22C43%22%2C%20%22phi43%22%2C%20%22C45%22%2C%20%22phi45%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%22C50%22%2C%20%22C52%22%2C%20%22phi52%22%2C%20%22C54%22%2C%20%22phi54%22%2C%20%22C56%22%2C%20%22phi56%22%2C%0A%20%20%20%20%20%20%20%20)%0A%0A%20%20%20%20%20%20%20%20_polar_aliases%20%3D%20%7B%0A%20%20%20%20%20%20%20%20%20%20%20%20%22defocus%22%3A%20%22C10%22%2C%20%22astigmatism%22%3A%20%22C12%22%2C%20%22astigmatism_angle%22%3A%20%22phi12%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%22coma%22%3A%20%22C21%22%2C%20%22coma_angle%22%3A%20%22phi21%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%22Cs%22%3A%20%22C30%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%22C5%22%3A%20%22C50%22%2C%0A%20%20%20%20%20%20%20%20%7D%0A%20%20%20%20%20%20%20%20%23%20fmt%3A%20on%0A%0A%20%20%20%20%20%20%20%20def%20__init__(%0A%20%20%20%20%20%20%20%20%20%20%20%20self%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20energy1%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20gpts%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20sampling%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20semiangle_cutoff%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20soft_aperture%3DTrue%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20parameters%3D%7B%7D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20**kwargs%2C%0A%20%20%20%20%20%20%20%20)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20self._energy%20%3D%20energy1%0A%20%20%20%20%20%20%20%20%20%20%20%20self._gpts%20%3D%20gpts%0A%20%20%20%20%20%20%20%20%20%20%20%20self._sampling%20%3D%20sampling%0A%20%20%20%20%20%20%20%20%20%20%20%20self._semiangle_cutoff%20%3D%20semiangle_cutoff%0A%20%20%20%20%20%20%20%20%20%20%20%20self._soft_aperture%20%3D%20soft_aperture%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20self._parameters%20%3D%20dict(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20zip(self._polar_symbols%2C%20%5B0.0%5D%20*%20len(self._polar_symbols))%0A%20%20%20%20%20%20%20%20%20%20%20%20)%0A%20%20%20%20%20%20%20%20%20%20%20%20parameters.update(kwargs)%0A%20%20%20%20%20%20%20%20%20%20%20%20self.set_parameters(parameters)%0A%20%20%20%20%20%20%20%20%20%20%20%20self._wavelength%20%3D%20energy2wavelength(self._energy)%0A%0A%20%20%20%20%20%20%20%20def%20set_parameters(self%2C%20parameters)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%22%22%22%20%22%22%22%0A%20%20%20%20%20%20%20%20%20%20%20%20for%20symbol%2C%20value%20in%20parameters.items()%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20if%20symbol%20in%20self._parameters.keys()%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20self._parameters%5Bsymbol%5D%20%3D%20value%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20elif%20symbol%20%3D%3D%20%22defocus%22%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20self._parameters%5Bself._polar_aliases%5Bsymbol%5D%5D%20%3D%20-value%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20elif%20symbol%20in%20self._polar_aliases.keys()%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20self._parameters%5Bself._polar_aliases%5Bsymbol%5D%5D%20%3D%20value%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20else%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20raise%20ValueError(%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%22%7B%7D%20not%20a%20recognized%20parameter%22.format(symbol)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20)%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20parameters%0A%0A%20%20%20%20%20%20%20%20def%20get_spatial_frequencies(self)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20tuple(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20np.fft.fftfreq(n%2C%20d)%20for%20n%2C%20d%20in%20zip(self._gpts%2C%20self._sampling)%0A%20%20%20%20%20%20%20%20%20%20%20%20)%0A%0A%20%20%20%20%20%20%20%20def%20get_scattering_angles(self)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20kx%2C%20ky%20%3D%20self.get_spatial_frequencies()%0A%20%20%20%20%20%20%20%20%20%20%20%20kx%2C%20ky%20%3D%20kx%20*%20self._wavelength%2C%20ky%20*%20self._wavelength%0A%20%20%20%20%20%20%20%20%20%20%20%20alpha%20%3D%20np.sqrt(kx%5B%3A%2C%20None%5D%20**%202%20%2B%20ky%5BNone%2C%20%3A%5D%20**%202)%0A%20%20%20%20%20%20%20%20%20%20%20%20phi%20%3D%20np.arctan2(ky%5BNone%2C%20%3A%5D%2C%20kx%5B%3A%2C%20None%5D)%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20alpha%2C%20phi%0A%0A%20%20%20%20%20%20%20%20def%20hard_aperture(self%2C%20alpha%2C%20semiangle_cutoff)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20alpha%20%3C%3D%20semiangle_cutoff%0A%0A%20%20%20%20%20%20%20%20def%20soft_aperture(self%2C%20alpha%2C%20semiangle_cutoff%2C%20angular_sampling)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20denominator%20%3D%20(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20np.sqrt(angular_sampling%5B0%5D%20**%202%20%2B%20angular_sampling%5B1%5D%20**%202)%20*%201e-3%0A%20%20%20%20%20%20%20%20%20%20%20%20)%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20np.clip((semiangle_cutoff%20-%20alpha)%20%2F%20denominator%20%2B%200.5%2C%200%2C%201)%0A%0A%20%20%20%20%20%20%20%20def%20evaluate_aperture(self%2C%20alpha%2C%20phi)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20if%20self._soft_aperture%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20return%20self.soft_aperture(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20alpha%2C%20self._semiangle_cutoff%20*%201e-3%2C%20self.angular_sampling%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20)%0A%20%20%20%20%20%20%20%20%20%20%20%20else%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20return%20self.hard_aperture(alpha%2C%20self._semiangle_cutoff%20*%201e-3)%0A%0A%20%20%20%20%20%20%20%20def%20evaluate_chi(self%2C%20alpha%2C%20phi)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20p%20%3D%20self._parameters%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20alpha2%20%3D%20alpha**2%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20array%20%3D%20np.zeros_like(alpha)%0A%20%20%20%20%20%20%20%20%20%20%20%20if%20any(%5Bp%5Bsymbol%5D%20!%3D%200.0%20for%20symbol%20in%20(%22C10%22%2C%20%22C12%22%2C%20%22phi12%22)%5D)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20array%20%2B%3D%20(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%201%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%2F%202%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20*%20alpha2%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20*%20(p%5B%22C10%22%5D%20%2B%20p%5B%22C12%22%5D%20*%20np.cos(2%20*%20(phi%20-%20p%5B%22phi12%22%5D)))%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20)%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20if%20any(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%5Bp%5Bsymbol%5D%20!%3D%200.0%20for%20symbol%20in%20(%22C21%22%2C%20%22phi21%22%2C%20%22C23%22%2C%20%22phi23%22)%5D%0A%20%20%20%20%20%20%20%20%20%20%20%20)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20array%20%2B%3D%20(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%201%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%2F%203%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20*%20alpha2%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20*%20alpha%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20*%20(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20p%5B%22C21%22%5D%20*%20np.cos(phi%20-%20p%5B%22phi21%22%5D)%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%2B%20p%5B%22C23%22%5D%20*%20np.cos(3%20*%20(phi%20-%20p%5B%22phi23%22%5D))%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20)%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20if%20any(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%5B%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20p%5Bsymbol%5D%20!%3D%200.0%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20for%20symbol%20in%20(%22C30%22%2C%20%22C32%22%2C%20%22phi32%22%2C%20%22C34%22%2C%20%22phi34%22)%0A%20%20%20%20%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%20)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20array%20%2B%3D%20(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%201%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%2F%204%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20*%20alpha2**2%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20*%20(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20p%5B%22C30%22%5D%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%2B%20p%5B%22C32%22%5D%20*%20np.cos(2%20*%20(phi%20-%20p%5B%22phi32%22%5D))%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%2B%20p%5B%22C34%22%5D%20*%20np.cos(4%20*%20(phi%20-%20p%5B%22phi34%22%5D))%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20)%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20if%20any(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%5B%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20p%5Bsymbol%5D%20!%3D%200.0%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20for%20symbol%20in%20(%22C41%22%2C%20%22phi41%22%2C%20%22C43%22%2C%20%22phi43%22%2C%20%22C45%22%2C%20%22phi41%22)%0A%20%20%20%20%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%20)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20array%20%2B%3D%20(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%201%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%2F%205%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20*%20alpha2**2%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20*%20alpha%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20*%20(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20p%5B%22C41%22%5D%20*%20np.cos((phi%20-%20p%5B%22phi41%22%5D))%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%2B%20p%5B%22C43%22%5D%20*%20np.cos(3%20*%20(phi%20-%20p%5B%22phi43%22%5D))%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%2B%20p%5B%22C45%22%5D%20*%20np.cos(5%20*%20(phi%20-%20p%5B%22phi45%22%5D))%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20)%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20if%20any(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%5B%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20p%5Bsymbol%5D%20!%3D%200.0%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20for%20symbol%20in%20(%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%22C50%22%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%22C52%22%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%22phi52%22%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%22C54%22%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%22phi54%22%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%22C56%22%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%22phi56%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20)%0A%20%20%20%20%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%20)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20array%20%2B%3D%20(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%201%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%2F%206%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20*%20alpha2**3%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20*%20(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20p%5B%22C50%22%5D%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%2B%20p%5B%22C52%22%5D%20*%20np.cos(2%20*%20(phi%20-%20p%5B%22phi52%22%5D))%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%2B%20p%5B%22C54%22%5D%20*%20np.cos(4%20*%20(phi%20-%20p%5B%22phi54%22%5D))%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%2B%20p%5B%22C56%22%5D%20*%20np.cos(6%20*%20(phi%20-%20p%5B%22phi56%22%5D))%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20)%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20array%20%3D%202%20*%20np.pi%20%2F%20self._wavelength%20*%20array%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20array%0A%0A%20%20%20%20%20%20%20%20def%20evaluate_aberrations(self%2C%20alpha%2C%20phi)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20np.exp(-1.0j%20*%20self.evaluate_chi(alpha%2C%20phi))%0A%0A%20%20%20%20%20%20%20%20def%20evaluate_ctf(self)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20alpha%2C%20phi%20%3D%20self.get_scattering_angles()%0A%20%20%20%20%20%20%20%20%20%20%20%20array%20%3D%20self.evaluate_aberrations(alpha%2C%20phi)%0A%20%20%20%20%20%20%20%20%20%20%20%20array%20*%3D%20self.evaluate_aperture(alpha%2C%20phi)%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20array%0A%0A%20%20%20%20%20%20%20%20def%20build(self)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20self._array_fourier%20%3D%20self.evaluate_ctf()%0A%20%20%20%20%20%20%20%20%20%20%20%20array%20%3D%20np.fft.ifft2(self._array_fourier)%0A%20%20%20%20%20%20%20%20%20%20%20%20array%20%2F%3D%20np.sqrt(np.sum(np.abs(array)%20**%202))%0A%20%20%20%20%20%20%20%20%20%20%20%20self._array%20%3D%20array%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20self%0A%0A%20%20%20%20%20%20%20%20%40property%0A%20%20%20%20%20%20%20%20def%20sampling(self)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20self._sampling%0A%0A%20%20%20%20%20%20%20%20%40property%0A%20%20%20%20%20%20%20%20def%20reciprocal_space_sampling(self)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20tuple(1%20%2F%20(n%20*%20s)%20for%20n%2C%20s%20in%20zip(self._gpts%2C%20self._sampling))%0A%0A%20%20%20%20%20%20%20%20%40property%0A%20%20%20%20%20%20%20%20def%20angular_sampling(self)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20tuple(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20dk%20*%20self._wavelength%20*%201e3%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20for%20dk%20in%20self.reciprocal_space_sampling%0A%20%20%20%20%20%20%20%20%20%20%20%20)%0A%0A%0A%20%20%20%20return%20ComplexProbe%2C%20energy2wavelength%0A%0A%0A%40app.cell%0Adef%20_(AnchoredSizeBar%2C%20cspace_convert%2C%20np%2C%20plt)%3A%0A%20%20%20%20%23%20Complex%20Plotting%20Utilities%0A%20%20%20%20def%20Complex2RGB(%0A%20%20%20%20%20%20%20%20complex_data%2C%20vmin%3DNone%2C%20vmax%3DNone%2C%20power%3DNone%2C%20chroma_boost%3D1%0A%20%20%20%20)%3A%0A%20%20%20%20%20%20%20%20%22%22%22%20%22%22%22%0A%20%20%20%20%20%20%20%20amp%20%3D%20np.abs(complex_data)%0A%20%20%20%20%20%20%20%20phase%20%3D%20np.angle(complex_data)%0A%0A%20%20%20%20%20%20%20%20if%20power%20is%20not%20None%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20amp%20%3D%20amp**power%0A%0A%20%20%20%20%20%20%20%20if%20np.isclose(np.max(amp)%2C%20np.min(amp))%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20if%20vmin%20is%20None%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20vmin%20%3D%200%0A%20%20%20%20%20%20%20%20%20%20%20%20if%20vmax%20is%20None%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20vmax%20%3D%20np.max(amp)%0A%20%20%20%20%20%20%20%20else%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20if%20vmin%20is%20None%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20vmin%20%3D%200.02%0A%20%20%20%20%20%20%20%20%20%20%20%20if%20vmax%20is%20None%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20vmax%20%3D%200.98%0A%20%20%20%20%20%20%20%20%20%20%20%20vals%20%3D%20np.sort(amp%5B~np.isnan(amp)%5D)%0A%20%20%20%20%20%20%20%20%20%20%20%20ind_vmin%20%3D%20np.round((vals.shape%5B0%5D%20-%201)%20*%20vmin).astype(%22int%22)%0A%20%20%20%20%20%20%20%20%20%20%20%20ind_vmax%20%3D%20np.round((vals.shape%5B0%5D%20-%201)%20*%20vmax).astype(%22int%22)%0A%20%20%20%20%20%20%20%20%20%20%20%20ind_vmin%20%3D%20np.max(%5B0%2C%20ind_vmin%5D)%0A%20%20%20%20%20%20%20%20%20%20%20%20ind_vmax%20%3D%20np.min(%5Blen(vals)%20-%201%2C%20ind_vmax%5D)%0A%20%20%20%20%20%20%20%20%20%20%20%20vmin%20%3D%20vals%5Bind_vmin%5D%0A%20%20%20%20%20%20%20%20%20%20%20%20vmax%20%3D%20vals%5Bind_vmax%5D%0A%0A%20%20%20%20%20%20%20%20amp%20%3D%20np.where(amp%20%3C%20vmin%2C%20vmin%2C%20amp)%0A%20%20%20%20%20%20%20%20amp%20%3D%20np.where(amp%20%3E%20vmax%2C%20vmax%2C%20amp)%0A%20%20%20%20%20%20%20%20amp%20%3D%20((amp%20-%20vmin)%20%2F%20vmax).clip(1e-16%2C%201)%0A%0A%20%20%20%20%20%20%20%20J%20%3D%20amp%20*%2061.5%20%20%23%20Note%20we%20restrict%20luminance%20to%20the%20monotonic%20chroma%20cutoff%0A%20%20%20%20%20%20%20%20C%20%3D%20np.minimum(chroma_boost%20*%2098%20*%20J%20%2F%20123%2C%20110)%0A%20%20%20%20%20%20%20%20h%20%3D%20np.rad2deg(phase)%20%2B%20180%0A%0A%20%20%20%20%20%20%20%20JCh%20%3D%20np.stack((J%2C%20C%2C%20h)%2C%20axis%3D-1)%0A%20%20%20%20%20%20%20%20rgb%20%3D%20cspace_convert(JCh%2C%20%22JCh%22%2C%20%22sRGB1%22).clip(0%2C%201)%0A%0A%20%20%20%20%20%20%20%20return%20rgb%0A%0A%0A%20%20%20%20def%20add_scalebar(ax%2C%20sampling%2C%20units)%3A%0A%20%20%20%20%20%20%20%20%22%22%22%20%22%22%22%0A%20%20%20%20%20%20%20%20bar%20%3D%20AnchoredSizeBar(%0A%20%20%20%20%20%20%20%20%20%20%20%20ax.transData%2C%0A%20%20%20%20%20%20%20%20%20%20%20%2020%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20f%22%7Bnp.round(sampling%2C1)*20%3A.0f%7D%20%7Bunits%7D%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%22lower%20right%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20pad%3D0.5%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20color%3D%22white%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20frameon%3DFalse%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20label_top%3DTrue%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20size_vertical%3D1%2C%0A%20%20%20%20%20%20%20%20)%0A%20%20%20%20%20%20%20%20ax.add_artist(bar)%0A%20%20%20%20%20%20%20%20return%20ax%0A%0A%0A%20%20%20%20def%20show_complex(%0A%20%20%20%20%20%20%20%20complex_data%2C%0A%20%20%20%20%20%20%20%20figax%3DNone%2C%0A%20%20%20%20%20%20%20%20vmin%3DNone%2C%0A%20%20%20%20%20%20%20%20vmax%3DNone%2C%0A%20%20%20%20%20%20%20%20power%3DNone%2C%0A%20%20%20%20%20%20%20%20ticks%3DTrue%2C%0A%20%20%20%20%20%20%20%20chroma_boost%3D1%2C%0A%20%20%20%20%20%20%20%20**kwargs%2C%0A%20%20%20%20)%3A%0A%20%20%20%20%20%20%20%20%22%22%22%20%22%22%22%0A%20%20%20%20%20%20%20%20rgb%20%3D%20Complex2RGB(%0A%20%20%20%20%20%20%20%20%20%20%20%20complex_data%2C%20vmin%2C%20vmax%2C%20power%3Dpower%2C%20chroma_boost%3Dchroma_boost%0A%20%20%20%20%20%20%20%20)%0A%0A%20%20%20%20%20%20%20%20figsize%20%3D%20kwargs.pop(%22figsize%22%2C%20(6%2C%206))%0A%20%20%20%20%20%20%20%20if%20figax%20is%20None%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20fig%2C%20ax%20%3D%20plt.subplots(figsize%3Dfigsize)%0A%20%20%20%20%20%20%20%20else%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20fig%2C%20ax%20%3D%20figax%0A%0A%20%20%20%20%20%20%20%20ax.imshow(rgb%2C%20**kwargs)%0A%20%20%20%20%20%20%20%20if%20ticks%20is%20False%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20ax.set_xticks(%5B%5D)%0A%20%20%20%20%20%20%20%20%20%20%20%20ax.set_yticks(%5B%5D)%0A%20%20%20%20%20%20%20%20return%20ax%0A%20%20%20%20return%20add_scalebar%2C%20show_complex%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%20%20%20%20from%20mpl_toolkits.axes_grid1.anchored_artists%20import%20AnchoredSizeBar%0A%20%20%20%20from%20colorspacious%20import%20cspace_convert%0A%20%20%20%20return%20AnchoredSizeBar%2C%20cspace_convert%2C%20mo%2C%20np%2C%20plt%0A%0A%0A%40app.cell%0Adef%20_()%3A%0A%20%20%20%20import%20typing%20as%20t%0A%20%20%20%20import%20matplotlib.pyplot%20as%20pyplot%0A%20%20%20%20import%20numpy%20as%20numpy%0A%20%20%20%20from%20numpy.typing%20import%20ArrayLike%2C%20NDArray%0A%20%20%20%20from%20matplotlib.colors%20import%20LinearSegmentedColormap%2C%20Normalize%2C%20Colormap%0A%20%20%20%20from%20sklearn.discriminant_analysis%20import%20QuadraticDiscriminantAnalysis%0A%20%20%20%20from%20matplotlib.axes%20import%20Axes%0A%20%20%20%20from%20matplotlib.figure%20import%20Figure%0A%20%20%20%20import%20matplotlib%20%0A%0A%0A%0A%0A%20%20%20%20class%20_FittedQDA%3A%0A%20%20%20%20%20%20%20%20def%20__init__(self%2C%20means%3A%20ArrayLike%2C%20rotations%3A%20t.Sequence%5BArrayLike%5D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20scalings%3A%20t.Sequence%5BArrayLike%5D%2C%20priors%3A%20t.Optional%5BArrayLike%5D%20%3D%20None)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20self.means_%20%3D%20numpy.asarray(means)%0A%20%20%20%20%20%20%20%20%20%20%20%20self.n_classes%20%3D%20self.means_.shape%5B0%5D%0A%20%20%20%20%20%20%20%20%20%20%20%20self.rotations_%20%3D%20list(map(numpy.asarray%2C%20rotations))%0A%20%20%20%20%20%20%20%20%20%20%20%20self.scalings_%20%3D%20list(map(numpy.asarray%2C%20scalings))%0A%20%20%20%20%20%20%20%20%20%20%20%20self.priors_%20%3D%20numpy.asarray(priors)%20if%20priors%20is%20not%20None%20else%20numpy.full(self.n_classes%2C%201%2Fself.n_classes)%0A%0A%20%20%20%20%20%20%20%20def%20predict_log_proba(self%2C%20X%3A%20ArrayLike)%20-%3E%20NDArray%5Bnumpy.floating%5D%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20X%20%3D%20numpy.asarray(X)%0A%20%20%20%20%20%20%20%20%20%20%20%20assert%20X.ndim%20%3D%3D%202%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20norm2%20%3D%20%5B%5D%0A%20%20%20%20%20%20%20%20%20%20%20%20for%20i%20in%20range(self.n_classes)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20X2%20%3D%20(X%20-%20self.means_%5Bi%5D)%20%40%20(self.rotations_%5Bi%5D%20*%20(self.scalings_%5Bi%5D%20**%20-0.5))%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20norm2.append(numpy.sum(X2**2%2C%20axis%3D-1))%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20norm2%20%3D%20numpy.array(norm2).T%0A%20%20%20%20%20%20%20%20%20%20%20%20u%20%3D%20numpy.asarray(%5Bnumpy.sum(numpy.log(s))%20for%20s%20in%20self.scalings_%5D)%0A%20%20%20%20%20%20%20%20%20%20%20%20scores%20%3D%20-0.5%20*%20(norm2%20%2B%20u)%20%2B%20numpy.log(self.priors_)%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20log_likelihood%20%3D%20scores%20-%20numpy.max(scores%2C%20axis%3D-1%2C%20keepdims%3DTrue)%0A%20%20%20%20%20%20%20%20%20%20%20%20log_likelihood%20-%3D%20numpy.log(numpy.sum(numpy.exp(log_likelihood)%2C%20axis%3D-1%2C%20keepdims%3DTrue))%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20log_likelihood%0A%0A%20%20%20%20%20%20%20%20def%20predict_proba(self%2C%20X%3A%20ArrayLike)%20-%3E%20NDArray%5Bnumpy.floating%5D%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20numpy.exp(self.predict_log_proba(X))%0A%0A%20%20%20%20%20%20%20%20def%20predict_prob_success(self%2C%20X%3A%20ArrayLike)%20-%3E%20NDArray%5Bnumpy.floating%5D%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20numpy.exp(self.predict_log_proba(X)%5B...%2C%201%5D)%0A%0A%20%20%20%20def%20_plot_probe_overlap(ax%3A%20'Axes'%2C%20metrics%3A%20t.Dict%5Bstr%2C%20float%5D)%3A%0A%20%20%20%20%20%20%20%20from%20matplotlib.patches%20import%20Circle%2C%20Rectangle%0A%0A%20%20%20%20%20%20%20%20probe_r%20%3D%20metrics%5B'probe_radius'%5D%0A%20%20%20%20%20%20%20%20scan_r%20%3D%20metrics%5B'scan_step'%5D%20%2F%202.0%0A%20%20%20%20%20%20%20%20box_r%20%3D%20metrics%5B'scan_step'%5D%20*%20metrics%5B'fund_samp'%5D%20%2F%202.%0A%0A%20%20%20%20%20%20%20%20ax_r%20%3D%20scan_r%20%2B%20max(probe_r%2C%20box_r)%20*%201.1%0A%0A%20%20%20%20%20%20%20%20ax.set_xlim(-ax_r%2C%20ax_r)%0A%20%20%20%20%20%20%20%20ax.set_ylim(-ax_r%2C%20ax_r)%0A%0A%20%20%20%20%20%20%20%20ax.set_axis_off()%0A%0A%20%20%20%20%20%20%20%20theta%20%3D%2020.0%20*%20numpy.pi%20%2F%20180.%0A%0A%20%20%20%20%20%20%20%20pos1%20%3D%20numpy.array(%5Bnumpy.cos(theta)%2C%20numpy.sin(theta)%5D)%20*%20scan_r%0A%0A%20%20%20%20%20%20%20%20_%20%3D%20%5Bax.add_patch(Circle(%0A%20%20%20%20%20%20%20%20%20%20%20%20t.cast(t.Tuple%5Bfloat%2C%20float%5D%2C%20tuple(pos))%2C%20probe_r%2C%20facecolor%3D'green'%2C%20alpha%3D0.6%2C%20edgecolor%3D'black'%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20linewidth%3D2.0%2C%0A%20%20%20%20%20%20%20%20))%20for%20pos%20in%20(pos1%2C%20-pos1)%5D%0A%0A%20%20%20%20%20%20%20%20ax.plot(%5Bpos1%5B0%5D%2C%20-pos1%5B0%5D%5D%2C%20%5Bpos1%5B1%5D%2C%20-pos1%5B1%5D%5D%2C%20'.-k'%2C%20linewidth%3D2.0)%0A%0A%20%20%20%20%20%20%20%20_%20%3D%20%5Bax.add_patch(Rectangle(%0A%20%20%20%20%20%20%20%20%20%20%20%20(pos%5B0%5D%20-%20box_r%2C%20pos%5B1%5D%20-%20box_r)%2C%202*box_r%2C%202*box_r%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20edgecolor%3D'black'%2C%20linewidth%3D2.0%2C%20fill%3DFalse%2C%0A%20%20%20%20%20%20%20%20))%20for%20pos%20in%20(pos1%2C%20-pos1)%5D%0A%0A%20%20%20%20def%20predict_recons_success(ronchi_mag%3A%20ArrayLike%2C%20areal_oversamp%3A%20ArrayLike)%20-%3E%20NDArray%5Bnumpy.floating%5D%3A%0A%20%20%20%20%20%20%20%20%22%22%22%0A%20%20%20%20%20%20%20%20Empirically%20predict%20the%20probability%20of%20reconstruction%20success%2C%20given%20the%20Ronchigram%20magnification%0A%20%20%20%20%20%20%20%20and%20areal%20oversampling.%0A%0A%20%20%20%20%20%20%20%20Broadcasts%20%60ronchi_mag%60%20and%20%60areal_oversamp%60%20together%2C%20and%20returns%20an%20array%20of%20the%20same%20shape%2C%0A%20%20%20%20%20%20%20%20with%20values%20indicating%20the%20estimated%20probability%20of%20success.%0A%0A%20%20%20%20%20%20%20%20Fitted%20on%20simulated%20Si%20data%2C%20using%20an%20intensity%20threshold%20of%2090%25%20to%20calculate%20the%20probe%20radius.%0A%20%20%20%20%20%20%20%20%22%22%22%0A%0A%20%20%20%20%20%20%20%20clf%20%3D%20_FittedQDA(%0A%20%20%20%20%20%20%20%20%20%20%20%20means%3D%5B%5B0.49027803%2C%201.82918678%5D%2C%20%5B0.80980859%2C%202.00158048%5D%5D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20rotations%3D%5B%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%5B%5B-0.2316652028331075%2C%200.972795576571098%5D%2C%20%5B0.972795576571098%2C%200.2316652028331075%5D%5D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%5B%5B-0.26985970170864165%2C%200.9628996528162853%5D%2C%20%5B0.9628996528162853%2C%200.26985970170864165%5D%5D%0A%20%20%20%20%20%20%20%20%20%20%20%20%5D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20scalings%3D%5B%5B0.56440416%2C%200.07769331%5D%2C%20%5B0.4110058%20%2C%200.06621369%5D%5D%2C%0A%20%20%20%20%20%20%20%20)%0A%0A%20%20%20%20%20%20%20%20ronchi_mag%2C%20areal_oversamp%20%3D%20numpy.broadcast_arrays(ronchi_mag%2C%20areal_oversamp)%0A%0A%20%20%20%20%20%20%20%20return%20clf.predict_prob_success(%0A%20%20%20%20%20%20%20%20%20%20%20%20numpy.log10(numpy.stack((ronchi_mag%2C%20areal_oversamp)%2C%20axis%3D-1).reshape(-1%2C%202))%0A%20%20%20%20%20%20%20%20).reshape(ronchi_mag.shape)%0A%0A%20%20%20%20def%20_plot_predicted_success(ax%3A%20'Axes'%2C%20metrics%3A%20t.Dict%5Bstr%2C%20float%5D)%3A%0A%0A%0A%20%20%20%20%20%20%20%20bwr_r%3A%20Colormap%20%3D%20matplotlib.colormaps%5B'bwr_r'%5D%20%20%23%20type%3A%20ignore%0A%0A%20%20%20%20%20%20%20%20gamma%20%3D%202.0%0A%20%20%20%20%20%20%20%20lin_norm%20%3D%20Normalize(0.5%2C%201.0)%0A%20%20%20%20%20%20%20%20prob_cmap%20%3D%20LinearSegmentedColormap.from_list(%0A%20%20%20%20%20%20%20%20%20%20%20%20'prob'%2C%20numpy.stack(%5B%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20bwr_r(lin_norm.inverse(numpy.abs(lin_norm(x))**gamma%20*%20numpy.sign(lin_norm(x))))%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20for%20x%20in%20numpy.linspace(0.%2C%201.%2C%20512%2C%20endpoint%3DTrue)%0A%20%20%20%20%20%20%20%20%20%20%20%20%5D%2C%20axis%3D0)%0A%20%20%20%20%20%20%20%20)%0A%0A%20%20%20%20%20%20%20%20ax.set_xscale('log')%0A%20%20%20%20%20%20%20%20ax.set_yscale('log')%0A%0A%20%20%20%20%20%20%20%20ax.set_xlabel(%22Ronchi.%20mag%20%5B%24%5C%5Cmathrm%7Bpx%2F%5C%5CAA%7D%24%5D%22)%0A%20%20%20%20%20%20%20%20ax.set_ylabel(%22Areal%20oversampling%22)%0A%0A%20%20%20%20%20%20%20%20ronchi_mag%20%3D%20metrics%5B'ronchi_mag'%5D%0A%20%20%20%20%20%20%20%20areal_oversamp%20%3D%20metrics%5B'areal_oversamp'%5D%0A%0A%20%20%20%20%20%20%20%20ylim%20%3D%20(%0A%20%20%20%20%20%20%20%20%20%20%20%20min(8e-1%2C%20areal_oversamp%20%2F%203)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20max(1e4%2C%20areal_oversamp%20*%203)%2C%0A%20%20%20%20%20%20%20%20)%0A%20%20%20%20%20%20%20%20xlim%20%3D%20(%0A%20%20%20%20%20%20%20%20%20%20%20%20min(8e-1%2C%20ronchi_mag%20%2F%201.5)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20max(1e2%2C%20ronchi_mag%20*%204)%2C%0A%20%20%20%20%20%20%20%20)%0A%0A%20%20%20%20%20%20%20%20yy%20%3D%20numpy.geomspace(ylim%5B1%5D%2C%20ylim%5B0%5D%2C%20100)%0A%20%20%20%20%20%20%20%20xx%20%3D%20numpy.geomspace(xlim%5B0%5D%2C%20xlim%5B1%5D%2C%20100)%0A%20%20%20%20%20%20%20%20yy%2C%20xx%20%3D%20numpy.meshgrid(yy%2C%20xx%2C%20indexing%3D'ij')%0A%0A%20%20%20%20%20%20%20%20pp%20%3D%20predict_recons_success(xx%2C%20yy)%0A%20%20%20%20%20%20%20%20ax.pcolormesh(xx%2C%20yy%2C%20pp%2C%20alpha%3D0.8%2C%20cmap%3Dprob_cmap%2C%20vmin%3D0.0%2C%20vmax%3D1.0)%0A%0A%20%20%20%20%20%20%20%20prob%20%3D%20predict_recons_success(ronchi_mag%2C%20areal_oversamp)%0A%20%20%20%20%20%20%20%20ax.scatter(%5Bronchi_mag%5D%2C%20%5Bareal_oversamp%5D%2C%20marker%3D'x'%2C%20s%3D80%2C%20c%3D'black')%0A%20%20%20%20%20%20%20%20ax.annotate(f%22%7Bprob%3A.1%25%7D%22%2C%20(ronchi_mag%2C%20areal_oversamp)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20(0.5%2C%200.7)%2C%20textcoords%3D'offset%20fontsize')%0A%0A%20%20%20%20def%20_plot_linear_metrics(ax%3A%20'Axes'%2C%20metrics%3A%20t.Dict%5Bstr%2C%20float%5D)%3A%0A%20%20%20%20%20%20%20%20%23%20mrad%0A%20%20%20%20%20%20%20%20diff_scale%20%3D%20metrics%5B'wavelength'%5D*1e3%2F(2.%20*%20metrics%5B'probe_radius'%5D)%0A%20%20%20%20%20%20%20%20step_scale%20%3D%202.*metrics%5B'probe_radius'%5D%0A%0A%20%20%20%20%20%20%20%20x%20%3D%20metrics%5B'scan_step'%5D%0A%20%20%20%20%20%20%20%20y%20%3D%20metrics%5B'diff_step'%5D%0A%0A%20%20%20%20%20%20%20%20ax.set_xlabel(%22Scan%20step%20%5B%24%5C%5Cmathrm%7B%5C%5CAA%2Fpx%7D%24%5D%22)%0A%20%20%20%20%20%20%20%20ax.set_xscale('log')%0A%20%20%20%20%20%20%20%20ax.set_ylabel(%22Diff.%20pixel%20size%20%5Bmrad%2Fpx%5D%22)%0A%20%20%20%20%20%20%20%20ax.set_yscale('log')%0A%0A%20%20%20%20%20%20%20%20ylim%20%3D%20(%0A%20%20%20%20%20%20%20%20%20%20%20%20min(y%2F1.5%2C%200.1*diff_scale)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20max(y*1.5%2C%2010.*diff_scale)%2C%0A%20%20%20%20%20%20%20%20)%0A%20%20%20%20%20%20%20%20xlim%20%3D%20(%0A%20%20%20%20%20%20%20%20%20%20%20%20min(x%2F1.5%2C%200.1*step_scale)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20max(x*1.5%2C%2010.*step_scale)%2C%0A%20%20%20%20%20%20%20%20)%0A%20%20%20%20%20%20%20%20yy%20%3D%20numpy.geomspace(*ylim%2C%20201%2C%20endpoint%3DTrue)%0A%20%20%20%20%20%20%20%20xx%20%3D%20numpy.geomspace(*xlim%2C%20201%2C%20endpoint%3DTrue)%0A%20%20%20%20%20%20%20%20(yy%2C%20xx)%20%3D%20numpy.meshgrid(yy%2C%20xx%2C%20indexing%3D'ij')%0A%0A%20%20%20%20%20%20%20%20fps%20%3D%201.0%20%2F%20(yy%20%2F%20diff_scale%20*%20xx%20%2F%20step_scale)%0A%0A%20%20%20%20%20%20%20%20ax.axhline(diff_scale%2C%20linestyle%3D'dashed'%2C%20color%3D'%234363d8')%20%23%20blue%0A%20%20%20%20%20%20%20%20ax.axvline(step_scale%2C%20linestyle%3D'dashdot'%2C%20color%3D'%23f58231')%20%23%20orange%0A%20%20%20%20%20%20%20%20ax.contour(xx%2C%20yy%2C%20fps%2C%20%5B1.0%5D%2C%20linestyles%3D%5B'solid'%5D%2C%20colors%3D%5B'%23e6194b'%5D)%20%23%20red%0A%0A%20%20%20%20%20%20%20%20ax.scatter(%5Bx%5D%2C%20%5By%5D%2C%20marker%3D'x'%2C%20color%3D'black'%2C%20s%3D80)%0A%0A%20%20%20%20%20%20%20%20ax.set_ylim(*ylim)%0A%20%20%20%20%20%20%20%20ax.set_xlim(*xlim)%0A%0A%20%20%20%20def%20plot_metrics(metrics%3A%20t.Dict%5Bstr%2C%20float%5D)%20-%3E%20'Figure'%3A%0A%20%20%20%20%20%20%20%20fig%2C%20(ax1%2C%20ax2%2C%20ax3)%20%3D%20pyplot.subplots(%0A%20%20%20%20%20%20%20%20%20%20%20%20ncols%3D3%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20gridspec_kw%3D%7B%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20'width_ratios'%3A%20%5B2.%2C%201.%2C%202.%5D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20'wspace'%3A%200.05%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%7D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20constrained_layout%3DTrue%0A%20%20%20%20%20%20%20%20)%0A%20%20%20%20%20%20%20%20fig.set_size_inches(12%2C%204)%0A%20%20%20%20%20%20%20%20ax1.set_box_aspect(1.)%0A%20%20%20%20%20%20%20%20ax2.set_aspect(1.)%0A%20%20%20%20%20%20%20%20ax3.set_box_aspect(1.)%0A%20%20%20%20%20%20%20%20ax1.set_title('Probe%20sampling%20vs.%20Diff.%20Oversampling')%0A%20%20%20%20%20%20%20%20ax2.set_title('Probe%20Overlap')%0A%20%20%20%20%20%20%20%20ax3.set_title('Heuristic%20success%20probability')%0A%0A%20%20%20%20%20%20%20%20_plot_linear_metrics(ax1%2C%20metrics)%0A%20%20%20%20%20%20%20%20_plot_probe_overlap(ax2%2C%20metrics)%0A%20%20%20%20%20%20%20%20_plot_predicted_success(ax3%2C%20metrics)%0A%0A%20%20%20%20%20%20%20%20return%20fig%0A%0A%0A%0A%0A%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_()%3A%0A%20%20%20%20return%0A%0A%0Aif%20__name__%20%3D%3D%20%22__main__%22%3A%0A%20%20%20%20app.run()%0A