getframes 2.2.0__tar.gz → 2.4.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.
- {getframes-2.2.0 → getframes-2.4.0}/CHANGELOG.md +91 -1
- {getframes-2.2.0 → getframes-2.4.0}/PKG-INFO +7 -5
- {getframes-2.2.0 → getframes-2.4.0}/README.md +4 -4
- {getframes-2.2.0 → getframes-2.4.0}/pyproject.toml +3 -0
- {getframes-2.2.0 → getframes-2.4.0}/src/getframes/__about__.py +1 -1
- {getframes-2.2.0 → getframes-2.4.0}/src/getframes/__init__.py +2 -1
- getframes-2.4.0/src/getframes/backend.py +391 -0
- {getframes-2.2.0 → getframes-2.4.0}/src/getframes/camera.py +48 -27
- {getframes-2.2.0 → getframes-2.4.0}/src/getframes/cli.py +12 -7
- {getframes-2.2.0 → getframes-2.4.0}/src/getframes/dataset.py +3 -2
- {getframes-2.2.0 → getframes-2.4.0}/src/getframes/frame.py +9 -6
- {getframes-2.2.0 → getframes-2.4.0}/src/getframes/noise.py +30 -19
- {getframes-2.2.0 → getframes-2.4.0}/src/getframes/scene/optics.py +3 -2
- {getframes-2.2.0 → getframes-2.4.0}/src/getframes/scene/psf.py +47 -34
- {getframes-2.2.0 → getframes-2.4.0}/src/getframes/scene/scene.py +19 -5
- {getframes-2.2.0 → getframes-2.4.0}/src/getframes/scene/thermal.py +2 -5
- getframes-2.4.0/tests/test_backend.py +242 -0
- {getframes-2.2.0 → getframes-2.4.0}/tests/test_cli.py +29 -0
- getframes-2.4.0/tests/test_conformance.py +149 -0
- {getframes-2.2.0 → getframes-2.4.0}/tests/test_gpu.py +58 -1
- {getframes-2.2.0 → getframes-2.4.0}/tests/test_scene.py +23 -0
- getframes-2.2.0/src/getframes/backend.py +0 -150
- {getframes-2.2.0 → getframes-2.4.0}/.gitignore +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/LICENSE +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/benchmarks/__init__.py +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/benchmarks/bench_detector_workspace.py +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/benchmarks/bench_devices.py +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/benchmarks/bench_fixed_map_dtype.py +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/benchmarks/detector-workspace-results.json +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/benchmarks/device-results.json +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/benchmarks/device-results.md +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/benchmarks/fixed-map-dtype-results.json +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/benchmarks/render_device_table.py +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/benchmarks/run.py +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/examples/01_basic_dark_frame.py +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/examples/02_custom_camera.py +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/examples/03_master_dark.py +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/examples/04_browse_presets.py +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/examples/05_visualise.py +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/examples/06_photon_transfer_curve.py +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/examples/07_star_field_exposure.py +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/examples/08_ao_limiting_magnitude.py +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/examples/09_transit_photometry.py +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/examples/10_detector_realism.py +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/examples/11_radiometry_and_ir.py +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/examples/12_ml_dataset.py +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/examples/13_crowded_field.py +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/examples/14_keck_lgs_ttf_trade_study.ipynb +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/examples/15_detector_characterization.py +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/examples/16_detector_showcase.py +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/examples/README.md +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/examples/_common.py +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/examples/detector_showcase.webp +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/src/getframes/analysis/__init__.py +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/src/getframes/analysis/apertures.py +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/src/getframes/analysis/characterize.py +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/src/getframes/analysis/nondestructive.py +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/src/getframes/analysis/ptc.py +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/src/getframes/calibrate.py +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/src/getframes/config.py +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/src/getframes/observation.py +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/src/getframes/presets/__init__.py +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/src/getframes/presets/data/__init__.py +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/src/getframes/presets/data/andor_cb1_0_5mp.toml +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/src/getframes/presets/data/andor_ikon_m934.toml +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/src/getframes/presets/data/andor_ixon_ultra_888.toml +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/src/getframes/presets/data/andor_marana_4_2b_11.toml +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/src/getframes/presets/data/andor_ocam2k.toml +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/src/getframes/presets/data/first_light_imaging_cred_one.toml +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/src/getframes/presets/data/generic_ccd.toml +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/src/getframes/presets/data/generic_cmos.toml +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/src/getframes/presets/data/generic_eapd.toml +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/src/getframes/presets/data/generic_emccd.toml +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/src/getframes/presets/data/generic_scmos.toml +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/src/getframes/presets/data/hamamatsu_orca_fusion.toml +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/src/getframes/presets/data/hamamatsu_orca_quest_2.toml +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/src/getframes/presets/data/leonardo_saphira.toml +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/src/getframes/presets/data/nuvu_hnu_128_omega.toml +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/src/getframes/presets/data/nuvu_hnu_240.toml +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/src/getframes/presets/data/photometrics_prime_95b.toml +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/src/getframes/presets/data/princeton_instruments_kuro_1200b.toml +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/src/getframes/presets/data/qhy530_pro_ii.toml +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/src/getframes/presets/data/scimeasure_little_joe_ccd39.toml +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/src/getframes/presets/data/tucsen_aries_6504_pro.toml +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/src/getframes/presets/data/zwo_asi2600mm.toml +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/src/getframes/py.typed +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/src/getframes/scene/__init__.py +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/src/getframes/scene/photometry.py +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/src/getframes/scene/sources.py +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/src/getframes/scene/wcs.py +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/src/getframes/spectral.py +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/tests/test_analysis.py +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/tests/test_benchmarks.py +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/tests/test_calibrate.py +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/tests/test_camera.py +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/tests/test_characterize.py +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/tests/test_config.py +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/tests/test_dataset.py +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/tests/test_detector.py +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/tests/test_frame.py +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/tests/test_gain.py +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/tests/test_noise.py +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/tests/test_nondestructive_analysis.py +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/tests/test_observation.py +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/tests/test_presets.py +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/tests/test_radiometry.py +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/tests/test_realism.py +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/tests/test_scale.py +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/tests/test_scene_enrich.py +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/tests/test_signal.py +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/tests/test_spectral.py +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/tests/test_validation.py +0 -0
- {getframes-2.2.0 → getframes-2.4.0}/tests/test_workspace.py +0 -0
|
@@ -6,6 +6,94 @@ to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
|
|
|
6
6
|
|
|
7
7
|
## [Unreleased]
|
|
8
8
|
|
|
9
|
+
## [2.4.0] - 2026-10-07
|
|
10
|
+
|
|
11
|
+
### Added
|
|
12
|
+
|
|
13
|
+
- **`device="gpu:N"` and `device="auto"`.** Every `device` argument (`Camera`,
|
|
14
|
+
`get_backend`, and the CLI's `[camera]` table) now speaks the AO stack's
|
|
15
|
+
vocabulary (aocore CONVENTIONS 8.1): `"cpu"`, `"gpu"`, `"gpu:N"` for CUDA
|
|
16
|
+
device `N`, and `"auto"`, which picks the GPU when CuPy is installed and sees
|
|
17
|
+
a device and the CPU otherwise. A `"gpu:N"` beyond the devices CuPy sees
|
|
18
|
+
raises a `ValueError` naming the count. The old spellings (`"numpy"`,
|
|
19
|
+
`"cuda"`, `"cupy"`) still work, now case-insensitively. A GPU camera is
|
|
20
|
+
pinned to its card: the fixed-pattern maps and the cuRAND streams are created
|
|
21
|
+
on it and every camera method runs with it current, so a `"gpu:1"` camera
|
|
22
|
+
works whichever device is current at the call, and `with_config` keeps it.
|
|
23
|
+
New: `Camera.device_id`, `ArrayBackend.device_id`, `ArrayBackend.spec`
|
|
24
|
+
(`"cpu"` or `"gpu:N"`) and `ArrayBackend.activate()` (the device context,
|
|
25
|
+
for calling the low-level `noise` functions on another card).
|
|
26
|
+
- **`precision="single"` / `"double"`.** The working precision takes the
|
|
27
|
+
shared names (aocore CONVENTIONS 8.2), with `"float32"`/`"float64"` kept as
|
|
28
|
+
aliases: on `Camera`, as a new `precision` keyword on
|
|
29
|
+
`Scene.photon_rate_map`/`photoelectron_rate_map` (beside `dtype`), and on the
|
|
30
|
+
`noise` functions that take a `float_dtype` (`simulate_frame`,
|
|
31
|
+
`fixed_pattern_maps`, `dark_signal_map`, `photo_signal_map`). A `dtype` and a
|
|
32
|
+
`precision` that disagree raise `ValueError`. `getframes.resolve_precision`
|
|
33
|
+
maps any of these names to the NumPy dtype. `Camera.precision` still reports
|
|
34
|
+
the dtype name (`"float32"`/`"float64"`) whichever spelling was passed.
|
|
35
|
+
`dataset.pairs(dtype=...)` is unchanged: it is the host *storage* type of the
|
|
36
|
+
finished arrays, not a working precision.
|
|
37
|
+
- The CLI's `[camera]` table takes a `device` key.
|
|
38
|
+
- **Conformance tests** for the device and precision vocabulary, including that
|
|
39
|
+
each precision name selects the same dtype as in aocore.
|
|
40
|
+
- **Edge-flux conformance for every PSF model.** `tests/test_conformance.py`
|
|
41
|
+
now uses aocore 0.1.3's image-builder checks (`check_point_source_centring`,
|
|
42
|
+
`check_point_source_flux`) instead of feeding analytic PSFs through the
|
|
43
|
+
OPD-driven checks with a dummy OPD, and adds `check_edge_flux_loss` for
|
|
44
|
+
Gaussian, Moffat, elliptical Gaussian, Airy and array PSFs, guarding the
|
|
45
|
+
2.3.0 fix. The `dev` extra pins `aocore>=0.1.3,<0.2`; the runtime
|
|
46
|
+
requirement is unchanged.
|
|
47
|
+
|
|
48
|
+
### Changed
|
|
49
|
+
|
|
50
|
+
- An unknown `device` string now raises `ValueError` listing the accepted words
|
|
51
|
+
(`'cpu', 'gpu', 'gpu:N' ... or 'auto'`), and a non-string `device` a
|
|
52
|
+
`TypeError`. `device="gpu"` with CuPy installed but no CUDA device raises
|
|
53
|
+
`RuntimeError` at construction rather than failing at the first frame.
|
|
54
|
+
- `Camera.__repr__` shows the GPU number (`device='gpu:0'`).
|
|
55
|
+
- The noise functions' `float_dtype` default is now `None` (still float64).
|
|
56
|
+
|
|
57
|
+
### Fixed
|
|
58
|
+
|
|
59
|
+
- `dataset.pairs` with a GPU camera failed on an implicit CuPy-to-NumPy
|
|
60
|
+
conversion; it now copies each frame to the host through `to_numpy`, as do
|
|
61
|
+
the CLI's `.npy`/`.npz` writers now that the CLI can select a GPU.
|
|
62
|
+
|
|
63
|
+
## [2.3.0] - 2026-10-07
|
|
64
|
+
|
|
65
|
+
### Fixed
|
|
66
|
+
|
|
67
|
+
- **`MoffatPSF`, `EllipticalGaussianPSF` and `AiryPSF` no longer put light that
|
|
68
|
+
falls off the frame back onto it.** They normalised their stamp *after*
|
|
69
|
+
clipping it to the frame, so a source on the edge column deposited its full
|
|
70
|
+
flux instead of the part that lands on the detector (a Moffat star centred on
|
|
71
|
+
column 0 kept 100% of its light, where `GaussianPSF` correctly keeps about
|
|
72
|
+
half). The stamp is now normalised before clipping, which follows the AO stack
|
|
73
|
+
convention that light lost at a detector edge is lost, not renormalised
|
|
74
|
+
(aocore CONVENTIONS 3.3). Sources wholly inside the frame are unchanged.
|
|
75
|
+
|
|
76
|
+
### Changed
|
|
77
|
+
|
|
78
|
+
- **Generic primitives now come from `aocore`.** getframes is part of an
|
|
79
|
+
adaptive-optics stack whose shared conventions and primitives live in
|
|
80
|
+
[`aocore`](https://github.com/jacotay7/aocore), now a core dependency
|
|
81
|
+
(`aocore>=0.1.2,<0.2`). `noise.block_sum` is `aocore.block_sum`, re-exported
|
|
82
|
+
so existing imports keep working; it accepts everything it did before and now
|
|
83
|
+
also takes a `(fy, fx)` factor and leading batch axes, and raises a clear
|
|
84
|
+
`ValueError` for a factor below 1 or a shape that does not divide (previously
|
|
85
|
+
a bare reshape error, or a `ZeroDivisionError` for a factor of 0).
|
|
86
|
+
`Frame.binned` sums through it. `AiryPSF` and `Thermal` take their radians per
|
|
87
|
+
arcsecond from `aocore.ARCSEC_TO_RAD`, and `Vignetting` its centred grid from
|
|
88
|
+
`aocore.coordinate_grid`. Every output is bit-for-bit unchanged.
|
|
89
|
+
|
|
90
|
+
### Added
|
|
91
|
+
|
|
92
|
+
- **Conformance tests** (`tests/test_conformance.py`) run the
|
|
93
|
+
`aocore.conformance` checks that apply to a detector package: every PSF model
|
|
94
|
+
deposits unit flux, a source at `(n - 1) / 2` and the vignetting pattern are
|
|
95
|
+
centred on the optical axis for odd and even windows.
|
|
96
|
+
|
|
9
97
|
## [2.2.0] - 2026-08-24
|
|
10
98
|
|
|
11
99
|
### Added
|
|
@@ -558,7 +646,9 @@ together in 1.0.
|
|
|
558
646
|
- Documentation, runnable examples, and CI (lint, type-check, test matrix, PyPI
|
|
559
647
|
release via Trusted Publishing).
|
|
560
648
|
|
|
561
|
-
[Unreleased]: https://github.com/jacotay7/getframes/compare/2.
|
|
649
|
+
[Unreleased]: https://github.com/jacotay7/getframes/compare/2.4.0...HEAD
|
|
650
|
+
[2.4.0]: https://github.com/jacotay7/getframes/compare/2.3.0...2.4.0
|
|
651
|
+
[2.3.0]: https://github.com/jacotay7/getframes/compare/2.2.0...2.3.0
|
|
562
652
|
[2.2.0]: https://github.com/jacotay7/getframes/compare/2.1.1...2.2.0
|
|
563
653
|
[2.1.1]: https://github.com/jacotay7/getframes/compare/2.1.0...2.1.1
|
|
564
654
|
[2.1.0]: https://github.com/jacotay7/getframes/compare/2.0.0...2.1.0
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.5
|
|
2
2
|
Name: getframes
|
|
3
|
-
Version: 2.
|
|
3
|
+
Version: 2.4.0
|
|
4
4
|
Summary: Generate physically realistic synthetic camera frames (CCD/CMOS/EMCCD/eAPD/sCMOS) — dark, bias, flat, and rendered star fields — with auditable noise physics for scientific imaging pipelines.
|
|
5
5
|
Project-URL: Homepage, https://github.com/jacotay7/getframes
|
|
6
6
|
Project-URL: Documentation, https://jacotay7.github.io/getframes/
|
|
@@ -23,11 +23,13 @@ Classifier: Topic :: Scientific/Engineering :: Astronomy
|
|
|
23
23
|
Classifier: Topic :: Scientific/Engineering :: Image Processing
|
|
24
24
|
Classifier: Typing :: Typed
|
|
25
25
|
Requires-Python: >=3.10
|
|
26
|
+
Requires-Dist: aocore<0.2,>=0.1.2
|
|
26
27
|
Requires-Dist: astropy>=5.0
|
|
27
28
|
Requires-Dist: numpy>=1.23
|
|
28
29
|
Requires-Dist: scipy>=1.10
|
|
29
30
|
Requires-Dist: tomli>=2.0; python_version < '3.11'
|
|
30
31
|
Provides-Extra: dev
|
|
32
|
+
Requires-Dist: aocore<0.2,>=0.1.3; extra == 'dev'
|
|
31
33
|
Requires-Dist: build>=1.0; extra == 'dev'
|
|
32
34
|
Requires-Dist: mypy>=1.8; extra == 'dev'
|
|
33
35
|
Requires-Dist: pytest-cov>=4.0; extra == 'dev'
|
|
@@ -72,7 +74,7 @@ switches to CUDA (via CuPy) with a single argument.
|
|
|
72
74
|
## Install
|
|
73
75
|
|
|
74
76
|
```bash
|
|
75
|
-
pip install getframes # CPU (NumPy + SciPy + astropy)
|
|
77
|
+
pip install getframes # CPU (NumPy + SciPy + astropy + aocore)
|
|
76
78
|
pip install 'getframes[gpu]' # + CuPy for CUDA 12.x
|
|
77
79
|
pip install -e '.[dev]' # from a clone, for development
|
|
78
80
|
```
|
|
@@ -105,7 +107,7 @@ frame = cam.with_config(resolution=(256, 256)).observe(scene, exposure=300.0, se
|
|
|
105
107
|
|
|
106
108
|
import cupy as cp # and the same path on a GPU
|
|
107
109
|
|
|
108
|
-
cam = gf.Camera.from_preset("andor_ocam2k", device="gpu", precision="
|
|
110
|
+
cam = gf.Camera.from_preset("andor_ocam2k", device="gpu", precision="single")
|
|
109
111
|
rate = cp.full(cam.resolution, 2.0e6, dtype=cp.float32) # photons/s/pixel
|
|
110
112
|
frame = cam.expose(rate, exposure=1.0e-3, seed=0) # CuPy ADU, no host copy
|
|
111
113
|
```
|
|
@@ -185,8 +187,8 @@ for the methodology.
|
|
|
185
187
|
- **Scale & datasets** — a float32 fast path, vectorised multi-source rendering,
|
|
186
188
|
a streaming raw+truth `dataset` generator and a `getframes` CLI; see
|
|
187
189
|
**[Scale & datasets](https://jacotay7.github.io/getframes/guides/datasets/)**.
|
|
188
|
-
- **GPU-optional** — every camera takes `device="gpu"`
|
|
189
|
-
detector path and truth arrays device-resident. CPU and GPU have independent
|
|
190
|
+
- **GPU-optional** — every camera takes `device="gpu"`, `"gpu:N"` or `"auto"`
|
|
191
|
+
(CuPy) and keeps the detector path and truth arrays device-resident. CPU and GPU have independent
|
|
190
192
|
RNG streams, so a `seed` repeats exactly on a fixed backend while parity across
|
|
191
193
|
backends means matching statistics, not identical pixels.
|
|
192
194
|
- **Reproducible and typed** — all randomness flows through a camera-owned seeded
|
|
@@ -27,7 +27,7 @@ switches to CUDA (via CuPy) with a single argument.
|
|
|
27
27
|
## Install
|
|
28
28
|
|
|
29
29
|
```bash
|
|
30
|
-
pip install getframes # CPU (NumPy + SciPy + astropy)
|
|
30
|
+
pip install getframes # CPU (NumPy + SciPy + astropy + aocore)
|
|
31
31
|
pip install 'getframes[gpu]' # + CuPy for CUDA 12.x
|
|
32
32
|
pip install -e '.[dev]' # from a clone, for development
|
|
33
33
|
```
|
|
@@ -60,7 +60,7 @@ frame = cam.with_config(resolution=(256, 256)).observe(scene, exposure=300.0, se
|
|
|
60
60
|
|
|
61
61
|
import cupy as cp # and the same path on a GPU
|
|
62
62
|
|
|
63
|
-
cam = gf.Camera.from_preset("andor_ocam2k", device="gpu", precision="
|
|
63
|
+
cam = gf.Camera.from_preset("andor_ocam2k", device="gpu", precision="single")
|
|
64
64
|
rate = cp.full(cam.resolution, 2.0e6, dtype=cp.float32) # photons/s/pixel
|
|
65
65
|
frame = cam.expose(rate, exposure=1.0e-3, seed=0) # CuPy ADU, no host copy
|
|
66
66
|
```
|
|
@@ -140,8 +140,8 @@ for the methodology.
|
|
|
140
140
|
- **Scale & datasets** — a float32 fast path, vectorised multi-source rendering,
|
|
141
141
|
a streaming raw+truth `dataset` generator and a `getframes` CLI; see
|
|
142
142
|
**[Scale & datasets](https://jacotay7.github.io/getframes/guides/datasets/)**.
|
|
143
|
-
- **GPU-optional** — every camera takes `device="gpu"`
|
|
144
|
-
detector path and truth arrays device-resident. CPU and GPU have independent
|
|
143
|
+
- **GPU-optional** — every camera takes `device="gpu"`, `"gpu:N"` or `"auto"`
|
|
144
|
+
(CuPy) and keeps the detector path and truth arrays device-resident. CPU and GPU have independent
|
|
145
145
|
RNG streams, so a `seed` repeats exactly on a fixed backend while parity across
|
|
146
146
|
backends means matching statistics, not identical pixels.
|
|
147
147
|
- **Reproducible and typed** — all randomness flows through a camera-owned seeded
|
|
@@ -36,6 +36,7 @@ classifiers = [
|
|
|
36
36
|
]
|
|
37
37
|
dynamic = ["version"]
|
|
38
38
|
dependencies = [
|
|
39
|
+
"aocore>=0.1.2,<0.2",
|
|
39
40
|
"numpy>=1.23",
|
|
40
41
|
"scipy>=1.10",
|
|
41
42
|
"astropy>=5.0",
|
|
@@ -44,6 +45,8 @@ dependencies = [
|
|
|
44
45
|
|
|
45
46
|
[project.optional-dependencies]
|
|
46
47
|
dev = [
|
|
48
|
+
# The test suite runs aocore's image-builder conformance checks (0.1.3+).
|
|
49
|
+
"aocore>=0.1.3,<0.2",
|
|
47
50
|
"pytest>=7.0",
|
|
48
51
|
"pytest-cov>=4.0",
|
|
49
52
|
"ruff>=0.6",
|
|
@@ -21,7 +21,7 @@ from .analysis import (
|
|
|
21
21
|
nondestructive_stack_statistics,
|
|
22
22
|
ramp_photon_transfer,
|
|
23
23
|
)
|
|
24
|
-
from .backend import ArrayBackend, get_array_module, get_backend, to_numpy
|
|
24
|
+
from .backend import ArrayBackend, get_array_module, get_backend, resolve_precision, to_numpy
|
|
25
25
|
from .calibrate import calibrate, combine
|
|
26
26
|
from .camera import Camera
|
|
27
27
|
from .config import CameraConfig, SensorType
|
|
@@ -108,5 +108,6 @@ __all__ = [
|
|
|
108
108
|
"load_preset",
|
|
109
109
|
"nondestructive_stack_statistics",
|
|
110
110
|
"ramp_photon_transfer",
|
|
111
|
+
"resolve_precision",
|
|
111
112
|
"to_numpy",
|
|
112
113
|
]
|
|
@@ -0,0 +1,391 @@
|
|
|
1
|
+
# SPDX-License-Identifier: MIT
|
|
2
|
+
"""Optional array backends for detector simulation.
|
|
3
|
+
|
|
4
|
+
NumPy remains the reference implementation. The CuPy backend is imported only
|
|
5
|
+
when requested, so installing and importing :mod:`getframes` stays CPU-only by
|
|
6
|
+
default. Detector physics receives an :class:`ArrayBackend` explicitly; this
|
|
7
|
+
keeps photon-rate, electron, truth, and ADU arrays on one device for the complete
|
|
8
|
+
signal chain.
|
|
9
|
+
"""
|
|
10
|
+
|
|
11
|
+
from __future__ import annotations
|
|
12
|
+
|
|
13
|
+
import contextlib
|
|
14
|
+
from collections.abc import Callable, Generator, Iterator
|
|
15
|
+
from dataclasses import dataclass
|
|
16
|
+
from functools import wraps
|
|
17
|
+
from inspect import isgeneratorfunction
|
|
18
|
+
from typing import Any, ParamSpec, TypeVar, cast
|
|
19
|
+
|
|
20
|
+
import numpy as np
|
|
21
|
+
from numpy.typing import DTypeLike
|
|
22
|
+
|
|
23
|
+
_P = ParamSpec("_P")
|
|
24
|
+
_R = TypeVar("_R")
|
|
25
|
+
|
|
26
|
+
_CPU_NAMES = frozenset({"cpu", "numpy"})
|
|
27
|
+
_GPU_NAMES = frozenset({"gpu", "cuda", "cupy"})
|
|
28
|
+
_DEVICE_HELP = (
|
|
29
|
+
"expected 'cpu', 'gpu', 'gpu:N' (CUDA device N) or 'auto' "
|
|
30
|
+
"(aliases: 'numpy' for 'cpu'; 'cuda'/'cupy' for 'gpu')"
|
|
31
|
+
)
|
|
32
|
+
|
|
33
|
+
# Working-precision vocabulary (aocore CONVENTIONS 8.2): the canonical names are
|
|
34
|
+
# "single"/"double"; "float32"/"float64" are aliases.
|
|
35
|
+
_PRECISIONS: dict[str, type[np.floating[Any]]] = {
|
|
36
|
+
"single": np.float32,
|
|
37
|
+
"double": np.float64,
|
|
38
|
+
"float32": np.float32,
|
|
39
|
+
"float64": np.float64,
|
|
40
|
+
}
|
|
41
|
+
|
|
42
|
+
|
|
43
|
+
def _cupy_seed(seed: Any) -> int | None:
|
|
44
|
+
"""Map NumPy-compatible seed input onto CuPy RandomState's uint32 seed."""
|
|
45
|
+
if seed is None:
|
|
46
|
+
return None
|
|
47
|
+
return int(np.random.SeedSequence(int(seed)).generate_state(1, dtype=np.uint32)[0])
|
|
48
|
+
|
|
49
|
+
|
|
50
|
+
class _CuPyGenerator:
|
|
51
|
+
"""Expose the NumPy Generator spellings over a fast CuPy RandomState."""
|
|
52
|
+
|
|
53
|
+
def __init__(self, generator: Any, xp: Any, float_dtype: Any) -> None:
|
|
54
|
+
self._generator = generator
|
|
55
|
+
self._xp = xp
|
|
56
|
+
self._float_dtype = float_dtype
|
|
57
|
+
|
|
58
|
+
def __getattr__(self, name: str) -> Any:
|
|
59
|
+
return getattr(self._generator, name)
|
|
60
|
+
|
|
61
|
+
def normal(self, loc: Any = 0.0, scale: Any = 1.0, size: Any = None) -> Any:
|
|
62
|
+
"""Draw a scaled normal variate directly in the working precision."""
|
|
63
|
+
return self._generator.normal(loc=loc, scale=scale, size=size, dtype=self._float_dtype)
|
|
64
|
+
|
|
65
|
+
def standard_normal(self, size: Any = None, dtype: Any = None) -> Any:
|
|
66
|
+
"""Draw a standard normal variate in the working precision."""
|
|
67
|
+
selected_dtype = self._float_dtype if dtype is None else dtype
|
|
68
|
+
return self._generator.standard_normal(size=size, dtype=selected_dtype)
|
|
69
|
+
|
|
70
|
+
def lognormal(self, mean: Any = 0.0, sigma: Any = 1.0, size: Any = None) -> Any:
|
|
71
|
+
"""Draw a log-normal variate without falling back to global RNG state."""
|
|
72
|
+
return self._xp.exp(self.normal(mean, sigma, size))
|
|
73
|
+
|
|
74
|
+
def gamma(self, shape: Any, scale: Any = 1.0, size: Any = None) -> Any:
|
|
75
|
+
"""Draw Gamma variates directly in the detector working precision."""
|
|
76
|
+
return self._generator.gamma(shape=shape, scale=scale, size=size, dtype=self._float_dtype)
|
|
77
|
+
|
|
78
|
+
def integers(self, low: Any, high: Any = None, size: Any = None) -> Any:
|
|
79
|
+
"""NumPy-Generator spelling for CuPy RandomState's ``randint``."""
|
|
80
|
+
return self._generator.randint(low, high=high, size=size)
|
|
81
|
+
|
|
82
|
+
def random(self, size: Any = None) -> Any:
|
|
83
|
+
"""NumPy-Generator spelling for CuPy RandomState's uniform sampler."""
|
|
84
|
+
return self._generator.random_sample(size=size)
|
|
85
|
+
|
|
86
|
+
def seed(self, seed: Any) -> None:
|
|
87
|
+
"""Reset this private per-call stream without rebuilding cuRAND state."""
|
|
88
|
+
self._generator.seed(_cupy_seed(seed))
|
|
89
|
+
|
|
90
|
+
|
|
91
|
+
@dataclass(frozen=True)
|
|
92
|
+
class ArrayBackend:
|
|
93
|
+
"""Array namespace and RNG factory for one detector execution device.
|
|
94
|
+
|
|
95
|
+
Attributes
|
|
96
|
+
----------
|
|
97
|
+
xp:
|
|
98
|
+
The array module, :mod:`numpy` or :mod:`cupy`.
|
|
99
|
+
device:
|
|
100
|
+
Device kind, ``"cpu"`` or ``"gpu"``.
|
|
101
|
+
device_id:
|
|
102
|
+
CUDA device number the GPU backend allocates on and launches kernels on
|
|
103
|
+
(``None`` for the CPU backend). :func:`get_backend` always fills it in for
|
|
104
|
+
a GPU backend, so a camera stays on one card whatever device is current
|
|
105
|
+
when it is later called.
|
|
106
|
+
"""
|
|
107
|
+
|
|
108
|
+
xp: Any
|
|
109
|
+
device: str
|
|
110
|
+
device_id: int | None = None
|
|
111
|
+
|
|
112
|
+
@property
|
|
113
|
+
def is_cpu(self) -> bool:
|
|
114
|
+
"""Whether arrays live in host NumPy storage."""
|
|
115
|
+
return self.device == "cpu"
|
|
116
|
+
|
|
117
|
+
@property
|
|
118
|
+
def spec(self) -> str:
|
|
119
|
+
"""The ``device`` string that selects this backend again (``"cpu"`` or ``"gpu:N"``)."""
|
|
120
|
+
if self.is_cpu or self.device_id is None:
|
|
121
|
+
return self.device
|
|
122
|
+
return f"{self.device}:{self.device_id}"
|
|
123
|
+
|
|
124
|
+
def activate(self) -> contextlib.AbstractContextManager[Any]:
|
|
125
|
+
"""Context that makes this backend's CUDA device current.
|
|
126
|
+
|
|
127
|
+
A no-op on the CPU backend. Camera methods enter it themselves; wrap direct
|
|
128
|
+
calls to the low-level :mod:`getframes.noise` functions in it when using a
|
|
129
|
+
GPU other than the current one.
|
|
130
|
+
"""
|
|
131
|
+
if self.is_cpu or self.device_id is None:
|
|
132
|
+
return contextlib.nullcontext()
|
|
133
|
+
return cast(contextlib.AbstractContextManager[Any], self.xp.cuda.Device(self.device_id))
|
|
134
|
+
|
|
135
|
+
def asarray(self, value: Any, *, dtype: Any | None = None) -> Any:
|
|
136
|
+
"""Convert ``value`` to an array on this backend."""
|
|
137
|
+
if self.is_cpu:
|
|
138
|
+
return self.xp.asarray(value, dtype=dtype)
|
|
139
|
+
with self.activate():
|
|
140
|
+
return self.xp.asarray(value, dtype=dtype)
|
|
141
|
+
|
|
142
|
+
def default_rng(self, seed: Any = None, *, float_dtype: Any = np.float64) -> Any:
|
|
143
|
+
"""Create a backend-native random generator (on this backend's device)."""
|
|
144
|
+
if self.is_cpu:
|
|
145
|
+
return self.xp.random.default_rng(seed)
|
|
146
|
+
# CuPy's Generator construction initializes device-side state and is much
|
|
147
|
+
# slower than RandomState for the per-exposure seed contract. RandomState
|
|
148
|
+
# still owns an independent, backend-native cuRAND stream and exposes all
|
|
149
|
+
# distributions used by the detector chain. The cuRAND state is created on
|
|
150
|
+
# the selected device, so it must be built inside the device context.
|
|
151
|
+
with self.activate():
|
|
152
|
+
state = self.xp.random.RandomState(_cupy_seed(seed))
|
|
153
|
+
return _CuPyGenerator(state, self.xp, float_dtype)
|
|
154
|
+
|
|
155
|
+
def convolve(self, array: Any, kernel: Any) -> Any:
|
|
156
|
+
"""Convolve with constant-zero boundary conditions on this backend."""
|
|
157
|
+
if self.is_cpu:
|
|
158
|
+
from scipy import ndimage
|
|
159
|
+
|
|
160
|
+
return ndimage.convolve(array, kernel, mode="constant", cval=0.0)
|
|
161
|
+
from cupyx.scipy import ndimage # pragma: no cover - optional CUDA dependency
|
|
162
|
+
|
|
163
|
+
with self.activate():
|
|
164
|
+
return ndimage.convolve(array, kernel, mode="constant", cval=0.0)
|
|
165
|
+
|
|
166
|
+
def scalar(self, value: Any) -> float:
|
|
167
|
+
"""Transfer one scalar to the host for validation or metadata."""
|
|
168
|
+
item = value.item() if hasattr(value, "item") else value
|
|
169
|
+
return float(item)
|
|
170
|
+
|
|
171
|
+
def to_numpy(self, value: Any) -> np.ndarray[Any, Any]:
|
|
172
|
+
"""Copy an array to host NumPy storage at an explicit boundary."""
|
|
173
|
+
if self.is_cpu:
|
|
174
|
+
return np.asarray(value)
|
|
175
|
+
return cast(np.ndarray[Any, Any], self.xp.asnumpy(value))
|
|
176
|
+
|
|
177
|
+
|
|
178
|
+
_CPU_BACKEND = ArrayBackend(np, "cpu")
|
|
179
|
+
|
|
180
|
+
|
|
181
|
+
def _import_cupy() -> Any:
|
|
182
|
+
"""Import CuPy lazily (raises :class:`ImportError` when it is not installed)."""
|
|
183
|
+
import cupy
|
|
184
|
+
|
|
185
|
+
return cupy
|
|
186
|
+
|
|
187
|
+
|
|
188
|
+
def _gpu_device_count(cupy: Any) -> int:
|
|
189
|
+
"""Number of CUDA devices CuPy can use (``0`` when the runtime is unusable)."""
|
|
190
|
+
try:
|
|
191
|
+
return int(cupy.cuda.runtime.getDeviceCount())
|
|
192
|
+
except Exception: # pragma: no cover - depends on the local CUDA driver/runtime
|
|
193
|
+
return 0
|
|
194
|
+
|
|
195
|
+
|
|
196
|
+
def _parse_device(device: str) -> tuple[str, int | None]:
|
|
197
|
+
"""Split a device string into ``("cpu" | "gpu" | "auto", index or None)``."""
|
|
198
|
+
if not isinstance(device, str):
|
|
199
|
+
raise TypeError(f"device must be a string, got {type(device).__name__}; {_DEVICE_HELP}.")
|
|
200
|
+
name = device.strip().lower()
|
|
201
|
+
base, sep, index = name.partition(":")
|
|
202
|
+
if base == "auto" and not sep:
|
|
203
|
+
return "auto", None
|
|
204
|
+
if base in _CPU_NAMES and not sep:
|
|
205
|
+
return "cpu", None
|
|
206
|
+
if base in _GPU_NAMES:
|
|
207
|
+
if not sep:
|
|
208
|
+
return "gpu", None
|
|
209
|
+
if index.isdigit():
|
|
210
|
+
return "gpu", int(index)
|
|
211
|
+
raise ValueError(f"unknown device {device!r}; {_DEVICE_HELP}.")
|
|
212
|
+
|
|
213
|
+
|
|
214
|
+
def _gpu_backend(device: str, index: int | None) -> ArrayBackend:
|
|
215
|
+
try:
|
|
216
|
+
cupy = _import_cupy()
|
|
217
|
+
except ImportError as exc:
|
|
218
|
+
raise ImportError(
|
|
219
|
+
f"device={device!r} requires CuPy; install getframes[gpu], "
|
|
220
|
+
"or use device='auto' to fall back to the CPU."
|
|
221
|
+
) from exc
|
|
222
|
+
count = _gpu_device_count(cupy)
|
|
223
|
+
if count < 1:
|
|
224
|
+
raise RuntimeError(
|
|
225
|
+
f"device={device!r} needs a CUDA device, but CuPy sees none; "
|
|
226
|
+
"use device='auto' to fall back to the CPU."
|
|
227
|
+
)
|
|
228
|
+
if index is None:
|
|
229
|
+
index = int(cupy.cuda.runtime.getDevice())
|
|
230
|
+
elif index >= count:
|
|
231
|
+
raise ValueError(
|
|
232
|
+
f"device={device!r} asks for CUDA device {index}, but CuPy sees {count} "
|
|
233
|
+
f"device(s) (numbered 0-{count - 1}; CUDA_VISIBLE_DEVICES controls the numbering)."
|
|
234
|
+
)
|
|
235
|
+
return ArrayBackend(cupy, "gpu", index)
|
|
236
|
+
|
|
237
|
+
|
|
238
|
+
def get_backend(device: str = "cpu") -> ArrayBackend:
|
|
239
|
+
"""Return the backend for ``device``.
|
|
240
|
+
|
|
241
|
+
Parameters
|
|
242
|
+
----------
|
|
243
|
+
device:
|
|
244
|
+
``"cpu"`` (NumPy, the reference), ``"gpu"`` (CuPy on the current CUDA
|
|
245
|
+
device), ``"gpu:N"`` (CuPy on CUDA device ``N``), or ``"auto"`` (the
|
|
246
|
+
current CUDA device when CuPy is installed and sees one, else the CPU).
|
|
247
|
+
Matching is case-insensitive; ``"numpy"`` is an alias of ``"cpu"`` and
|
|
248
|
+
``"cuda"``/``"cupy"`` of ``"gpu"`` (including ``"cuda:N"``).
|
|
249
|
+
|
|
250
|
+
Raises
|
|
251
|
+
------
|
|
252
|
+
ValueError
|
|
253
|
+
For an unknown device string, or a GPU number CuPy does not see.
|
|
254
|
+
ImportError
|
|
255
|
+
For ``"gpu"``/``"gpu:N"`` without CuPy installed.
|
|
256
|
+
RuntimeError
|
|
257
|
+
For ``"gpu"``/``"gpu:N"`` when CuPy is installed but sees no CUDA device.
|
|
258
|
+
|
|
259
|
+
Notes
|
|
260
|
+
-----
|
|
261
|
+
CuPy is an optional dependency and is imported lazily, only for a GPU (or
|
|
262
|
+
``"auto"``) device.
|
|
263
|
+
"""
|
|
264
|
+
kind, index = _parse_device(device)
|
|
265
|
+
if kind == "cpu":
|
|
266
|
+
return _CPU_BACKEND
|
|
267
|
+
if kind == "auto":
|
|
268
|
+
try:
|
|
269
|
+
cupy = _import_cupy()
|
|
270
|
+
except ImportError:
|
|
271
|
+
return _CPU_BACKEND
|
|
272
|
+
if _gpu_device_count(cupy) < 1:
|
|
273
|
+
return _CPU_BACKEND
|
|
274
|
+
return _gpu_backend(device, index)
|
|
275
|
+
|
|
276
|
+
|
|
277
|
+
def resolve_precision(precision: str | DTypeLike) -> np.dtype[Any]:
|
|
278
|
+
"""Return the floating-point working dtype named by ``precision``.
|
|
279
|
+
|
|
280
|
+
Parameters
|
|
281
|
+
----------
|
|
282
|
+
precision:
|
|
283
|
+
``"single"`` or ``"double"`` (the shared AO-stack vocabulary), their
|
|
284
|
+
aliases ``"float32"``/``"float64"`` (case-insensitive), or a NumPy
|
|
285
|
+
``float32``/``float64`` dtype.
|
|
286
|
+
|
|
287
|
+
Raises
|
|
288
|
+
------
|
|
289
|
+
ValueError
|
|
290
|
+
For any other precision.
|
|
291
|
+
"""
|
|
292
|
+
if isinstance(precision, str):
|
|
293
|
+
selected = _PRECISIONS.get(precision.strip().lower())
|
|
294
|
+
if selected is not None:
|
|
295
|
+
return np.dtype(selected)
|
|
296
|
+
elif precision is not None: # np.dtype(None) would silently mean float64
|
|
297
|
+
try:
|
|
298
|
+
dtype = np.dtype(precision)
|
|
299
|
+
except TypeError:
|
|
300
|
+
dtype = None
|
|
301
|
+
if dtype is not None and dtype in (np.dtype(np.float32), np.dtype(np.float64)):
|
|
302
|
+
return dtype
|
|
303
|
+
raise ValueError(
|
|
304
|
+
f"precision must be 'single' or 'double' (aliases 'float32'/'float64'), got {precision!r}."
|
|
305
|
+
)
|
|
306
|
+
|
|
307
|
+
|
|
308
|
+
def _working_dtype(
|
|
309
|
+
dtype: DTypeLike | None, precision: str | None, *, name: str = "dtype"
|
|
310
|
+
) -> np.dtype[Any]:
|
|
311
|
+
"""Merge a legacy ``dtype`` argument with the ``precision`` vocabulary.
|
|
312
|
+
|
|
313
|
+
``None`` for both gives ``float64``. Both may be given only when they agree.
|
|
314
|
+
"""
|
|
315
|
+
from_dtype = None if dtype is None else np.dtype(dtype)
|
|
316
|
+
if precision is None:
|
|
317
|
+
return np.dtype(np.float64) if from_dtype is None else from_dtype
|
|
318
|
+
from_precision = resolve_precision(precision)
|
|
319
|
+
if from_dtype is not None and from_dtype != from_precision:
|
|
320
|
+
raise ValueError(
|
|
321
|
+
f"{name}={from_dtype.name!r} conflicts with precision={precision!r}; pass one of them."
|
|
322
|
+
)
|
|
323
|
+
return from_precision
|
|
324
|
+
|
|
325
|
+
|
|
326
|
+
def _on_device(method: Callable[_P, _R]) -> Callable[_P, _R]:
|
|
327
|
+
"""Run a method of an object with a ``_backend`` inside that backend's device context.
|
|
328
|
+
|
|
329
|
+
Generator methods enter the context for every step, so lazily produced frames
|
|
330
|
+
are also computed on the object's device.
|
|
331
|
+
"""
|
|
332
|
+
if isgeneratorfunction(method):
|
|
333
|
+
|
|
334
|
+
@wraps(method)
|
|
335
|
+
def generator_wrapper(*args: _P.args, **kwargs: _P.kwargs) -> Any:
|
|
336
|
+
backend: ArrayBackend = args[0]._backend # type: ignore[attr-defined]
|
|
337
|
+
iterator = cast(Generator[Any, None, Any], method(*args, **kwargs))
|
|
338
|
+
return _stepped_on(backend, iterator)
|
|
339
|
+
|
|
340
|
+
return cast(Callable[_P, _R], generator_wrapper)
|
|
341
|
+
|
|
342
|
+
@wraps(method)
|
|
343
|
+
def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R:
|
|
344
|
+
backend: ArrayBackend = args[0]._backend # type: ignore[attr-defined]
|
|
345
|
+
with backend.activate():
|
|
346
|
+
return method(*args, **kwargs)
|
|
347
|
+
|
|
348
|
+
return wrapper
|
|
349
|
+
|
|
350
|
+
|
|
351
|
+
def _stepped_on(backend: ArrayBackend, iterator: Generator[Any, None, Any]) -> Iterator[Any]:
|
|
352
|
+
"""Advance ``iterator`` one item at a time inside ``backend``'s device context."""
|
|
353
|
+
while True:
|
|
354
|
+
with backend.activate():
|
|
355
|
+
try:
|
|
356
|
+
item = next(iterator)
|
|
357
|
+
except StopIteration:
|
|
358
|
+
return
|
|
359
|
+
yield item
|
|
360
|
+
|
|
361
|
+
|
|
362
|
+
def _array_device(value: Any) -> contextlib.AbstractContextManager[Any]:
|
|
363
|
+
"""Context that makes ``value``'s CUDA device current (a no-op for host arrays)."""
|
|
364
|
+
if type(value).__module__.split(".", 1)[0] == "cupy":
|
|
365
|
+
return cast(contextlib.AbstractContextManager[Any], value.device)
|
|
366
|
+
return contextlib.nullcontext()
|
|
367
|
+
|
|
368
|
+
|
|
369
|
+
def get_array_module(value: Any) -> Any:
|
|
370
|
+
"""Return NumPy or CuPy for an existing array without copying it."""
|
|
371
|
+
module = type(value).__module__.split(".", 1)[0]
|
|
372
|
+
if module == "cupy":
|
|
373
|
+
return _import_cupy()
|
|
374
|
+
return np
|
|
375
|
+
|
|
376
|
+
|
|
377
|
+
def to_numpy(value: Any) -> np.ndarray[Any, Any]:
|
|
378
|
+
"""Return ``value`` in host NumPy storage, copying device arrays explicitly."""
|
|
379
|
+
module = type(value).__module__.split(".", 1)[0]
|
|
380
|
+
if module == "cupy":
|
|
381
|
+
return cast(np.ndarray[Any, Any], _import_cupy().asnumpy(value))
|
|
382
|
+
return np.asarray(value)
|
|
383
|
+
|
|
384
|
+
|
|
385
|
+
__all__ = [
|
|
386
|
+
"ArrayBackend",
|
|
387
|
+
"get_array_module",
|
|
388
|
+
"get_backend",
|
|
389
|
+
"resolve_precision",
|
|
390
|
+
"to_numpy",
|
|
391
|
+
]
|