xptycho 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
xptycho-0.1.0/LICENSE ADDED
@@ -0,0 +1,29 @@
1
+ BSD 3-Clause License
2
+
3
+ Copyright (c) 2026, Charles A Bouman
4
+ All rights reserved.
5
+
6
+ Redistribution and use in source and binary forms, with or without
7
+ modification, are permitted provided that the following conditions are met:
8
+
9
+ 1. Redistributions of source code must retain the above copyright notice, this
10
+ list of conditions and the following disclaimer.
11
+
12
+ 2. Redistributions in binary form must reproduce the above copyright notice,
13
+ this list of conditions and the following disclaimer in the documentation
14
+ and/or other materials provided with the distribution.
15
+
16
+ 3. Neither the name of the copyright holder nor the names of its
17
+ contributors may be used to endorse or promote products derived from
18
+ this software without specific prior written permission.
19
+
20
+ THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
21
+ AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
22
+ IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
23
+ DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
24
+ FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
25
+ DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
26
+ SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
27
+ CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
28
+ OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
29
+ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
xptycho-0.1.0/PKG-INFO ADDED
@@ -0,0 +1,63 @@
1
+ Metadata-Version: 2.4
2
+ Name: xptycho
3
+ Version: 0.1.0
4
+ Summary: Ptychographic reconstruction with PMACE in PyTorch
5
+ Author: Charles A. Bouman, Brendt Wohlberg, Gregery T. Buzzard
6
+ License: BSD-3-Clause
7
+ Project-URL: Homepage, https://github.com/cabouman/xptycho
8
+ Project-URL: Documentation, https://xptycho.readthedocs.io
9
+ Requires-Python: >=3.11
10
+ Description-Content-Type: text/markdown
11
+ License-File: LICENSE
12
+ Requires-Dist: numpy
13
+ Requires-Dist: scipy
14
+ Requires-Dist: torch>=2.2
15
+ Requires-Dist: h5py
16
+ Requires-Dist: tifffile
17
+ Requires-Dist: matplotlib
18
+ Provides-Extra: test
19
+ Requires-Dist: pytest; extra == "test"
20
+ Provides-Extra: docs
21
+ Requires-Dist: sphinx; extra == "docs"
22
+ Requires-Dist: sphinx_book_theme; extra == "docs"
23
+ Requires-Dist: sphinx-copybutton; extra == "docs"
24
+ Requires-Dist: sphinxext-opengraph; extra == "docs"
25
+ Requires-Dist: sphinx-design; extra == "docs"
26
+ Dynamic: license-file
27
+
28
+ # xptycho
29
+
30
+ xptycho: ptychographic reconstruction with projected multi-agent consensus
31
+ equilibrium (PMACE) using [PyTorch](https://pytorch.org/).
32
+
33
+ Features include:
34
+ * Reconstruction of the complex object image, magnitude and phase, from far-field diffraction frames.
35
+ * Estimation of the probe together with the object, with one or more probe modes.
36
+ * Refinement of the probe positions.
37
+ * Preprocessing of raw detector frames: dark subtraction, outlier removal, centering, and cropping.
38
+ * One HDF5 file format for scans, samples, and reconstructions.
39
+ * Demos on simulated data and on measured data, from raw file to image.
40
+ * Seamless operation on 1 or more GPUs, Mac MPS, or CPU.
41
+
42
+ Full documentation at [https://xptycho.readthedocs.io/](https://xptycho.readthedocs.io/)
43
+
44
+ Design pages at [https://cabouman.github.io/xptycho/](https://cabouman.github.io/xptycho/)
45
+
46
+ Install:
47
+ ```bash
48
+ pip install xptycho
49
+ ```
50
+
51
+ Reconstruct in a few lines:
52
+ ```python
53
+ import xptycho as xpt
54
+ scan = xpt.Scan.load('scan.h5')
55
+ model = xpt.PtychoModel.from_scan(scan)
56
+ recon = model.recon(scan)
57
+ xpt.view_sample(recon)
58
+ ```
59
+
60
+ xptycho implements the PMACE method of Qiuchen Zhai, Gregery T. Buzzard,
61
+ Kevin Mertes, Brendt Wohlberg, and Charles A. Bouman. For the papers to
62
+ cite, the source of the demo data, and the funding support, see
63
+ [Credits](https://xptycho.readthedocs.io/en/latest/credits.html).
@@ -0,0 +1,36 @@
1
+ # xptycho
2
+
3
+ xptycho: ptychographic reconstruction with projected multi-agent consensus
4
+ equilibrium (PMACE) using [PyTorch](https://pytorch.org/).
5
+
6
+ Features include:
7
+ * Reconstruction of the complex object image, magnitude and phase, from far-field diffraction frames.
8
+ * Estimation of the probe together with the object, with one or more probe modes.
9
+ * Refinement of the probe positions.
10
+ * Preprocessing of raw detector frames: dark subtraction, outlier removal, centering, and cropping.
11
+ * One HDF5 file format for scans, samples, and reconstructions.
12
+ * Demos on simulated data and on measured data, from raw file to image.
13
+ * Seamless operation on 1 or more GPUs, Mac MPS, or CPU.
14
+
15
+ Full documentation at [https://xptycho.readthedocs.io/](https://xptycho.readthedocs.io/)
16
+
17
+ Design pages at [https://cabouman.github.io/xptycho/](https://cabouman.github.io/xptycho/)
18
+
19
+ Install:
20
+ ```bash
21
+ pip install xptycho
22
+ ```
23
+
24
+ Reconstruct in a few lines:
25
+ ```python
26
+ import xptycho as xpt
27
+ scan = xpt.Scan.load('scan.h5')
28
+ model = xpt.PtychoModel.from_scan(scan)
29
+ recon = model.recon(scan)
30
+ xpt.view_sample(recon)
31
+ ```
32
+
33
+ xptycho implements the PMACE method of Qiuchen Zhai, Gregery T. Buzzard,
34
+ Kevin Mertes, Brendt Wohlberg, and Charles A. Bouman. For the papers to
35
+ cite, the source of the demo data, and the funding support, see
36
+ [Credits](https://xptycho.readthedocs.io/en/latest/credits.html).
@@ -0,0 +1,38 @@
1
+ [build-system]
2
+ requires = ["setuptools>=64"]
3
+ build-backend = "setuptools.build_meta"
4
+
5
+ [project]
6
+ name = "xptycho"
7
+ description = "Ptychographic reconstruction with PMACE in PyTorch"
8
+ readme = "README.md"
9
+ license = { text = "BSD-3-Clause" }
10
+ requires-python = ">=3.11"
11
+ dynamic = ["version"]
12
+ authors = [
13
+ { name = "Charles A. Bouman" },
14
+ { name = "Brendt Wohlberg" },
15
+ { name = "Gregery T. Buzzard" },
16
+ ]
17
+ dependencies = [
18
+ "numpy",
19
+ "scipy",
20
+ "torch>=2.2",
21
+ "h5py",
22
+ "tifffile",
23
+ "matplotlib",
24
+ ]
25
+
26
+ [project.optional-dependencies]
27
+ test = ["pytest"]
28
+ docs = ["sphinx", "sphinx_book_theme", "sphinx-copybutton", "sphinxext-opengraph", "sphinx-design"]
29
+
30
+ [project.urls]
31
+ Homepage = "https://github.com/cabouman/xptycho"
32
+ Documentation = "https://xptycho.readthedocs.io"
33
+
34
+ [tool.setuptools.packages.find]
35
+ include = ["xptycho*"]
36
+
37
+ [tool.setuptools.dynamic]
38
+ version = { attr = "xptycho.__version__" }
@@ -0,0 +1,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+
@@ -0,0 +1,34 @@
1
+ """A scan and a sample are groups of one HDF5 file and read back unchanged."""
2
+ import numpy as np
3
+
4
+ import xptycho as xpt
5
+ from tiny import PIXEL_PITCH, make_model, make_object, make_probe
6
+
7
+
8
+ def test_one_file_holds_a_scan_and_a_sample(tmp_path):
9
+ path = str(tmp_path / 'tiny.h5')
10
+ truth = xpt.Sample(make_object(), make_probe(), PIXEL_PITCH, name='tiny')
11
+ model = make_model()
12
+ scan = model.simulate(truth)
13
+ truth.save(path)
14
+ scan.save(path) # the same file; the sample is kept
15
+
16
+ sample, scan_back = xpt.Sample.load(path), xpt.Scan.load(path)
17
+ assert np.array_equal(sample.object, truth.object) and np.array_equal(sample.probe, truth.probe)
18
+ assert sample.pixel_pitch == PIXEL_PITCH and sample.name == 'tiny' and sample.run is None
19
+ assert np.array_equal(scan_back.frames, scan.frames) and scan_back.det_pixel_pitch == scan.det_pixel_pitch
20
+
21
+
22
+ def test_a_reconstruction_keeps_the_record_of_its_run(tmp_path):
23
+ path = str(tmp_path / 'recon.h5')
24
+ model = make_model()
25
+ probe = make_probe()
26
+ scan = model.simulate(xpt.Sample(make_object(), probe, PIXEL_PITCH))
27
+ recon = model.recon(scan, probe=probe, iterations=3, verbose=0)
28
+ recon.save(path)
29
+ back = xpt.Sample.load(path)
30
+ assert np.array_equal(back.object, recon.object) and back.origin == recon.origin
31
+ assert back.run.iterations == 3 and back.run.data_error == recon.run.data_error
32
+ assert np.array_equal(back.run.coverage, recon.run.coverage)
33
+ assert [row['name'] for row in back.run.parameters] == [row['name'] for row in recon.run.parameters]
34
+ assert np.array_equal(back.scanned_region(), recon.scanned_region())
@@ -0,0 +1,6 @@
1
+ import xptycho
2
+
3
+
4
+ def test_version():
5
+ assert isinstance(xptycho.__version__, str)
6
+ assert xptycho.__version__
@@ -0,0 +1,32 @@
1
+ """The preprocessing steps on small arrays with known answers."""
2
+ import numpy as np
3
+
4
+ import xptycho.preprocess as xpp
5
+
6
+
7
+ def test_subtract_dark_clips_at_zero():
8
+ frames = np.array([[[5.0, 1.0]], [[3.0, 9.0]]])
9
+ dark = np.array([[[2.0, 2.0]], [[4.0, 4.0]]]) # mean dark is 3
10
+ assert np.array_equal(xpp.subtract_dark(frames, dark), [[[2.0, 0.0]], [[0.0, 6.0]]])
11
+
12
+
13
+ def test_find_outlier_frames():
14
+ frames = np.ones((20, 4, 4))
15
+ frames[7] *= 3
16
+ outlier = xpp.find_outlier_frames(frames, threshold=2)
17
+ assert outlier.tolist() == [j == 7 for j in range(20)]
18
+
19
+
20
+ def test_center_and_crop_put_the_pattern_at_the_frame_center():
21
+ frames = np.zeros((3, 21, 21))
22
+ frames[:, 12, 8] = 1.0 # every pattern at row 12, column 8
23
+ center = xpp.diffraction_center(frames)
24
+ assert center == (12.0, 8.0)
25
+ cropped = xpp.crop_frames(frames, center, 8)
26
+ assert cropped.shape == (3, 8, 8) and np.all(cropped[:, 4, 4] == 1) and cropped.sum() == 3
27
+
28
+
29
+ def test_tukey_window():
30
+ window = xpp.tukey_window(64, 0.5)
31
+ assert window.shape == (64, 64) and window[32, 32] == 1 and window[0, 0] == 0
32
+ assert np.allclose(window, window.T) and window.min() >= 0 and window.max() == 1
@@ -0,0 +1,22 @@
1
+ """Position refinement finds whole-pixel errors in the positions."""
2
+ import numpy as np
3
+
4
+ import xptycho as xpt
5
+ from tiny import PIXEL_PITCH, make_model, make_object, make_probe
6
+
7
+
8
+ def test_refine_positions_recovers_shifts():
9
+ truth, probe = make_object(), make_probe()
10
+ model = make_model()
11
+ scan = model.simulate(xpt.Sample(truth, probe, PIXEL_PITCH))
12
+ true_positions = model.probe_positions.copy()
13
+
14
+ wrong = true_positions.copy()
15
+ wrong[10] += (PIXEL_PITCH, 0) # one pixel down
16
+ wrong[30] += (0, -PIXEL_PITCH) # one pixel left
17
+ model.set_params(probe_positions=wrong)
18
+ refined, misfit = model.refine_probe_positions(scan, truth, probe, max_shift=1)
19
+
20
+ assert np.allclose(refined, true_positions, atol=PIXEL_PITCH / 100)
21
+ assert misfit[10] > 0.5 and misfit[30] > 0.5
22
+ assert misfit[0] == 0
@@ -0,0 +1,44 @@
1
+ """The new code against the old ptycho_pmace code on the tiny problem,
2
+ with a known probe, on one device and divided among several."""
3
+ import json
4
+ import os
5
+
6
+ import numpy as np
7
+ import pytest
8
+
9
+ import xptycho as xpt
10
+ from tiny import PIXEL_PITCH, make_model, make_object, make_probe
11
+
12
+ BASELINE = json.load(open(os.path.join(os.path.dirname(__file__), 'data', 'tiny_baseline.json')))
13
+ TOLERANCE = 1e-4 # relative; single precision agrees to about 1e-6
14
+
15
+
16
+ def numbers(recon, truth):
17
+ est = xpt.match_scale(recon.object, truth)
18
+ return dict(frobenius=np.linalg.norm(est), sum_real=est.real.sum(), sum_imag=est.imag.sum(),
19
+ max_abs=np.abs(est).max(), nrmse_obj=xpt.nrmse(recon.object, truth),
20
+ nrmse_meas=recon.run.data_error[-1])
21
+
22
+
23
+ @pytest.mark.parametrize('devices', [('cpu',), ('cpu', 'cpu'), ('cpu',) * 5])
24
+ @pytest.mark.parametrize('iterations', [1, 5, 20])
25
+ def test_matches_old_code(devices, iterations):
26
+ truth, probe = make_object(), make_probe()
27
+ model = make_model(devices, object_data_fit=BASELINE['obj_data_fit_prm'], relaxation=BASELINE['rho'],
28
+ probe_weight_exponent=BASELINE['probe_exp'])
29
+ scan = model.simulate(xpt.Sample(truth, probe, PIXEL_PITCH))
30
+ recon = model.recon(scan, probe=probe, iterations=iterations, verbose=0)
31
+ mine, old = numbers(recon, truth), BASELINE['iterates'][str(iterations)]
32
+ for name in mine:
33
+ assert mine[name] == pytest.approx(old[name], rel=TOLERANCE), name
34
+
35
+
36
+ def test_devices_agree():
37
+ """Five devices put block boundaries inside scan lines."""
38
+ truth, probe = make_object(), make_probe()
39
+ results = []
40
+ for devices in [('cpu',), ('cpu',) * 5]:
41
+ model = make_model(devices)
42
+ scan = model.simulate(xpt.Sample(truth, probe, PIXEL_PITCH))
43
+ results.append(model.recon(scan, probe=probe, iterations=10, verbose=0).object)
44
+ assert np.abs(results[0] - results[1]).max() < 1e-4 * np.abs(results[0]).max()
@@ -0,0 +1,70 @@
1
+ """The new code against the old ptycho_pmace code on the tiny problem
2
+ with two probe modes, the probe estimated and the second mode added at
3
+ iteration 2, on one device and divided among several."""
4
+ import json
5
+ import os
6
+
7
+ import numpy as np
8
+ import pytest
9
+
10
+ import xptycho as xpt
11
+ from tiny import PIXEL_PITCH, make_model, make_object, make_probe, make_second_mode, numbers
12
+
13
+ BASELINE = json.load(open(os.path.join(os.path.dirname(__file__), 'data', 'tiny_blind_baseline.json')))
14
+ TOLERANCE = 1e-3 # relative; the first iterations agree to about 1e-5
15
+
16
+
17
+ def setup(devices):
18
+ truth = make_object()
19
+ modes = np.stack([make_probe(), make_second_mode(scale=BASELINE['second_mode_scale'])])
20
+ model = make_model(devices, object_data_fit=BASELINE['obj_data_fit_prm'],
21
+ probe_data_fit=BASELINE['probe_data_fit_prm'], relaxation=BASELINE['rho'],
22
+ probe_weight_exponent=BASELINE['probe_exp'], mode_schedule=BASELINE['add_mode'],
23
+ mode_energy_fraction=BASELINE['energy_ratio'])
24
+ model.set_params(probe_fresnel_radius_pixels=np.sqrt(model.wavelength * BASELINE['propagation_dist'])
25
+ / model.sample_pixel_pitch)
26
+ model.num_probe_modes = 2
27
+ return model, model.simulate(xpt.Sample(truth, modes, PIXEL_PITCH)), truth
28
+
29
+
30
+ def close(mine, old):
31
+ for name in old:
32
+ assert mine[name] == pytest.approx(old[name], rel=TOLERANCE, abs=1e-4), name
33
+
34
+
35
+ def test_start_matches_old_code():
36
+ model, scan, _ = setup(('cpu',))
37
+ probe = model.initial_probe(scan)
38
+ close(numbers(probe[0]), BASELINE['init_probe'])
39
+ close(numbers(model.initial_object(scan, probe)), BASELINE['init_obj'])
40
+
41
+
42
+ @pytest.mark.parametrize('devices', [('cpu',), ('cpu',) * 3])
43
+ @pytest.mark.parametrize('iterations', [1, 3, 6])
44
+ def test_matches_old_code(devices, iterations):
45
+ model, scan, truth = setup(devices)
46
+ recon = model.recon(scan, iterations=iterations, verbose=0)
47
+ old = BASELINE['iterates'][str(iterations)]
48
+ close(numbers(xpt.match_scale(recon.object, truth)), old['object'])
49
+ assert len(recon.probe) == len(old['probe'])
50
+ for mine, theirs in zip(recon.probe, old['probe']):
51
+ close(numbers(mine), theirs)
52
+ assert recon.run.data_error[-1] == pytest.approx(old['nrmse_meas'], rel=TOLERANCE)
53
+
54
+
55
+ def test_orthogonalize_modes_keeps_the_fit():
56
+ """Making the modes orthogonal changes the modes, not the intensities
57
+ they predict."""
58
+ model, scan, _ = setup(('cpu',))
59
+ plain = model.recon(scan, iterations=BASELINE['add_mode'][0] - 1, verbose=0) # stop before the addition
60
+
61
+ def after_addition(orthogonalize):
62
+ model.set_params(orthogonalize_modes=orthogonalize)
63
+ return model.recon(scan, init=plain, iterations=1, verbose=0)
64
+
65
+ model.set_params(mode_schedule=[1]) # add the mode in the first iteration of the continued run
66
+ off, on = after_addition(False), after_addition(True)
67
+ inner = abs(np.vdot(on.probe[0], on.probe[1])) / (np.linalg.norm(on.probe[0]) * np.linalg.norm(on.probe[1]))
68
+ inner_off = abs(np.vdot(off.probe[0], off.probe[1])) / (np.linalg.norm(off.probe[0]) * np.linalg.norm(off.probe[1]))
69
+ assert len(on.probe) == 2 and inner < inner_off
70
+ assert np.linalg.norm(on.probe[0]) >= np.linalg.norm(on.probe[1])
@@ -0,0 +1,24 @@
1
+ """xptycho: ptychographic reconstruction with PMACE in PyTorch.
2
+
3
+ xptycho reconstructs the complex transmittance image of a thin object
4
+ from a ptychographic scan: far-field diffraction patterns recorded
5
+ while a focused X-ray probe steps across the object in overlapping
6
+ positions. A :class:`Scan` holds the measurement, a :class:`PtychoModel`
7
+ holds the parameters and the forward model, and
8
+ :meth:`PtychoModel.recon` returns a :class:`Sample`: the object and the
9
+ probe.
10
+ """
11
+
12
+ __version__ = '0.1.0'
13
+
14
+ from .scan import Scan, energy_to_wavelength
15
+ from .model import PtychoModel
16
+ from .sample import Sample, RunRecord
17
+ from .metrics import nrmse, match_scale
18
+ from .view import view_scan, view_sample, save_figures
19
+ from .datasets import download, scan_positions
20
+ from . import operators
21
+ from . import preprocess
22
+
23
+ __all__ = ['Scan', 'PtychoModel', 'Sample', 'RunRecord', 'download', 'scan_positions', 'view_scan', 'view_sample', 'save_figures', 'nrmse',
24
+ 'match_scale', 'energy_to_wavelength', 'operators', 'preprocess']
@@ -0,0 +1,124 @@
1
+ """How a reconstruction is divided among devices, and the two operations
2
+ that move data between them. See the Computing page of the design.
3
+
4
+ Three terms:
5
+
6
+ block
7
+ The positions are sorted by the first row of their patch and cut into
8
+ one block per device, of nearly equal count.
9
+ band
10
+ The rows of the object image are cut into one band per device, of
11
+ nearly equal height. A device owns its band.
12
+ window
13
+ The range of image rows that the patches of a device's block touch.
14
+ """
15
+ import numpy as np
16
+ import torch
17
+
18
+
19
+ class Layout:
20
+ """The blocks, bands, and windows of one run.
21
+
22
+ Args:
23
+ devices (list of torch.device): one entry per shard. The same
24
+ device may appear more than once, which is how the tests run
25
+ several shards on one CPU.
26
+ starts (ndarray): int ``(J, 2)``, the row and column of the first
27
+ pixel of each patch on the object grid.
28
+ frame_size (int): the patches are ``frame_size`` square.
29
+ object_shape (tuple of int): ``(rows, cols)`` of the object grid.
30
+
31
+ Attributes:
32
+ blocks (list of ndarray): for each device, the indices of its
33
+ positions in the scan's own order.
34
+ bands (list of tuple): for each device, the ``(first, stop)`` rows
35
+ it owns.
36
+ windows (list of tuple): for each device, the ``(first, stop)``
37
+ rows its patches touch.
38
+ """
39
+
40
+ def __init__(self, devices, starts, frame_size, object_shape):
41
+ self.devices = [torch.device(d) for d in devices]
42
+ n = len(self.devices)
43
+ rows = int(object_shape[0])
44
+ if n > len(starts):
45
+ raise ValueError('{} devices were asked for but the scan has only {} positions'.format(n, len(starts)))
46
+ if n > rows:
47
+ raise ValueError('{} devices were asked for but the object has only {} rows'.format(n, rows))
48
+ order = np.argsort(starts[:, 0], kind='stable')
49
+ self.blocks = np.array_split(order, n)
50
+ edges = np.linspace(0, rows, n + 1).round().astype(int)
51
+ self.bands = [(int(edges[g]), int(edges[g + 1])) for g in range(n)]
52
+ self.windows = [(int(starts[b, 0].min()), int(starts[b, 0].max()) + frame_size) for b in self.blocks]
53
+ self.object_shape = (rows, int(object_shape[1]))
54
+
55
+ def __len__(self):
56
+ return len(self.devices)
57
+
58
+ def _overlap(self, window, band):
59
+ first, stop = max(window[0], band[0]), min(window[1], band[1])
60
+ return (first, stop) if stop > first else None
61
+
62
+ def sum_to_owner(self, parts):
63
+ """Add the devices' partial sums onto the owners of the rows.
64
+
65
+ Args:
66
+ parts (list of Tensor): for each device, its partial sum over
67
+ its window, shape ``(..., window rows, cols)``.
68
+
69
+ Returns:
70
+ list of Tensor: for each device, the complete sum over its
71
+ band, shape ``(..., band rows, cols)``. The pieces are added
72
+ in device order.
73
+ """
74
+ out = []
75
+ for g, band in enumerate(self.bands):
76
+ lead = parts[g].shape[:-2]
77
+ total = torch.zeros(lead + (band[1] - band[0], self.object_shape[1]),
78
+ dtype=parts[g].dtype, device=self.devices[g])
79
+ for h, window in enumerate(self.windows):
80
+ rows = self._overlap(window, band)
81
+ if rows is not None:
82
+ piece = parts[h][..., rows[0] - window[0]:rows[1] - window[0], :]
83
+ total[..., rows[0] - band[0]:rows[1] - band[0], :] += piece.to(self.devices[g])
84
+ out.append(total)
85
+ return out
86
+
87
+ def copy_to_window(self, bands):
88
+ """Give each device the rows of its window, from their owners.
89
+
90
+ Args:
91
+ bands (list of Tensor): for each device, an array over its band.
92
+
93
+ Returns:
94
+ list of Tensor: for each device, the same array over its window.
95
+ """
96
+ out = []
97
+ for g, window in enumerate(self.windows):
98
+ pieces = []
99
+ for h, band in enumerate(self.bands):
100
+ rows = self._overlap(window, band)
101
+ if rows is not None:
102
+ pieces.append(bands[h][..., rows[0] - band[0]:rows[1] - band[0], :].to(self.devices[g]))
103
+ out.append(torch.cat(pieces, dim=-2))
104
+ return out
105
+
106
+ def sum_small(self, parts):
107
+ """Add small partial sums on the first device, in device order."""
108
+ total = parts[0].clone()
109
+ for part in parts[1:]:
110
+ total += part.to(self.devices[0])
111
+ return total
112
+
113
+ def broadcast(self, x):
114
+ """A copy of a small array on every device."""
115
+ return [x.to(d) for d in self.devices]
116
+
117
+ def gather(self, bands):
118
+ """Join the bands into one image on the host, as numpy."""
119
+ return np.concatenate([b.cpu().numpy() for b in bands], axis=-2)
120
+
121
+ def split(self, image):
122
+ """Cut a host image into bands and place each on its owner."""
123
+ image = torch.as_tensor(image)
124
+ return [image[..., first:stop, :].to(d) for (first, stop), d in zip(self.bands, self.devices)]
@@ -0,0 +1,57 @@
1
+ """Tools for demos: a file download and simple scan patterns."""
2
+ import os
3
+ import subprocess
4
+ import urllib.error
5
+ import urllib.request
6
+
7
+ import numpy as np
8
+
9
+
10
+ def download(url, directory):
11
+ """Download a file into a directory, unless it is already there.
12
+
13
+ Args:
14
+ url (str): the address of the file.
15
+ directory (str): the local directory. It is created if needed.
16
+
17
+ Returns:
18
+ str: the path of the local file, named as in the address.
19
+ """
20
+ path = os.path.join(directory, os.path.basename(url))
21
+ if os.path.isfile(path):
22
+ return path
23
+ os.makedirs(directory, exist_ok=True)
24
+ print('downloading {} to {}'.format(url, path))
25
+ partial = path + '.part'
26
+ try:
27
+ urllib.request.urlretrieve(url, partial)
28
+ except urllib.error.URLError:
29
+ # Some servers refuse Python's downloader or do not send their
30
+ # intermediate certificate, which Python cannot verify; curl works.
31
+ result = subprocess.run(['curl', '-L', '--fail', '-sS', '-o', partial, url], capture_output=True, text=True)
32
+ if result.returncode != 0:
33
+ raise RuntimeError('could not download {}: {}'.format(url, result.stderr.strip()))
34
+ os.replace(partial, path)
35
+ return path
36
+
37
+
38
+ def scan_positions(grid, step, max_offset=0.0, seed=0):
39
+ """The positions of a rectangular scan, with a random offset at each.
40
+
41
+ Args:
42
+ grid (tuple of int): the number of positions along rows and columns.
43
+ step (float): meters between neighboring positions.
44
+ max_offset (float, optional): meters. Each position is moved by a
45
+ uniform random amount up to this along each axis. Defaults to 0.
46
+ seed (int, optional): the random seed of the offsets.
47
+
48
+ Returns:
49
+ ndarray: ``(grid[0] * grid[1], 2)``, row then column, in meters,
50
+ centered on zero.
51
+ """
52
+ rows = (np.arange(grid[0]) - (grid[0] - 1) / 2) * step
53
+ cols = (np.arange(grid[1]) - (grid[1] - 1) / 2) * step
54
+ probe_positions = np.stack(np.meshgrid(rows, cols, indexing='ij'), axis=-1).reshape(-1, 2)
55
+ if max_offset:
56
+ probe_positions = probe_positions + np.random.default_rng(seed).uniform(-max_offset, max_offset, probe_positions.shape)
57
+ return probe_positions
@@ -0,0 +1,39 @@
1
+ """Error metrics."""
2
+ import numpy as np
3
+
4
+
5
+ def match_scale(image, truth, region=None):
6
+ """Return ``image`` times the one complex number that brings it
7
+ closest to ``truth``. A ptychographic reconstruction is determined
8
+ only up to a complex constant, so this is applied before comparing.
9
+
10
+ Args:
11
+ image (ndarray): complex, the reconstruction.
12
+ truth (ndarray): complex, the same shape.
13
+ region (ndarray, optional): only pixels where this is positive
14
+ are used; the result is zero elsewhere.
15
+ """
16
+ if region is not None:
17
+ mask = np.asarray(region) > 0
18
+ image, truth = image * mask, truth * mask
19
+ scale = np.sum(np.conj(image) * truth) / np.sum(np.abs(image) ** 2)
20
+ return (scale * image).astype(np.complex64)
21
+
22
+
23
+ def nrmse(image, truth, region=None):
24
+ """The normalized root mean square error between a reconstruction and
25
+ its truth, after :func:`match_scale`.
26
+
27
+ Args:
28
+ image (ndarray): complex, the reconstruction.
29
+ truth (ndarray): complex, the same shape.
30
+ region (ndarray, optional): only pixels where this is positive
31
+ count, for example a reconstruction's ``scanned_region()``.
32
+
33
+ Returns:
34
+ float: ``||c * image - truth|| / ||truth||`` over the counted pixels.
35
+ """
36
+ scaled = match_scale(image, truth, region)
37
+ if region is not None:
38
+ truth = truth * (np.asarray(region) > 0)
39
+ return float(np.sqrt(np.sum(np.abs(scaled - truth) ** 2) / np.sum(np.abs(truth) ** 2)))