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 +29 -0
- xptycho-0.1.0/PKG-INFO +63 -0
- xptycho-0.1.0/README.md +36 -0
- xptycho-0.1.0/pyproject.toml +38 -0
- xptycho-0.1.0/setup.cfg +4 -0
- xptycho-0.1.0/tests/test_files.py +34 -0
- xptycho-0.1.0/tests/test_import.py +6 -0
- xptycho-0.1.0/tests/test_preprocess.py +32 -0
- xptycho-0.1.0/tests/test_refine_positions.py +22 -0
- xptycho-0.1.0/tests/test_tiny_baseline.py +44 -0
- xptycho-0.1.0/tests/test_tiny_blind.py +70 -0
- xptycho-0.1.0/xptycho/__init__.py +24 -0
- xptycho-0.1.0/xptycho/_sharding.py +124 -0
- xptycho-0.1.0/xptycho/datasets.py +57 -0
- xptycho-0.1.0/xptycho/metrics.py +39 -0
- xptycho-0.1.0/xptycho/model.py +484 -0
- xptycho-0.1.0/xptycho/operators.py +80 -0
- xptycho-0.1.0/xptycho/pmace.py +296 -0
- xptycho-0.1.0/xptycho/preprocess/__init__.py +12 -0
- xptycho-0.1.0/xptycho/preprocess/cxi.py +47 -0
- xptycho-0.1.0/xptycho/preprocess/utilities.py +99 -0
- xptycho-0.1.0/xptycho/sample.py +172 -0
- xptycho-0.1.0/xptycho/scan.py +114 -0
- xptycho-0.1.0/xptycho/view.py +249 -0
- xptycho-0.1.0/xptycho.egg-info/PKG-INFO +63 -0
- xptycho-0.1.0/xptycho.egg-info/SOURCES.txt +27 -0
- xptycho-0.1.0/xptycho.egg-info/dependency_links.txt +1 -0
- xptycho-0.1.0/xptycho.egg-info/requires.txt +16 -0
- xptycho-0.1.0/xptycho.egg-info/top_level.txt +1 -0
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).
|
xptycho-0.1.0/README.md
ADDED
|
@@ -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__" }
|
xptycho-0.1.0/setup.cfg
ADDED
|
@@ -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,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)))
|