Metadata-Version: 2.4
Name: fast-acbf
Version: 0.9.0
Summary: GPU-based Bright-field Ptychography Solver for tcBF and acBF, compatible with PtyRAD
Author-email: Chia-Hao Lee <cl2696@cornell.edu>
License-Expression: LGPL-3.0
Project-URL: Repository, https://github.com/chiahao3/fast-acbf
Project-URL: Changelog, https://github.com/chiahao3/fast-acbf/blob/main/CHANGELOG.md
Project-URL: PtyRAD, https://github.com/chiahao3/ptyrad
Keywords: Ptychography
Classifier: Development Status :: 4 - Beta
Classifier: Programming Language :: Python :: 3
Classifier: Operating System :: OS Independent
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: h5py
Requires-Dist: matplotlib
Requires-Dist: numpy
Requires-Dist: psutil
Requires-Dist: scipy
Requires-Dist: torch>=2.4
Requires-Dist: torchvision
Provides-Extra: dev
Requires-Dist: pytest>=8.0; extra == "dev"
Requires-Dist: pytest-cov>=5.0; extra == "dev"
Requires-Dist: zarr>=3.0; extra == "dev"
Dynamic: license-file

# GPU-accelerated tcBF and acBF

This repository provides a fast GPU implementation of aberration-corrected bright-field 4D-STEM (tcBF/acBF) reconstruction.

It is designed to serve as the seed for:

1. Future native integration of tcBF and acBF into [PtyRAD](https://github.com/chiahao3/ptyrad) pipeline
2. Real-time visualization during 4D-STEM data acquisition

The underlying imaging theory and initial CPU implementation of **acBF** were developed by Dr. Desheng Ma and Dr. Steven Zeltmann [1, 2], while the **tcBF** method has a longer history; readers are encouraged to read this paper by Dr. Yue Yu [3].

This implementation (**fast-acBF**) was developed independently focusing on GPU acceleration and integration with the PtyRAD reconstruction framework. It is shared as a working research implementation; interfaces may change as development continues.

## Installation guide

**Major Dependencies:**
- python >=3.10
- pytorch >=2.4

fast-acbf is a standalone library and does not need PtyRAD. Its aberration notation
(`Aberrations`, ported from PtyRAD and kept in sync with it) and output frames match
PtyRAD's, so exported aberrations and probes can be passed to PtyRAD unchanged. To use
both, install PtyRAD separately (`pip install ptyrad`); fast-acbf does not pull it in.

### 1. Get the fast-acBF code from GitHub
You can either download the repository as a .zip file and extract it, or use the following command if you have `git` installed.

```bash
git clone https://github.com/chiahao3/fast-acbf
```

### 2. Create and Activate the Python Environment
Assuming you're using conda, you can create an independent environment and install the packages with these commands:

```bash
# Enter the commands one by one
conda create -n fast-acbf python=3.12 -y
conda activate fast-acbf
cd fast-acbf
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu126
pip install -e .
```

If you prefer a legacy version of PyTorch, or a different version of CUDA runtime other than CUDA 12.6, see instruction [here](https://pytorch.org/get-started/previous-versions/).

## Get Started

1. Download the demo tBL-WSe2 data "Figure 4.zip" from the [Zenodo link](https://doi.org/10.5281/zenodo.15283331)
2. Run the `get_acBF.ipynb` Jupyter notebook to reconstruct tcBF / acBF images

## Data Pipeline Defaults

`BFSolver` is the main user entry point. By default it uses
`pipeline="balanced"` with automatic policy resolution:

```python
from fast_acbf import BFSolver, Dataset4D

dataset = Dataset4D.from_hdf5("scan.h5", key="array")  # lazy by default
# EMPAD .raw has no shape metadata, so pass it (gap=0 for gapless raw files):
# dataset = Dataset4D.from_raw("scan.raw", scan_shape=(128, 128), detector_shape=(128, 128), gap=1024)

solver = BFSolver(
    dataset=dataset,
    max_alpha=25.0,
    scan_step_size=0.43,
    dk=0.04,
    wavelength=0.04176,
    device="cuda",
)
```

The solver consumes `ImageFFT` chunks internally. The policy controls where that
FFT cache lives, when it is filled, and how virtual BF images are extracted from
the raw 4D data.

| Pipeline | Intended use | Default behavior |
| --- | --- | --- |
| `speed` | Small data that can afford temporary materialization | Prefer device ImageFFT and, when raw 4D can be materialized and fit in VRAM with vBF/ImageFFT, do a whole-pass `device_mask` precompute. |
| `balanced` | General default | Cache ImageFFT on device if it fits, otherwise host RAM, otherwise stream. Avoids raw-on-device extraction for lazy disk data unless explicitly requested. |
| `memory` | Largest data / lowest persistent memory | Keep `imagefft_storage="none"` and compute on the fly, using disk extraction for lazy data. |

Advanced users can override the resolved policy:

```python
solver = BFSolver(
    dataset=dataset,
    ...,
    pipeline="balanced",
    imagefft_storage="host",     # auto | device | host | none
    imagefft_fill="precompute",  # auto | precompute | lazy | on_the_fly
    extractor_strategy="auto",   # auto | device_mask | host_mask | disk_*
)
```

Impossible or wasteful combinations are rejected early with memory/path guidance.
For example, `extractor_strategy="device_mask"` requires persistent ImageFFT
storage and `imagefft_fill="precompute"` because it is a whole-pass route:
temporarily move raw 4D to the compute device, extract all vBF images, precompute
ImageFFT, then release raw/vBF intermediates.

## References

[1] Ma, Desheng, et al. "Information in 4D-STEM: Where it is, and How to Use it." Ultramicroscopy (2026). https://doi.org/10.1016/j.ultramic.2026.114351

[2] Ma, Desheng, David A. Muller, and Steven E. Zeltmann. "Using Aberrations to Improve Dose-Efficient Tilt-corrected 4D-STEM Imaging." Microscopy and Microanalysis (2026). https://doi.org/10.1093/mam/ozag008

[3] Yu, Yue, et al. "Dose-efficient cryo-electron microscopy for thick samples using tilt-corrected scanning transmission electron microscopy." Nature Methods (2025). https://doi.org/10.1038/s41592-025-02834-9

## Relevant Repositories

- [tcBF-STEM](https://github.com/yyu2017/tcBFSTEM)
- [acBF-STEM](https://github.com/dsmagiya/acBF-STEM)
- [PtyRAD](https://github.com/chiahao3/ptyrad)
- [py4DSTEM](https://github.com/py4dstem/py4DSTEM)
- [quantem](https://github.com/electronmicroscopy/quantem)

## Author 

Chia-Hao Lee (cl2696@cornell.edu)

Developed at the Muller Group, Cornell University.


