phaserEM 0.1__tar.gz → 0.2__tar.gz

{phaserem-0.1 → phaserem-0.2}/PKG-INFO

@@ -1,11 +1,12 @@
 Metadata-Version: 2.4
 Name: phaserEM
-Version: 0.1
+Version: 0.2
 Summary: Weapon of choice for electron ptychographic reconstructions
 Author-email: Colin Gilgenbach <hexane@mit.edu>
 License-Expression: MPL-2.0
 Project-URL: Repository, https://github.com/hexane360/phaser.git
 Project-URL: Documentation, https://hexane360.github.io/phaser/dev
+Keywords: science,atoms,materials,chemistry
 Classifier: Development Status :: 4 - Beta
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3 :: Only
@@ -13,32 +14,37 @@ Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
+Classifier: Intended Audience :: Science/Research
+Classifier: Topic :: Scientific/Engineering :: Physics
 Classifier: Typing :: Typed
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE.txt
-Requires-Dist: numpy<2.5,>=1.22
-Requires-Dist: scipy<1.18,>=1.7.0
+Requires-Dist: numpy<2.6,>=1.22
+Requires-Dist: scipy<1.19,>=1.7.0
 Requires-Dist: matplotlib~=3.8
 Requires-Dist: h5py~=3.8
 Requires-Dist: pyyaml>=5.3.1
 Requires-Dist: click~=8.1.0
-Requires-Dist: rich~=12.0.0
+Requires-Dist: rich<15,>=12.0.0
 Requires-Dist: tifffile>=2023.8.25
-Requires-Dist: py-pane==0.11.2
+Requires-Dist: py-pane==0.11.3
 Requires-Dist: typing_extensions~=4.7
 Provides-Extra: dev
 Requires-Dist: pytest>=6.2.4; extra == "dev"
 Requires-Dist: pytest-cov>=3.0.0; extra == "dev"
 Requires-Dist: coveragepy-lcov>=0.1.1; extra == "dev"
-Provides-Extra: cupy
-Requires-Dist: cupy-wheel; extra == "cupy"
-Requires-Dist: pynvml>=11.0.0; extra == "cupy"
+Provides-Extra: cupy11
+Requires-Dist: cupy-cuda11x>=12.0.0; extra == "cupy11"
+Requires-Dist: pynvml>=11.0.0; extra == "cupy11"
+Provides-Extra: cupy12
+Requires-Dist: cupy-cuda12x>=12.0.0; extra == "cupy12"
+Requires-Dist: pynvml>=11.0.0; extra == "cupy12"
 Provides-Extra: jax
-Requires-Dist: jax<0.6,>=0.4.25; extra == "jax"
+Requires-Dist: jax<0.8,>=0.4.25; extra == "jax"
 Requires-Dist: optax>=0.2.2; extra == "jax"
 Provides-Extra: web
-Requires-Dist: Quart>=0.19.5; extra == "web"
+Requires-Dist: Quart>=0.20.0; extra == "web"
 Requires-Dist: backoff==2.2.1; extra == "web"
 Requires-Dist: requests>=2.31.0; extra == "web"
 Provides-Extra: docs
@@ -52,9 +58,11 @@ Dynamic: license-file
 
 phaser: The weapon of choice for ptychographic reconstructions
 ---
-[![][ci-badge]][ci-url] [![][commit-badge]][commit-url] [![][docs-dev-badge]][docs-dev-url]
+[![][ci-badge]][ci-url] [![][commit-badge]][commit-url] [![][docs-dev-badge]][docs-dev-url] [![][binder-badge]][binder-url]
+
 
 `phaser` is a fully-featured package for multislice electron ptychography.
+See our [arXiv paper](https://arxiv.org/abs/2505.14372) for more details.
 
 ## Features
 
@@ -64,11 +72,18 @@ phaser: The weapon of choice for ptychographic reconstructions
 - Multiple incoherent probe modes
 - Probe position correction
 - Upsampled (sPIE) and segmented ptychography (work in progress)
-- Adaptive propagator correction (work in progress)
+- Adaptive propagator correction (contributed by M Zhu)
+
+## Documentation
+
+Documentation on `phaser` can be found here: https://hexane360.github.io/phaser/dev/
+
+Documentation is still very much a work in progress, so please feel free to open an issue or email me if you have any questions!
 
 ## Installation
 
-To install, first clone the repository from github:
+To install, first clone the repository from github.
+This can be done from GitHub Desktop, or from the git command line:
 
 ```sh
 $ git clone https://github.com/hexane360/phaser
@@ -76,9 +91,29 @@ $ git clone https://github.com/hexane360/phaser
 $ cd phaser
 ```
 
-If you're using [`cupy`][cupy] or [`jax`][jax] with a GPU, follow the installation instructions for those packages.
+We recommend using a conda environment or Python virtual environment to keep things clean, although this is not mandatory.
+
+`phaser` supports multiple computational backends. The simplest (and slowest) is `numpy`. `cupy` can be used for CUDA-accelerated. `jax` supports CPU and GPU acceleration, and is the only backend which supports the gradient descent engine.
+If you're unsure what engines to use, we recommend installing the `jax` engine.
+
+If you're using [`cupy`][cupy] or [`jax`][jax] with a GPU, start by following the installation instructions for those packages.
+Jax can be installed with or without CUDA support, if you're using CUDA make sure you install the correct version. Currently, Jax does not support CUDA on Windows.
 
-Then, install with `pip`:
+Before moving on to installing `phaser`, make sure those packages you've installed work:
+```sh
+$ python
+>>> import jax
+>>> jax.default_backend()
+'gpu'  # should be 'gpu' on cuda, 'cpu' otherwise
+>>> jax.numpy.array([1, 2, 3, 4]) + 1   # test a basic operation
+Array([2, 3, 4, 5], dtype=int32)
+
+>>> import cupy
+>>> cupy.array([1, 2, 3, 4]) + 1
+array([2, 3, 4, 5])
+```
+
+Then, install `phaser` using `pip`:
 
 ```sh
 $ python -m pip install -e .
@@ -87,7 +122,22 @@ $ python -m pip install -e .
 For the jax or cupy backend, or for the optional webserver, install with the corresponding options:
 
 ```sh
-$ python -m pip install -e '.[cupy,web]' # or '.[jax,web]'
+$ python -m pip install -e ".[jax,cupy12,web]" # for the 'jax', 'cupy12', and 'web' options
+```
+
+Depending on your command line, you may need to put double quotes around the options (as shown).
+
+Here are the supported installation options:
+
+- `jax`: For the [`jax`][jax] backend (required for the gradient descent engine)
+- `cupy11`: `cupy` for CUDA toolkit 11.x
+- `cupy12`: `cupy` for CUDA toolkit 12.x
+- `web`: For the web interface
+
+For [Optuna](https://optuna.org/) hyperoptimization, install it as well:
+
+```sh
+$ pip install optuna
 ```
 
 ## Running
@@ -100,6 +150,20 @@ To run the webserver, call `phaser serve`. By default, the server serves on http
 
 To run a worker, call `phaser worker <url>`, where `url` is the URL of a running job server.
 
+## Sample data & Examples
+
+Sample data can be downloaded from the following dropbox link: https://www.dropbox.com/scl/fo/txm3k88ubrzvt541v23ir/AL-l_m6VnGlFxzHWZSSc0TA?rlkey=8qxtwnc8cwhpff6jpr5s40y6i&st=x9pbwke0&dl=0
+
+Copy the `sample_data` directory into the root code folder.
+
+```sh
+$ curl --output sample_data.zip -L 'https://www.dropbox.com/scl/fo/txm3k88ubrzvt541v23ir/AL-l_m6VnGlFxzHWZSSc0TA?rlkey=8qxtwnc8cwhpff6jpr5s40y6i&st=x9pbwke0&dl=1'
+$ unzip sample_data.zip -x / -d sample_data
+```
+
+Sample data includes simulated and experimental MoS2 data, simulated and experimental Si data, and experimental PrScO3 data.
+After the data is downloaded, any of the example reconstructions can be run as `phaser run examples/mos2_grad.yaml` (for example).
+
 ## Alternatives
 
 Other notable ptychography packages:
@@ -108,6 +172,7 @@ Other notable ptychography packages:
  - [`py4DSTEM`](https://github.com/py4dstem/py4DSTEM)
  - [`PtyLab.m/py/jl`](https://opg.optica.org/oe/fulltext.cfm?uri=oe-31-9-13763&id=529026)
  - [PtyPy](https://ptycho.github.io/ptypy/)
+ - [PtyRAD](https://github.com/chiahao3/ptyrad)
 
 [numpy]: https://numpy.org/
 [cupy]: https://cupy.dev/
@@ -119,3 +184,5 @@ Other notable ptychography packages:
 [docs-dev-url]: https://hexane360.github.io/phaser/dev/
 [commit-badge]: https://img.shields.io/github/last-commit/hexane360/phaser
 [commit-url]: https://github.com/hexane360/phaser/commits
+[binder-badge]: https://mybinder.org/badge_logo.svg
+[binder-url]: https://mybinder.org/v2/gh/hexane360/phaser/HEAD

phaserem-0.2/README.md

@@ -0,0 +1,130 @@
+phaser: The weapon of choice for ptychographic reconstructions
+---
+[![][ci-badge]][ci-url] [![][commit-badge]][commit-url] [![][docs-dev-badge]][docs-dev-url] [![][binder-badge]][binder-url]
+
+
+`phaser` is a fully-featured package for multislice electron ptychography.
+See our [arXiv paper](https://arxiv.org/abs/2505.14372) for more details.
+
+## Features
+
+- ePIE, LSQ-MLs, and gradient descent algorithms
+- [`numpy`][numpy], [`cupy`][cupy], and [`jax`][jax] backends.
+- Single and multislice ptychography
+- Multiple incoherent probe modes
+- Probe position correction
+- Upsampled (sPIE) and segmented ptychography (work in progress)
+- Adaptive propagator correction (contributed by M Zhu)
+
+## Documentation
+
+Documentation on `phaser` can be found here: https://hexane360.github.io/phaser/dev/
+
+Documentation is still very much a work in progress, so please feel free to open an issue or email me if you have any questions!
+
+## Installation
+
+To install, first clone the repository from github.
+This can be done from GitHub Desktop, or from the git command line:
+
+```sh
+$ git clone https://github.com/hexane360/phaser
+# enter phaser directory
+$ cd phaser
+```
+
+We recommend using a conda environment or Python virtual environment to keep things clean, although this is not mandatory.
+
+`phaser` supports multiple computational backends. The simplest (and slowest) is `numpy`. `cupy` can be used for CUDA-accelerated. `jax` supports CPU and GPU acceleration, and is the only backend which supports the gradient descent engine.
+If you're unsure what engines to use, we recommend installing the `jax` engine.
+
+If you're using [`cupy`][cupy] or [`jax`][jax] with a GPU, start by following the installation instructions for those packages.
+Jax can be installed with or without CUDA support, if you're using CUDA make sure you install the correct version. Currently, Jax does not support CUDA on Windows.
+
+Before moving on to installing `phaser`, make sure those packages you've installed work:
+```sh
+$ python
+>>> import jax
+>>> jax.default_backend()
+'gpu'  # should be 'gpu' on cuda, 'cpu' otherwise
+>>> jax.numpy.array([1, 2, 3, 4]) + 1   # test a basic operation
+Array([2, 3, 4, 5], dtype=int32)
+
+>>> import cupy
+>>> cupy.array([1, 2, 3, 4]) + 1
+array([2, 3, 4, 5])
+```
+
+Then, install `phaser` using `pip`:
+
+```sh
+$ python -m pip install -e .
+```
+
+For the jax or cupy backend, or for the optional webserver, install with the corresponding options:
+
+```sh
+$ python -m pip install -e ".[jax,cupy12,web]" # for the 'jax', 'cupy12', and 'web' options
+```
+
+Depending on your command line, you may need to put double quotes around the options (as shown).
+
+Here are the supported installation options:
+
+- `jax`: For the [`jax`][jax] backend (required for the gradient descent engine)
+- `cupy11`: `cupy` for CUDA toolkit 11.x
+- `cupy12`: `cupy` for CUDA toolkit 12.x
+- `web`: For the web interface
+
+For [Optuna](https://optuna.org/) hyperoptimization, install it as well:
+
+```sh
+$ pip install optuna
+```
+
+## Running
+
+After installation, the `phaser` command should be available. Phaser can be run from the command line, or through a job server.
+
+To run a single reconstruction on the command line, call `phaser run <file>`, where `file` is the path to a reconstruction plan file.
+
+To run the webserver, call `phaser serve`. By default, the server serves on https://localhost:5050/, so navigate there in a web browser. The server interface can be used to start workers and schedule reconstruction jobs.
+
+To run a worker, call `phaser worker <url>`, where `url` is the URL of a running job server.
+
+## Sample data & Examples
+
+Sample data can be downloaded from the following dropbox link: https://www.dropbox.com/scl/fo/txm3k88ubrzvt541v23ir/AL-l_m6VnGlFxzHWZSSc0TA?rlkey=8qxtwnc8cwhpff6jpr5s40y6i&st=x9pbwke0&dl=0
+
+Copy the `sample_data` directory into the root code folder.
+
+```sh
+$ curl --output sample_data.zip -L 'https://www.dropbox.com/scl/fo/txm3k88ubrzvt541v23ir/AL-l_m6VnGlFxzHWZSSc0TA?rlkey=8qxtwnc8cwhpff6jpr5s40y6i&st=x9pbwke0&dl=1'
+$ unzip sample_data.zip -x / -d sample_data
+```
+
+Sample data includes simulated and experimental MoS2 data, simulated and experimental Si data, and experimental PrScO3 data.
+After the data is downloaded, any of the example reconstructions can be run as `phaser run examples/mos2_grad.yaml` (for example).
+
+## Alternatives
+
+Other notable ptychography packages:
+
+ - [`fold_slice`](https://github.com/yijiang1/fold_slice) branch of [PtychoShelves](https://www.psi.ch/en/sls/csaxs/software#coming-soon-ptychoshelves-a-versatile-high-level-framework-for-high-performance-analysis-of)
+ - [`py4DSTEM`](https://github.com/py4dstem/py4DSTEM)
+ - [`PtyLab.m/py/jl`](https://opg.optica.org/oe/fulltext.cfm?uri=oe-31-9-13763&id=529026)
+ - [PtyPy](https://ptycho.github.io/ptypy/)
+ - [PtyRAD](https://github.com/chiahao3/ptyrad)
+
+[numpy]: https://numpy.org/
+[cupy]: https://cupy.dev/
+[jax]: https://docs.jax.dev/en/latest/
+
+[ci-badge]: https://github.com/hexane360/phaser/workflows/Tests/badge.svg
+[ci-url]: https://github.com/hexane360/phaser/actions/workflows/ci.yaml
+[docs-dev-badge]: https://img.shields.io/badge/docs-dev-blue
+[docs-dev-url]: https://hexane360.github.io/phaser/dev/
+[commit-badge]: https://img.shields.io/github/last-commit/hexane360/phaser
+[commit-url]: https://github.com/hexane360/phaser/commits
+[binder-badge]: https://mybinder.org/badge_logo.svg
+[binder-url]: https://mybinder.org/v2/gh/hexane360/phaser/HEAD

{phaserem-0.1 → phaserem-0.2}/phaser/engines/common/noise_models.py

@@ -93,12 +93,11 @@ class PoissonNoiseModel(NoiseModel[None]):
     ) -> t.Tuple[Float, None]:
         xp = get_array_module(model_wave, model_intensity, exp_patterns, mask)
         patterns = xp.maximum(exp_patterns, 0.0)
-        #intensity - patterns * xp.log(intensity + self.offset)
 
         loss = xp.sum(mask * (
-            #model_intensity - patterns * xp.log(model_intensity + self.eps)
-            model_intensity + patterns * (xp.log(patterns + self.eps) - xp.log(model_intensity + self.eps) - 1.0)
-            #patterns - (model_intensity + self.offset) * xp.log(patterns)
+            model_intensity + self.eps + patterns * (
+                xp.log(patterns + self.eps) - xp.log(model_intensity + self.eps) - 1.0
+            )
         )).astype(exp_patterns.dtype)
         return (loss, state)
 

{phaserem-0.1 → phaserem-0.2}/phaser/engines/common/output.py

@@ -1,9 +1,9 @@
 from functools import partial
+import logging
 from pathlib import Path
 import typing as t
 
 import numpy
-from numpy.typing import NDArray
 import tifffile
 
 from phaser.utils.num import to_numpy, abs2, fft2, get_array_module
@@ -18,9 +18,10 @@ def output_images(state: ReconsState, out_dir: Path, options: SaveOptions):
         if ty not in _SAVE_FUNCS:
             raise ValueError(f"Unknown image type '{ty}'")
 
+        ext = options.plot_ext if ty in _PLOT_FUNCS else 'tiff'
         try:
             out_name = options.img_fmt.format(
-                type=ty, iter=state.iter,
+                type=ty, iter=state.iter, ext=ext
             )
             out_path = out_dir / out_name
         except KeyError as e:
@@ -122,7 +123,7 @@ def _save_object_phase(state: ReconsState, out_path: Path, options: SaveOptions,
     )
 
     if crop:
-        obj_phase = obj_phase[..., *state.object.sampling.get_region_crop()]
+        obj_phase = obj_phase[(Ellipsis, *state.object.sampling.get_region_crop())]
         mask = xp.ones(obj_phase.shape[-2:], dtype=numpy.bool_)
     else:
         # include whole image, but only scale based on ROI
@@ -158,7 +159,7 @@ def _save_object_mag(state: ReconsState, out_path: Path, options: SaveOptions, s
     xp = get_array_module(state.object.data)
     obj_mag = abs2(state.object.data)
     if crop:
-        obj_mag = obj_mag[..., *state.object.sampling.get_region_crop()]
+        obj_mag = obj_mag[(Ellipsis, *state.object.sampling.get_region_crop())]
         mask = numpy.ones(obj_mag.shape[-2:], dtype=numpy.bool_)
     else:
         # include whole image, but only scale based on ROI
@@ -176,6 +177,65 @@ def _save_object_mag(state: ReconsState, out_path: Path, options: SaveOptions, s
         w.write(obj_mag, **write_opts)
 
 
+def _plot_scan(state: ReconsState, out_path: Path, options: SaveOptions):
+    from matplotlib import pyplot
+    fig, ax = pyplot.subplots(figsize=(4, 4), dpi=options.plot_dpi, constrained_layout=True)
+
+    ax.set_aspect(1.)
+    [left, right, bottom, top] = state.object.sampling.mpl_extent()
+    ax.set_xlim(left, right)
+    ax.set_ylim(bottom, top)
+
+    scan = to_numpy(state.scan)
+    i = numpy.arange(scan[..., 0].size)
+    ax.scatter(scan[..., 1].ravel(), scan[..., 0].ravel(), c=i, cmap='plasma', s=0.5, edgecolors='none')
+
+    fig.savefig(out_path)
+    pyplot.close(fig)
+
+
+def _plot_tilt(state: ReconsState, out_path: Path, options: SaveOptions):
+    from matplotlib import pyplot
+
+    if state.tilt is None:
+        logger = logging.getLogger(__name__)
+        logger.warning("Tilt map (`state.tilt`) is missing, skipping `plot_tilt`")
+        return
+
+    fig, ax = pyplot.subplots(figsize=(4, 4), dpi=options.plot_dpi, constrained_layout=True)
+
+    ax.set_aspect(1.)
+    [left, right, bottom, top] = state.object.sampling.mpl_extent()
+    ax.set_xlim(left, right)
+    ax.set_ylim(bottom, top)
+
+    scan = to_numpy(state.scan)
+    tilt = to_numpy(state.tilt)
+    tilt = tilt[..., 1] + tilt[..., 0]*1.j
+    max_tilt = max(numpy.max(numpy.abs(tilt)), 1.0)  # at least 1 mrad
+    c = colorize_complex(tilt.ravel() / max_tilt, amp=True, rescale=False)
+    ax.scatter(scan[..., 1].ravel(), scan[..., 0].ravel(), c=c, s=0.5, edgecolors='none')
+
+    fig.draw_without_rendering()
+    trans = ax.transAxes + fig.transFigure.inverted()
+    legend_ax_max = trans.transform([0.95, 0.02])
+    legend_ax_size = (0.1, 0.1)
+    legend_ax = fig.add_axes((legend_ax_max[0] - legend_ax_size[0], legend_ax_max[1], *legend_ax_size), projection='polar')
+
+    legend_ax.set_rmax(max_tilt)  # type: ignore
+    legend_ax.set_theta_direction(-1)  # type: ignore
+    legend_ax.set_axis_off()
+
+    thetas = numpy.linspace(0., 2*numpy.pi, 70)
+    rs = numpy.concatenate([[0.0], numpy.geomspace(0.1, 1.0, 30)])
+    rr, tt = numpy.meshgrid(rs, thetas, indexing='ij')
+    c2 = colorize_complex(rr * numpy.exp(1.j * tt), rescale=False)
+    legend_ax.pcolormesh(tt, rr * max_tilt, c2)
+    legend_ax.text(-numpy.pi/2., max_tilt * 1.05, f"{max_tilt:.1f} mrad", ha='center', va='bottom', size='small')
+
+    fig.savefig(out_path)
+    pyplot.close(fig)
+
 
 _SAVE_FUNCS: t.Dict[str, t.Callable[[ReconsState, Path, SaveOptions], t.Any]] = {
     'probe': _save_probe,
@@ -186,4 +246,8 @@ _SAVE_FUNCS: t.Dict[str, t.Callable[[ReconsState, Path, SaveOptions], t.Any]] =
     'object_phase_sum': partial(_save_object_phase, stack=False),
     'object_mag_stack': partial(_save_object_mag, stack=True),
     'object_mag_sum': partial(_save_object_mag, stack=False),
-}
+    'scan': _plot_scan,
+    'tilt': _plot_tilt,
+}
+# save functions with special handling of file extensions
+_PLOT_FUNCS: t.Set[str] = {'scan', 'tilt'}

{phaserem-0.1 → phaserem-0.2}/phaser/engines/common/regularizers.py

@@ -6,14 +6,14 @@ import numpy
 from numpy.typing import NDArray
 
 from phaser.utils.num import (
-    get_array_module, get_scipy_module, Float,
-    jit, fft2, ifft2, abs2, xp_is_jax, to_real_dtype
+    get_array_module, get_scipy_module, Float, unstack,
+    jit, fft2, ifft2, abs2, xp_is_jax, to_real_dtype, to_numpy
 )
 from phaser.state import ReconsState
 from phaser.hooks.regularization import (
     ClampObjectAmplitudeProps, LimitProbeSupportProps,
-    RegularizeLayersProps, ObjLowPassProps,
-    CostRegularizerProps, TVRegularizerProps
+    RegularizeLayersProps, ObjLowPassProps, GaussianProps,
+    CostRegularizerProps, TVRegularizerProps, UnstructuredGaussianProps
 )
 
 
@@ -161,6 +161,37 @@ class ObjLowPass:
         return (sim, state)
 
 
+class ObjGaussian:
+    def __init__(self, args: None, props: GaussianProps):
+        self.sigma = props.sigma
+        self.weight = props.weight
+
+    def init_state(self, sim: ReconsState) -> NDArray[numpy.floating]:
+        samp = sim.object.sampling
+        xp = get_array_module(sim.object.data)
+        dtype = to_real_dtype(sim.object.data.dtype)
+
+        ky = xp.fft.fftfreq(samp.shape[0], samp.sampling[0]).astype(dtype)
+        kx = xp.fft.fftfreq(samp.shape[1], samp.sampling[1]).astype(dtype)
+        (ky, kx) = xp.meshgrid(ky, kx, indexing='ij')
+        k2 = ky**2 + kx**2
+
+        filt = xp.exp(-(2 * numpy.pi**2 * self.sigma**2) * k2).astype(dtype)
+        return (1. - self.weight * (1. - filt)).astype(dtype)
+
+    def apply_group(
+        self, group: NDArray[numpy.integer], sim: ReconsState, state: NDArray[numpy.floating]
+    ) -> t.Tuple[ReconsState, NDArray[numpy.floating]]:
+        return self.apply_iter(sim, state)
+
+    def apply_iter(
+        self, sim: ReconsState, state: NDArray[numpy.floating]
+    ) -> t.Tuple[ReconsState, NDArray[numpy.floating]]:
+        # TODO: should this be done in-place?
+        sim.object.data = ifft2(state * fft2(sim.object.data))
+        return (sim, state)
+
+
 class ObjL1:
     def __init__(self, args: None, props: CostRegularizerProps):
         self.cost: float = props.cost
@@ -395,6 +426,77 @@ class ProbeRecipTotalVariation:
         return (cost * cost_scale * self.cost, state)
 
 
+class UnstructuredGaussian:
+    def __init__(self, args: None, props: UnstructuredGaussianProps):
+        self.weight = props.weight
+        self.sigma = props.sigma
+        self.attr_path = props.attr_path
+
+    def init_state(self, sim: ReconsState) -> NDArray[numpy.floating]:
+        xp = get_array_module(sim.scan)
+        try:
+            self.getattr_nested(sim, self.attr_path)
+        except AttributeError as e:
+            raise AttributeError(f"Can't get path '{self.attr_path}' in reconstruction state") from e
+
+        # precompute Gaussian filter
+        obj_samp = sim.object.sampling
+        ky = xp.fft.fftfreq(obj_samp.shape[0], obj_samp.sampling[0])
+        kx = xp.fft.fftfreq(obj_samp.shape[1], obj_samp.sampling[1])
+        ky, kx = xp.meshgrid(ky, kx, indexing='ij')
+        k2 = ky**2 + kx**2
+        return xp.exp(- (2 * numpy.pi**2 * self.sigma**2) * k2)
+
+    def getattr_nested(self, obj: t.Any, attr_path: str) -> t.Any:
+        for attr in attr_path.split('.'):
+            obj = getattr(obj, attr)
+        return obj
+
+    def setattr_nested(self, obj: t.Any, attr_path: str, value: t.Any):
+        *parents, last = attr_path.split('.')
+        for attr in parents:
+            obj = getattr(obj, attr)
+        setattr(obj, last, value)
+
+    def apply_iter(self, sim: ReconsState, state: NDArray[numpy.floating]) -> t.Tuple[ReconsState, NDArray[numpy.floating]]:
+        from scipy.spatial import KDTree
+        obj_samp = sim.object.sampling
+        scan_flat = sim.scan.reshape(-1, 2)
+        scan_ndim = sim.scan.ndim - 1
+
+        attr = self.getattr_nested(sim, self.attr_path)
+        vals = t.cast(NDArray[numpy.inexact], getattr(attr, 'data', attr))  # Extract raw array
+        xp = get_array_module(vals)
+
+        is_complex = xp.iscomplexobj(vals)
+
+        if is_complex:
+            vals = xp.stack([vals.real, vals.imag], axis=-1)
+
+        tree = KDTree(to_numpy(scan_flat))
+        obj_pts = numpy.stack(obj_samp.grid(xp=numpy), axis=-1)
+        # val_img should be on GPU
+        val_img = vals.reshape((-1, *vals.shape[scan_ndim:]))[xp.array(tree.query(obj_pts)[1])]
+
+        # blur in Fourier space
+        # HACK: the transposes are so we FFT over the first two axes
+        val_img_blur = ifft2(fft2(val_img.T) * state.T).real.T
+
+        # then sample back to the scan positions
+        idxs = xp.round((scan_flat - obj_samp.corner) / obj_samp.sampling).astype(numpy.int_)
+        blur_vals = val_img_blur[unstack(idxs, axis=-1)].reshape(vals.shape)
+
+        new_vals = self.weight * blur_vals + (1 - self.weight) * vals
+
+        if is_complex:
+            new_vals = new_vals[..., 0] + new_vals[..., 1] * 1.j
+
+        # TODO: this should be expressed as immutable
+        self.setattr_nested(sim, self.attr_path, new_vals)
+
+        return (sim, state)
+
+
 def img_grad(img: numpy.ndarray) -> t.Tuple[numpy.ndarray, numpy.ndarray]:
     xp = get_array_module(img)
     return (

{phaserem-0.1 → phaserem-0.2}/phaser/engines/common/simulation.py

@@ -8,7 +8,7 @@ from typing_extensions import Self
 
 from phaser.utils.num import (
     get_array_module, to_real_dtype, to_complex_dtype,
-    fft2, ifft2, is_jax, to_numpy, block_until_ready,
+    fft2, ifft2, is_jax, to_numpy, block_until_ready, ufunc_outer
 )
 from phaser.utils.misc import FloatKey, jax_dataclass, create_compact_groupings, create_sparse_groupings, shuffled
 from phaser.utils.optics import fresnel_propagator, fourier_shift_filter
@@ -190,6 +190,38 @@ def make_propagators(state: ReconsState, bwlim_frac: t.Optional[float] = 2/3) ->
     )
 
 
+def tilt_propagators(
+    ky: NDArray[numpy.floating], kx: NDArray[numpy.floating],
+    state: ReconsState, 
+    props: t.Optional[NDArray[numpy.complexfloating]],  # shape: (Nz-1, Ny, Nx)
+    tilts: t.Optional[NDArray[numpy.floating]]          # shape: (..., 2), in mrad
+) -> t.Optional[NDArray[numpy.complexfloating]]:
+    """
+    Applies tilt and slice-dependent propagation phase shifts to props.
+    -------
+    NDArray[complex] or None
+        Tilted propagators of shape (n_layers-1, ..., Ny, Nx), or None if no slices.
+    """
+    if props is None:
+        return None
+    if tilts is None:
+        return props[:, None, ...]
+
+    xp = get_array_module(state.probe.data)
+    dtype = to_real_dtype(state.probe.data.dtype)
+    complex_dtype = to_complex_dtype(dtype)
+    delta_zs = state.object.thicknesses[:-1]
+
+    tilt_ramps = xp.exp(  # (n_layers-1, batch, Ny, Nx)
+        2.j * xp.pi * ufunc_outer(xp.multiply, delta_zs, (
+            ufunc_outer(xp.multiply, xp.tan(tilts[..., 0] * 1e-3), ky) +
+            ufunc_outer(xp.multiply, xp.tan(tilts[..., 1] * 1e-3), kx)
+        ))
+    )
+
+    return props[(slice(None), *(None,)*(tilts.ndim - 1), Ellipsis)] * tilt_ramps.astype(complex_dtype)
+
+
 @t.overload
 def cutout_group(
     ky: NDArray[numpy.floating], kx: NDArray[numpy.floating],
@@ -236,16 +268,20 @@ def slice_forwards(
     if props is None:
         return f(0, None, state)
 
-    n_slices = len(props) + 1
+    n_slices = len(props) + 1  # props shape: (Nz-1, batch, Ny, Nx)
 
     if is_jax(props):
         import jax
-        state = jax.lax.fori_loop(0, n_slices - 1, lambda slice_i, state: f(slice_i, props[slice_i], state), state, unroll=False)
+        def step_fn(carry, slice_i):
+            new_state = f(slice_i, props[slice_i], carry)
+            return new_state, None
+
+        state, _ = jax.lax.scan(step_fn, state, jax.numpy.arange(n_slices - 1))
         return f(n_slices - 1, None, state)
 
+    # fallback numpy mode
     for slice_i in range(n_slices - 1):
         state = f(slice_i, props[slice_i], state)
-
     return f(n_slices - 1, None, state)
 
 
@@ -267,4 +303,4 @@ def slice_backwards(
     for slice_i in range(n_slices - 1, 0, -1):
         state = f(slice_i, props[slice_i - 1], state)
 
-    return f(0, None, state)
+    return f(0, None, state)