diff --git a/Demos/nxs/nxs_generator.py b/Demos/nxs/nxs_generator.py deleted file mode 100644 index cdf76c8d..00000000 --- a/Demos/nxs/nxs_generator.py +++ /dev/null @@ -1,135 +0,0 @@ -#!/usr/bin/env python3 -# -*- coding: utf-8 -*- -import argparse -import h5py -import numpy as np -import pathlib -import psutil -import tomophantom - -""" -This script uses TomoPhantom to generate an NXtomo-compliant NeXus file, with HDF5 as the underlying storage format. Users can specify the model number, phantom shape, and angular range. - -The script generates the phantom incrementally in chunks, based on the available system memory, and writes each chunk directly to the output file. This approach enables the creation of extremely large phantom datasets, even on systems with limited CPU memory. - -Because the resulting NeXus file conforms to the NXtomo data standard, it can be read by multiple data-processing and reconstruction frameworks. This work was carried out to support testing of high-resolution datasets within the HTTomo big-data processing framework developed at Diamond Light Source: -https://github.com/DiamondLightSource/httomo - -Example to generate (128 x 256 x 362) data file from the model no. 13 using the command line: python nxs_generator.py -m 13 -s 128 256 362 -""" - - -def main(args: argparse.Namespace): - model = args.model_number - path = pathlib.Path(tomophantom.__file__).parent - path_library3D = str(path / "phantomlib" / "Phantom3DLibrary.dat") - - sinogram_dtype = np.uint16() - Horiz_det = args.sinogram_shape[2] - Vert_det = args.sinogram_shape[0] - angles_num = args.sinogram_shape[1] - square_phantom_slice_width = int(Horiz_det / np.sqrt(2)) - - angles = np.linspace(*args.angle_range, angles_num, dtype="float32") # in degrees - - available_memory_bytes = psutil.virtual_memory().available - slice_memory_bytes = Horiz_det * angles_num * sinogram_dtype.itemsize - if slice_memory_bytes > available_memory_bytes: - print( - f"A signle slice ({slice_memory_bytes} bytes) is would be bigger than available memory ({available_memory_bytes} bytes)." - ) - return - - chunk_size = available_memory_bytes // slice_memory_bytes - chunk_count = int(np.ceil(Vert_det / chunk_size)) - print( - f"Creating phantom in {chunk_count} number of chunks of at max {chunk_size} slices per chunk." - ) - - with h5py.File(args.output_path, "w") as file: - file.attrs["default"] = "entry" - entry_group = file.create_group("entry") - entry_group.attrs["NX_class"] = "NXentry" - entry_group.attrs["default"] = "data" - entry_group.attrs["NX_application"] = "NXtomo" - - instrument_group = entry_group.create_group("instrument") - detector_group = instrument_group.create_group("detector") - detector_group.create_dataset( - "image_key", data=np.zeros_like(angles, dtype=np.int8) - ) - - entry_group.create_dataset("definition", data="NXtomo") - data_group = entry_group.create_group("data") - sinogram_dataset = data_group.create_dataset( - "data", (angles_num, Vert_det, Horiz_det), sinogram_dtype - ) - data_group.create_dataset("rotation_angle", data=angles) - - file.create_dataset(args.output_path.name, data=[0, 0]) - - data_dims_group = file.create_group("data_dims") - data_dims_group.create_dataset("detector_x_y", data=[Horiz_det, Vert_det]) - - for i in range(chunk_count): - chunk_start = i * chunk_size - chunk_end = min((i + 1) * chunk_size, Vert_det) - - projData3D_analyt = tomophantom.TomoP3D.ModelSinoSub( - model, - square_phantom_slice_width, - Horiz_det, - Vert_det, - (chunk_start, chunk_end), - angles, - path_library3D, - ) - - f_min, f_max = projData3D_analyt.min(), projData3D_analyt.max() - projData3D_analyt = ( - (projData3D_analyt - f_min) / (f_max - f_min) * 65535 - ).astype(np.uint16) - - swapped_projData3D_analyt = np.swapaxes(projData3D_analyt, 0, 1) - sinogram_dataset[:angles_num, chunk_start:chunk_end, :Horiz_det] = ( - swapped_projData3D_analyt - ) - - print( - f"#slices: {Vert_det}, #projections: {angles_num}, detector width: {Horiz_det}" - ) - - -if __name__ == "__main__": - parser = argparse.ArgumentParser( - description="TomoPhantom phantom .nxs generator. Output can be used with httomo, e.g. 'httomo run phantom.nxs pipeline.yaml output'." - ) - parser.add_argument( - "-m", - "--model-number", - type=int, - default=13, # Shepp-Logan - help="Model number in TomoPhantom's model library. See https://github.com/dkazanc/TomoPhantom/tree/master/tomophantom/phantomlib.", - ) - parser.add_argument( - "-s", - "--sinogram-shape", - nargs=3, - metavar=("detector height", "projection count", "detector width"), - type=int, - default=(2368, 256, 4416), - ) - parser.add_argument( - "-a", - "--angle-range", - nargs=2, - metavar=("start", "end"), - type=float, - default=(0.0, 179.9), - ) - parser.add_argument( - "-o", "--output-path", type=pathlib.Path, default="./phantom.nxs" - ) - - args = parser.parse_args() - main(args) diff --git a/scripts/README.md b/scripts/README.md new file mode 100644 index 00000000..9bf98a82 --- /dev/null +++ b/scripts/README.md @@ -0,0 +1,150 @@ +# TomoPhantom NXtomo generator: realistic data and optional artefacts + +This is a replacement for `Demos/nxs/nxs_generator.py`, with the original +`-m`, `-s`, `-a`, and `-o` arguments retained. The default shape is deliberately +small: `(128, 256, 362)` in **detector height, projection count, detector width** +order. Install TomoPhantom in your environment; this bundle does not contain it. + +```bash +conda install -c httomo -c conda-forge tomophantom h5py psutil +``` + +Some artefacts, particularly subpixel shifts, also require `scikit-image`. +No GPU, ASTRA or reconstruction package is needed to run this generator. + +## Generate data + +Realistic background, flats and darks: + +```bash +python nxs_generator.py --realistic -m 13 -s 128 256 362 \ + --flats 20 --darks 10 --source-intensity 50000 \ + --flat-settings flat_settings.json --seed 1 -o realistic.nxs +``` + +Add stripes and zingers to the simulated detector signal: + +```bash +python nxs_generator.py --realistic -m 13 -s 128 256 362 \ + --artefacts artefacts.json --artefacts-stage raw \ + --flat-settings flat_settings.json --seed 1 -o artefacts.nxs +``` + +The example JSON files pass the documented keyword arguments to +`synth_flats` and `artefacts_mix`. Unknown keys are rejected; the implementation's +stripe keyword is `stripes_type`, not the singular spelling in its docstring. +The `detectors_miscallibration` and `fresnel_wavelenght` spellings are upstream API names. + +For blur, Fresnel propagation, or line-integral noise experiments, use +`--artefacts-stage line-integrals`. For example, save this as `blur.json`: + +```json +{"pve_strength": 1, "verbose": false} +``` + +```bash +python nxs_generator.py --realistic -s 128 256 362 \ + --artefacts blur.json --artefacts-stage line-integrals -o blurred.nxs +``` + +To simulate motion, use one positive `datashifts_maxamplitude_pixel` or +`datashifts_maxamplitude_subpixel` setting. Returned shifts are saved under +`/entry/generator/artefact_shifts`, in the API's native component order. The +`projection_frame_indices` dataset maps these to the combined frame stack. + +Clean line-integral data without calibration frames: + +```bash +python nxs_generator.py -m 13 -s 128 256 362 --chunk-rows 16 -o clean.nxs +``` + +Use `python nxs_generator.py --help` for all options. Existing output is protected +unless `--overwrite` is supplied. Output is published only after writing completes. + +## Data layout and image keys + +The frame order is **darks, flats, projections**. All are stored in one uint16 +dataset; calibration images are selected by their image keys, not by separate files. + +| Content | HDF5 path | +| --- | --- | +| Combined frame data | `/entry/data/data` | +| Image keys | `/entry/instrument/detector/image_key` | +| Frame-aligned rotation angles | `/entry/data/rotation_angle` | +| Detector data (same hard-linked dataset) | `/entry/instrument/detector/data` | +| Image keys (same hard-linked dataset) | `/entry/data/image_key` | +| Sample angles (same hard-linked dataset) | `/entry/sample/rotation_angle` | +| Configuration and provenance | `/entry/generator` | + +| Image key | Frame type | Count in the default realistic example | +| --- | --- | --- | +| `2` | Dark field | 10 | +| `1` | Flat field | 20 | +| `0` | Sample projection | 256 | + +For that example, stored shape is `(286, 128, 362)`: **frame, detector y, detector x**. +The key and angle arrays each have length 286. Calibration frames have the initial +stage angle; use `image_key == 0` to obtain the 256 projection angles. +An inclusive `linspace` uses the supplied angular endpoints, as in the original script. + +## Realistic generation model + +1. Generate non-negative analytical projections with `TomoP3D.ModelSinoSub`. +2. Optionally call `artefacts_mix` in the line-integral domain. +3. Call `synth_flats` once for the entire detector volume, returning raw projections, + simulated flat frames and a detector-response map. The map is not stored. +4. Optionally call `artefacts_mix` on the raw signal, scaled to approximately `[0, 1]`. +5. Apply the same signal gain to projections and flats, add a dark-current/read-noise + model, then round, clip and convert to uint16. Generate dark frames with no beam signal. + +For raw projections/flats, before rounding and clipping: + +`stored = signal_scale * synth_flats_signal + Poisson(dark_mean) + Gaussian(0, read_noise)` + +Dark frames use the same expression with the beam signal set to zero. Defaults are +`signal_scale=0.9`, `dark_mean=100` and `read_noise=2` standard deviation. +This is a simple additive detector-dark model, not measured detector calibration. +The original realistic example uses `darks=None`; `synth_flats` does not return darks. + +Both projections and flats use a common added gain, and no new maximum rescaling is +performed after raw artefacts. Out-of-range values are clipped, not wrapped; the +clipped voxel count is reported and stored. `synth_flats` itself normalises projections +and flats internally, so its output is synthetic detector signal rather than absolutely +calibrated photon counts. `source_intensity` controls its noise model, not the stored ADU maximum. + +`artefacts_mix` has its own fixed application order (PVE, Fresnel, zingers, stripes, +shifts, noise), regardless of JSON key order. Artefacts affect projections, not +calibration frames. Gaussian `noise_amplitude` follows upstream behaviour: it is the +standard deviation relative to that function's normalised data. + +The Poisson branch in `artefacts_mix` exponentiates its input; it must not receive +already-raw detector counts. Raw-stage Poisson/Fresnel and `noise_prelog=True` are +therefore rejected. Realistic synthesis already adds its own noise, so explicitly +enabling line-integral noise adds another noise contribution. + +## Memory, scaling and repeatability + +- Clean mode without artefacts remains chunked. It uses a first pass to find a + **global maximum**, then a second pass to write uint16 data. This removes the original + per-chunk min/max scaling, which could cause seams and inconsistent attenuation. +- Realistic mode or any artefact configuration requires a **full-volume working array**. + The chunk-rows setting limits analytic generation and writing, not realistic synthesis. + Independently synthesising chunks would reset flat-field background, normalisation + and motion/PVE context at each boundary. This version does not claim arbitrary-size + out-of-core realistic generation; that requires a stateful streaming synthesis API. +- A heuristic working-memory check uses 50% of available RAM by default. Use + `--memory-fraction` or `--max-memory-gb` to adjust the budget. It is not a hard process + memory limit; reduce the dimensions if the check fails. +- Clean data are scaled line integrals, not raw transmission: do not apply dark/flat + correction and minus-log to clean-mode output. The storage scale is recorded. +- Use the same seed, parameters, chunk-rows and library version for reproducibility. + Different chunk sizes can change the added dark/read-noise realisation. Upstream + routines also contain their own RNG behaviour; cross-version bitwise reproducibility + is not guaranteed. + +Sources: +- [Original generator](https://github.com/dkazanc/TomoPhantom/blob/master/Demos/nxs/nxs_generator.py) +- [Realistic example](https://github.com/dkazanc/TomoPhantom/blob/afadc7c522919e08cdae5aeef6db41cded0fc213/Demos/3D/ReconASTRA3D_realistic.py) +- [synth_flats](https://github.com/dkazanc/TomoPhantom/blob/afadc7c522919e08cdae5aeef6db41cded0fc213/tomophantom/flatsgen.py) +- [artefacts_mix](https://github.com/dkazanc/TomoPhantom/blob/afadc7c522919e08cdae5aeef6db41cded0fc213/tomophantom/artefacts.py) +- [NXtomo application definition](https://manual.nexusformat.org/classes/applications/NXtomo.html) diff --git a/scripts/artefacts.json b/scripts/artefacts.json new file mode 100644 index 00000000..cc54b454 --- /dev/null +++ b/scripts/artefacts.json @@ -0,0 +1,10 @@ +{ + "zingers_percentage": 0.05, + "zingers_modulus": 10, + "stripes_percentage": 3.0, + "stripes_maxthickness": 3, + "stripes_intensity": 0.05, + "stripes_type": "full", + "stripes_variability": 0.002, + "verbose": false +} diff --git a/scripts/flat_settings.json b/scripts/flat_settings.json new file mode 100644 index 00000000..1a60c6b7 --- /dev/null +++ b/scripts/flat_settings.json @@ -0,0 +1,9 @@ +{ + "detectors_miscallibration": 0.05, + "variations_number": 3, + "arguments_Bessel": [1, 25], + "specklesize": 2, + "kbar": 2, + "sigmasmooth": 3, + "jitter_projections": 0.0 +} diff --git a/scripts/nxs_generator.py b/scripts/nxs_generator.py new file mode 100644 index 00000000..d1222c7a --- /dev/null +++ b/scripts/nxs_generator.py @@ -0,0 +1,475 @@ +#!/usr/bin/env python3 +"""Generate TomoPhantom projections with optional realistic detector data/artefacts. + +Adapted from TomoPhantom's Demos/nxs/nxs_generator.py (Apache-2.0). +Axes inside TomoPhantom: (detector_y, angle, detector_x). +Axes on disk: (frame, detector_y, detector_x). + +Clean generation is streamed in detector-row chunks with ONE global scale. +Realistic synthesis and artefacts require a full volume: synth_flats has global +normalisation/background state, and shifts/PVE need neighbouring detector rows. +The memory estimate is a conservative planning heuristic, not an OOM guarantee. +""" + +from __future__ import annotations + +import argparse +from datetime import datetime, timezone +import importlib.metadata +import json +import math +from pathlib import Path +import random +import tempfile + +import h5py +import numpy as np +import psutil + +ARTEFACT_KEYS = { + "noise_type", + "noise_amplitude", + "noise_seed", + "noise_prelog", + "zingers_percentage", + "zingers_modulus", + "stripes_percentage", + "stripes_maxthickness", + "stripes_intensity", + "stripes_type", + "stripes_variability", + "datashifts_maxamplitude_pixel", + "datashifts_maxamplitude_subpixel", + "pve_strength", + "fresnel_dist_observation", + "fresnel_scale_factor", + "fresnel_wavelenght", + "verbose", +} + + +def parser(): + p = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + p.add_argument("-m", "--model-number", type=int, default=13) + p.add_argument( + "-s", + "--sinogram-shape", + type=int, + nargs=3, + metavar=("HEIGHT", "ANGLES", "WIDTH"), + default=(128, 256, 362), + ) + p.add_argument("-a", "--angle-range", type=float, nargs=2, default=(0, 179.9)) + p.add_argument("-o", "--output-path", type=Path, default=Path("phantom.nxs")) + p.add_argument( + "--realistic", action="store_true", help="Use synth_flats; add darks." + ) + p.add_argument("--flats", type=int, default=20) + p.add_argument("--darks", type=int, default=10) + p.add_argument("--source-intensity", type=int, default=50000) + p.add_argument("--flat-settings", type=Path, help="JSON overrides for synth_flats.") + p.add_argument( + "--dark-mean", type=float, default=100.0, help="Dark current in ADU." + ) + p.add_argument("--read-noise", type=float, default=2.0, help="Gaussian SD in ADU.") + p.add_argument( + "--signal-scale", + type=float, + default=0.9, + help="Common gain on synth_flats uint16 signals; reserves headroom.", + ) + p.add_argument( + "--artefacts", type=Path, help="JSON keyword arguments to artefacts_mix." + ) + p.add_argument( + "--artefacts-stage", + choices=["raw", "line-integrals"], + default="raw", + help="raw: before dark offset, on normalised detector counts.", + ) + p.add_argument("--seed", type=int, default=1) + p.add_argument( + "--chunk-rows", + type=int, + default=32, + help="Rows per analytic-generation/write chunk (not synth_flats).", + ) + p.add_argument("--memory-fraction", type=float, default=0.5) + p.add_argument( + "--max-memory-gb", type=float, help="Additional working-memory cap (GiB)." + ) + p.add_argument("--compression", choices=["none", "gzip", "lzf"], default="none") + p.add_argument("--overwrite", action="store_true") + return p + + +def read_json(path): + if path is None: + return {} + obj = json.loads(path.read_text()) + if not isinstance(obj, dict): + raise ValueError(f"{path}: expected a JSON object") + return obj + + +def settings(args): + import inspect + from tomophantom.flatsgen import synth_flats + + if min(args.sinogram_shape) < 2: + raise ValueError("Each shape dimension must be at least 2.") + if args.chunk_rows < 1 or not 0 < args.memory_fraction <= 1: + raise ValueError( + "chunk-rows must be positive; memory-fraction must be in (0,1]." + ) + if ( + not np.isfinite(args.angle_range).all() + or args.angle_range[1] <= args.angle_range[0] + ): + raise ValueError("The angle range must be finite and increasing.") + if args.max_memory_gb is not None and ( + not math.isfinite(args.max_memory_gb) or args.max_memory_gb <= 0 + ): + raise ValueError("max-memory-gb must be finite and positive.") + if not 0 <= args.seed < 2**32: + raise ValueError("seed must be between 0 and 2**32-1.") + if not np.isfinite([args.dark_mean, args.read_noise, args.signal_scale]).all(): + raise ValueError("Detector settings must be finite.") + if args.dark_mean < 0 or args.read_noise < 0 or not 0 < args.signal_scale <= 1: + raise ValueError("Invalid dark mean, read noise or signal scale.") + if args.realistic and (min(args.flats, args.darks, args.source_intensity) < 1): + raise ValueError( + "Realistic generation requires positive flat/dark counts and intensity." + ) + flat = read_json(args.flat_settings) + allowed = set(inspect.signature(synth_flats).parameters) - { + "projData3D_clean", + "source_intensity", + "flatsnum", + } + if set(flat) - allowed: + raise ValueError(f"Unsupported flat settings: {sorted(set(flat) - allowed)}") + if flat and not args.realistic: + raise ValueError("--flat-settings requires --realistic.") + flat = dict(source_intensity=args.source_intensity, flatsnum=args.flats, **flat) + if not 1 <= flat.get("variations_number", 3) <= 3: + raise ValueError("synth_flats supports 1 to 3 detector-response variations.") + art = read_json(args.artefacts) + if set(art) - ARTEFACT_KEYS: + raise ValueError(f"Unknown artefact keys: {sorted(set(art) - ARTEFACT_KEYS)}") + if art and not args.realistic and args.artefacts_stage == "raw": + raise ValueError("Use --realistic or --artefacts-stage line-integrals.") + if art.get("noise_prelog"): + raise ValueError( + "noise_prelog is unsupported; --realistic provides raw data and flats." + ) + if art.get("noise_type") not in (None, "Gaussian", "Poisson"): + raise ValueError("noise_type must be Gaussian or Poisson.") + if args.artefacts_stage == "raw" and ( + art.get("noise_type") == "Poisson" + or art.get("fresnel_dist_observation") is not None + ): + raise ValueError( + "Poisson/Fresnel in artefacts_mix expects line-integral input. " + "Use --artefacts-stage line-integrals. synth_flats already adds noise." + ) + shifts = [k for k in art if k.startswith("datashifts_") and art[k] is not None] + if len(shifts) > 1 or any(art[k] <= 0 for k in shifts): + raise ValueError( + "Specify at most one shift type, with strictly positive amplitude." + ) + return flat, art + + +def checked(data, shape, label): + if data.shape != shape or not np.isfinite(data).all(): + raise ValueError( + f"{label}: invalid shape or non-finite data; expected {shape}." + ) + return data + + +def apply_artefacts(data, config, seed): + """Use the API's own order; preserve shifts when it returns [data, shifts].""" + from tomophantom.artefacts import artefacts_mix + + if not config: + return data, None + config = dict(config) + noise_seed = config.pop("noise_seed", seed) + # The cited API checks `seed is True`, not an arbitrary integer seed. + random.seed(seed) + np.random.seed(seed if noise_seed is None else int(noise_seed)) + config.update(noise_seed=False, noise_prelog=False) + config.setdefault("verbose", False) + result = artefacts_mix(np.array(data, dtype=np.float32, copy=True), **config) + shifts = None + if isinstance(result, (tuple, list)): + result, shifts = result + checked(result, data.shape, "artefacts_mix") + return result, shifts + + +def nxgroup(parent, name, nxclass): + group = parent.create_group(name) + group.attrs["NX_class"] = nxclass + return group + + +def initialise(file, args, angles, metadata): + height, count, width = args.sinogram_shape + nd, nf = (args.darks, args.flats) if args.realistic else (0, 0) + keys = np.concatenate([np.full(nd, 2), np.full(nf, 1), np.zeros(count)]).astype( + "int8" + ) + # Calibration frames are taken at the initial stage angle and identified by key. + frame_angles = np.concatenate([np.full(nd + nf, angles[0]), angles]).astype( + "float32" + ) + file.attrs.update( + default="entry", + file_name=args.output_path.name, + file_time=datetime.now(timezone.utc).isoformat(), + ) + entry = nxgroup(file, "entry", "NXentry") + entry.attrs["default"] = "data" + entry.create_dataset("definition", data="NXtomo") + entry.create_dataset("title", data=f"TomoPhantom model {args.model_number}") + entry.create_dataset("start_time", data=datetime.now(timezone.utc).isoformat()) + instrument = nxgroup(entry, "instrument", "NXinstrument") + detector = nxgroup(instrument, "detector", "NXdetector") + data = detector.create_dataset( + "data", + shape=(len(keys), height, width), + dtype="uint16", + chunks=(1, min(height, args.chunk_rows), min(width, 512)), + compression=None if args.compression == "none" else args.compression, + ) + data.attrs.update( + units="counts" if args.realistic else "arbitrary", + interpretation="image", + axis_order="frame,detector_y,detector_x", + ) + key = detector.create_dataset("image_key", data=keys) + key.attrs["meaning"] = "0=projection; 1=flat; 2=dark" + sample = nxgroup(entry, "sample", "NXsample") + sample.create_dataset("name", data=f"Synthetic model {args.model_number}") + rotation = sample.create_dataset("rotation_angle", data=frame_angles) + rotation.attrs["units"] = "deg" + nxdata = nxgroup(entry, "data", "NXdata") + nxdata.attrs["signal"] = "data" + nxdata.attrs["axes"] = np.asarray( + ["rotation_angle", ".", "."], dtype=h5py.string_dtype() + ) + nxdata["data"], nxdata["image_key"], nxdata["rotation_angle"] = data, key, rotation + process = nxgroup(entry, "generator", "NXprocess") + process.create_dataset("program", data="nxs_generator.py") + process.create_dataset("version", data="2.0") + process.create_dataset("date", data=datetime.now(timezone.utc).isoformat()) + process.create_dataset("sequence_index", data=1) + process.create_dataset( + "configuration", data=json.dumps(metadata, default=str, sort_keys=True) + ) + process.create_dataset( + "projection_frame_indices", data=np.arange(nd + nf, len(keys)) + ) + return data, process, nd, nf + + +def main(args): + import tomophantom + from tomophantom import TomoP3D + from tomophantom.flatsgen import synth_flats + + flat_config, art = settings(args) + args.output_path = Path(args.output_path) + if args.output_path.exists() and not args.overwrite: + raise FileExistsError( + f"{args.output_path} exists; use --overwrite to replace it." + ) + height, count, width = args.sinogram_shape + rows = min(height, args.chunk_rows) + angles = np.linspace(*args.angle_range, count, dtype="float32") + library = str( + Path(tomophantom.__file__).parent / "phantomlib" / "Phantom3DLibrary.dat" + ) + full = args.realistic or bool(art) + budget = int(psutil.virtual_memory().available * args.memory_fraction) + if args.max_memory_gb is not None: + budget = min(budget, int(args.max_memory_gb * 2**30)) + # Includes float work arrays, calibration stacks, maps and projection scratch. + estimate = ( + height * count * width * 64 + height * width * (8 * args.flats + 256) + if full + else rows * count * width * 16 + ) + 64 * 2**20 + if estimate > budget: + raise MemoryError( + f"Estimated working memory {estimate / 2**30:.2f} GiB exceeds " + f"budget {budget / 2**30:.2f} GiB. Reduce shape" + + ( + "; realistic/artefact modes require a full volume." + if full + else " or --chunk-rows." + ) + ) + print( + f"Estimated working memory: {estimate / 2**30:.2f} GiB; " + f'mode: {"full-volume synthesis" if full else "streamed clean projections"}' + ) + random.seed(args.seed) + np.random.seed(args.seed) + rng = np.random.default_rng(args.seed) + + def analytic(start, stop): + data = TomoP3D.ModelSinoSub( + args.model_number, + int(width / np.sqrt(2)), + width, + height, + (start, stop), + angles, + library, + ) + checked(data, (stop - start, count, width), "ModelSinoSub") + if data.min() < 0: + raise ValueError( + "This generator requires non-negative attenuation projections." + ) + return data + + projections = None + if full: + projections = np.empty((height, count, width), dtype="float32") + for start in range(0, height, rows): + stop = min(start + rows, height) + projections[start:stop] = analytic(start, stop) + clean_max = float(projections.max()) + else: + # First pass gives one common scaling for all chunks, avoiding seams. + clean_max = max( + float(analytic(s, min(s + rows, height)).max()) + for s in range(0, height, rows) + ) + if clean_max <= 0: + raise ValueError("The selected phantom has no positive projection data.") + shifts = None + if art and args.artefacts_stage == "line-integrals": + projections, shifts = apply_artefacts(projections, art, args.seed) + np.maximum(projections, 0, out=projections) + flats = None + synthesis_peak = None + if args.realistic: + synthesis_peak = float(projections.max()) + if synthesis_peak <= 0: + raise ValueError( + "Artefacts removed all positive signal before synth_flats." + ) + projections, flats, response_map = synth_flats(projections, **flat_config) + checked(projections, (height, count, width), "synth_flats projections") + checked(flats, (height, args.flats, width), "synth_flats flats") + if not np.isfinite(response_map).all(): + raise ValueError("synth_flats produced an invalid detector-response map.") + del response_map + if art and args.artefacts_stage == "raw": + projections, shifts = apply_artefacts( + projections.astype("float32") / 65535, art, args.seed + ) + projections *= 65535 # No new max normalisation after corruption. + clean_scale = float(projections.max()) if full and not args.realistic else clean_max + clean_scale = max(clean_scale, np.finfo("float32").tiny) + try: + version = importlib.metadata.version("tomophantom") + except importlib.metadata.PackageNotFoundError: + version = getattr(tomophantom, "__version__", "source checkout / unknown") + metadata = dict( + arguments=vars(args), + flat_settings=flat_config if args.realistic else {}, + artefacts=art, + tomophantom_version=version, + clean_projection_max=clean_max, + clean_storage_scale=None if args.realistic else clean_scale, + synth_flats_input_max=synthesis_peak, + data_domain=( + "raw_detector_counts" if args.realistic else "scaled_line_integrals" + ), + frame_order="darks, flats, projections", + estimated_working_bytes=estimate, + ) + args.output_path.parent.mkdir(parents=True, exist_ok=True) + # Write beside destination, then publish only a complete file. + temp = tempfile.NamedTemporaryFile( + dir=args.output_path.parent, + prefix=args.output_path.name + ".", + suffix=".partial", + delete=False, + ) + temporary = Path(temp.name) + temp.close() + clipped = 0 + + def quantise(block, detector_signal=False): + nonlocal clipped + values = np.array(block, dtype="float32", copy=True) + if detector_signal: + values *= args.signal_scale + values += rng.poisson(args.dark_mean, size=values.shape).astype("float32") + if args.read_noise: + values += rng.normal(0, args.read_noise, size=values.shape).astype( + "float32" + ) + if not np.isfinite(values).all(): + raise ValueError("Non-finite values encountered before uint16 conversion.") + clipped += int(np.count_nonzero((values < 0) | (values > 65535))) + return np.rint(np.clip(values, 0, 65535)).astype("uint16") + + try: + with h5py.File(temporary, "w") as f: + output, process, nd, nf = initialise(f, args, angles, metadata) + # One shared dark-current distribution for projections, flats and darks. + for frame in range(nd): + output[frame] = quantise(np.zeros((height, width), "float32"), True) + for frame in range(nf): + output[nd + frame] = quantise(flats[:, frame, :], True) + for start in range(0, height, rows): + stop = min(start + rows, height) + block = projections[start:stop] if full else analytic(start, stop) + if not args.realistic: + block = block / clean_scale * 65535 + output[nd + nf :, start:stop] = np.swapaxes( + quantise(block, args.realistic), 0, 1 + ) + if shifts is not None: + ds = process.create_dataset("artefact_shifts", data=shifts) + ds.attrs.update( + units="pixel", + description="As returned by TomoPhantom; projection frames only", + ) + process.create_dataset("clipped_voxels", data=clipped) + f["entry"].create_dataset( + "end_time", data=datetime.now(timezone.utc).isoformat() + ) + if args.overwrite: + temporary.replace(args.output_path) + else: + # Atomic no-clobber publication on the same filesystem. + import os + + os.link(temporary, args.output_path) + temporary.unlink() + finally: + temporary.unlink(missing_ok=True) + print( + f"Wrote {args.output_path}: {nd} darks + {nf} flats + {count} projections; " + f"frame shape {height} x {width}; clipped voxels: {clipped}" + ) + + +if __name__ == "__main__": + p = parser() + try: + main(p.parse_args()) + except (ValueError, MemoryError, FileExistsError, OSError) as exc: + p.exit(2, f"Error: {exc}\n") diff --git a/scripts/test_nxs_generator.py b/scripts/test_nxs_generator.py new file mode 100644 index 00000000..51e9e220 --- /dev/null +++ b/scripts/test_nxs_generator.py @@ -0,0 +1,158 @@ +"""Run with pytest in an environment with the real TomoPhantom library installed.""" + +import json + +import h5py +import numpy as np +import pytest + +import nxs_generator as gen + + +def args_for(tmp_path, *extra): + return gen.parser().parse_args( + [ + "-s", + "64", + "24", + "64", + "--flats", + "4", + "--darks", + "3", + "-o", + str(tmp_path / "phantom.nxs"), + *extra, + ] + ) + + +@pytest.mark.parametrize( + "stage,config", + [ + ("raw", {}), + ( + "raw", + { + "zingers_percentage": 0.02, + "stripes_percentage": 2, + "stripes_maxthickness": 1, + "stripes_intensity": 0.05, + }, + ), + ( + "line-integrals", + {"pve_strength": 1, "noise_type": "Gaussian", "noise_amplitude": 0.01}, + ), + ("raw", {"datashifts_maxamplitude_pixel": 1}), + ], +) +def test_realistic_frames_and_links(tmp_path, stage, config): + config_file = tmp_path / "artefacts.json" + config_file.write_text(json.dumps(config)) + a = args_for( + tmp_path, + "--realistic", + "--artefacts", + str(config_file), + "--artefacts-stage", + stage, + ) + gen.main(a) + with h5py.File(a.output_path) as f: + detector = f["entry/instrument/detector"] + data = f["entry/data/data"] + assert data.shape == (31, 64, 64) + assert data.dtype == np.dtype("uint16") + assert data.id == detector["data"].id + assert f["entry/data/image_key"].id == detector["image_key"].id + assert f["entry/data/rotation_angle"].id == f["entry/sample/rotation_angle"].id + np.testing.assert_array_equal( + detector["image_key"][:], [2] * 3 + [1] * 4 + [0] * 24 + ) + angles = f["entry/sample/rotation_angle"][:] + np.testing.assert_allclose( + angles[7:], np.linspace(0, 179.9, 24, dtype="float32") + ) + assert len(angles) == len(detector["image_key"]) + assert f["entry/sample/rotation_angle"].attrs["units"] == "deg" + assert data[:3].mean() < data[3:7].mean() + assert data[3:7].mean() > data[7:].mean() + assert data.maxshape == data.shape + if "datashifts_maxamplitude_pixel" in config: + assert f["entry/generator/artefact_shifts"].shape == (24, 2) + np.testing.assert_array_equal( + f["entry/generator/projection_frame_indices"][:], np.arange(7, 31) + ) + + +def test_clean_scaling_is_independent_of_chunk_rows(tmp_path): + a = args_for(tmp_path, "--chunk-rows", "7") + gen.main(a) + with h5py.File(a.output_path) as f: + expected = f["entry/data/data"][:] + assert np.all(f["entry/data/image_key"][:] == 0) + a.chunk_rows = 64 + a.output_path = tmp_path / "whole.nxs" + gen.main(a) + with h5py.File(a.output_path) as f: + np.testing.assert_array_equal(expected, f["entry/data/data"][:]) + + +def test_same_seed_repeats(tmp_path): + a = args_for(tmp_path, "--realistic") + gen.main(a) + with h5py.File(a.output_path) as f: + first = f["entry/data/data"][:] + a.output_path = tmp_path / "repeat.nxs" + gen.main(a) + with h5py.File(a.output_path) as f: + np.testing.assert_array_equal(first, f["entry/data/data"][:]) + + +def test_memory_guard_leaves_no_file(tmp_path): + a = args_for(tmp_path, "--realistic", "--max-memory-gb", ".001") + with pytest.raises(MemoryError): + gen.main(a) + assert not a.output_path.exists() + + +def test_existing_file_is_preserved(tmp_path): + a = args_for(tmp_path) + a.output_path.write_bytes(b"existing file") + with pytest.raises(FileExistsError): + gen.main(a) + assert a.output_path.read_bytes() == b"existing file" + + +@pytest.mark.parametrize( + "config", + [ + {"stripe_type": "full"}, # Actual API keyword is stripes_type. + { + "noise_type": "Poisson" + }, # Cannot use line-integral Poisson API on raw counts. + {"noise_prelog": True}, + {"datashifts_maxamplitude_pixel": 0}, + ], +) +def test_invalid_artefact_config(tmp_path, config): + path = tmp_path / "bad.json" + path.write_text(json.dumps(config)) + with pytest.raises(ValueError): + gen.main(args_for(tmp_path, "--realistic", "--artefacts", str(path))) + assert not (tmp_path / "phantom.nxs").exists() + + +def test_failure_does_not_replace_existing_output(tmp_path, monkeypatch): + a = args_for(tmp_path, "--overwrite") + a.output_path.write_bytes(b"keep this") + + def fail(*args): + raise RuntimeError("simulated writer failure") + + monkeypatch.setattr(gen, "initialise", fail) + with pytest.raises(RuntimeError): + gen.main(a) + assert a.output_path.read_bytes() == b"keep this" + assert not list(tmp_path.glob("*.partial"))