diff --git a/AGENTS.md b/AGENTS.md index f37a4cb..97ad735 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -21,6 +21,12 @@ Read the most local `AGENTS.md` for the files you touch: user-facing unit boundaries. - `orgui/datautils/xrayutils/AGENTS.md`: scientific core, detector geometry, reciprocal-space math, CTR calculations, and related tests. +- `orgui/reconstruction_*.py`, `orgui/backend/`: no nested `AGENTS.md` yet. + Follow this file plus `orgui/datautils/xrayutils/AGENTS.md` — the + reconstruction pipeline is scientific/geometry code built on + `orgui/datautils/xrayutils/reconstruction.py`, and `orgui/backend/` handles + the beamline-specific metadata normalization called out in + `orgui/app/AGENTS.md`. - `doc/AGENTS.md`: Sphinx documentation source, the changelog/release-notes workflow, and what counts as a user-facing change worth documenting. @@ -79,6 +85,7 @@ it: - `pytest orgui/datautils/xrayutils/test/test_HKLcalc.py` - `pytest orgui/datautils/xrayutils/test/test_DetectorCalibration.py` - `pytest orgui/datautils/xrayutils/test/test_CTRcalc.py` +- `pytest orgui/datautils/xrayutils/test/test_reconstruction*.py` When changing config-loading or unit-conversion behavior, inspect `examples/config_minimal` and the other `examples/config_*` files. diff --git a/CHANGELOG.md b/CHANGELOG.md index 475379a..cb97ed9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -195,6 +195,20 @@ ESRF ID31 beamline support and reciprocal-space display: raised before showing the dialog. Subscans such as ``1.10`` are no longer mistaken for the ``1.1`` fast-counter subscan. +Reciprocal-space reconstruction: + +- Added a centralized, out-of-core reciprocal-space reconstruction pipeline. + A new "Reconstruct reciprocal space" dialog (Configuration menu) defines + HKL/Q output grids, previews coverage and storage cost, and prepares, + runs, and resumes jobs. Jobs can run locally, as SGE/Slurm cluster batch + arrays with parallel submap reduction, or from the command line via a new + ``reconstruction_cli`` entry point. +- Added a shared HDF5 output-settings dialog (chunk shape, compression) + reachable from both the reconstruction dialog and the Configuration menu. +- Exposure-time normalization and monitor-counter corrections now live in + the reconstruction dialog rather than the shared scan-options panel, since + they affect reconstruction output only, not ROI/CTR image integration. + ## [1.5.0] (2026-06-07) diff --git a/CLAUDE.md b/CLAUDE.md new file mode 100644 index 0000000..43c994c --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1 @@ +@AGENTS.md diff --git a/benchmarks/benchmark_reconstruction_alignment.py b/benchmarks/benchmark_reconstruction_alignment.py new file mode 100644 index 0000000..14acdc6 --- /dev/null +++ b/benchmarks/benchmark_reconstruction_alignment.py @@ -0,0 +1,149 @@ +"""Measure native reconstruction sensitivity to input-buffer alignment.""" + +from __future__ import annotations + +import argparse +from dataclasses import replace +import gc +import json +from pathlib import Path +from statistics import median +from time import perf_counter + +import numpy as np + +from orgui.datautils.xrayutils.reconstruction import ( + _detector_corner_rays, + _kernel_for_grid, +) +from orgui.reconstruction_job import ( + _correction_pipeline, + _load_assets, + read_job, +) + + +def _aligned_copy(values, address_mod_64): + values = np.ascontiguousarray(values) + storage = np.empty(values.nbytes + 63, dtype=np.uint8) + offset = (address_mod_64 - storage.ctypes.data) % 64 + result = storage[offset : offset + values.nbytes].view(values.dtype) + result = result.reshape(values.shape) + np.copyto(result, values) + if result.ctypes.data % 64 != address_mod_64: + raise RuntimeError("Failed to construct requested input alignment") + return result + + +def _arguments(): + parser = argparse.ArgumentParser() + parser.add_argument("job", type=Path) + parser.add_argument("--frame", type=int, default=0) + parser.add_argument("--tile-size", type=int, default=1024) + parser.add_argument("--threads", type=int, default=4) + parser.add_argument("--depth", type=int, default=2) + parser.add_argument("--repeats", type=int, default=5) + parser.add_argument( + "--offsets", + type=int, + nargs="+", + default=tuple(range(0, 64, 8)), + ) + return parser.parse_args() + + +def main(): + """Benchmark naturally aligned and deliberately offset native inputs.""" + arguments = _arguments() + job = read_job(arguments.job) + config = job.config_data + scan = job.scan + payload = scan.get_raw_img(arguments.frame) + image = np.asarray(payload.img) + tile = ( + 0, + min(arguments.tile_size, image.shape[0]), + 0, + min(arguments.tile_size, image.shape[1]), + ) + selection = np.s_[tile[0] : tile[1], tile[2] : tile[3]] + correction = _correction_pipeline( + config, + scan, + _load_assets(job), + {}, + ) + corrected = correction.correct_frame(payload, image, arguments.frame) + inputs = [ + np.ascontiguousarray(corrected[0][selection], dtype=np.float64), + np.ascontiguousarray(corrected[1][selection], dtype=np.float64), + np.ascontiguousarray(corrected[2][selection], dtype=bool), + _detector_corner_rays(config.detector, tile), + ] + bounds = scan.exposure_angle_bounds( + config, + fallback=job.angle_fallback, + )[arguments.frame] + angles_start = np.ascontiguousarray(bounds[0], dtype=np.float64) + angles_end = np.ascontiguousarray(bounds[1], dtype=np.float64) + spec = replace( + job.internal_spec(), + max_depth=arguments.depth, + ) + kernel = _kernel_for_grid( + spec, + spec.grids[0], + config.ub_calculator, + threads=arguments.threads, + memory_budget_bytes=spec.memory_budget_bytes, + ) + + offsets = tuple(arguments.offsets) + if any(offset < 0 or offset >= 64 for offset in offsets): + raise ValueError("Alignment offsets must be in [0, 64)") + aligned_inputs = { + offset: tuple(_aligned_copy(values, offset) for values in inputs) + for offset in offsets + } + kernel.accumulate( + *aligned_inputs[0], + angles_start, + angles_end, + ) + timings = {offset: [] for offset in offsets} + gc.disable() + try: + for repeat in range(arguments.repeats): + order = offsets if repeat % 2 == 0 else tuple(reversed(offsets)) + for offset in order: + started = perf_counter() + kernel.accumulate( + *aligned_inputs[offset], + angles_start, + angles_end, + ) + timings[offset].append(perf_counter() - started) + finally: + gc.enable() + + baseline = median(timings[0]) + result = { + "tile": tile, + "threads": arguments.threads, + "depth": arguments.depth, + "repeats": arguments.repeats, + "results": [ + { + "address_mod_64": offset, + "median_seconds": median(timings[offset]), + "relative_to_aligned": median(timings[offset]) / baseline, + "samples": timings[offset], + } + for offset in offsets + ], + } + print(json.dumps(result, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/benchmark_reconstruction_mapping.py b/benchmarks/benchmark_reconstruction_mapping.py new file mode 100644 index 0000000..6ba6da3 --- /dev/null +++ b/benchmarks/benchmark_reconstruction_mapping.py @@ -0,0 +1,392 @@ +"""Benchmark reciprocal-space mapping layouts on a prepared orGUI job.""" + +from __future__ import annotations + +import argparse +from concurrent.futures import ThreadPoolExecutor +from dataclasses import replace +import gc +import json +import math +from pathlib import Path +from time import perf_counter + +import numpy as np + +from orgui.datautils.xrayutils.reconstruction import ( + _detector_corner_rays, + _kernel_for_grid, + _map_frame_range, + _reduce_batches, + _xxh3_128, +) +import orgui.datautils.xrayutils.reconstruction as reconstruction +from orgui.reconstruction_job import ( + _correction_pipeline, + _load_assets, + read_job, +) + + +def _arguments(): + parser = argparse.ArgumentParser() + parser.add_argument("job", type=Path) + parser.add_argument("--frame", type=int, default=0) + parser.add_argument("--threads", type=int, nargs="+", default=[4, 6]) + parser.add_argument( + "--layouts", + choices=("current", "equal"), + nargs="+", + default=["current", "equal"], + ) + parser.add_argument("--tile-sizes", type=int, nargs="+") + parser.add_argument("--work-blocks", type=int, nargs="+") + parser.add_argument("--depths", type=int, nargs="+") + parser.add_argument("--parallel-frames", type=int, default=0) + parser.add_argument("--native-profile", action="store_true") + parser.add_argument("--profile-load-frames", type=int, default=0) + parser.add_argument("--profile-load-workers", type=int, default=1) + parser.add_argument("--output", type=Path) + return parser.parse_args() + + +def _factor_grid(count, rows, columns): + candidates = [] + for row_count in range(1, count + 1): + if count % row_count: + continue + column_count = count // row_count + tile_rows = rows / row_count + tile_columns = columns / column_count + distortion = abs(math.log(tile_rows / tile_columns)) + candidates.append((distortion, row_count, column_count)) + _, row_count, column_count = min(candidates) + return row_count, column_count + + +def _equal_tiles(rows, columns, count): + row_count, column_count = _factor_grid(count, rows, columns) + row_edges = [(index * rows) // row_count for index in range(row_count + 1)] + column_edges = [ + (index * columns) // column_count + for index in range(column_count + 1) + ] + return tuple( + ( + row_edges[row], + row_edges[row + 1], + column_edges[column], + column_edges[column + 1], + ) + for row in range(row_count) + for column in range(column_count) + ) + + +def _fixed_tiles(rows, columns, tile_rows, tile_columns): + return tuple( + ( + row, + min(row + tile_rows, rows), + column, + min(column + tile_columns, columns), + ) + for row in range(0, rows, tile_rows) + for column in range(0, columns, tile_columns) + ) + + +def _run_case( + *, + name, + tiles, + threads, + spec, + job, + config, + payload, + image, + angles, + assets, + native_profile=False, +): + provenance = {} + correct = _correction_pipeline(config, job.scan, assets, provenance) + rays = { + tile: _detector_corner_rays(config.detector, tile) + for tile in tiles + } + kernel = _kernel_for_grid( + spec, + spec.grids[0], + config.ub_calculator, + threads=threads, + memory_budget_bytes=spec.memory_budget_bytes, + ) + batches = [] + native_profiles = [] + correction_seconds = 0.0 + native_seconds = 0.0 + started = perf_counter() + correction_started = perf_counter() + corrected_frame = correct.correct_frame(payload, image, 0) + correction_seconds = perf_counter() - correction_started + for tile in tiles: + row_start, row_stop, column_start, column_stop = tile + selection = np.s_[ + row_start:row_stop, column_start:column_stop + ] + intensity, variance, mask = ( + values[selection] for values in corrected_frame + ) + native_started = perf_counter() + batch = kernel.accumulate( + np.ascontiguousarray(intensity, dtype=np.float64), + np.ascontiguousarray(variance, dtype=np.float64), + np.ascontiguousarray(mask, dtype=bool), + rays[tile], + angles[0], + angles[1], + profile=native_profile, + ) + if native_profile: + native_profiles.append(batch.pop("_profile")) + batches.append(batch) + native_seconds += perf_counter() - native_started + reduction_started = perf_counter() + reduced = _reduce_batches(batches) + reduction_seconds = perf_counter() - reduction_started + elapsed = perf_counter() - started + rows, columns = image.shape + result = { + "name": name, + "threads": threads, + "tiles": len(tiles), + "largest_tile_pixels": max( + (tile[1] - tile[0]) * (tile[3] - tile[2]) + for tile in tiles + ), + "seconds": elapsed, + "megapixels_per_second": rows * columns / elapsed / 1e6, + "correction_seconds": correction_seconds, + "native_seconds": native_seconds, + "reduction_seconds": reduction_seconds, + "records": int(reduced["chunk_id"].size), + } + if native_profiles: + result["native_profile"] = { + name: ( + max(profile[name] for profile in native_profiles) + if name == "maximum_weights_per_pixel" + else sum(profile[name] for profile in native_profiles) + ) + for name in native_profiles[0] + } + result["result_xxh3_128"] = { + name: _xxh3_128(values) + for name, values in reduced.items() + } + del batches, reduced, kernel, rays + gc.collect() + return result + + +def main(): + """Benchmark current square tiling against equal-area thread-count tiling.""" + arguments = _arguments() + job = read_job(arguments.job) + config = job.config_data + scan = job.scan + if arguments.profile_load_frames: + frame_count = min(len(scan), arguments.profile_load_frames) + started = perf_counter() + total_bytes = 0 + def load(frame): + return np.asarray(scan.get_raw_img(frame).img).nbytes + + with ThreadPoolExecutor( + max_workers=arguments.profile_load_workers + ) as executor: + total_bytes = sum(executor.map(load, range(frame_count))) + elapsed = perf_counter() - started + print( + json.dumps( + { + "frames": frame_count, + "workers": arguments.profile_load_workers, + "seconds": elapsed, + "frames_per_second": frame_count / elapsed, + "megabytes_per_second": total_bytes / elapsed / 1e6, + }, + sort_keys=True, + ) + ) + return + payload = scan.get_raw_img(arguments.frame) + image = np.asarray(payload.img) + bounds = scan.exposure_angle_bounds( + config, + fallback=job.angle_fallback, + ) + angles = np.ascontiguousarray(bounds[arguments.frame], dtype=np.float64) + assets = _load_assets(job) + spec = job.internal_spec() + rows, columns = image.shape + cases = [] + for threads in arguments.threads: + tile_sizes = arguments.tile_sizes or [1024] + work_blocks = arguments.work_blocks or [spec.work_block_pixels] + depths = arguments.depths or [spec.max_depth] + for tile_size in tile_sizes: + if "current" in arguments.layouts: + for work_block in work_blocks: + for depth in depths: + cases.append( + ( + f"fixed-{tile_size}/{threads}/" + f"block-{work_block}/depth-{depth}", + _fixed_tiles( + rows, + columns, + tile_size, + tile_size, + ), + threads, + work_block, + depth, + ) + ) + if "equal" in arguments.layouts and arguments.tile_sizes is None: + for work_block in work_blocks: + for depth in depths: + cases.append( + ( + f"equal-{threads}/{threads}/" + f"block-{work_block}/depth-{depth}", + _equal_tiles(rows, columns, threads), + threads, + work_block, + depth, + ) + ) + if arguments.parallel_frames: + frame_indices = list( + range( + arguments.frame, + min(len(scan), arguments.frame + arguments.parallel_frames), + ) + ) + payloads = { + frame: scan.get_raw_img(frame) + for frame in frame_indices + } + for payload in payloads.values(): + np.asarray(payload.img) + original_write = reconstruction._write_parquet + original_checksum = reconstruction._uri_checksum_and_size + reconstruction._write_parquet = lambda *args, **kwargs: None + reconstruction._uri_checksum_and_size = lambda *args, **kwargs: ("", 0) + try: + parallel_results = [] + for name, tiles, threads, work_block, depth in cases: + case_spec = replace( + spec, + threads=24, + work_block_pixels=work_block, + max_depth=depth, + ) + rays = { + tile: _detector_corner_rays(config.detector, tile) + for tile in tiles + } + fingerprints = { + tile: _xxh3_128(values) + for tile, values in rays.items() + } + worker_count = max(1, 24 // threads) + provenance = {} + correct = _correction_pipeline( + config, scan, assets, provenance + ) + started = perf_counter() + + def map_frame(frame): + return _map_frame_range( + case_spec, + scan, + config.detector, + config.ub_calculator, + (frame, frame + 1), + tiles, + bounds[frame : frame + 1], + Path("."), + correction_pipeline=correct, + job_digest=job.digest, + image_payloads={frame: payloads[frame]}, + corner_rays=rays, + corner_rays_fingerprints=fingerprints, + kernel_threads=threads, + kernel_memory_budget_bytes=spec.memory_budget_bytes, + accumulation_budget_bytes=spec.memory_budget_bytes, + ) + + with ThreadPoolExecutor(max_workers=worker_count) as executor: + manifests = list(executor.map(map_frame, frame_indices)) + elapsed = perf_counter() - started + result = { + "name": name, + "threads": threads, + "image_workers": worker_count, + "tiles": len(tiles), + "frames": len(frame_indices), + "seconds": elapsed, + "frames_per_second": len(frame_indices) / elapsed, + "records": sum( + sum(partition.rows for partition in manifest.partitions) + for manifest in manifests + ), + } + parallel_results.append(result) + print(json.dumps(result, sort_keys=True), flush=True) + del manifests, rays + gc.collect() + if arguments.output: + arguments.output.write_text( + json.dumps(parallel_results, indent=2, sort_keys=True), + encoding="utf-8", + ) + return + finally: + reconstruction._write_parquet = original_write + reconstruction._uri_checksum_and_size = original_checksum + results = [] + for name, tiles, threads, work_block, depth in cases: + result = _run_case( + name=name, + tiles=tiles, + threads=threads, + spec=replace( + spec, + threads=24, + work_block_pixels=work_block, + max_depth=depth, + ), + job=job, + config=config, + payload=payload, + image=image, + angles=angles, + assets=assets, + native_profile=arguments.native_profile, + ) + results.append(result) + print(json.dumps(result, sort_keys=True), flush=True) + if arguments.output: + arguments.output.write_text( + json.dumps(results, indent=2, sort_keys=True), + encoding="utf-8", + ) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/benchmark_reconstruction_reduce.py b/benchmarks/benchmark_reconstruction_reduce.py new file mode 100644 index 0000000..cc909b8 --- /dev/null +++ b/benchmarks/benchmark_reconstruction_reduce.py @@ -0,0 +1,139 @@ +"""Benchmark bounded reciprocal-space reduction on an existing job.""" + +from __future__ import annotations + +import argparse +from collections import defaultdict +import json +import math +from pathlib import Path +from time import perf_counter + +from orgui.datautils.xrayutils import reconstruction + + +def _arguments(): + parser = argparse.ArgumentParser() + parser.add_argument("job", type=Path) + parser.add_argument("--chunks", type=int, default=64) + parser.add_argument("--memory-gib", type=float, default=16.0) + return parser.parse_args() + + +def main(): + """Stream and reduce representative chunks from the largest map bucket.""" + arguments = _arguments() + job = json.loads(arguments.job.read_text(encoding="utf-8")) + manifests = [ + reconstruction._read_manifest(path) + for path in job["map_manifests"] + ] + grouped = defaultdict(list) + for manifest in manifests: + for partition in manifest.partitions: + grouped[(partition.grid_name, partition.bucket)].append(partition) + group_key, partitions = max( + grouped.items(), + key=lambda item: sum(partition.rows for partition in item[1]), + ) + grid_name, bucket = group_key + spec = reconstruction._ReconstructionSpec.from_dict(manifests[0].spec) + grid = next(grid for grid in spec.grids if grid.grid_name == grid_name) + memory_bytes = int(arguments.memory_gib * 1024**3) + batch_rows = max( + 4096, + min( + 131072, + memory_bytes // (max(1, len(partitions)) * 48 * 4), + ), + ) + readers = [ + reconstruction._ParquetRangeReader( + partition.uri, + batch_size=batch_rows, + ) + for partition in sorted(partitions, key=lambda item: item.uri) + ] + chunk_grid = tuple( + math.ceil(size / chunk) + for size, chunk in zip(grid.shape, grid.chunk_shape) + ) + chunk_start = bucket * spec.partition_chunk_span + chunk_stop = min( + math.prod(chunk_grid), + chunk_start + spec.partition_chunk_span, + ) + metadata_minima = [] + for reader in readers: + for row_group in range(reader.parquet.metadata.num_row_groups): + statistics = reader.parquet.metadata.row_group(row_group).column( + reader.chunk_column + ).statistics + if statistics is not None and statistics.has_min_max: + metadata_minima.append(int(statistics.min)) + if metadata_minima: + chunk_start = max(chunk_start, min(metadata_minima)) + chunk_stop = min(chunk_stop, chunk_start + arguments.chunks) + + input_rows = 0 + output_rows = 0 + started = perf_counter() + for ordinal, chunk_id in enumerate( + range(chunk_start, chunk_stop), + start=1, + ): + coordinates = reconstruction._chunk_coordinates(chunk_id, grid) + local_shape = tuple( + min(chunk, size - coordinate * chunk) + for coordinate, chunk, size in zip( + coordinates, + grid.chunk_shape, + grid.shape, + ) + ) + local_stop = ( + (local_shape[0] - 1) + * grid.chunk_shape[1] + * grid.chunk_shape[2] + + (local_shape[1] - 1) * grid.chunk_shape[2] + + local_shape[2] + ) + levels = [] + for reader in readers: + batch = reader.read(chunk_id, 0, local_stop) + input_rows += int(batch["chunk_id"].size) + if not batch["chunk_id"].size: + continue + level = 0 + while level < len(levels) and levels[level] is not None: + batch = reconstruction._merge_sorted_batches( + levels[level], batch + ) + levels[level] = None + level += 1 + if level == len(levels): + levels.append(batch) + else: + levels[level] = batch + reduced = reconstruction._empty_batch() + for batch in reversed(levels): + if batch is not None: + reduced = reconstruction._merge_sorted_batches(reduced, batch) + output_rows += int(reduced["chunk_id"].size) + if ordinal % 8 == 0 or chunk_id + 1 == chunk_stop: + elapsed = perf_counter() - started + print( + f"{ordinal}/{chunk_stop - chunk_start} chunks; " + f"{input_rows / elapsed:,.0f} input rows/s; " + f"{input_rows * 48 / elapsed / 1024**2:,.1f} MiB/s" + ) + elapsed = perf_counter() - started + print( + f"bucket={bucket}, partitions={len(partitions)}, " + f"batch_rows={batch_rows}, input_rows={input_rows}, " + f"output_rows={output_rows}, seconds={elapsed:.3f}" + ) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/benchmark_reconstruction_storage.py b/benchmarks/benchmark_reconstruction_storage.py new file mode 100644 index 0000000..c359662 --- /dev/null +++ b/benchmarks/benchmark_reconstruction_storage.py @@ -0,0 +1,171 @@ +"""Benchmark reconstruction map-task and Parquet storage controls.""" + +from __future__ import annotations + +import argparse +from concurrent.futures import ThreadPoolExecutor +from dataclasses import replace +import json +from pathlib import Path +from tempfile import TemporaryDirectory +from time import perf_counter + +import numpy as np + +from orgui.datautils.xrayutils.reconstruction import ( + _detector_corner_rays, + _map_frame_range, + _xxh3_128, +) +from orgui.reconstruction_job import ( + _correction_pipeline, + _load_assets, + read_job, +) + +from .benchmark_reconstruction_mapping import _fixed_tiles + + +def _arguments(): + parser = argparse.ArgumentParser() + parser.add_argument("job", type=Path) + parser.add_argument("--frames", type=int, default=16) + parser.add_argument("--threads", type=int, default=2) + parser.add_argument("--tile-size", type=int, default=1024) + parser.add_argument("--work-block", type=int, default=65536) + parser.add_argument("--output", type=Path) + return parser.parse_args() + + +def _cases(): + return ( + ("baseline", 4, 1024, 2000), + ("frame-batch-1", 1, 1024, 2000), + ("frame-batch-16", 16, 1024, 2000), + ("accumulation-64", 16, 64, 2000), + ("accumulation-256", 16, 256, 2000), + ("accumulation-2048", 16, 2048, 2000), + ("partition-span-256", 4, 1024, 256), + ("partition-span-1000", 4, 1024, 1000), + ("partition-span-4000", 4, 1024, 4000), + ) + + +def main(): + """Run controlled map/Parquet cases while retaining 24 total threads.""" + arguments = _arguments() + job = read_job(arguments.job) + config = job.config_data + scan = job.scan + assets = _load_assets(job) + frame_count = min(arguments.frames, len(scan)) + frame_indices = list(range(frame_count)) + payloads = {frame: scan.get_raw_img(frame) for frame in frame_indices} + for payload in payloads.values(): + np.asarray(payload.img) + bounds = scan.exposure_angle_bounds( + config, + fallback=job.angle_fallback, + ) + rows, columns = config.detector.detector.shape + tiles = _fixed_tiles( + rows, + columns, + arguments.tile_size, + arguments.tile_size, + ) + rays = { + tile: _detector_corner_rays(config.detector, tile) + for tile in tiles + } + fingerprints = { + tile: _xxh3_128(values) + for tile, values in rays.items() + } + root = arguments.output or Path("build/reconstruction-storage-benchmark") + root.mkdir(parents=True, exist_ok=True) + results = [] + for name, frame_batch, accumulation_mib, partition_span in _cases(): + ranges = [ + (start, min(start + frame_batch, frame_count)) + for start in range(0, frame_count, frame_batch) + ] + spec = replace( + job.internal_spec(), + threads=24, + work_block_pixels=arguments.work_block, + partition_chunk_span=partition_span, + ) + provenance = {} + correct = _correction_pipeline( + config, + scan, + assets, + provenance, + ) + started = perf_counter() + with TemporaryDirectory(prefix=f"{name}-", dir=root) as output: + + def map_range(frame_range): + return _map_frame_range( + spec, + scan, + config.detector, + config.ub_calculator, + frame_range, + tiles, + bounds[frame_range[0] : frame_range[1]], + output, + correction_pipeline=correct, + job_digest=job.digest, + image_payloads={ + frame: payloads[frame] + for frame in range(*frame_range) + }, + corner_rays=rays, + corner_rays_fingerprints=fingerprints, + kernel_threads=arguments.threads, + kernel_memory_budget_bytes=spec.memory_budget_bytes, + accumulation_budget_bytes=accumulation_mib * 1024**2, + ) + + workers = min( + len(ranges), + max(1, 24 // arguments.threads), + ) + with ThreadPoolExecutor(max_workers=workers) as executor: + manifests = list(executor.map(map_range, ranges)) + elapsed = perf_counter() - started + partitions = [ + partition + for manifest in manifests + for partition in manifest.partitions + ] + result = { + "name": name, + "frame_batch": frame_batch, + "accumulation_MiB": accumulation_mib, + "partition_span": partition_span, + "workers": workers, + "seconds": elapsed, + "frames_per_second": frame_count / elapsed, + "parquet_files": len(partitions), + "parquet_MiB": sum( + partition.size_bytes or 0 for partition in partitions + ) + / 1024**2, + "segments": sum( + int(manifest.metadata["accumulation_segments"]) + for manifest in manifests + ), + } + results.append(result) + print(json.dumps(result, sort_keys=True), flush=True) + (root / "results.json").write_text( + json.dumps(results, indent=2, sort_keys=True), + encoding="utf-8", + ) + + +if __name__ == "__main__": + main() diff --git a/doc/source/index.rst b/doc/source/index.rst index 76082ac..73b1c6f 100644 --- a/doc/source/index.rst +++ b/doc/source/index.rst @@ -19,6 +19,7 @@ large 2D detectors. geometry ub_matrix image_integration + reciprocal_space_reconstruction acceleration_backends entry_points benchmarks/ctr_accel_backends diff --git a/doc/source/reciprocal_space_reconstruction.rst b/doc/source/reciprocal_space_reconstruction.rst new file mode 100644 index 0000000..c8db2ec --- /dev/null +++ b/doc/source/reciprocal_space_reconstruction.rst @@ -0,0 +1,714 @@ +Reciprocal-Space Reconstruction +=============================== + +The reciprocal-space reconstruction workflow maps corrected detector images +into one or more regular three-dimensional HKL or momentum-transfer grids. It +is designed for data sets that are much larger than memory. Pixel geometry, +angle transforms, adaptive footprint splitting, and local accumulation run in +a C++17 extension. Python coordinates scan loading, corrections, resumable +tasks, Parquet scratch records, and final NeXus/HDF5 output. + +The experiment configuration is not entered a second time. Reconstruction uses +the active orGUI scan, detector calibration, crystal, UB matrix, mask, +background, correction selections, exclusions, CPU limit, memory limit, and +HDF5 filter registry. Preparing a job freezes those settings into a +checksummed, resumable snapshot. + +Installation +------------ + +The native reconstruction extension is built with orGUI. PyArrow is required +for the out-of-core scratch format, and ``hdf5plugin`` makes the optional HDF5 +filters registered by orGUI available: + +.. code-block:: bash + + pip install "orGUI[reconstruction]" + +The reconstruction feature does not use Numba, OpenMP, or TBB. The C++ kernel +uses a bounded ``std::thread`` worker pool. + +Opening the Workflow +-------------------- + +Load the scan and experiment configuration normally, then open +``Reciprocal space -> Reconstruct reciprocal space``. The dialog contains: + +* ``Experiment``: active scientific state, corrections, exclusions, and + exposure-bound policy; +* ``Output grids``: coordinate systems, bounds, steps, interval counts, HDF5 + settings, and geometry-matched step estimation; +* ``Performance``: footprint accuracy, CPU and memory allocation, and advanced + task layout; +* ``Job and output``: job descriptor, scratch directory, and final HDF5 path; +* ``Preview and status``: estimates, prepared JSON, status, results, and + user-facing errors. + +Preparing Versus Running +~~~~~~~~~~~~~~~~~~~~~~~~ + +``Preview`` reads the current live orGUI state and estimates grid sizes and +execution layout without freezing anything. + +``Prepare Job`` writes the JSON descriptor and immutable HDF5 asset bundle. +Large masks, backgrounds, background variances, and repair geometry are stored +in the asset bundle rather than duplicated in JSON. The JSON and asset bundle +are checksummed. + +``Run Locally`` prepares the current settings and executes them. ``Open Job`` +opens an existing descriptor, while ``Resume`` verifies its sources and +checksums and continues incomplete work. A completed result is registered in +the active orGUI database as an external result; the large arrays remain in the +standalone file. + +Experiment Parameters +--------------------- + +Active Experiment +~~~~~~~~~~~~~~~~~ + +The read-only overview shows the active scan name, frame count, detector shape, +X-ray energy in keV, UB matrix, correction state, and global CPU and memory +limits. Active masks and backgrounds are shown as resolved inputs. Asset paths +such as ``/mask`` refer to datasets that will be created in the prepared asset +bundle. + +``Refresh from active orGUI state`` + Re-reads the central state. If no grid exists, it also derives the initial + HKL grid and chooses writable default paths under the current working + directory. + +Corrections and Exclusions +~~~~~~~~~~~~~~~~~~~~~~~~~~ + +``Mask and repair`` + Opens the shared orGUI mask and pixel-repair settings. Reconstruction uses + the same active detector mask as image integration. A matching mask is + required when mask correction or repair is enabled. + +``Background`` + Opens the shared background-image control. The active background and + optional background variance are frozen into the job asset bundle. + +``Excluded frames`` + Opens the shared frame-exclusion editor. Excluded frames are omitted from + mapping, task layout, geometry-step estimation, and coverage accounting. + +``Missing exact angle bounds`` + Selects the policy used when a scan backend has no exact encoder start/end + positions. ``Stationary exposure`` uses identical start and end positions. + ``Midpoint inference`` explicitly estimates boundaries halfway between + adjacent nominal positions. Exact backend bounds always take precedence. + +``User note`` + Optional free text stored in the job descriptor and final provenance. + +``Normalize by exposure time`` + Divides each frame by its exposure time when the active scan backend + provides one. Enabled by default. This setting is specific to + reconstruction and does not affect ROI/CTR image integration. + +``Monitor corrections`` + Comma-separated scan counter names applied as divisive monitor + normalizations, each with uncertainty propagated when the backend + exposes a matching ``_variance`` counter. This setting is specific + to reconstruction and does not affect ROI/CTR image integration. + +Output Grid Parameters +---------------------- + +Each row describes one independent output coordinate system. + +``Name`` + Optional HDF5 group name. Empty names are generated from the frame name and + made unique. + +``Frame`` + Coordinate system. ``hkl`` is in reciprocal lattice units. Q frames are in + :math:`\mathrm{\mathring{A}}^{-1}`: + + ``lab`` + Laboratory axes. + + ``alpha`` + Frame after the alpha rotation. + + ``omega`` + Frame after alpha and omega rotations. + + ``chi`` + Frame after alpha, omega, and chi rotations. + + ``phi`` + Frame after the full sample rotation. + + ``crystal`` + Crystal-oriented Cartesian Q coordinates using the orientation matrix. + +``Min 1``, ``Min 2``, ``Min 3`` + Lower voxel edges. Units are r.l.u. for HKL and + :math:`\mathrm{\mathring{A}}^{-1}` for Q. + +``Max 1``, ``Max 2``, ``Max 3`` + Exclusive upper grid bounds in the same units as the minima. + +``Step 1``, ``Step 2``, ``Step 3`` + Voxel widths. Editing a step recalculates its interval count. + +``Intervals 1``, ``Intervals 2``, ``Intervals 3`` + Editable numbers of voxels along each axis. Editing an interval count + recalculates the corresponding step from the current bounds. + +``Est. size`` + Dense uncompressed payload estimate: + + .. math:: + + N_1 N_2 N_3 \times 32\ \mathrm{bytes}. + + The 32 bytes comprise float64 intensity, variance, and weight plus uint64 + contributors. Sparse HDF5 chunk allocation and compression usually reduce + payload size; axes, metadata, and HDF5 structures add a smaller overhead. + +Adding and Removing Grids +~~~~~~~~~~~~~~~~~~~~~~~~~ + +``Add derived HKL grid`` + Adds a grid covering the active detector and scan in HKL. + +``Add derived Q grid`` + Prompts for a Q reference frame and adds its derived coverage. + +``Remove selected grid`` + Removes every row containing a selected table cell. + +Automatic Coverage and Initial Steps +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Coverage derivation transforms all four detector corners at the start and end +of every exposure. Per-axis minima and maxima bound those transformed points. +The initial sampling count is: + +.. code-block:: python + + max(64, min(512, max(detector_rows, detector_columns, frame_count))) + +The initial step is the axis extent divided by that scalar count. One step of +padding is added on both sides, so the resulting grid has approximately two +more intervals. The rule is a coverage heuristic, not an instrumental +resolution model. + +Geometry-Matched Step Estimator +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +``Estimate geometry-matched steps`` applies to selected grid rows, or all rows +when nothing is selected. It samples a 5 by 5 set of detector pixels and up to +32 included scan frames. At each point it evaluates finite-difference +generators for one detector-row pixel, one detector-column pixel, and the +exposure sweep. A stationary exposure uses the adjacent-frame center +displacement for the third direction. + +With those generators as columns of a local Jacobian :math:`J`, uniform pixel +and scan-cell widths give: + +.. math:: + + \Sigma_y = \frac{J J^\mathsf{T}}{12}. + +The local axis estimates are +:math:`\sqrt{\operatorname{diag}(\Sigma_y)}`. The estimator dialog exposes the +percentile applied independently to the sampled axis estimates. The default is +10 percent. Lower percentiles select finer steps and protect more locally +high-resolution regions; higher percentiles produce smaller grids. + +This is a geometry-matched sampling estimate. It deliberately excludes beam +divergence, energy bandwidth, detector point-spread, sample mosaicity, and +calibration covariance. + +HDF5 Settings +------------- + +``HDF5 settings`` is available from the grid tab and the main Configuration +menu. These values are shared by all reconstruction grids. + +``Chunk shape (voxels)`` + Three spatial HDF5 chunk lengths. A chunk is the unit written by the + finalizer. Every populated chunk is written once; untouched chunks remain + unallocated. One float64 dataset chunk must remain below the HDF5 4 GiB + limit. + +``Active database filter`` + Compression used by normal orGUI database writes and, by default, by the + reconstruction. + +``Override for reconstruction output`` + When enabled, ``Reconstruction filter`` replaces the active database filter + only for prepared reconstruction output. The selected registered filter is + frozen into the job. + +Accuracy Parameters +------------------- + +Each detector pixel and exposure sweep is treated as a parametric cell. +Adaptive splitting tests transformed cell corners against voxel boundaries. +A leaf wholly inside one voxel contributes its full weight. Otherwise it is +subdivided until it becomes single-voxel or reaches the selected depth, where +its centroid is assigned. Total source-pixel weight is conserved. + +``Center only (depth 0)`` + Assign one centroid per pixel. Fastest, without finite-footprint splitting. + +``Low (depth 1)`` + One adaptive subdivision level. + +``Balanced (depth 2)`` + Default compromise between boundary accuracy and runtime. + +``High (depth 3)`` + Finer boundary subdivision at substantially higher cost. + +``Very high (depth 4)`` + Intended for validation or demanding coarse-voxel boundaries. + +``Maximum (depth 5)`` + Highest UI setting. Benchmark memory and runtime before large production + jobs. + +Stationary cells split into four spatial children per level. Swept cells split +in detector row, detector column, and exposure position and can create eight +children per level. Consequently, worst-case work grows rapidly with depth. + +Performance Parameters +---------------------- + +Detected values appear in ``Detected execution layout``. An unchecked +``Override`` editor displays the detected value but does not freeze it into the +job. + +``Total thread budget`` + Optional per-job replacement for ``orGUI.numberthreads``. It is divided + between concurrent image workers and native threads within each image. + +``Native threads per image`` + C++ threads assigned to one image. The default is 4. Remaining threads can + run other images concurrently. The actual value is bounded by the total + thread and memory budgets. + +``Total memory budget`` + Optional per-job replacement for ``orGUI.maxMemory``, in MiB. It bounds + image workers, native working memory, retained records, sorting, and Arrow + conversion. + +``Accumulation per worker`` + Maximum reduced-record memory retained by each image worker before records + are merged and flushed as a larger Parquet segment. A larger value reduces + small files and repeated merges but consumes more RAM. The global budget + reserves transient merge and conversion memory in addition to this value. + +``Frames per task`` + Consecutive included frames in one deterministic resumable map task. + Automatic sizing targets several tasks per concurrent image worker and caps + a task at 64 frames. + +``Tile rows`` and ``Tile columns`` + Detector dimensions passed to one native-kernel call. Tiles bound image, + variance, mask, and corner-ray arrays and may end at arbitrary pixel + boundaries. All tiles of one loaded image are processed without reloading + or reapplying full-image corrections. + +``Native block pixels`` + Flattened pixels in one C++ scheduling and local-reduction block inside a + detector tile. It is not the tile area. For example, a 1024 by 1024 tile + and a 65,536-pixel block create 16 blocks. Multiple blocks are necessary + for native threads to share a tile; very small blocks add scheduling, + sorting, and merge overhead. + +``Parquet chunk span`` + Number of consecutive HDF5 spatial chunk IDs grouped into a coarse Parquet + partition range. Larger spans create fewer, larger files; smaller spans + narrow reducer reads but increase object count. + +The detected summary also reports frame tasks, detector tiles, total map tasks, +concurrent image workers, memory per image, and effective accumulation memory. + +Correction and Variance Model +----------------------------- + +Corrections are evaluated once for each loaded full image, before detector +tiling: + +1. Use variance supplied by the image backend. If unavailable, use the flagged + approximation :math:`v=\max(I,0)` from the current raw image. +2. Subtract the active background. Add background variance when supplied; + otherwise record that the background was treated as deterministic. +3. Apply the active detector mask and optional pixel repair. +4. Apply inverse solid-angle and inverse polarization factors. +5. Normalize by exposure time when enabled and available. +6. Apply explicitly selected divisive monitor counters. +7. Mask non-finite corrected intensity or variance. + +For a deterministic multiplicative factor :math:`c`: + +.. math:: + + I' = cI,\qquad v' = c^2v. + +When factor variance :math:`v_c` is available: + +.. math:: + + v' = c^2v + I^2v_c. + +The constant detector mask is analyzed once. Repairable components and their +neighbors are stored in a native repair plan reused for all images. Isolated +pixels use the local valid-neighbor median; accepted small components use +inverse-distance weights. Interpolation variance uses the squared weights. +Repair introduces spatial covariance that is not stored in the final file. + +Voxel Accumulation +------------------ + +Leaves from the same source pixel entering the same voxel are combined before +global accumulation, preserving their correlation. For independent source +pixels the kernel accumulates: + +.. math:: + + S_I = \sum wI,\qquad + S_v = \sum w^2v,\qquad + S_w = \sum w. + +Final voxel datasets are: + +.. math:: + + I_\mathrm{voxel} = \frac{S_I}{S_w},\qquad + v_\mathrm{voxel} = \frac{S_v}{S_w^2}. + +Empty intensity and variance voxels are NaN. ``weight`` stores :math:`S_w`; +``contributors`` stores the number of independent source pixels. Adaptive +footprint splitting creates cross-voxel covariance, so ``variance`` contains +marginal variances only. + +Out-of-Core Execution +--------------------- + +The execution stages are: + +1. Verify the job JSON, scan reference, source fingerprint, and asset checksum. +2. Precompute detector corner-ray lattices for bounded tiles. +3. Load and correct images in parallel. +4. Map tiles in the native kernel, releasing the Python GIL. +5. Sort and locally reduce compact records keyed by HDF5 chunk and voxel IDs. +6. Accumulate records in worker memory until its byte budget is reached. +7. Write immutable Zstandard-compressed Parquet partitions and manifests. +8. Reduce partitions in bounded external merge passes into finalized HDF5 + chunk shards. Contiguous shard ranges are reduced concurrently with private + Parquet readers; the global memory budget is divided across reducer workers. +9. Create the standalone HDF5 file and write each populated spatial chunk once. +10. Validate the output checksum, register it as an external result, then + remove map partitions, reduced scratch data, and the asset bundle. + +Interrupted jobs retain scratch data. Cleanup occurs only after successful +HDF5 close, validation, and checksum calculation. Cleanup errors are recorded +but do not invalidate a completed scientific result. + +Reduction uses up to the global thread budget. The effective worker count is +also limited to the number of pending shards and to one worker per 64 MiB of +configured memory. Each worker owns a contiguous range of shard plans and +private forward-only Parquet readers. This permits parallel reduction within a +single large bucket without sharing stateful readers or exceeding the global +memory budget. Mapping-partition verification is parallelized by the same +limit. Checkpoints and progress callbacks remain serialized on the coordinating +thread. + +Paths +----- + +``Job JSON`` + Resumable descriptor containing frozen scientific settings, paths, build + metadata, task state, and checksums. + +``Scratch directory`` + Writable per-job directory for the asset bundle, map partitions, reduction + shards, and manifests. It should normally be on fast local storage. + +``Standalone HDF5`` + Final NeXus-style output file. Defaults are created under the current + working directory because raw-data locations are often read-only. + +Command-Line Execution +---------------------- + +The CLI consumes the exact job prepared by the UI; it does not provide a +parallel experiment-configuration interface: + +.. code-block:: bash + + orGUI rsmap run JOB.json + orGUI rsmap resume JOB.json + orGUI rsmap status JOB.json + +The direct alias is equivalent: + +.. code-block:: bash + + orGUI-rsmap run JOB.json + orGUI-rsmap resume JOB.json + orGUI-rsmap status JOB.json + +``run`` and ``resume`` both verify and continue deterministic task state. +``status`` is read-only and prints the current descriptor status as JSON. + +Cluster Batch Execution +----------------------- + +The **Cluster** tab generates an SGE or Slurm batch bundle from the same +prepared job JSON used for local execution. It does not introduce another +experiment configuration. The bundle contains: + +``*-map.sge`` or ``*-map.slurm`` + A job array with one element per deterministic frame-range map task. Each + element reads the immutable job and asset bundle and writes only its own + Parquet partitions and map manifest. Array elements never update the shared + job JSON, so they may execute concurrently. + +``*-finalize.sge`` or ``*-finalize.slurm`` + A single dependent job. It verifies that every expected array manifest is + complete and checksummed, records them in the job JSON, reduces shards with + its own CPU and memory allocation, and creates the final HDF5 file. + +``*-submit.sh`` + A convenience submission wrapper. For SGE it captures ``qsub -terse`` and + submits the finalizer with ``-hold_jid``. For Slurm it captures + ``sbatch --parsable`` and uses ``afterok`` on the complete array. The + finalizer always verifies every map task, including on SGE installations + where a job hold records completion rather than successful exit. + +The scripts require the job JSON, scratch directory, immutable assets, scan +source, and final output directory to be reachable at the same absolute paths +from every compute node. Scratch should normally reside on a high-throughput +shared filesystem or node-local storage explicitly staged by site-specific +setup commands. + +The generated worker commands are also available for inspection and manual +scheduler integration: + +.. code-block:: bash + + orGUI rsmap cluster-map JOB.json --task-index 0 --cpus 4 --memory-gib 16 + orGUI rsmap cluster-finalize JOB.json --cpus 24 --memory-gib 64 + orGUI rsmap cluster-scripts JOB.json --output-directory batch + +Scheduler Parameters +~~~~~~~~~~~~~~~~~~~~ + +``Scheduler`` + ``SGE`` (default) or ``Slurm``. SGE arrays are one-based and converted to + orGUI's zero-based task index. Slurm arrays are generated zero-based. + +``Job name`` + Scheduler-safe base name for the map array and finalizer. Letters, numbers, + periods, underscores, and hyphens are accepted. + +``Queue / partition`` + Optional SGE ``-q`` queue or Slurm ``--partition``. + +``Project / account`` + Optional SGE ``-P`` project or Slurm ``--account``. + +``Script directory`` + Destination for the three generated scripts. It defaults beneath the + current writable working directory. + +``Working directory`` + Shared directory selected with ``cd`` before environment setup and Python + execution. + +``Python executable`` + Python command used to invoke ``orgui.reconstruction_cli`` after environment + setup. ``python`` is the portable default; an absolute cluster-environment + path can be supplied. + +``Setup commands`` + Verbatim shell lines placed after ``set -euo pipefail`` and ``cd``. Use + these for module loading and conda or virtual-environment activation. + +``CPUs / slots per mapping task`` + Native C++ threads used by each array element. This allocation is + independent of the number of simultaneously running array elements. + +``Memory per mapping task`` + Total RAM budget passed to one mapping process. Slurm receives it through + ``--mem``. SGE memory complexes are normally per-slot, so orGUI divides the + total by the slot count and rounds upward. + +``Mapping wall time`` + Per-element SGE ``h_rt`` or Slurm ``--time`` limit. + +``Maximum concurrent tasks`` + Optional array throttle: SGE ``-tc`` or the Slurm ``%N`` array suffix. + Zero leaves concurrency to site policy. + +``Reduction CPUs / slots`` + Independent worker capacity for parallel checksum verification and + contiguous shard reduction. It may be larger or smaller than the mapping + task allocation. + +``Reduction memory`` + Total reducer/finalizer RAM budget. The reducer divides it among active + workers and retains bounded shard buffers. + +``Reduction wall time`` + SGE ``h_rt`` or Slurm ``--time`` for the dependent job. + +``SGE parallel environment`` + Name requested through ``-pe``; ``smp`` is the default. It must match the + target site's configured shared-memory parallel environment. + +``SGE memory resource`` + Per-slot consumable complex used in ``-l`` requests. ``h_vmem`` is the + default, but sites may require ``mem_free`` or another name. + +``Extra map/finalizer directives`` + Optional scheduler-specific header lines. Each non-empty line must start + with ``#$`` for SGE or ``#SBATCH`` for Slurm. + +See the `Grid Engine qsub reference +`_ and the `Slurm job +array reference `_ for scheduler +semantics. + +Job Descriptor Reference +------------------------ + +Users normally create this file through the UI. The fields are documented for +inspection, scheduling, and provenance: + +``schema_version`` + Exact reconstruction job schema. Unsupported schemas are rejected. + +``config`` + Authoritative central ``ConfigData`` snapshot containing detector + calibration, crystal, UB, diffractometer values, and correction state. + +``scan_reference`` + Serializable reference for the active scan, including supported slices, + interlaced scans, manual scans, simulations, and external backend files. + +``grids`` + List of output grid dictionaries: minimum, maximum, step, frame, name, and + chunk shape. + +``scratch_path``, ``output_path`` + Absolute local paths described above. + +``compression`` + Frozen name from the central HDF5 filter registry. + +``assets_path``, ``assets_sha256`` + Immutable job asset bundle and SHA-256 checksum. + +``source_fingerprint_sha256`` + Digest of the serialized scan reference used to detect changed sources. + +``build_metadata`` + orGUI, NumPy, h5py, compiler, and native build information. + +``runtime_threads``, ``runtime_memory_bytes`` + Global orGUI defaults captured when the job was prepared. + +``thread_override``, ``memory_override_bytes`` + Optional per-job replacements for the captured defaults. + +``threads_per_image`` + Requested native threads per concurrent image. + +``accumulation_budget_bytes`` + Optional retained-record bytes per image worker. + +``angle_fallback`` + ``stationary`` or explicitly selected ``midpoint`` inference. + +``accuracy`` + Named footprint-depth preset. + +``advanced_depth`` + Internal legacy field in the current descriptor schema. UI-prepared jobs + leave it null and use ``accuracy``. + +``frame_batch`` + Optional frames-per-task override. + +``tile_shape`` + Optional detector tile rows and columns. + +``work_block_pixels`` + Optional native scheduling-block override. + +``partition_chunk_span`` + Optional Parquet chunk-range override. + +``user_note`` + Optional free-text provenance. + +``expected_map_tasks`` + Number of deterministic map tasks in the prepared layout. + +``status`` + Current job state, such as ``prepared``, ``mapping``, ``reducing``, + ``finalizing``, or ``complete``. + +``map_manifests``, ``reduction_manifest`` + Completed scratch task manifests used for retry and resume. + +``output_sha256`` + Final standalone-file checksum after verified completion. + +``correction_provenance`` + Recorded variance source, factor-uncertainty assumptions, repair-plan + configuration, and image-processing provenance. + +``cleanup_errors`` + Nonfatal failures encountered while deleting verified scratch outputs. + +``cluster_settings`` + Frozen scheduler, environment, map-array resource, reduction-resource, and + optional directive settings used to regenerate the batch bundle. + +Output File Layout +------------------ + +The final file contains an ``NXentry`` at ``/entry``. The reconstruction +``NXprocess`` stores the frozen compute configuration, scientific context, +provenance, marginal-variance warning, and central orGUI scan configuration. + +Each selected coordinate system is an ``NXdata`` group below +``/entry/reconstruction/results``. It contains: + +* voxel-center axes ``h``, ``k``, ``l`` in r.l.u. or ``qx``, ``qy``, ``qz`` in + :math:`\mathrm{\mathring{A}}^{-1}`; +* float64 ``intensity``; +* float64 marginal ``variance``; +* float64 ``weight``; +* uint64 ``contributors``; +* coordinate-frame, signal, axes, and units attributes. + +All scientific arrays default to float64 except the contributor count. The +file is conventional HDF5 and can be read with h5py, silx, and NeXus-aware +tools. + +Practical Guidance +------------------ + +* Start with ``Balanced`` depth and the 10-percent geometry estimator. +* Check interval counts and the dense size estimate before preparing. +* Put scratch data on fast local SSD when possible. +* Increase accumulation memory to reduce small Parquet files only when the + total memory budget has sufficient headroom. +* Keep enough native blocks per tile to occupy all native threads. +* Prefer automatic tiling and task sizing until representative benchmarks + justify overrides. +* Use ``status`` before resuming a job copied from another machine. +* Treat the variance arrays as marginal uncertainties, not a complete + covariance representation. diff --git a/doc/source/release_notes.rst b/doc/source/release_notes.rst index 623797d..c219a82 100644 --- a/doc/source/release_notes.rst +++ b/doc/source/release_notes.rst @@ -56,6 +56,12 @@ ESRF ID31 beamline support and reciprocal-space display: - **BREAKING CHANGE:** HDF5 file locking is now disabled by default (``HDF5_USE_FILE_LOCKING=False``), so that files still open for writing by the acquisition system can be read. A manually set environment variable still wins, and ``--hdflocking`` / ``-l`` restores the previous behavior. - **BREAKING CHANGE:** the segmented ("interlaced") scan loader no longer guesses a file's layout from a hardcoded list of beamtime ids. Backends can now answer directly with the new optional ``Scan.listScans`` classmethod, which returns the scan identifiers -- numbers such as ``[1, 2, 10]``, names such as ``["ascan_12", "dscan_3"]``, or ``(identifier, label)`` pairs to label the rows of the selection dialog. Backends that do not implement it are handled by applying their own ``parse_h5_node`` to every entry of the file root, which requires no backend change but relies on ``parse_h5_node`` raising for entries that are not scans. Either way the loader can no longer disagree with how the same backend opens a single scan. This fixes the ``id31_default_p4`` backend, which the old list did not cover, and makes custom backends work regardless of what their class is called -- the example backend under ``examples/backend/ID31_EBS_p4_backend.py`` was itself affected. Backends addressed by a name rather than by a number are now supported here as well, which makes segmented scans work for the legacy ``ch5523`` beamtime for the first time: it was listed as ID31-style, so the loader looked for a ``"."`` name its files never had and raised before showing the dialog. Subscans such as ``1.10`` are no longer mistaken for the ``1.1`` fast-counter subscan. +Reciprocal-space reconstruction: + +- Added a centralized, out-of-core reciprocal-space reconstruction pipeline. A new "Reconstruct reciprocal space" dialog (Configuration menu) defines HKL/Q output grids, previews coverage and storage cost, and prepares, runs, and resumes jobs. Jobs can run locally, as SGE/Slurm cluster batch arrays with parallel submap reduction, or from the command line via a new ``reconstruction_cli`` entry point. +- Added a shared HDF5 output-settings dialog (chunk shape, compression) reachable from both the reconstruction dialog and the Configuration menu. +- Exposure-time normalization and monitor-counter corrections now live in the reconstruction dialog rather than the shared scan-options panel, since they affect reconstruction output only, not ROI/CTR image integration. + 1.5.0 (2026-06-07) ------------------ diff --git a/meson.build b/meson.build index 54b800b..fbefdcd 100644 --- a/meson.build +++ b/meson.build @@ -209,6 +209,14 @@ py.extension_module( subdir: 'orgui/datautils/xrayutils' ) +py.extension_module( + '_reciprocal_reconstruction_cpp', + files('orgui/datautils/xrayutils/cpp/reciprocal_reconstruction_cpp.cpp'), + dependencies: [pybind11_dep, xxhash_dep], + install: true, + subdir: 'orgui/datautils/xrayutils' +) + py.extension_module( '_roi_sum_cpp', files('orgui/app/cpp/roi_sum_cpp.cpp'), diff --git a/orgui/app/AGENTS.md b/orgui/app/AGENTS.md index 2e5637a..5ffe761 100644 --- a/orgui/app/AGENTS.md +++ b/orgui/app/AGENTS.md @@ -9,6 +9,9 @@ This directory contains GUI components and user workflows: - `database.py`: config loading, detector/crystal setup, Nexus/HDF5 handling. - `QScanSelector.py`, `ROIutils.py`, `orGUI.py`: integration and image workflows. +- `ReconstructionDialog.py`: GUI front-end for the out-of-core reciprocal-space + reconstruction pipeline (`orgui/reconstruction_job.py`, + `reconstruction_cli.py`, `reconstruction_cluster.py`). - Dialogs and tests under this directory cover user-facing behavior and GUI regressions. diff --git a/orgui/app/HDF5SettingsDialog.py b/orgui/app/HDF5SettingsDialog.py new file mode 100644 index 0000000..efc3347 --- /dev/null +++ b/orgui/app/HDF5SettingsDialog.py @@ -0,0 +1,147 @@ +"""Shared HDF5 output settings dialog.""" + +from math import prod + +import numpy as np +from silx.gui import qt + +from .database import FILTERS + + +def compression_filter_name(value): + """Return the registered name for an HDF5 compression filter.""" + for name, available in FILTERS.items(): + try: + if value is available or value == available: + return name + except Exception: + continue + raise ValueError("The active HDF5 compression filter is not registered") + + +class HDF5SettingsDialog(qt.QDialog): + """Edit shared reciprocal-space HDF5 storage settings.""" + + def __init__( + self, + active_compression, + chunk_shape=(64, 64, 64), + compression_override=None, + parent=None, + ): + super().__init__(parent) + self.setWindowTitle("HDF5 settings") + layout = qt.QVBoxLayout(self) + + chunk_group = qt.QGroupBox("Spatial chunks") + chunk_form = qt.QFormLayout(chunk_group) + chunk_row = qt.QWidget() + chunk_layout = qt.QHBoxLayout(chunk_row) + chunk_layout.setContentsMargins(0, 0, 0, 0) + chunk_tooltip = ( + "HDF5 spatial chunk dimensions in reciprocal-space voxels, shared " + "by every output grid." + ) + self.chunk_editors = [] + for axis, value in zip("XYZ", chunk_shape): + label = qt.QLabel(axis) + label.setToolTip(chunk_tooltip) + editor = qt.QSpinBox() + editor.setRange(1, 1_000_000) + editor.setValue(int(value)) + editor.setToolTip( + f"HDF5 chunk length along reciprocal-space axis {axis}." + ) + chunk_layout.addWidget(label) + chunk_layout.addWidget(editor) + self.chunk_editors.append(editor) + chunk_layout.addStretch(1) + chunk_label = qt.QLabel("Chunk shape (voxels):") + chunk_label.setToolTip(chunk_tooltip) + chunk_row.setToolTip(chunk_tooltip) + chunk_form.addRow(chunk_label, chunk_row) + layout.addWidget(chunk_group) + + compression_group = qt.QGroupBox("Compression") + compression_form = qt.QFormLayout(compression_group) + self.database_compression = qt.QComboBox() + self.database_compression.addItems(FILTERS) + self.database_compression.setCurrentText( + compression_filter_name(active_compression) + ) + self.database_compression.setToolTip( + "Compression used by normal orGUI database writes." + ) + database_label = qt.QLabel("Active database filter:") + database_label.setToolTip(self.database_compression.toolTip()) + compression_form.addRow(database_label, self.database_compression) + self.override_compression = qt.QCheckBox( + "Override for reconstruction output" + ) + self.override_compression.setToolTip( + "When disabled, reconstruction output uses the active orGUI " + "database compression filter." + ) + compression_form.addRow(self.override_compression) + self.output_compression = qt.QComboBox() + self.output_compression.addItems(FILTERS) + selected = ( + compression_override + if compression_override is not None + else compression_filter_name(active_compression) + ) + self.output_compression.setCurrentText(selected) + self.output_compression.setEnabled(compression_override is not None) + self.override_compression.setChecked( + compression_override is not None + ) + self.override_compression.toggled.connect( + self.output_compression.setEnabled + ) + self.output_compression.setToolTip( + "Compression filter stored in the prepared job and used for its " + "final HDF5 datasets." + ) + output_label = qt.QLabel("Reconstruction filter:") + output_label.setToolTip(self.output_compression.toolTip()) + compression_form.addRow(output_label, self.output_compression) + layout.addWidget(compression_group) + + buttons = qt.QDialogButtonBox( + qt.QDialogButtonBox.Ok | qt.QDialogButtonBox.Cancel + ) + buttons.accepted.connect(self.accept) + buttons.rejected.connect(self.reject) + layout.addWidget(buttons) + + @property + def chunk_shape(self): + """Return the selected three-dimensional chunk shape in voxels.""" + return tuple(editor.value() for editor in self.chunk_editors) + + @property + def compression_override(self): + """Return the selected override name, or ``None`` for the active filter.""" + return ( + self.output_compression.currentText() + if self.override_compression.isChecked() + else None + ) + + @property + def database_compression_name(self): + """Return the selected active database compression-filter name.""" + return self.database_compression.currentText() + + def accept(self): + """Validate settings and close the dialog.""" + chunk_bytes = prod(self.chunk_shape) * np.dtype(np.float64).itemsize + if chunk_bytes >= 2**32: + qt.QMessageBox.warning( + self, + "Invalid HDF5 chunk shape", + "One float64 dataset chunk must be smaller than 4 GiB. " + f"The selected shape requires {chunk_bytes / 1024**3:.2f} GiB.", + ) + return + super().accept() diff --git a/orgui/app/QScanSelector.py b/orgui/app/QScanSelector.py index 684b0a1..eb313b8 100644 --- a/orgui/app/QScanSelector.py +++ b/orgui/app/QScanSelector.py @@ -50,7 +50,6 @@ from . import qutils from .QReflectionSelector import QReflectionAnglesDialog from .QHKLDialog import HKLDialog -import runpy from contextlib import contextmanager @@ -1061,25 +1060,7 @@ def _onLoadBackend(self): # qutils.warning_detailed_message(self, "Cannot load backend", "Cannot load backend", traceback.format_exc()) # noqa: E501 def loadBackendFile(self, filename): - backend_file = runpy.run_path(filename) - found_backends = [] - for e in backend_file: - try: - if ( - issubclass(backend_file[e], scans.Scan) - and backend_file[e] != scans.Scan - ): - found_backends.append((e, backend_file[e])) - except Exception: - pass - # traceback.print_exc() - if not found_backends: - raise ValueError(f"Found no backend in file {filename}") - if len(found_backends) > 1: - raise ValueError( - f"Found more than one Scan class in backend file {filename}. Only one is permitted" # noqa: E501 - ) - name, scancls = found_backends[0] + name, scancls = scans.load_scan_backend_file(filename) self.btid.addItem(name) backends.fscans[name] = scancls self.btid.setCurrentText(name) diff --git a/orgui/app/ReconstructionDialog.py b/orgui/app/ReconstructionDialog.py new file mode 100644 index 0000000..33575a3 --- /dev/null +++ b/orgui/app/ReconstructionDialog.py @@ -0,0 +1,1698 @@ +"""GUI editor for centralized reciprocal-space reconstruction jobs.""" + +from __future__ import annotations + +import json +import logging +from pathlib import Path + +import numpy as np +from silx.gui import qt + +from .. import logger_utils +from ..reconstruction_cluster import ( + ClusterSettings, + generate_cluster_scripts, +) +from ..reconstruction_job import ( + ACCURACY_DEPTHS, + ReconstructionGrid, + derive_grid, + estimate_geometry_steps, + job_status, + prepare_job, + read_job, + reconstruction_execution_settings, + run_job, +) +from .config_data import ConfigData +from .database import FILTERS +from .HDF5SettingsDialog import ( + HDF5SettingsDialog, + compression_filter_name, +) + + +_GRID_COLUMNS = ( + "Name", + "Frame", + "Min 1", + "Min 2", + "Min 3", + "Max 1", + "Max 2", + "Max 3", + "Step 1", + "Step 2", + "Step 3", + "Intervals 1", + "Intervals 2", + "Intervals 3", + "Est. size", +) +_FRAMES = ("hkl", "lab", "alpha", "omega", "chi", "phi", "crystal") +logger = logging.getLogger(__name__) + + +class _GridNumberItem(qt.QTableWidgetItem): + """Display a compact number while retaining its exact editable value.""" + + def __init__(self, value): + super().__init__() + self.setData(qt.Qt.EditRole, float(value)) + + def data(self, role): + if role == qt.Qt.DisplayRole: + value = super().data(qt.Qt.EditRole) + try: + return format(float(value), ".6g") + except (TypeError, ValueError): + return value + return super().data(role) + + +class _GeometryResolutionDialog(qt.QDialog): + """Choose robust sampling for the local-Jacobian step estimate.""" + + def __init__(self, parent=None): + super().__init__(parent) + self.setWindowTitle("Geometry-matched grid steps") + layout = qt.QVBoxLayout(self) + explanation = qt.QLabel( + "Local detector-row, detector-column, and scan-direction " + "Jacobians are sampled across the experiment. A lower percentile " + "selects finer steps and protects more high-resolution regions." + ) + explanation.setWordWrap(True) + layout.addWidget(explanation) + form = qt.QFormLayout() + self.percentile_editor = qt.QDoubleSpinBox() + self.percentile_editor.setRange(0.1, 100.0) + self.percentile_editor.setDecimals(1) + self.percentile_editor.setSingleStep(1.0) + self.percentile_editor.setValue(10.0) + self.percentile_editor.setSuffix(" %") + tooltip = ( + "Percentile of sampled local one-sigma axis resolutions. Lower " + "values produce finer grids; 10% is a conservative default." + ) + self.percentile_editor.setToolTip(tooltip) + percentile_label = qt.QLabel("Resolution percentile:") + percentile_label.setToolTip(tooltip) + form.addRow(percentile_label, self.percentile_editor) + layout.addLayout(form) + buttons = qt.QDialogButtonBox( + qt.QDialogButtonBox.Ok | qt.QDialogButtonBox.Cancel + ) + buttons.accepted.connect(self.accept) + buttons.rejected.connect(self.reject) + layout.addWidget(buttons) + + @property + def percentile(self): + """Return the selected local-resolution percentile.""" + return self.percentile_editor.value() + + +def _format_size(size_bytes): + for suffix in ("B", "KiB", "MiB", "GiB", "TiB", "PiB"): + if size_bytes < 1024 or suffix == "PiB": + return f"{size_bytes:.3g} {suffix}" + size_bytes /= 1024 + + +class _ReconstructionCancelled(RuntimeError): + """Signal cancellation between reconstruction tasks.""" + + +class ReconstructionDialog(qt.QDialog): + """Configure, prepare, run, and resume centralized reconstruction jobs.""" + + def __init__(self, orgui, parent=None): + super().__init__(parent or orgui) + self.orgui = orgui + if not hasattr(self.orgui, "reconstruction_chunk_shape"): + self.orgui.reconstruction_chunk_shape = (64, 64, 64) + if not hasattr(self.orgui, "reconstruction_compression_override"): + self.orgui.reconstruction_compression_override = None + if not hasattr(self.orgui, "reconstruction_normalize_exposure"): + self.orgui.reconstruction_normalize_exposure = True + if not hasattr(self.orgui, "reconstruction_monitor_corrections"): + self.orgui.reconstruction_monitor_corrections = () + self.setWindowTitle("Reciprocal-space reconstruction") + self.resize(1100, 760) + layout = qt.QVBoxLayout(self) + self.tabs = qt.QTabWidget() + self.tabs.addTab(self._data_tab(), "Experiment") + self.tabs.addTab(self._grid_tab(), "Output grids") + self.tabs.addTab(self._performance_tab(), "Performance") + self.tabs.addTab(self._cluster_tab(), "Cluster") + self.tabs.addTab(self._paths_tab(), "Job and output") + self.output_tab = qt.QWidget() + output_layout = qt.QVBoxLayout(self.output_tab) + self.preview_output = qt.QPlainTextEdit() + self.preview_output.setReadOnly(True) + self.preview_output.setToolTip( + "Preview estimates, prepared-job JSON, execution results, and " + "status or error details." + ) + output_layout.addWidget(self.preview_output) + self.tabs.addTab(self.output_tab, "Preview and status") + layout.addWidget(self.tabs, stretch=1) + buttons = qt.QDialogButtonBox(qt.QDialogButtonBox.Close) + for label, slot in ( + ("Preview", self.preview), + ("Prepare Job", self.prepare), + ("Run Locally", self.run_local), + ("Create Cluster Scripts", self.create_cluster_scripts), + ("Open Job", self.open_job), + ("Resume", self.resume), + ): + button = buttons.addButton(label, qt.QDialogButtonBox.ActionRole) + button.clicked.connect(slot) + button.setToolTip( + { + "Preview": "Estimate coverage, storage, and execution layout.", + "Prepare Job": "Freeze the current settings into a resumable job.", + "Run Locally": "Prepare and execute the configured job locally.", + "Create Cluster Scripts": ( + "Prepare the job and create an SGE or Slurm map array " + "with a dependent reduction/finalization job." + ), + "Open Job": "Select an existing reconstruction job JSON file.", + "Resume": "Verify and continue the selected prepared job.", + }[label] + ) + buttons.rejected.connect(self.reject) + layout.addWidget(buttons) + self.refresh_live_state() + + def _data_tab(self): + widget = qt.QWidget() + layout = qt.QVBoxLayout(widget) + + experiment_group = qt.QGroupBox("Active experiment") + experiment_layout = qt.QVBoxLayout(experiment_group) + self.experiment_summary = qt.QPlainTextEdit() + self.experiment_summary.setReadOnly(True) + self.experiment_summary.setToolTip( + "Live summary of the active scan, geometry, UB matrix, corrections, " + "and global runtime limits." + ) + experiment_layout.addWidget(self.experiment_summary) + refresh = qt.QPushButton("Refresh from active orGUI state") + refresh.setToolTip( + "Reload the summary and default grid from the current orGUI state." + ) + refresh.clicked.connect(self.refresh_live_state) + experiment_layout.addWidget(refresh) + layout.addWidget(experiment_group) + + correction_group = qt.QGroupBox("Corrections and exclusions") + shared = qt.QHBoxLayout(correction_group) + mask = qt.QPushButton("Mask and repair") + mask.setToolTip( + "Open the shared detector-mask and masked-pixel repair settings." + ) + mask.clicked.connect( + lambda: self._invoke_ui_action( + "Cannot open mask settings", + self.orgui._onShowMaskConfig, + ) + ) + background = qt.QPushButton("Background") + background.setToolTip( + "Select or clear the shared background image and its variance." + ) + background.clicked.connect( + lambda: self._invoke_ui_action( + "Cannot open background settings", + self.orgui.backgroundImageAct.trigger, + ) + ) + exclusions = qt.QPushButton("Excluded frames") + exclusions.setToolTip( + "Choose scan frames that must not contribute to reconstruction." + ) + exclusions.clicked.connect( + lambda: self._invoke_ui_action( + "Cannot open excluded-frame settings", + self.orgui.excludedImagesDialog.show, + ) + ) + shared.addWidget(mask) + shared.addWidget(background) + shared.addWidget(exclusions) + layout.addWidget(correction_group) + + metadata_group = qt.QGroupBox("Exposure and job metadata") + form = qt.QFormLayout(metadata_group) + self.angle_fallback = qt.QComboBox() + self.angle_fallback.addItem("Stationary exposure", "stationary") + self.angle_fallback.addItem( + "Midpoint inference (explicit fallback)", "midpoint" + ) + self._add_form_row( + form, + "Missing exact angle bounds:", + self.angle_fallback, + "Choose how exposure bounds are represented when the scan backend " + "does not provide exact encoder start and end positions.", + ) + self.user_note = qt.QLineEdit() + self.user_note.setPlaceholderText("Optional note stored with the job") + self._add_form_row( + form, + "User note:", + self.user_note, + "Optional free-text note embedded in the job descriptor and output.", + ) + layout.addWidget(metadata_group) + + normalization_group = qt.QGroupBox("Exposure and monitor normalization") + normalization_form = qt.QFormLayout(normalization_group) + self.normalize_exposure = qt.QCheckBox("Normalize by exposure time") + self.normalize_exposure.setChecked( + bool(self.orgui.reconstruction_normalize_exposure) + ) + self.normalize_exposure.toggled.connect( + self._on_normalize_exposure_changed + ) + self._add_form_row( + normalization_form, + "", + self.normalize_exposure, + "Divide each frame by its exposure time when the scan backend " + "provides one. Applies only to reciprocal-space reconstruction.", + ) + self.monitor_corrections = qt.QLineEdit() + self.monitor_corrections.setPlaceholderText( + "Optional scan counters, comma-separated" + ) + self.monitor_corrections.setText( + ", ".join(self.orgui.reconstruction_monitor_corrections) + ) + self.monitor_corrections.editingFinished.connect( + self._on_monitor_corrections_changed + ) + self._add_form_row( + normalization_form, + "Monitor corrections:", + self.monitor_corrections, + "Counters applied as divisive monitor normalizations. Applies " + "only to reciprocal-space reconstruction.", + ) + layout.addWidget(normalization_group) + return widget + + def _on_normalize_exposure_changed(self, checked): + self.orgui.reconstruction_normalize_exposure = bool(checked) + + def _on_monitor_corrections_changed(self): + self.orgui.reconstruction_monitor_corrections = tuple( + value.strip() + for value in self.monitor_corrections.text().split(",") + if value.strip() + ) + + def _grid_tab(self): + widget = qt.QWidget() + layout = qt.QVBoxLayout(widget) + grid_group = qt.QGroupBox("Grid definitions") + grid_layout = qt.QVBoxLayout(grid_group) + self.grid_table = qt.QTableWidget(0, len(_GRID_COLUMNS)) + self.grid_table.setHorizontalHeaderLabels(_GRID_COLUMNS) + self.grid_table.setToolTip( + "Edit coordinate bounds and voxel spacing for each output " + "coordinate system." + ) + header_tooltips = ( + "Unique output-grid name.", + "HKL or momentum-transfer reference frame.", + "Lower bound of coordinate axis 1.", + "Lower bound of coordinate axis 2.", + "Lower bound of coordinate axis 3.", + "Upper bound of coordinate axis 1.", + "Upper bound of coordinate axis 2.", + "Upper bound of coordinate axis 3.", + "Voxel spacing along coordinate axis 1.", + "Voxel spacing along coordinate axis 2.", + "Voxel spacing along coordinate axis 3.", + "Editable number of voxel intervals along coordinate axis 1.", + "Editable number of voxel intervals along coordinate axis 2.", + "Editable number of voxel intervals along coordinate axis 3.", + "Estimated dense, uncompressed size of intensity, variance, " + "weight, and contributor datasets. Sparse allocation and " + "compression normally reduce the final file size.", + ) + for column, tooltip in enumerate(header_tooltips): + self.grid_table.horizontalHeaderItem(column).setToolTip(tooltip) + self.grid_table.horizontalHeader().setSectionResizeMode( + qt.QHeaderView.ResizeToContents + ) + self.grid_table.cellChanged.connect(self._on_grid_cell_changed) + grid_layout.addWidget(self.grid_table) + buttons = qt.QHBoxLayout() + add_hkl = qt.QPushButton("Add derived HKL grid") + add_hkl.setToolTip( + "Estimate editable HKL bounds and spacing from the active scan." + ) + add_hkl.clicked.connect(lambda: self.add_derived_grid("hkl")) + add_q = qt.QPushButton("Add derived Q grid") + add_q.setToolTip( + "Add an editable momentum-transfer grid in a selected frame." + ) + add_q.clicked.connect(self._add_q_grid) + remove = qt.QPushButton("Remove selected grid") + remove.setToolTip("Remove every grid row containing a selected cell.") + remove.clicked.connect(self._remove_grid) + estimate_steps = qt.QPushButton("Estimate geometry-matched steps") + estimate_steps.setToolTip( + "Estimate axis steps from local detector and scan geometry " + "Jacobians, then apply them to selected grids or all grids." + ) + estimate_steps.clicked.connect(self._estimate_geometry_steps) + hdf5_settings = qt.QPushButton("HDF5 settings") + hdf5_settings.setToolTip( + "Set one chunk shape for all output grids and optionally override " + "the active database compression." + ) + hdf5_settings.clicked.connect(self._edit_hdf5_settings) + buttons.addWidget(add_hkl) + buttons.addWidget(add_q) + buttons.addWidget(remove) + buttons.addWidget(estimate_steps) + buttons.addWidget(hdf5_settings) + buttons.addStretch(1) + grid_layout.addLayout(buttons) + self.hdf5_summary = qt.QLabel() + self.hdf5_summary.setToolTip( + "Shared HDF5 chunk shape and compression selection for every grid." + ) + grid_layout.addWidget(self.hdf5_summary) + self._refresh_hdf5_summary() + layout.addWidget(grid_group) + return widget + + def _refresh_hdf5_summary(self): + chunk = self.orgui.reconstruction_chunk_shape + override = self.orgui.reconstruction_compression_override + if override is None: + database = getattr(self.orgui, "database", None) + compression = ( + f"active database filter " + f"({compression_filter_name(database.compression)})" + if database is not None + else "active database filter" + ) + else: + compression = f"override ({override})" + self.hdf5_summary.setText( + f"All grids: chunks {chunk[0]} × {chunk[1]} × {chunk[2]} voxels; " + f"compression: {compression}." + ) + + def _edit_hdf5_settings(self): + try: + database = self.orgui.database + dialog = HDF5SettingsDialog( + database.compression, + self.orgui.reconstruction_chunk_shape, + self.orgui.reconstruction_compression_override, + self, + ) + if dialog.exec() != qt.QDialog.Accepted: + return + database.compression = FILTERS[dialog.database_compression_name] + self.orgui.reconstruction_chunk_shape = dialog.chunk_shape + self.orgui.reconstruction_compression_override = ( + dialog.compression_override + ) + self._refresh_hdf5_summary() + except Exception as error: + self._report_failure("Cannot update HDF5 settings", error) + + def _estimate_geometry_steps(self): + try: + if self.orgui.fscan is None: + raise ValueError( + "Load a scan before estimating geometry-matched steps" + ) + rows = sorted( + {index.row() for index in self.grid_table.selectedIndexes()} + ) + if not rows: + rows = list(range(self.grid_table.rowCount())) + if not rows: + raise ValueError("Add an output grid before estimating steps") + dialog = _GeometryResolutionDialog(self) + if dialog.exec() != qt.QDialog.Accepted: + return + self._apply_geometry_steps(rows, dialog.percentile) + except Exception as error: + self._report_failure( + "Cannot estimate geometry-matched grid steps", error + ) + + def _apply_geometry_steps(self, rows, percentile): + config = ConfigData.from_gui(self.orgui) + estimates = {} + with qt.QSignalBlocker(self.grid_table): + for row in rows: + frame = self.grid_table.item(row, 1).text().strip() + if frame not in estimates: + estimates[frame] = estimate_geometry_steps( + config, + self.orgui.fscan, + frame=frame, + percentile=percentile, + ) + for axis, step in enumerate(estimates[frame]): + self.grid_table.item(row, 8 + axis).setData( + qt.Qt.EditRole, step + ) + self._update_grid_row(row) + + def _performance_tab(self): + widget = qt.QWidget() + layout = qt.QVBoxLayout(widget) + + accuracy_group = qt.QGroupBox("Accuracy") + accuracy_form = qt.QFormLayout(accuracy_group) + self.accuracy = qt.QComboBox() + self.accuracy.addItem("Center only (depth 0)", "center") + self.accuracy.addItem("Low (depth 1)", "low") + self.accuracy.addItem("Balanced (depth 2)", "balanced") + self.accuracy.addItem("High (depth 3)", "high") + self.accuracy.addItem("Very high (depth 4)", "very_high") + self.accuracy.addItem("Maximum (depth 5)", "maximum") + self.accuracy.setCurrentIndex( + self.accuracy.findData("balanced") + ) + self._add_form_row( + accuracy_form, + "Footprint preset:", + self.accuracy, + "Select the adaptive pixel-footprint subdivision depth. Higher " + "depths resolve voxel boundaries more accurately but require " + "substantially more computation.", + ) + layout.addWidget(accuracy_group) + + execution_group = qt.QGroupBox("Parallel execution and memory") + execution_form = qt.QFormLayout(execution_group) + thread_tooltip = ( + "Total CPU-thread budget shared by concurrent images and the " + "native threads working on each image." + ) + self.thread_override = self._optional_spin( + 1, 4096, thread_tooltip + ) + self._add_form_row( + execution_form, + "Total thread budget:", + self.thread_override[0], + thread_tooltip, + ) + self.threads_per_image = qt.QSpinBox() + self.threads_per_image.setRange(1, 4096) + self.threads_per_image.setValue(4) + self._add_form_row( + execution_form, + "Native threads per image:", + self.threads_per_image, + "Native C++ threads assigned to one image; concurrent image workers " + "use the remaining total thread and memory budgets.", + ) + memory_tooltip = ( + "Maximum total RAM available to this reconstruction job." + ) + self.memory_override = self._optional_spin( + 1, 1024 * 1024, memory_tooltip, suffix=" MiB" + ) + self._add_form_row( + execution_form, + "Total memory budget:", + self.memory_override[0], + memory_tooltip, + ) + accumulation_tooltip = ( + "Maximum reduced-record RAM retained by each image worker before " + "it writes a larger Parquet segment. The global memory budget " + "also reserves space for sorting and Arrow conversion." + ) + self.accumulation_memory = self._optional_spin( + 1, 1024 * 1024, accumulation_tooltip, suffix=" MiB" + ) + self._add_form_row( + execution_form, + "Accumulation per worker:", + self.accumulation_memory[0], + accumulation_tooltip, + ) + layout.addWidget(execution_group) + + advanced_group = qt.QGroupBox("Advanced settings") + advanced_form = qt.QFormLayout(advanced_group) + frame_tooltip = ( + "Number of consecutive scan frames assigned to each map task." + ) + self.frame_batch = self._optional_spin(1, 100000, frame_tooltip) + self._add_form_row( + advanced_form, + "Frames per task:", + self.frame_batch[0], + frame_tooltip, + ) + tile_tooltip = ( + "Rectangular detector tiles may end at any pixel boundary; " + "neighboring tiles share the same detector corner rays." + ) + self.tile_rows = self._optional_spin(1, 100000, tile_tooltip) + self._add_form_row( + advanced_form, "Tile rows:", self.tile_rows[0], tile_tooltip + ) + self.tile_columns = self._optional_spin(1, 100000, tile_tooltip) + self._add_form_row( + advanced_form, + "Tile columns:", + self.tile_columns[0], + tile_tooltip, + ) + block_tooltip = ( + "Fixed number of detector pixels scheduled as one native work block." + ) + self.work_block = self._optional_spin( + 1, 100000000, block_tooltip + ) + self._add_form_row( + advanced_form, + "Native block pixels:", + self.work_block[0], + block_tooltip, + ) + partition_tooltip = ( + "Number of consecutive HDF5 chunk IDs grouped into one Parquet " + "partition range." + ) + self.partition_span = self._optional_spin( + 1, 100000000, partition_tooltip + ) + self._add_form_row( + advanced_form, + "Parquet chunk span:", + self.partition_span[0], + partition_tooltip, + ) + layout.addWidget(advanced_group) + + detected_group = qt.QGroupBox("Detected execution layout") + detected_layout = qt.QVBoxLayout(detected_group) + self.performance_summary = qt.QPlainTextEdit() + self.performance_summary.setReadOnly(True) + self.performance_summary.setMaximumHeight(190) + self.performance_summary.setToolTip( + "Resolved task, tiling, thread, and memory values for this job." + ) + detected_layout.addWidget(self.performance_summary) + note = qt.QLabel( + "Unset values are derived from the active CPU count, orGUI memory " + "budget, detector shape, and output chunk geometry. Unchecked " + "editors display the detected values; enable Override to freeze " + "an edited value into this job." + ) + note.setWordWrap(True) + note.setToolTip( + "Enable Override beside a setting only when its detected value " + "should be replaced for this job." + ) + detected_layout.addWidget(note) + layout.addWidget(detected_group) + layout.addStretch(1) + return widget + + def _paths_tab(self): + widget = qt.QWidget() + layout = qt.QVBoxLayout(widget) + self.job_path = qt.QLineEdit() + self.scratch_path = qt.QLineEdit() + self.output_path = qt.QLineEdit() + + job_group = qt.QGroupBox("Job descriptor") + job_form = qt.QFormLayout(job_group) + job_tooltip = ( + "Resumable JSON descriptor containing the frozen experiment and " + "execution settings." + ) + self._add_form_row( + job_form, + "Job JSON:", + self._path_row( + self.job_path, + "job", + save=True, + tooltip=job_tooltip, + ), + job_tooltip, + ) + layout.addWidget(job_group) + + storage_group = qt.QGroupBox("Storage") + storage_form = qt.QFormLayout(storage_group) + scratch_tooltip = ( + "Writable directory for immutable Parquet partitions and the " + "checksummed job asset bundle." + ) + self._add_form_row( + storage_form, + "Scratch directory:", + self._path_row( + self.scratch_path, + "directory", + tooltip=scratch_tooltip, + ), + scratch_tooltip, + ) + output_tooltip = ( + "Final standalone NeXus/HDF5 reciprocal-space reconstruction file." + ) + self._add_form_row( + storage_form, + "Standalone HDF5:", + self._path_row( + self.output_path, + "hdf5", + save=True, + tooltip=output_tooltip, + ), + output_tooltip, + ) + note = qt.QLabel( + "Scratch data are retained after interruption and removed " + "automatically only after successful HDF5 validation." + ) + note.setWordWrap(True) + note.setToolTip( + "Interrupted jobs keep their scratch data so they can be resumed." + ) + storage_form.addRow(note) + layout.addWidget(storage_group) + layout.addStretch(1) + return widget + + def _cluster_tab(self): + widget = qt.QWidget() + outer = qt.QVBoxLayout(widget) + scroll = qt.QScrollArea() + scroll.setWidgetResizable(True) + contents = qt.QWidget() + layout = qt.QVBoxLayout(contents) + + scheduler_group = qt.QGroupBox("Scheduler") + scheduler_form = qt.QFormLayout(scheduler_group) + self.cluster_scheduler = qt.QComboBox() + self.cluster_scheduler.addItem("Sun/Grid Engine (SGE)", "sge") + self.cluster_scheduler.addItem("Slurm", "slurm") + self._add_form_row( + scheduler_form, + "Scheduler:", + self.cluster_scheduler, + "Generate SGE qsub scripts or Slurm sbatch scripts. SGE is the " + "default and primary target.", + ) + self.cluster_job_name = qt.QLineEdit("orgui-rsmap") + self._add_form_row( + scheduler_form, + "Job name:", + self.cluster_job_name, + "Scheduler-safe base name used for the map array and finalizer.", + ) + self.cluster_queue = qt.QLineEdit() + self._add_form_row( + scheduler_form, + "Queue / partition:", + self.cluster_queue, + "Optional SGE queue (-q) or Slurm partition (--partition).", + ) + self.cluster_account = qt.QLineEdit() + self._add_form_row( + scheduler_form, + "Project / account:", + self.cluster_account, + "Optional SGE project (-P) or Slurm account (--account).", + ) + layout.addWidget(scheduler_group) + + environment_group = qt.QGroupBox("Python environment") + environment_form = qt.QFormLayout(environment_group) + self.cluster_script_directory = qt.QLineEdit() + self._add_form_row( + environment_form, + "Script directory:", + self._path_row( + self.cluster_script_directory, + "directory", + tooltip=( + "Directory receiving map, finalizer, and submission scripts." + ), + ), + "Directory receiving map, finalizer, and submission scripts.", + ) + self.cluster_working_directory = qt.QLineEdit(str(Path.cwd())) + self._add_form_row( + environment_form, + "Working directory:", + self._path_row( + self.cluster_working_directory, + "directory", + tooltip="Shared working directory selected before Python starts.", + ), + "Shared working directory selected before Python starts.", + ) + self.cluster_python = qt.QLineEdit("python") + self._add_form_row( + environment_form, + "Python executable:", + self.cluster_python, + "Python command available after environment setup, normally " + "'python' from the activated orGUI environment.", + ) + self.cluster_environment = qt.QPlainTextEdit() + self.cluster_environment.setPlaceholderText( + "module load ...\nsource /path/to/venv/bin/activate" + ) + self.cluster_environment.setMaximumHeight(90) + self._add_form_row( + environment_form, + "Setup commands:", + self.cluster_environment, + "Shell commands run before every task, for example module loads " + "and conda or virtual-environment activation.", + ) + layout.addWidget(environment_group) + + map_group = qt.QGroupBox("Mapping array") + map_form = qt.QFormLayout(map_group) + self.cluster_array_cpus = qt.QSpinBox() + self.cluster_array_cpus.setRange(1, 4096) + self.cluster_array_cpus.setValue(4) + self._add_form_row( + map_form, + "CPUs / slots per task:", + self.cluster_array_cpus, + "Native C++ threads used by each independent map-array task.", + ) + self.cluster_array_memory = qt.QDoubleSpinBox() + self.cluster_array_memory.setRange(0.25, 1024 * 1024) + self.cluster_array_memory.setDecimals(2) + self.cluster_array_memory.setValue(16.0) + self.cluster_array_memory.setSuffix(" GiB") + self._add_form_row( + map_form, + "Memory per task:", + self.cluster_array_memory, + "Total RAM budget for one array task. For SGE, the generated " + "h_vmem request is divided across the requested slots.", + ) + self.cluster_array_walltime = qt.QLineEdit("24:00:00") + self._add_form_row( + map_form, + "Wall time:", + self.cluster_array_walltime, + "Scheduler wall-time limit for each mapping task.", + ) + self.cluster_array_concurrency = qt.QSpinBox() + self.cluster_array_concurrency.setRange(0, 1000000) + self.cluster_array_concurrency.setValue(0) + self._add_form_row( + map_form, + "Maximum concurrent tasks:", + self.cluster_array_concurrency, + "Optional array throttle (-tc or %N). Zero leaves the scheduler " + "limit unchanged.", + ) + self.cluster_summary = qt.QPlainTextEdit() + self.cluster_summary.setReadOnly(True) + self.cluster_summary.setMaximumHeight(75) + self._add_form_row( + map_form, + "Detected array:", + self.cluster_summary, + "One array element is created for each detected deterministic " + "map task. Prepare the job to update this count.", + ) + layout.addWidget(map_group) + + reduce_group = qt.QGroupBox("Reduction and finalization") + reduce_form = qt.QFormLayout(reduce_group) + self.cluster_reduce_cpus = qt.QSpinBox() + self.cluster_reduce_cpus.setRange(1, 4096) + self.cluster_reduce_cpus.setValue(24) + self._add_form_row( + reduce_form, + "CPUs / slots:", + self.cluster_reduce_cpus, + "Independent CPU count for parallel shard reduction; it need not " + "match the map-array task size.", + ) + self.cluster_reduce_memory = qt.QDoubleSpinBox() + self.cluster_reduce_memory.setRange(0.25, 1024 * 1024) + self.cluster_reduce_memory.setDecimals(2) + self.cluster_reduce_memory.setValue(64.0) + self.cluster_reduce_memory.setSuffix(" GiB") + self._add_form_row( + reduce_form, + "Memory:", + self.cluster_reduce_memory, + "Total RAM budget divided among parallel reducer workers.", + ) + self.cluster_reduce_walltime = qt.QLineEdit("24:00:00") + self._add_form_row( + reduce_form, + "Wall time:", + self.cluster_reduce_walltime, + "Wall-time limit for reduction and final HDF5 creation.", + ) + layout.addWidget(reduce_group) + + advanced_group = qt.QGroupBox("Scheduler-specific settings") + advanced_form = qt.QFormLayout(advanced_group) + self.cluster_sge_pe = qt.QLineEdit("smp") + self._add_form_row( + advanced_form, + "SGE parallel environment:", + self.cluster_sge_pe, + "SGE parallel environment requested with -pe, commonly 'smp'.", + ) + self.cluster_sge_memory = qt.QLineEdit("h_vmem") + self._add_form_row( + advanced_form, + "SGE memory resource:", + self.cluster_sge_memory, + "SGE per-slot complex used for memory requests, commonly h_vmem " + "or mem_free.", + ) + self.cluster_array_directives = qt.QPlainTextEdit() + self.cluster_array_directives.setMaximumHeight(70) + self._add_form_row( + advanced_form, + "Extra map directives:", + self.cluster_array_directives, + "Optional scheduler directives, one '#$' or '#SBATCH' line each.", + ) + self.cluster_reduce_directives = qt.QPlainTextEdit() + self.cluster_reduce_directives.setMaximumHeight(70) + self._add_form_row( + advanced_form, + "Extra finalizer directives:", + self.cluster_reduce_directives, + "Optional scheduler directives for reduction/finalization.", + ) + layout.addWidget(advanced_group) + layout.addStretch(1) + scroll.setWidget(contents) + outer.addWidget(scroll) + return widget + + def _cluster_settings(self): + return ClusterSettings( + scheduler=self.cluster_scheduler.currentData(), + job_name=self.cluster_job_name.text().strip(), + script_directory=self.cluster_script_directory.text().strip(), + working_directory=self.cluster_working_directory.text().strip(), + python_executable=self.cluster_python.text().strip(), + environment_setup=self.cluster_environment.toPlainText(), + queue=self.cluster_queue.text().strip(), + account=self.cluster_account.text().strip(), + array_cpus=self.cluster_array_cpus.value(), + array_memory_gib=self.cluster_array_memory.value(), + array_walltime=self.cluster_array_walltime.text().strip(), + array_concurrency=self.cluster_array_concurrency.value(), + reduce_cpus=self.cluster_reduce_cpus.value(), + reduce_memory_gib=self.cluster_reduce_memory.value(), + reduce_walltime=self.cluster_reduce_walltime.text().strip(), + sge_parallel_environment=self.cluster_sge_pe.text().strip(), + sge_memory_resource=self.cluster_sge_memory.text().strip(), + extra_array_directives=( + self.cluster_array_directives.toPlainText() + ), + extra_reduce_directives=( + self.cluster_reduce_directives.toPlainText() + ), + ) + + def _set_cluster_settings(self, values): + settings = ClusterSettings.from_dict(values) + self.cluster_scheduler.setCurrentIndex( + self.cluster_scheduler.findData(settings.scheduler) + ) + self.cluster_job_name.setText(settings.job_name) + self.cluster_script_directory.setText(settings.script_directory) + self.cluster_working_directory.setText(settings.working_directory) + self.cluster_python.setText(settings.python_executable) + self.cluster_environment.setPlainText(settings.environment_setup) + self.cluster_queue.setText(settings.queue) + self.cluster_account.setText(settings.account) + self.cluster_array_cpus.setValue(settings.array_cpus) + self.cluster_array_memory.setValue(settings.array_memory_gib) + self.cluster_array_walltime.setText(settings.array_walltime) + self.cluster_array_concurrency.setValue( + settings.array_concurrency + ) + self.cluster_reduce_cpus.setValue(settings.reduce_cpus) + self.cluster_reduce_memory.setValue(settings.reduce_memory_gib) + self.cluster_reduce_walltime.setText(settings.reduce_walltime) + self.cluster_sge_pe.setText(settings.sge_parallel_environment) + self.cluster_sge_memory.setText(settings.sge_memory_resource) + self.cluster_array_directives.setPlainText( + settings.extra_array_directives + ) + self.cluster_reduce_directives.setPlainText( + settings.extra_reduce_directives + ) + + @staticmethod + def _set_control_tooltip(control, tooltip): + widgets = control if isinstance(control, tuple) else (control,) + for child in widgets: + child.setToolTip(tooltip) + + def _add_form_row(self, form, label, control, tooltip): + label_widget = qt.QLabel(label) + label_widget.setToolTip(tooltip) + self._set_control_tooltip(control, tooltip) + form.addRow(label_widget, control) + + def _optional_spin(self, minimum, maximum, tooltip, suffix=""): + container = qt.QWidget() + layout = qt.QHBoxLayout(container) + layout.setContentsMargins(0, 0, 0, 0) + enabled = qt.QCheckBox("Override") + editor = qt.QSpinBox() + editor.setRange(minimum, maximum) + editor.setSuffix(suffix) + editor.setEnabled(False) + enabled.toggled.connect(editor.setEnabled) + layout.addWidget(enabled) + layout.addWidget(editor) + layout.addStretch(1) + control = (container, enabled, editor) + self._set_control_tooltip(control, tooltip) + return control + + def _path_row(self, editor, kind, save=False, tooltip=""): + widget = qt.QWidget() + layout = qt.QHBoxLayout(widget) + layout.setContentsMargins(0, 0, 0, 0) + widget.setToolTip(tooltip) + editor.setToolTip(tooltip) + layout.addWidget(editor) + button = qt.QPushButton("Browse…") + button.setToolTip(tooltip) + button.clicked.connect( + lambda: self._browse(editor, kind=kind, save=save) + ) + layout.addWidget(button) + return widget + + def _browse(self, editor, *, kind, save): + try: + directory = str(self._history_directory()) + if kind == "directory": + value = qt.QFileDialog.getExistingDirectory( + self, "Select scratch directory", directory + ) + else: + file_filter = ( + "JSON job (*.json)" + if kind == "job" + else "HDF5/NeXus (*.h5 *.hdf5 *.nxs)" + ) + dialog = ( + qt.QFileDialog.getSaveFileName + if save + else qt.QFileDialog.getOpenFileName + ) + value, _ = dialog(self, "Select file", directory, file_filter) + if value: + editor.setText(value) + self.orgui.filedialogdir = str(Path(value).parent) + return True + return False + except Exception as error: + self._report_failure("Cannot select reconstruction path", error) + return False + + def _history_directory(self): + history = Path(str(self.orgui.filedialogdir)).expanduser() + try: + if history.is_dir(): + return history + if history.parent.is_dir(): + return history.parent + except OSError: + pass + return Path.cwd() + + def _set_default_paths(self, stem): + base = Path.cwd() + if not self.job_path.text(): + self.job_path.setText(str(base / f"{stem}-rsmap.json")) + if not self.scratch_path.text(): + self.scratch_path.setText(str(base / f".{stem}-rsmap-scratch")) + if not self.output_path.text(): + self.output_path.setText(str(base / f"{stem}-rsmap.h5")) + if not self.cluster_script_directory.text(): + self.cluster_script_directory.setText( + str(base / f"{stem}-rsmap-cluster") + ) + + def _show_output(self, description): + self.preview_output.setPlainText(description) + self.tabs.setCurrentWidget(self.output_tab) + + def _report_message(self, title, description): + self._show_output(description) + logger.warning( + title, + extra={ + "title": title, + "description": description, + "show_dialog": True, + "dialog_level": logging.WARNING, + "parent": self, + }, + ) + + def _report_failure(self, title, error): + description = str(error) or type(error).__name__ + self._show_output(description) + logger.warning( + title, + exc_info=True, + extra={ + "title": title, + "description": description, + "show_dialog": True, + "dialog_level": logging.WARNING, + "parent": self, + }, + ) + + def _invoke_ui_action(self, title, action): + try: + action() + except Exception as error: + self._report_failure(title, error) + + def _live_correction_summary(self, config, detector_shape): + correction = config.corrections + summary = correction.to_dict() + for asset_field in ( + "mask_asset", + "background_asset", + "background_variance_asset", + ): + summary.pop(asset_field, None) + + mask = ( + self.orgui.get_detector_mask(detector_shape) + if correction.use_mask + else None + ) + if not correction.use_mask: + mask_status = {"status": "disabled"} + elif mask is None: + mask_status = {"status": "enabled but unavailable"} + else: + mask_status = { + "status": "active", + "masked_pixels": int(np.count_nonzero(mask)), + "job_asset_path": "/mask", + } + + background = ( + getattr(self.orgui, "background_image", None) + if correction.use_background + else None + ) + if not correction.use_background: + background_status = {"status": "disabled"} + elif background is None: + background_status = {"status": "enabled but unavailable"} + else: + background_status = { + "status": "active", + "job_asset_path": "/background", + "variance": ( + "available" + if getattr(self.orgui, "background_variance", None) + is not None + else "deterministic" + ), + } + summary["active_inputs"] = { + "mask": mask_status, + "background": background_status, + } + return summary + + def refresh_live_state(self): + """Refresh the experiment and correction summary from orGUI.""" + try: + if self.orgui.fscan is None: + self.experiment_summary.setPlainText("No active scan.") + return + config = ConfigData.from_gui(self.orgui) + detector_shape = self.orgui.ubcalc.detectorCal.detector.shape + summary = { + "scan": getattr( + self.orgui.fscan, + "name", + getattr(self.orgui.fscan, "title", ""), + ), + "frames": len(self.orgui.fscan), + "detector_shape": detector_shape, + "energy_keV": self.orgui.ubcalc.ubCal.getEnergy(), + "UB": self.orgui.ubcalc.ubCal.getUB().tolist(), + "corrections": self._live_correction_summary( + config, detector_shape + ), + "threads": self.orgui.numberthreads, + "memory_MiB": self.orgui.maxMemory, + } + self.experiment_summary.setPlainText( + json.dumps(summary, indent=2, sort_keys=True) + ) + stem = "".join( + character if character.isalnum() or character in "-_" else "_" + for character in str(summary["scan"]) + ).strip("_") or "scan" + self._set_default_paths(stem) + if self.grid_table.rowCount() == 0: + self.add_derived_grid("hkl") + except Exception as error: + self._report_failure("Cannot refresh reconstruction settings", error) + + def _add_q_grid(self): + try: + frame, accepted = qt.QInputDialog.getItem( + self, + "Q coordinate frame", + "Frame:", + list(_FRAMES[1:]), + 0, + False, + ) + if accepted: + self.add_derived_grid(frame) + except Exception as error: + self._report_failure("Cannot add Q output grid", error) + + def add_derived_grid(self, frame): + """Add a grid derived from current scan coverage.""" + if self.orgui.fscan is None: + self._report_message( + "No scan loaded", + "Load a scan before deriving reciprocal-space grid coverage.", + ) + return + try: + config = ConfigData.from_gui(self.orgui) + grid = derive_grid(config, self.orgui.fscan, frame=frame) + self._append_grid(grid) + except Exception as error: + self._report_failure("Cannot derive output grid", error) + + def _append_grid(self, grid): + row = self.grid_table.rowCount() + with qt.QSignalBlocker(self.grid_table): + self.grid_table.insertRow(row) + values = [ + grid.name or "", + grid.frame, + *grid.minimum, + *grid.maximum, + *grid.step, + ] + for column, value in enumerate(values): + item = ( + qt.QTableWidgetItem(str(value)) + if column < 2 + else _GridNumberItem(value) + ) + self.grid_table.setItem(row, column, item) + self._update_grid_row(row) + + def _on_grid_cell_changed(self, row, column): + try: + self._update_grid_row(row, changed_column=column) + except (TypeError, ValueError, OverflowError): + size_item = self.grid_table.item(row, 14) + if size_item is not None: + with qt.QSignalBlocker(self.grid_table): + size_item.setText("invalid") + + def _update_grid_row(self, row, changed_column=None): + with qt.QSignalBlocker(self.grid_table): + intervals = [] + for axis in range(3): + minimum = float( + self.grid_table.item(row, 2 + axis).data(qt.Qt.EditRole) + ) + maximum = float( + self.grid_table.item(row, 5 + axis).data(qt.Qt.EditRole) + ) + extent = maximum - minimum + if not np.isfinite(extent) or extent <= 0: + raise ValueError("Grid extent must be finite and positive") + interval_column = 11 + axis + if changed_column == interval_column: + value = float( + self.grid_table.item(row, interval_column).data( + qt.Qt.EditRole + ) + ) + count = int(value) + if count < 1 or value != count: + raise ValueError( + "Grid interval counts must be positive integers" + ) + step = np.nextafter(extent / count, np.inf) + self.grid_table.item(row, 8 + axis).setData( + qt.Qt.EditRole, step + ) + else: + step = float( + self.grid_table.item(row, 8 + axis).data( + qt.Qt.EditRole + ) + ) + if not np.isfinite(step) or step <= 0: + raise ValueError( + "Grid steps must be finite and positive" + ) + count = int(np.ceil(extent / step)) + item = self.grid_table.item(row, interval_column) + if item is None: + item = qt.QTableWidgetItem() + self.grid_table.setItem(row, interval_column, item) + item.setData(qt.Qt.EditRole, count) + intervals.append(count) + + estimated_bytes = int(np.prod(intervals, dtype=object)) * 32 + size_item = self.grid_table.item(row, 14) + if size_item is None: + size_item = qt.QTableWidgetItem() + size_item.setFlags( + size_item.flags() & ~qt.Qt.ItemIsEditable + ) + self.grid_table.setItem(row, 14, size_item) + size_item.setText(f"{_format_size(estimated_bytes)} uncompressed") + + def _remove_grid(self): + try: + rows = sorted( + {index.row() for index in self.grid_table.selectedIndexes()}, + reverse=True, + ) + for row in rows: + self.grid_table.removeRow(row) + except Exception as error: + self._report_failure("Cannot remove output grid", error) + + def _grids(self): + grids = [] + for row in range(self.grid_table.rowCount()): + items = [ + self.grid_table.item(row, column) + for column in range(self.grid_table.columnCount()) + ] + if any(item is None for item in items): + raise ValueError(f"Grid row {row + 1} contains an empty cell") + values = [item.text().strip() for item in items[:2]] + if values[1] not in _FRAMES: + raise ValueError( + f"Grid row {row + 1} has an unknown frame: {values[1]}" + ) + grid = ReconstructionGrid( + minimum=tuple( + float(item.data(qt.Qt.EditRole)) for item in items[2:5] + ), + maximum=tuple( + float(item.data(qt.Qt.EditRole)) for item in items[5:8] + ), + step=tuple( + float(item.data(qt.Qt.EditRole)) for item in items[8:11] + ), + chunk_shape=tuple( + self.orgui.reconstruction_chunk_shape + ), + frame=values[1], + name=values[0] or None, + ) + intervals = tuple( + int(item.data(qt.Qt.EditRole)) for item in items[11:14] + ) + if grid.to_spec().shape != intervals: + raise ValueError( + f"Grid row {row + 1} interval counts are inconsistent " + "with its bounds and steps" + ) + if any(step <= 0 for step in grid.step): + raise ValueError(f"Grid row {row + 1} steps must be positive") + if any( + upper <= lower + for lower, upper in zip(grid.minimum, grid.maximum) + ): + raise ValueError( + f"Grid row {row + 1} maxima must exceed its minima" + ) + grids.append(grid) + if not grids: + raise ValueError("At least one output grid is required") + return grids + + @staticmethod + def _optional_value(control): + _, enabled, editor = control + return editor.value() if enabled.isChecked() else None + + @staticmethod + def _set_detected_value(control, value): + _, enabled, editor = control + if enabled.isChecked() or value is None: + return + editor.setValue( + max(editor.minimum(), min(editor.maximum(), int(round(value)))) + ) + + def _show_execution_settings(self, job, *, scan=None, config=None): + settings = reconstruction_execution_settings( + job, scan=scan, config=config + ) + tile_rows, tile_columns = settings["detector_tile_shape"] + for control, value in ( + (self.thread_override, settings["thread_budget"]), + (self.memory_override, settings["memory_budget_MiB"]), + ( + self.accumulation_memory, + settings["accumulation_budget_MiB_per_worker"], + ), + (self.frame_batch, settings["frames_per_task"]), + (self.tile_rows, tile_rows), + (self.tile_columns, tile_columns), + (self.work_block, settings["native_work_block_pixels"]), + (self.partition_span, settings["parquet_chunk_span"]), + ): + self._set_detected_value(control, value) + self.threads_per_image.setValue( + settings["native_threads_per_image"] + ) + self.performance_summary.setPlainText( + json.dumps(settings, indent=2, sort_keys=True) + ) + self.cluster_summary.setPlainText( + json.dumps( + { + "array_tasks": settings["map_tasks"], + "cpus_per_task": self.cluster_array_cpus.value(), + "maximum_concurrent_tasks": ( + self.cluster_array_concurrency.value() or "scheduler" + ), + }, + indent=2, + ) + ) + return settings + + def _prepare(self): + if self.orgui.fscan is None: + raise ValueError("Load a scan before preparing a reconstruction job") + tile_rows = self._optional_value(self.tile_rows) + tile_columns = self._optional_value(self.tile_columns) + if (tile_rows is None) != (tile_columns is None): + raise ValueError("Override both detector tile dimensions together") + job_path = self.job_path.text().strip() + scratch_path = self.scratch_path.text().strip() + output_path = self.output_path.text().strip() + if not job_path: + raise ValueError("Select a job JSON path") + if not scratch_path: + raise ValueError("Select a scratch directory") + if not output_path: + raise ValueError("Select an output HDF5 path") + job = prepare_job( + self.orgui, + job_path, + grids=self._grids(), + scratch_path=scratch_path, + output_path=output_path, + accuracy=self.accuracy.currentData(), + compression_override=( + self.orgui.reconstruction_compression_override + ), + angle_fallback=self.angle_fallback.currentData(), + user_note=self.user_note.text(), + thread_override=self._optional_value(self.thread_override), + memory_override_bytes=( + None + if self._optional_value(self.memory_override) is None + else self._optional_value(self.memory_override) * 1024 * 1024 + ), + frame_batch=self._optional_value(self.frame_batch), + tile_shape=( + None + if tile_rows is None + else (tile_rows, tile_columns) + ), + work_block_pixels=self._optional_value(self.work_block), + partition_chunk_span=self._optional_value(self.partition_span), + threads_per_image=self.threads_per_image.value(), + accumulation_budget_bytes=( + None + if self._optional_value(self.accumulation_memory) is None + else self._optional_value(self.accumulation_memory) + * 1024 + * 1024 + ), + cluster_settings=self._cluster_settings().to_dict(), + ) + self._show_execution_settings( + job, + scan=self.orgui.fscan, + config=job.config_data, + ) + return job + + @qt.Slot() + def preview(self): + """Show live grid and resource estimates without freezing state.""" + try: + if self.orgui.fscan is None: + raise ValueError( + "Load a scan before previewing a reconstruction job" + ) + grids = self._grids() + depth = ACCURACY_DEPTHS[self.accuracy.currentData()] + grid_rows = [] + final_bytes = 0 + chunk_count = 0 + for grid in grids: + shape = tuple( + int(np.ceil((upper - lower) / step)) + for lower, upper, step in zip( + grid.minimum, grid.maximum, grid.step + ) + ) + voxels = int(np.prod(shape, dtype=np.int64)) + chunks = int( + np.prod( + [ + np.ceil(size / chunk) + for size, chunk in zip(shape, grid.chunk_shape) + ] + ) + ) + final_bytes += voxels * 32 + chunk_count += chunks + grid_rows.append({**grid.__dict__, "shape": shape}) + corrections = ConfigData.from_gui(self.orgui).corrections + included_frames = len(self.orgui.fscan) - len( + { + frame + for frame in corrections.excluded_frames + if 0 <= frame < len(self.orgui.fscan) + } + ) + detector_pixels = int( + np.prod(self.orgui.ubcalc.detectorCal.detector.shape) + ) + result = { + "grids": grid_rows, + "frames": included_frames, + "threads": self._optional_value(self.thread_override) + or self.orgui.numberthreads, + "native_threads_per_image": self.threads_per_image.value(), + "memory_MiB": self._optional_value(self.memory_override) + or self.orgui.maxMemory, + "accumulation_MiB_per_worker": self._optional_value( + self.accumulation_memory + ) + or "automatic", + "estimated_spatial_chunks": chunk_count, + "uncompressed_final_GiB": final_bytes / 1024**3, + "footprint_leaf_upper_bound": ( + included_frames + * detector_pixels + * min(8**depth, 4096) + ), + } + self._show_output( + json.dumps(result, indent=2, sort_keys=True) + ) + except Exception as error: + self._report_failure("Cannot preview reconstruction", error) + + @qt.Slot() + def prepare(self): + """Freeze current state and save a prepared job.""" + try: + job = self._prepare() + self._show_output( + json.dumps(job.to_dict(), indent=2, sort_keys=True) + ) + except Exception as error: + self._report_failure("Cannot prepare reconstruction job", error) + + def _run_path(self, path): + path = str(path).strip() + if not path: + raise ValueError("Select a prepared job JSON path") + progress = None + + def update(value, maximum, message): + progress.total = maximum + if hasattr(progress, "dialog"): + progress.dialog.setMaximum(maximum) + progress.update(value, message) + if progress.wasCanceled(): + raise _ReconstructionCancelled( + "Reconstruction cancelled between tasks" + ) + + try: + progress = logger_utils.create_progress_logger( + self, 1, "Reciprocal-space reconstruction" + ) + result = run_job( + path, + progress=update, + ) + try: + self._register_external_result(result) + except Exception as error: + self._report_failure( + "Reconstruction completed but result registration failed", + error, + ) + self._show_output( + json.dumps(result, indent=2, sort_keys=True) + ) + finally: + if progress is not None: + try: + progress.finish() + except Exception: + logger.warning( + "Cannot close reconstruction progress reporting.", + exc_info=True, + ) + + @qt.Slot() + def run_local(self): + """Freeze live state and execute the resulting job locally.""" + try: + self._prepare() + self._run_path(self.job_path.text()) + except _ReconstructionCancelled as error: + self._report_message("Reconstruction cancelled", str(error)) + except Exception as error: + self._report_failure("Cannot run reconstruction", error) + + @qt.Slot() + def create_cluster_scripts(self): + """Prepare the current job and create scheduler batch scripts.""" + try: + if self.orgui.fscan is None: + job_path = self.job_path.text().strip() + if not job_path: + raise ValueError("Select a prepared job JSON path") + job = read_job(job_path) + else: + job = self._prepare() + job_path = self.job_path.text().strip() + result = generate_cluster_scripts(job_path, job) + self._show_output( + json.dumps(result, indent=2, sort_keys=True) + ) + except Exception as error: + self._report_failure("Cannot create cluster scripts", error) + + @qt.Slot() + def open_job(self): + """Open and display an existing prepared job.""" + try: + if not self._browse(self.job_path, kind="job", save=False): + return + job_path = self.job_path.text().strip() + job = read_job(job_path) + self._show_output( + json.dumps(job_status(job_path), indent=2) + ) + self.output_path.setText(job.output_path) + self.scratch_path.setText(job.scratch_path) + chunk_shapes = { + tuple(values["chunk_shape"]) for values in job.grids + } + if len(chunk_shapes) != 1: + raise ValueError( + "The prepared job contains different per-grid HDF5 chunk " + "shapes and cannot be represented by the global setting." + ) + self.orgui.reconstruction_chunk_shape = chunk_shapes.pop() + active_compression = compression_filter_name( + self.orgui.database.compression + ) + self.orgui.reconstruction_compression_override = ( + None + if job.compression == active_compression + else job.compression + ) + self._refresh_hdf5_summary() + self._set_cluster_settings(job.cluster_settings) + self.grid_table.setRowCount(0) + for values in job.grids: + self._append_grid(ReconstructionGrid(**values)) + self._show_execution_settings(job) + except Exception as error: + self._report_failure("Cannot open reconstruction job", error) + + @qt.Slot() + def resume(self): + """Verify and resume the selected prepared job.""" + try: + self._run_path(self.job_path.text()) + except _ReconstructionCancelled as error: + self._report_message("Reconstruction cancelled", str(error)) + except Exception as error: + self._report_failure("Cannot resume reconstruction", error) + + def _register_external_result(self, result): + self.orgui.database.register_external_result( + result["output_path"], + result["output_sha256"], + result["grids"], + result["status"], + result["job_sha256"], + ) diff --git a/orgui/app/_roi_sum_accel.py b/orgui/app/_roi_sum_accel.py index 688e2c2..83ae614 100644 --- a/orgui/app/_roi_sum_accel.py +++ b/orgui/app/_roi_sum_accel.py @@ -35,7 +35,12 @@ def _import_cpp_backend(): return importlib.import_module("orgui.app._roi_sum_cpp") except ModuleNotFoundError as package_error: repo_root = Path(__file__).resolve().parents[2] - candidates = sorted((repo_root / "build").glob("cp*/_roi_sum_cpp*.so")) + candidates = sorted( + ( + *(repo_root / "build").glob("cp*/_roi_sum_cpp*.so"), + *(repo_root / "build").glob("cp*/_roi_sum_cpp*.pyd"), + ) + ) if not candidates: raise package_error extension_path = candidates[-1] @@ -167,6 +172,9 @@ def set_accel_backend(backend): _cpp_backend = None HAS_CPP_BACKEND = _cpp_backend is not None +PixelRepairPlan = ( + None if _cpp_backend is None else getattr(_cpp_backend, "PixelRepairPlan", None) +) HAS_NUMBA_BACKEND = False HAS_ACCEL_BACKEND = ROI_ACCEL_BACKEND != "numpy" _bind_backend(ROI_ACCEL_BACKEND, _counter_backend) diff --git a/orgui/app/config_data.py b/orgui/app/config_data.py index 212045c..395f449 100644 --- a/orgui/app/config_data.py +++ b/orgui/app/config_data.py @@ -19,7 +19,9 @@ import configparser import datetime +import json from dataclasses import dataclass, field +from typing import Any import h5py import numpy as np @@ -30,6 +32,67 @@ SCHEMA_VERSION = 1 +def _json_value(value): + if isinstance(value, dict): + return {str(key): _json_value(item) for key, item in value.items()} + if isinstance(value, list | tuple): + return [_json_value(item) for item in value] + if isinstance(value, np.ndarray): + return value.tolist() + if isinstance(value, np.generic): + return value.item() + if isinstance(value, bytes): + return value.decode() + return value + + +@dataclass +class CorrectionState: + """Persist the correction selections shared by orGUI workflows. + + Large mask, background, and variance arrays are stored in a job asset + bundle. The corresponding fields here identify datasets in that bundle. + """ + + use_mask: bool = False + use_background: bool = False + use_solid_angle: bool = False + use_polarization: bool = False + repair_masked_pixels: bool = False + repair_max_component_pixels: int | None = None + repair_max_span: int | None = None + repair_radius: int | None = None + repair_min_valid_neighbors: int | None = None + repair_use_pyfai_gaps: bool = True + repair_gap_size_px: int = 1 + normalize_exposure: bool = True + monitor_corrections: tuple[str, ...] = () + excluded_frames: tuple[int, ...] = () + mask_asset: str | None = None + background_asset: str | None = None + background_variance_asset: str | None = None + uncertainty_provenance: dict[str, Any] = field(default_factory=dict) + + def to_dict(self) -> dict[str, Any]: + """Return a JSON-compatible correction-state dictionary.""" + result = _json_value(self.__dict__) + result["monitor_corrections"] = list(self.monitor_corrections) + result["excluded_frames"] = list(self.excluded_frames) + return result + + @classmethod + def from_dict(cls, values): + """Build correction state from a JSON-compatible dictionary.""" + values = dict(values or {}) + values["monitor_corrections"] = tuple( + values.get("monitor_corrections", ()) + ) + values["excluded_frames"] = tuple( + int(value) for value in values.get("excluded_frames", ()) + ) + return cls(**values) + + def _as_text(value): if isinstance(value, bytes): return value.decode() @@ -53,7 +116,8 @@ def _string_array(strings): def _read_string_array(data): array = np.asarray(data) - if array.dtype == np.uint8 and array.ndim == 2: + if array.dtype.kind in "iu" and array.ndim == 2: + array = array.astype(np.uint8, copy=False) return [bytes(row[row != 0]).decode("utf-8") for row in array] return [_as_text(value) for value in data] @@ -208,6 +272,7 @@ class ConfigData: phi: float = 0.0 refraction_index: float = 1.0 reference_reflections: list = field(default_factory=list) + corrections: CorrectionState = field(default_factory=CorrectionState) orgui: dict = field(default_factory=dict) @classmethod @@ -284,6 +349,49 @@ def from_gui(cls, gui): cell = unit_cell_from_nxdict(unit_cell_to_nxdict(ub_widget.crystal)) ub_calculator = HKLVlieg.UBCalculator(cell, ub_widget.ubCal.getEnergy()) ub_calculator.setU(ub_widget.ubCal.getU()) + corrections = CorrectionState() + if hasattr(gui, "scanSelector"): + options = gui.scanSelector.get_integration_options() + repair = getattr(getattr(gui, "maskManager", None), "settings", None) + repair = getattr(repair, "pixel_repair", None) + repair_enabled = bool(getattr(repair, "enabled", False)) + excluded = getattr(gui, "excludedImagesDialog", None) + excluded = () if excluded is None else excluded.getData() + corrections = CorrectionState( + use_mask=bool(options.get("mask", False)) or repair_enabled, + use_background=getattr(gui, "background_image", None) + is not None, + use_solid_angle=bool(options.get("solidAngle", False)), + use_polarization=bool(options.get("polarization", False)), + repair_masked_pixels=repair_enabled, + repair_max_component_pixels=getattr( + repair, "max_component_pixels", None + ), + repair_max_span=getattr(repair, "max_span", None), + repair_radius=getattr(repair, "radius", None), + repair_min_valid_neighbors=getattr( + repair, "min_valid_neighbors", None + ), + repair_use_pyfai_gaps=bool( + getattr(repair, "use_pyfai_gaps", True) + ), + repair_gap_size_px=int( + getattr(repair, "gap_size_px", 1) + ), + normalize_exposure=bool( + getattr(gui, "reconstruction_normalize_exposure", True) + ), + monitor_corrections=tuple( + getattr(gui, "reconstruction_monitor_corrections", ()) + ), + excluded_frames=tuple( + sorted( + int(value) + for value in np.asarray(excluded).ravel() + if int(value) >= 0 + ) + ), + ) return cls( detector=ub_widget.detectorCal, unit_cell=cell, @@ -293,6 +401,7 @@ def from_gui(cls, gui): phi=getattr(ub_widget, "phi", 0.0), refraction_index=getattr(ub_widget, "n", 1.0), reference_reflections=reflections, + corrections=corrections, ) def apply_to_gui(self, gui): @@ -326,6 +435,44 @@ def apply_to_gui(self, gui): ) if hasattr(gui, "reflectionSel"): gui.reflectionSel.setReflections(self.reference_reflections) + if hasattr(gui, "scanSelector"): + gui.scanSelector.set_integration_options( + { + "mask": self.corrections.use_mask, + "solidAngle": self.corrections.use_solid_angle, + "polarization": self.corrections.use_polarization, + } + ) + gui.reconstruction_normalize_exposure = self.corrections.normalize_exposure + gui.reconstruction_monitor_corrections = self.corrections.monitor_corrections + if ( + hasattr(gui, "maskManager") + and self.corrections.repair_max_component_pixels is not None + ): + from .mask_config import PixelRepairSettings + + gui.maskManager.set_pixel_repair_settings( + PixelRepairSettings( + enabled=self.corrections.repair_masked_pixels, + max_component_pixels=( + self.corrections.repair_max_component_pixels + ), + max_span=self.corrections.repair_max_span, + radius=self.corrections.repair_radius, + min_valid_neighbors=( + self.corrections.repair_min_valid_neighbors + ), + use_pyfai_gaps=( + self.corrections.repair_use_pyfai_gaps + ), + gap_size_px=self.corrections.repair_gap_size_px, + ) + ) + if hasattr(gui, "excludedImagesDialog"): + excluded = np.asarray( + self.corrections.excluded_frames or (-1,), dtype=np.int64 + ) + gui.excludedImagesDialog.updateArrayData(excluded) if hasattr(ub_widget, "updateReflectionMismatch"): ub_widget.updateReflectionMismatch() @@ -364,6 +511,12 @@ def to_nxdict(self, role="scan", source=None): "@wavelength_unit": "Angstrom", }, "refraction_index": self.refraction_index, + "integration_corrections": { + "@NX_class": "NXcollection", + "json": json.dumps( + self.corrections.to_dict(), sort_keys=True + ), + }, **self.orgui, }, } @@ -383,6 +536,16 @@ def from_nxdict(cls, nxdict): ub_calculator = HKLVlieg.UBCalculator(unit_cell, energy) ub_calculator.setU(np.asarray(nxdict["sample"]["orientation_matrix"])) diffrac = nxdict.get("orgui", {}).get("diffractometer", {}) + correction_json = ( + nxdict.get("orgui", {}) + .get("integration_corrections", {}) + .get("json", "{}") + ) + correction_json = np.asarray(correction_json) + if correction_json.shape == (): + correction_json = correction_json.item() + if isinstance(correction_json, bytes): + correction_json = correction_json.decode() return cls( detector=detector, unit_cell=unit_cell, @@ -394,8 +557,18 @@ def from_nxdict(cls, nxdict): nxdict.get("orgui", {}).get("refraction_index", 1.0) ), reference_reflections=reflections_from_nxdict(nxdict), + corrections=CorrectionState.from_dict(json.loads(correction_json)), ) + def to_json_dict(self) -> dict[str, Any]: + """Serialize this configuration through the central NeXus schema.""" + return _json_value(self.to_nxdict(role="reconstruction")) + + @classmethod + def from_json_dict(cls, values): + """Deserialize a central configuration JSON dictionary.""" + return cls.from_nxdict(dict(values)) + class ConfigHandler: """Read and write orGUI config groups in an HDF5 database.""" diff --git a/orgui/app/cpp/roi_sum_cpp.cpp b/orgui/app/cpp/roi_sum_cpp.cpp index c3bf1fa..5a00c1e 100644 --- a/orgui/app/cpp/roi_sum_cpp.cpp +++ b/orgui/app/cpp/roi_sum_cpp.cpp @@ -5,6 +5,7 @@ #include #include #include +#include #include #include @@ -359,6 +360,369 @@ void repair_masked_pixels_inplace( } } +struct PlannedRepairCandidate { + py::ssize_t index; + double inverse_distance2; + unsigned char sides; +}; + +struct PlannedRepairComponent { + std::vector targets; + std::vector candidates; + bool median; +}; + +int side_count(const unsigned char sides) { + return static_cast(sides & 1U) + + static_cast((sides >> 1U) & 1U) + + static_cast((sides >> 2U) & 1U) + + static_cast((sides >> 3U) & 1U); +} + +class PixelRepairPlan { +public: + PixelRepairPlan( + const BoolArray2D mask, + const IntArray2D row_gaps, + const IntArray2D col_gaps, + const int max_component_pixels, + const int max_span, + const int radius, + const int min_valid_neighbors + ) + : max_component_pixels_(max_component_pixels), + max_span_(max_span), + radius_(radius), + min_valid_neighbors_(min_valid_neighbors) { + const auto mask_info = mask.request(); + const auto row_gap_info = row_gaps.request(); + const auto col_gap_info = col_gaps.request(); + require_ndim(mask_info, 2, "mask"); + require_intervals_shape(row_gap_info, "row_gaps"); + require_intervals_shape(col_gap_info, "col_gaps"); + if (max_component_pixels < 1 || max_span < 1 || radius < 1 + || min_valid_neighbors < 1) { + throw py::value_error("repair settings must be positive"); + } + height_ = mask_info.shape[0]; + width_ = mask_info.shape[1]; + size_ = mask_info.size; + const bool *mask_data = bool_ptr(mask_info); + const std::int32_t *row_gap_data = int32_ptr(row_gap_info); + const std::int32_t *col_gap_data = int32_ptr(col_gap_info); + mask_.resize(static_cast(size_)); + { + py::gil_scoped_release release; + for (py::ssize_t index = 0; index < size_; ++index) { + mask_[static_cast(index)] = + static_cast(mask_data[index]); + } + build( + row_gap_data, + row_gap_info, + col_gap_data, + col_gap_info + ); + } + } + + py::tuple apply_inplace(Array2D intensity, Array2D variance) const { + const auto intensity_info = intensity.request(); + const auto variance_info = variance.request(); + require_ndim(intensity_info, 2, "intensity"); + require_same_image_shape(intensity_info, variance_info, "variance"); + if ( + intensity_info.shape[0] != height_ + || intensity_info.shape[1] != width_ + ) { + throw py::value_error( + "intensity and variance must match the repair-plan shape" + ); + } + BoolArray2D remaining({height_, width_}); + BoolArray2D repaired({height_, width_}); + const auto remaining_info = remaining.request(); + const auto repaired_info = repaired.request(); + double *intensity_data = mutable_double_ptr(intensity_info); + double *variance_data = mutable_double_ptr(variance_info); + bool *remaining_data = static_cast(remaining_info.ptr); + bool *repaired_data = static_cast(repaired_info.ptr); + { + py::gil_scoped_release release; + for (py::ssize_t index = 0; index < size_; ++index) { + remaining_data[index] = + mask_[static_cast(index)] != 0; + repaired_data[index] = false; + } + std::vector valid; + std::vector> ordered; + for (const PlannedRepairComponent &component : components_) { + valid.clear(); + unsigned char sides = 0; + for (const PlannedRepairCandidate &candidate + : component.candidates) { + if (!std::isfinite(intensity_data[candidate.index])) { + continue; + } + valid.push_back(&candidate); + sides = static_cast( + sides | candidate.sides + ); + } + if ( + static_cast(valid.size()) < min_valid_neighbors_ + || side_count(sides) < 2 + ) { + continue; + } + double value = 0.0; + double value_variance = 0.0; + if (component.median) { + ordered.clear(); + ordered.reserve(valid.size()); + for (const PlannedRepairCandidate *candidate : valid) { + ordered.emplace_back( + intensity_data[candidate->index], + candidate->index + ); + } + std::stable_sort( + ordered.begin(), + ordered.end(), + [](const auto &left, const auto &right) { + return left.first < right.first; + } + ); + const std::size_t middle = ordered.size() / 2; + if (ordered.size() % 2 == 0) { + const auto &lower = ordered[middle - 1]; + const auto &upper = ordered[middle]; + value = 0.5 * (lower.first + upper.first); + value_variance = 0.25 * ( + variance_data[lower.second] + + variance_data[upper.second] + ); + } else { + const auto &selected = ordered[middle]; + value = selected.first; + value_variance = variance_data[selected.second]; + } + } else { + double total_weight = 0.0; + for (const PlannedRepairCandidate *candidate : valid) { + total_weight += candidate->inverse_distance2; + } + for (const PlannedRepairCandidate *candidate : valid) { + const double weight = + candidate->inverse_distance2 / total_weight; + value += intensity_data[candidate->index] * weight; + value_variance += ( + variance_data[candidate->index] * weight * weight + ); + } + } + for (const py::ssize_t target : component.targets) { + intensity_data[target] = value; + variance_data[target] = value_variance; + remaining_data[target] = false; + repaired_data[target] = true; + } + } + } + return py::make_tuple(remaining, repaired); + } + + py::dict configuration() const { + py::dict result; + result["shape"] = py::make_tuple(height_, width_); + result["components"] = components_.size(); + result["repairable_pixels"] = repairable_pixels_; + result["max_component_pixels"] = max_component_pixels_; + result["max_span"] = max_span_; + result["radius"] = radius_; + result["min_valid_neighbors"] = min_valid_neighbors_; + return result; + } + +private: + void build( + const std::int32_t *row_gaps, + const py::buffer_info &row_gap_info, + const std::int32_t *col_gaps, + const py::buffer_info &col_gap_info + ) { + std::vector visited( + static_cast(size_), + 0 + ); + for (py::ssize_t initial = 0; initial < size_; ++initial) { + if ( + visited[static_cast(initial)] + || mask_[static_cast(initial)] == 0 + ) { + continue; + } + std::vector component; + std::queue pending; + visited[static_cast(initial)] = 1; + pending.push(initial); + while (!pending.empty()) { + const py::ssize_t current = pending.front(); + pending.pop(); + component.push_back(current); + const py::ssize_t row = current / width_; + const py::ssize_t column = current % width_; + const py::ssize_t neighbors[4][2] = { + {row - 1, column}, + {row + 1, column}, + {row, column - 1}, + {row, column + 1}, + }; + for (const auto &neighbor : neighbors) { + const py::ssize_t next_row = neighbor[0]; + const py::ssize_t next_column = neighbor[1]; + if ( + next_row < 0 || next_column < 0 + || next_row >= height_ || next_column >= width_ + ) { + continue; + } + const py::ssize_t next = + next_row * width_ + next_column; + if ( + !visited[static_cast(next)] + && mask_[static_cast(next)] != 0 + ) { + visited[static_cast(next)] = 1; + pending.push(next); + } + } + } + if ( + component.empty() + || static_cast(component.size()) + > max_component_pixels_ + || touches_gap( + component, + width_, + row_gaps, + row_gap_info, + col_gaps, + col_gap_info + ) + ) { + continue; + } + py::ssize_t row_min = height_; + py::ssize_t row_max = -1; + py::ssize_t column_min = width_; + py::ssize_t column_max = -1; + for (const py::ssize_t target : component) { + const py::ssize_t row = target / width_; + const py::ssize_t column = target % width_; + row_min = std::min(row_min, row); + row_max = std::max(row_max, row); + column_min = std::min(column_min, column); + column_max = std::max(column_max, column); + } + if ( + row_max - row_min + 1 > max_span_ + || column_max - column_min + 1 > max_span_ + ) { + continue; + } + PlannedRepairComponent plan; + plan.targets = component; + plan.median = component.size() == 1; + unsigned char available_sides = 0; + for ( + py::ssize_t row = std::max( + 0, row_min - radius_ + ); + row <= std::min( + height_ - 1, row_max + radius_ + ); + ++row + ) { + for ( + py::ssize_t column = std::max( + 0, column_min - radius_ + ); + column <= std::min( + width_ - 1, column_max + radius_ + ); + ++column + ) { + const py::ssize_t index = row * width_ + column; + if (mask_[static_cast(index)] != 0) { + continue; + } + double distance2 = + std::numeric_limits::infinity(); + for (const py::ssize_t target : component) { + const double delta_row = static_cast( + row - target / width_ + ); + const double delta_column = static_cast( + column - target % width_ + ); + distance2 = std::min( + distance2, + delta_row * delta_row + + delta_column * delta_column + ); + } + if ( + distance2 <= 0.0 + || distance2 > radius_ * radius_ + ) { + continue; + } + unsigned char sides = 0; + if (column < column_min) { + sides = static_cast(sides | 1U); + } + if (column > column_max) { + sides = static_cast(sides | 2U); + } + if (row < row_min) { + sides = static_cast(sides | 4U); + } + if (row > row_max) { + sides = static_cast(sides | 8U); + } + available_sides = static_cast( + available_sides | sides + ); + plan.candidates.push_back( + {index, 1.0 / distance2, sides} + ); + } + } + if ( + static_cast(plan.candidates.size()) + < min_valid_neighbors_ + || side_count(available_sides) < 2 + ) { + continue; + } + repairable_pixels_ += plan.targets.size(); + components_.push_back(std::move(plan)); + } + } + + py::ssize_t height_ = 0; + py::ssize_t width_ = 0; + py::ssize_t size_ = 0; + int max_component_pixels_; + int max_span_; + int radius_; + int min_valid_neighbors_; + std::size_t repairable_pixels_ = 0; + std::vector mask_; + std::vector components_; +}; + RoiSums sum_roi_with_mask( const double *image_data, const double *correction_data, @@ -1685,6 +2049,35 @@ void calcMaxSum_bg( PYBIND11_MODULE(_roi_sum_cpp, module) { module.doc() = "CPU C++ acceleration kernels for ROI image summing."; + py::class_(module, "PixelRepairPlan") + .def( + py::init< + const BoolArray2D, + const IntArray2D, + const IntArray2D, + int, + int, + int, + int + >(), + py::arg("mask"), + py::arg("row_gaps"), + py::arg("column_gaps"), + py::arg("max_component_pixels") = 4, + py::arg("max_span") = 3, + py::arg("radius") = 2, + py::arg("min_valid_neighbors") = 6 + ) + .def( + "apply_inplace", + &PixelRepairPlan::apply_inplace, + py::arg("intensity"), + py::arg("variance") + ) + .def( + "configuration", + &PixelRepairPlan::configuration + ); module.def("processImage_Carr", &processImage_Carr); module.def("processImage_bg_Carr", &processImage_bg_Carr); module.def( diff --git a/orgui/app/database.py b/orgui/app/database.py index 2a20b96..d9be339 100644 --- a/orgui/app/database.py +++ b/orgui/app/database.py @@ -40,6 +40,7 @@ import h5py import datetime +import json import os import traceback import time @@ -53,6 +54,18 @@ logger = logging.getLogger(__name__) +def config_data_to_json(config): + """Convert :class:`ConfigData` to the central JSON representation.""" + if not isinstance(config, ConfigData): + raise TypeError("config must be a ConfigData instance") + return config.to_json_dict() + + +def config_data_from_json(values): + """Create :class:`ConfigData` from the central JSON representation.""" + return ConfigData.from_json_dict(values) + + class DBCloseError(IOError): pass @@ -719,6 +732,26 @@ def add_nxdict(self, nxentry, update_mode="add", h5path="/"): self.hdf5model.synchronizeH5pyObject(nxfile) self.view.expandToDepth(0) + def register_external_result( + self, path, checksum, grids, status, job_digest + ): + """Register a standalone reconstruction without copying its datasets.""" + self.add_nxdict( + { + "external_reconstructions": { + job_digest[:16]: { + "@NX_class": "NXnote", + "@orgui_meta": "external_reconstruction", + "path": os.path.abspath(path), + "sha256": checksum, + "status": status, + "grids_json": json.dumps(grids, sort_keys=True), + "job_sha256": job_digest, + } + } + } + ) + def write_scan_config(self, scan_name, config=None): """Write the default config group for a scan.""" if config is None: diff --git a/orgui/app/mask_config.py b/orgui/app/mask_config.py index 0bcb898..439bd68 100644 --- a/orgui/app/mask_config.py +++ b/orgui/app/mask_config.py @@ -151,6 +151,184 @@ def parse_mask_settings(config, base_dir=None): return MaskSettings(mask=mask_path, pixel_repair=repair) +def repair_intensity_variance( + intensity, + variance, + mask, + *, + max_component_pixels, + max_span, + radius, + min_valid_neighbors, + row_gaps=(), + column_gaps=(), +): + """Repair masked defects and propagate interpolation-weight variance. + + This follows the existing ROI repair policy: isolated pixels use the + median of surrounding valid pixels, while larger accepted components use + inverse-distance weighting. Returned masks remain ``True`` for defects + that cannot be repaired. + + :returns: + Repaired intensity, marginal variance, remaining mask, and repaired + pixel mask. + """ + intensity = np.asarray(intensity, dtype=np.float64).copy() + variance = np.asarray(variance, dtype=np.float64).copy() + mask = np.asarray(mask, dtype=bool) + if intensity.shape != variance.shape or intensity.shape != mask.shape: + raise ValueError("intensity, variance, and mask shapes must match") + height, width = mask.shape + visited = np.zeros(mask.shape, dtype=bool) + repaired = np.zeros(mask.shape, dtype=bool) + remaining = mask.copy() + row_gaps = tuple(tuple(map(int, interval)) for interval in row_gaps) + column_gaps = tuple(tuple(map(int, interval)) for interval in column_gaps) + + def touches_gap(component): + for row, column in component: + if any(start - 1 <= row <= stop for start, stop in row_gaps): + return True + if any(start - 1 <= column <= stop for start, stop in column_gaps): + return True + return False + + for initial_row, initial_column in zip(*np.nonzero(mask & ~visited)): + stack = [(int(initial_row), int(initial_column))] + visited[initial_row, initial_column] = True + component = [] + while stack: + row, column = stack.pop() + component.append((row, column)) + for next_row, next_column in ( + (row - 1, column), + (row + 1, column), + (row, column - 1), + (row, column + 1), + ): + if ( + 0 <= next_row < height + and 0 <= next_column < width + and mask[next_row, next_column] + and not visited[next_row, next_column] + ): + visited[next_row, next_column] = True + stack.append((next_row, next_column)) + rows = [item[0] for item in component] + columns = [item[1] for item in component] + if ( + len(component) > max_component_pixels + or max(rows) - min(rows) + 1 > max_span + or max(columns) - min(columns) + 1 > max_span + or touches_gap(component) + ): + continue + neighbors = [] + sides = set() + for row in range( + max(0, min(rows) - radius), + min(height, max(rows) + radius + 1), + ): + for column in range( + max(0, min(columns) - radius), + min(width, max(columns) + radius + 1), + ): + if mask[row, column] or not np.isfinite(intensity[row, column]): + continue + distance2 = min( + (row - item_row) ** 2 + (column - item_column) ** 2 + for item_row, item_column in component + ) + if distance2 <= 0 or distance2 > radius**2: + continue + neighbors.append((row, column, 1.0 / distance2)) + if column < min(columns): + sides.add("left") + if column > max(columns): + sides.add("right") + if row < min(rows): + sides.add("top") + if row > max(rows): + sides.add("bottom") + if len(neighbors) < min_valid_neighbors or len(sides) < 2: + continue + if len(component) == 1: + ordered = sorted( + neighbors, key=lambda item: intensity[item[0], item[1]] + ) + middle = len(ordered) // 2 + selected = ( + [(ordered[middle][0], ordered[middle][1], 1.0)] + if len(ordered) % 2 + else [ + (ordered[middle - 1][0], ordered[middle - 1][1], 0.5), + (ordered[middle][0], ordered[middle][1], 0.5), + ] + ) + else: + total_weight = sum(item[2] for item in neighbors) + selected = [ + (row, column, weight / total_weight) + for row, column, weight in neighbors + ] + value = sum( + intensity[row, column] * weight + for row, column, weight in selected + ) + value_variance = sum( + variance[row, column] * weight**2 + for row, column, weight in selected + ) + for row, column in component: + intensity[row, column] = value + variance[row, column] = value_variance + repaired[row, column] = True + remaining[row, column] = False + return intensity, variance, remaining, repaired + + +def create_pixel_repair_plan( + mask, + *, + max_component_pixels, + max_span, + radius, + min_valid_neighbors, + row_gaps=(), + column_gaps=(), +): + """Create a reusable native repair plan for a constant detector mask. + + The connected components, admissible neighbors, interpolation distances, + and detector-gap exclusions are computed once. Applying the returned plan + to an image and variance array modifies both arrays in place and releases + the Python GIL for the numerical work. + + :raises RuntimeError: + If the native ROI extension is unavailable. + """ + from ._roi_sum_accel import PixelRepairPlan + + if PixelRepairPlan is None: + raise RuntimeError( + "Native pixel repair is required for reciprocal-space mapping" + ) + row_gaps = np.ascontiguousarray(row_gaps, dtype=np.int32).reshape(-1, 2) + column_gaps = np.ascontiguousarray( + column_gaps, dtype=np.int32 + ).reshape(-1, 2) + return PixelRepairPlan( + np.ascontiguousarray(mask, dtype=bool), + row_gaps, + column_gaps, + int(max_component_pixels), + int(max_span), + int(radius), + int(min_valid_neighbors), + ) + + class MaskManager: """Manage detector mask state without depending on Qt widgets.""" diff --git a/orgui/app/orGUI.py b/orgui/app/orGUI.py index 6bb0f5b..a399f83 100644 --- a/orgui/app/orGUI.py +++ b/orgui/app/orGUI.py @@ -67,6 +67,8 @@ from .peak1Dintegr import RockingPeakIntegrator from .ArrayTableDialog import ArrayTableDialog from .MaskConfigDialog import MaskConfigDialog +from .HDF5SettingsDialog import HDF5SettingsDialog +from .ReconstructionDialog import ReconstructionDialog from .bgroi import RectangleBgROI from .database import DataBase, FILTERS from .mask_config import MaskManager @@ -164,6 +166,8 @@ def __init__(self, configfile, parent=None): ) self.maxMemory = MAX_MEMORY self.maxROIs = MAX_ROIS_DISPLAY + self.reconstruction_chunk_shape = (64, 64, 64) + self.reconstruction_compression_override = None self.filedialogdir = os.getcwd() @@ -443,7 +447,7 @@ def __init__(self, configfile, parent=None): self.backgroundImageAct.setCheckable(True) self.backgroundImageAct.setChecked(False) - self.dbCompressionAct = qt.QAction("Database compression", self) + self.dbCompressionAct = qt.QAction("HDF5 settings", self) self.dbCompressionAct.triggered.connect(self._onChangeDBCompression) self.maskConfigAct = qt.QAction("Mask", self) self.maskConfigAct.triggered.connect(self._onShowMaskConfig) @@ -549,6 +553,8 @@ def __init__(self, configfile, parent=None): rs = menu_bar.addMenu("&Reciprocal space") rs.addAction(calcCTRsAvailableAct) rs.addAction(editUAct) + reconstructionAct = rs.addAction("Reconstruct reciprocal space") + reconstructionAct.triggered.connect(self._onShowReconstruction) simul = menu_bar.addMenu("&Simulation") @@ -588,6 +594,29 @@ def _loadMaskConfig(self, configfile): np.ascontiguousarray(mask, dtype=np.uint8) ) + # GUI-only: user-triggered reciprocal-space reconstruction dialog. + def _onShowReconstruction(self): + """Open the reciprocal-space reconstruction workflow.""" + try: + if not hasattr(self, "_reconstruction_dialog"): + self._reconstruction_dialog = ReconstructionDialog(self) + self._reconstruction_dialog.show() + self._reconstruction_dialog.raise_() + except Exception: + logger.warning( + "Cannot open reciprocal-space reconstruction.", + exc_info=True, + extra={ + "title": "Cannot open reconstruction", + "description": ( + "The reconstruction dialog could not be initialized." + ), + "show_dialog": True, + "dialog_level": logging.WARNING, + "parent": self, + }, + ) + # GUI-only: user-triggered mask configuration dialog. def _onShowMaskConfig(self): """Open the detector mask configuration dialog.""" @@ -2613,23 +2642,21 @@ def _onSelectROIcount(self): self.maxROIs = rois def _onChangeDBCompression(self): - """GUI-only: ask the user for the database compression filter.""" - filter_names = list(FILTERS.keys()) - currentCompression = self.database.compression - for fn in filter_names: - if currentCompression == FILTERS[fn]: - break - idx = filter_names.index(fn) - selection, success = qt.QInputDialog.getItem( + """GUI-only: edit shared HDF5 reconstruction output settings.""" + dialog = HDF5SettingsDialog( + self.database.compression, + self.reconstruction_chunk_shape, + self.reconstruction_compression_override, self, - "Data compression settings", - "Available data base compression methods:\n(See discussion under https://github.com/tifuchs/orGUI/issues/16)\nRecommended: Blosc-lz4-Shuffle-5", # noqa: E501 - filter_names, - idx, - False, ) - if success: - self.database.compression = FILTERS[selection] + if dialog.exec() == qt.QDialog.Accepted: + self.database.compression = FILTERS[ + dialog.database_compression_name + ] + self.reconstruction_chunk_shape = dialog.chunk_shape + self.reconstruction_compression_override = ( + dialog.compression_override + ) def calcBraggRefl(self): """Calculate and display available Bragg reflections for the scan. diff --git a/orgui/app/test/test_config_data.py b/orgui/app/test/test_config_data.py index 547fdf0..212e8a2 100644 --- a/orgui/app/test/test_config_data.py +++ b/orgui/app/test/test_config_data.py @@ -2,10 +2,13 @@ import numpy as np from silx.io.dictdump import dicttonx, nxtodict import pytest +from types import SimpleNamespace from orgui.app.QReflectionSelector import HKLReflection -from orgui.app.config_data import ConfigData, ConfigHandler +from orgui.app.config_data import CorrectionState, ConfigData, ConfigHandler +from orgui.app.database import config_data_from_json, config_data_to_json from orgui.datautils.xrayutils import CTRcalc, DetectorCalibration, HKLVlieg +from orgui.reconstruction_job import _snapshot_assets def _make_config(): @@ -50,6 +53,15 @@ def _make_config(): phi=0.3, refraction_index=0.999, reference_reflections=reflections, + corrections=CorrectionState( + use_mask=True, + use_background=True, + use_solid_angle=True, + normalize_exposure=False, + monitor_corrections=("mondio",), + excluded_frames=(2, 7), + uncertainty_provenance={"background": "measured"}, + ), ) @@ -111,6 +123,114 @@ def test_config_data_round_trips_through_nexus_dict(tmp_path): ] assert np.allclose(loaded.reference_reflections[0].hkl, [1.0, 0.0, 0.0]) assert np.allclose(loaded.reference_reflections[1].xy, [21.0, 22.0]) + assert loaded.corrections == config.corrections + + +def test_config_data_round_trips_through_database_json(): + config = _make_config() + loaded = config_data_from_json(config_data_to_json(config)) + + assert loaded.corrections == config.corrections + assert np.allclose(loaded.ub_calculator.getUB(), config.ub_calculator.getUB()) + + +def test_enabled_pixel_repair_implies_mask_correction(): + config = _make_config() + repair = SimpleNamespace( + enabled=True, + max_component_pixels=4, + max_span=3, + radius=2, + min_valid_neighbors=6, + use_pyfai_gaps=True, + gap_size_px=1, + ) + gui = SimpleNamespace( + ubcalc=SimpleNamespace( + detectorCal=config.detector, + crystal=config.unit_cell, + ubCal=config.ub_calculator, + mu=config.mu, + chi=config.chi, + phi=config.phi, + n=config.refraction_index, + ), + scanSelector=SimpleNamespace( + get_integration_options=lambda: { + "mask": False, + "solidAngle": False, + "polarization": False, + } + ), + maskManager=SimpleNamespace( + settings=SimpleNamespace(pixel_repair=repair) + ), + excludedImagesDialog=SimpleNamespace( + getData=lambda: np.empty(0, dtype=np.int64) + ), + reconstruction_normalize_exposure=False, + reconstruction_monitor_corrections=("mondio",), + ) + + captured = ConfigData.from_gui(gui) + + assert captured.corrections.repair_masked_pixels is True + assert captured.corrections.use_mask is True + assert captured.corrections.normalize_exposure is False + assert captured.corrections.monitor_corrections == ("mondio",) + + +def test_apply_to_gui_sets_reconstruction_normalization_attributes(): + config = _make_config() + config.corrections = CorrectionState( + normalize_exposure=False, monitor_corrections=("mondio",) + ) + gui = SimpleNamespace( + ubcalc=SimpleNamespace( + detectorCal=config.detector, + crystal=config.unit_cell, + ubCal=config.ub_calculator, + mu=config.mu, + chi=config.chi, + phi=config.phi, + n=config.refraction_index, + ), + ) + + config.apply_to_gui(gui) + + assert gui.reconstruction_normalize_exposure is False + assert gui.reconstruction_monitor_corrections == ("mondio",) + + +def test_snapshot_assets_serializes_active_mask(tmp_path): + config = _make_config() + config.corrections = CorrectionState(use_mask=True) + mask = np.zeros(config.detector.detector.shape, dtype=bool) + mask[3, 4] = True + gui = SimpleNamespace( + ubcalc=SimpleNamespace(detectorCal=config.detector), + get_detector_mask=lambda shape: mask, + ) + assets = tmp_path / "job-assets.nxs" + + _snapshot_assets(gui, config, assets) + + assert config.corrections.mask_asset == "/mask" + with h5py.File(assets, "r") as h5file: + np.testing.assert_array_equal(h5file["mask"][()], mask) + + +def test_snapshot_assets_rejects_missing_enabled_mask(tmp_path): + config = _make_config() + config.corrections = CorrectionState(use_mask=True) + gui = SimpleNamespace( + ubcalc=SimpleNamespace(detectorCal=config.detector), + get_detector_mask=lambda shape: None, + ) + + with pytest.raises(ValueError, match="no active mask matches"): + _snapshot_assets(gui, config, tmp_path / "job-assets.nxs") def test_config_handler_writes_scan_and_integration_config(tmp_path): diff --git a/orgui/app/test/test_mask_config.py b/orgui/app/test/test_mask_config.py index ce237c2..fd7eeed 100644 --- a/orgui/app/test/test_mask_config.py +++ b/orgui/app/test/test_mask_config.py @@ -3,7 +3,13 @@ import numpy as np import pytest -from orgui.app.mask_config import MaskManager, parse_mask_settings +from orgui.app.mask_config import ( + MaskManager, + create_pixel_repair_plan, + parse_mask_settings, + repair_intensity_variance, +) +from orgui.app._roi_sum_accel import PixelRepairPlan def test_mask_without_pixel_repair_disables_repair(tmp_path): @@ -143,3 +149,63 @@ def test_mask_manager_shape_mismatch_returns_none(caplog): assert manager.get_mask((3, 2)) is None assert "does not match image shape" in caplog.text + + +@pytest.mark.skipif( + PixelRepairPlan is None, + reason="The active native ROI extension has no PixelRepairPlan", +) +def test_native_repair_plan_matches_single_pixel_reference(): + intensity = np.arange(49, dtype=np.float64).reshape(7, 7) + variance = np.linspace(1.0, 2.0, 49).reshape(7, 7) + mask = np.zeros((7, 7), dtype=bool) + mask[3, 3] = True + settings = { + "max_component_pixels": 4, + "max_span": 3, + "radius": 2, + "min_valid_neighbors": 6, + } + expected = repair_intensity_variance( + intensity, variance, mask, **settings + ) + plan = create_pixel_repair_plan(mask, **settings) + actual_intensity = intensity.copy() + actual_variance = variance.copy() + + remaining, repaired = plan.apply_inplace( + actual_intensity, actual_variance + ) + + np.testing.assert_allclose(actual_intensity, expected[0]) + np.testing.assert_allclose(actual_variance, expected[1]) + np.testing.assert_array_equal(remaining, expected[2]) + np.testing.assert_array_equal(repaired, expected[3]) + assert plan.configuration()["repairable_pixels"] == 1 + + +@pytest.mark.skipif( + PixelRepairPlan is None, + reason="The active native ROI extension has no PixelRepairPlan", +) +def test_native_repair_plan_reuses_component_geometry(): + mask = np.zeros((9, 9), dtype=bool) + mask[4, 4:6] = True + plan = create_pixel_repair_plan( + mask, + max_component_pixels=4, + max_span=3, + radius=2, + min_valid_neighbors=6, + ) + + for scale in (1.0, 2.0): + intensity = ( + np.arange(81, dtype=np.float64).reshape(9, 9) * scale + ) + variance = np.ones((9, 9), dtype=np.float64) + remaining, repaired = plan.apply_inplace(intensity, variance) + + assert repaired[4, 4:6].all() + assert not remaining[4, 4:6].any() + assert intensity[4, 4] == intensity[4, 5] diff --git a/orgui/app/test/test_reconstruction_dialog.py b/orgui/app/test/test_reconstruction_dialog.py new file mode 100644 index 0000000..04faa3f --- /dev/null +++ b/orgui/app/test/test_reconstruction_dialog.py @@ -0,0 +1,464 @@ +"""Regression tests for reconstruction-dialog error containment.""" + +import logging +from types import SimpleNamespace +from unittest.mock import Mock + +import numpy as np +from silx.gui import qt + +import orgui.app.ReconstructionDialog as reconstruction_dialog_module +from orgui.app.ReconstructionDialog import ReconstructionDialog +from orgui.app.config_data import CorrectionState +from orgui.app.database import FILTERS +from orgui.app.HDF5SettingsDialog import HDF5SettingsDialog + + +def _dialog(tmp_path): + app = qt.QApplication.instance() or qt.QApplication([]) + parent = qt.QWidget() + orgui = SimpleNamespace( + fscan=None, + filedialogdir=str(tmp_path), + numberthreads=1, + maxMemory=128, + database=SimpleNamespace(compression=FILTERS["Raw"]), + _onShowMaskConfig=lambda: None, + backgroundImageAct=SimpleNamespace(trigger=lambda: None), + excludedImagesDialog=SimpleNamespace(show=lambda: None), + ) + dialog = ReconstructionDialog(orgui, parent=parent) + dialog._test_app = app + dialog._test_parent = parent + return dialog + + +def test_no_scan_actions_are_reported_without_raising(tmp_path, caplog): + """User actions without an active scan must remain non-fatal.""" + dialog = _dialog(tmp_path) + caplog.set_level(logging.WARNING) + + dialog.add_derived_grid("hkl") + dialog.preview() + dialog.prepare() + dialog.run_local() + dialog.create_cluster_scripts() + dialog.resume() + dialog.refresh_live_state() + + assert "No active scan" in dialog.experiment_summary.toPlainText() + records = [ + record for record in caplog.records if hasattr(record, "show_dialog") + ] + assert len(records) >= 6 + assert all(record.show_dialog is True for record in records) + assert all(record.dialog_level == logging.WARNING for record in records) + + dialog.close() + dialog._test_parent.close() + + +def test_auxiliary_callback_failures_are_contained( + tmp_path, monkeypatch, caplog +): + """File, shared-setting, and grid callback failures must not escape.""" + dialog = _dialog(tmp_path) + caplog.set_level(logging.WARNING) + + monkeypatch.setattr( + qt.QFileDialog, + "getExistingDirectory", + Mock(side_effect=RuntimeError("file dialog failed")), + ) + assert ( + dialog._browse( + dialog.scratch_path, + kind="directory", + save=False, + ) + is False + ) + + dialog._invoke_ui_action( + "Cannot open test settings", + Mock(side_effect=RuntimeError("settings failed")), + ) + + dialog.grid_table.insertRow(0) + dialog.orgui.fscan = object() + dialog.preview() + + records = [ + record for record in caplog.records if hasattr(record, "show_dialog") + ] + assert len(records) == 3 + assert "empty cell" in dialog.preview_output.toPlainText() + + dialog.close() + dialog._test_parent.close() + + +def test_default_paths_use_working_directory(tmp_path, monkeypatch): + """Generated writable paths must not default beside input data.""" + dialog = _dialog(tmp_path) + history_file = tmp_path / "La3Ni2O7" + history_file.touch() + dialog.orgui.filedialogdir = str(history_file) + working_directory = tmp_path / "working" + working_directory.mkdir() + monkeypatch.chdir(working_directory) + + dialog._set_default_paths("39_1") + + assert dialog.job_path.text() == str( + working_directory / "39_1-rsmap.json" + ) + assert dialog.scratch_path.text() == str( + working_directory / ".39_1-rsmap-scratch" + ) + assert dialog.output_path.text() == str( + working_directory / "39_1-rsmap.h5" + ) + + dialog.close() + dialog._test_parent.close() + + +def test_detected_performance_values_are_visible_without_becoming_overrides( + tmp_path, monkeypatch +): + """Detected execution values populate disabled advanced editors.""" + dialog = _dialog(tmp_path) + detected = { + "thread_budget": 16, + "native_threads_per_image": 4, + "memory_budget_MiB": 2048.0, + "accumulation_budget_MiB_per_worker": 256.0, + "frames_per_task": 1, + "detector_tile_shape": (1024, 768), + "native_work_block_pixels": 65536, + "parquet_chunk_span": 256, + "frame_tasks": 20, + "detector_tiles": 12, + "map_tasks": 240, + "parallel_layouts": [ + { + "exposure": "stationary", + "concurrent_image_workers": 8, + "native_threads_per_image": 1, + "tiles_per_image": 12, + "memory_per_image_MiB": 256.0, + "accumulation_MiB_per_worker": 256.0, + } + ], + } + monkeypatch.setattr( + reconstruction_dialog_module, + "reconstruction_execution_settings", + lambda job, scan=None, config=None: detected, + ) + + dialog._show_execution_settings(object()) + + assert dialog.thread_override[2].value() == 16 + assert dialog.threads_per_image.value() == 4 + assert dialog.memory_override[2].value() == 2048 + assert dialog.accumulation_memory[2].value() == 256 + assert dialog.tile_rows[2].value() == 1024 + assert dialog.tile_columns[2].value() == 768 + assert not dialog.thread_override[1].isChecked() + assert not dialog.tile_rows[1].isChecked() + assert '"concurrent_image_workers": 8' in ( + dialog.performance_summary.toPlainText() + ) + + dialog.close() + dialog._test_parent.close() + + +def test_settings_are_grouped_and_have_tooltips(tmp_path): + """Reconstruction settings use concise grouped labels and help text.""" + dialog = _dialog(tmp_path) + + group_titles = { + group.title() for group in dialog.findChildren(qt.QGroupBox) + } + assert { + "Active experiment", + "Corrections and exclusions", + "Exposure and job metadata", + "Grid definitions", + "Accuracy", + "Parallel execution and memory", + "Advanced settings", + "Detected execution layout", + "Scheduler", + "Python environment", + "Mapping array", + "Reduction and finalization", + "Scheduler-specific settings", + "Job descriptor", + "Storage", + } <= group_titles + + labels = [ + label.text() for label in dialog.findChildren(qt.QLabel) + ] + assert not any(label.startswith("Advanced ") for label in labels) + assert not any(label.startswith("One-job ") for label in labels) + assert [ + dialog.accuracy.itemText(index) + for index in range(dialog.accuracy.count()) + ] == [ + "Center only (depth 0)", + "Low (depth 1)", + "Balanced (depth 2)", + "High (depth 3)", + "Very high (depth 4)", + "Maximum (depth 5)", + ] + assert [ + reconstruction_dialog_module.ACCURACY_DEPTHS[ + dialog.accuracy.itemData(index) + ] + for index in range(dialog.accuracy.count()) + ] == list(range(6)) + assert dialog.accuracy.currentData() == "balanced" + + setting_controls = [ + dialog.angle_fallback, + dialog.user_note, + dialog.grid_table, + dialog.accuracy, + dialog.thread_override[0], + dialog.threads_per_image, + dialog.memory_override[0], + dialog.accumulation_memory[0], + dialog.frame_batch[0], + dialog.tile_rows[0], + dialog.tile_columns[0], + dialog.work_block[0], + dialog.partition_span[0], + dialog.performance_summary, + dialog.cluster_scheduler, + dialog.cluster_job_name, + dialog.cluster_script_directory, + dialog.cluster_working_directory, + dialog.cluster_python, + dialog.cluster_environment, + dialog.cluster_array_cpus, + dialog.cluster_array_memory, + dialog.cluster_array_walltime, + dialog.cluster_array_concurrency, + dialog.cluster_summary, + dialog.cluster_reduce_cpus, + dialog.cluster_reduce_memory, + dialog.cluster_reduce_walltime, + dialog.cluster_sge_pe, + dialog.cluster_sge_memory, + dialog.cluster_array_directives, + dialog.cluster_reduce_directives, + dialog.job_path, + dialog.scratch_path, + dialog.output_path, + ] + assert all(control.toolTip() for control in setting_controls) + assert all( + dialog.grid_table.horizontalHeaderItem(column).toolTip() + for column in range(dialog.grid_table.columnCount()) + ) + + dialog.close() + dialog._test_parent.close() + + +def test_cluster_tab_uses_sge_defaults_and_separate_reduce_resources(tmp_path): + """Cluster settings default to SGE and preserve independent resources.""" + dialog = _dialog(tmp_path) + + dialog.cluster_array_cpus.setValue(4) + dialog.cluster_array_memory.setValue(16) + dialog.cluster_reduce_cpus.setValue(32) + dialog.cluster_reduce_memory.setValue(128) + settings = dialog._cluster_settings() + + assert settings.scheduler == "sge" + assert settings.array_cpus == 4 + assert settings.array_memory_gib == 16 + assert settings.reduce_cpus == 32 + assert settings.reduce_memory_gib == 128 + assert any( + dialog.tabs.tabText(index) == "Cluster" + for index in range(dialog.tabs.count()) + ) + + dialog.close() + dialog._test_parent.close() + + +def test_json_output_has_its_own_tab_and_is_selected_when_updated(tmp_path): + """JSON and status output must not compress the settings tabs.""" + dialog = _dialog(tmp_path) + + assert dialog.tabs.indexOf(dialog.output_tab) >= 0 + assert dialog.tabs.tabText( + dialog.tabs.indexOf(dialog.output_tab) + ) == "Preview and status" + assert dialog.preview_output.parentWidget() is dialog.output_tab + + dialog._show_output('{"status": "prepared"}') + + assert dialog.tabs.currentWidget() is dialog.output_tab + assert dialog.preview_output.toPlainText() == '{"status": "prepared"}' + + dialog.close() + dialog._test_parent.close() + + +def test_grid_numbers_are_compact_without_losing_values_and_chunks_are_global( + tmp_path, +): + """The table display is compact while its editable values remain exact.""" + dialog = _dialog(tmp_path) + dialog.orgui.reconstruction_chunk_shape = (32, 64, 128) + grid = reconstruction_dialog_module.ReconstructionGrid( + minimum=(0.1234567890123, -1.234567890123, 2.345678901234), + maximum=(1.234567890123, 2.345678901234, 3.456789012345), + step=(0.001234567890123, 0.002345678901234, 0.003456789012345), + frame="hkl", + ) + + dialog._append_grid(grid) + + assert dialog.grid_table.columnCount() == 15 + assert dialog.grid_table.item(0, 2).text() == "0.123457" + assert ( + dialog.grid_table.item(0, 2).data(qt.Qt.EditRole) + == grid.minimum[0] + ) + rebuilt = dialog._grids()[0] + assert rebuilt.minimum == grid.minimum + assert rebuilt.maximum == grid.maximum + assert rebuilt.step == grid.step + assert rebuilt.chunk_shape == (32, 64, 128) + assert all( + int(dialog.grid_table.item(0, column).data(qt.Qt.EditRole)) > 0 + for column in range(11, 14) + ) + assert "uncompressed" in dialog.grid_table.item(0, 14).text() + + dialog.grid_table.item(0, 11).setData(qt.Qt.EditRole, 100) + + rebuilt = dialog._grids()[0] + assert rebuilt.to_spec().shape[0] == 100 + assert rebuilt.step[0] == np.nextafter( + (grid.maximum[0] - grid.minimum[0]) / 100, np.inf + ) + + dialog.close() + dialog._test_parent.close() + + +def test_hdf5_settings_expose_global_chunks_and_compression_override(tmp_path): + """Shared HDF5 settings cover chunking and optional output compression.""" + app = qt.QApplication.instance() or qt.QApplication([]) + dialog = HDF5SettingsDialog( + FILTERS["Raw"], + chunk_shape=(32, 64, 128), + compression_override="GZip", + ) + dialog._test_app = app + + assert dialog.chunk_shape == (32, 64, 128) + assert dialog.database_compression_name == "Raw" + assert dialog.compression_override == "GZip" + dialog.override_compression.setChecked(False) + assert dialog.compression_override is None + + dialog.close() + + +def test_geometry_resolution_dialog_exposes_percentile(tmp_path): + """The local-Jacobian estimator exposes its robust percentile.""" + app = qt.QApplication.instance() or qt.QApplication([]) + dialog = reconstruction_dialog_module._GeometryResolutionDialog() + dialog._test_app = app + + assert dialog.percentile == 10.0 + dialog.percentile_editor.setValue(25.0) + assert dialog.percentile == 25.0 + + dialog.close() + + +def test_geometry_step_estimate_updates_steps_intervals_and_size( + tmp_path, monkeypatch +): + """Applying a geometry estimate keeps all derived grid fields consistent.""" + dialog = _dialog(tmp_path) + dialog.orgui.fscan = object() + dialog._append_grid( + reconstruction_dialog_module.ReconstructionGrid( + minimum=(0.0, 0.0, 0.0), + maximum=(1.0, 2.0, 3.0), + step=(0.5, 0.5, 0.5), + frame="hkl", + ) + ) + captured = {} + + monkeypatch.setattr( + reconstruction_dialog_module.ConfigData, + "from_gui", + lambda gui: object(), + ) + + def estimate(config, scan, *, frame, percentile): + captured.update(frame=frame, percentile=percentile) + return (0.1, 0.2, 0.3) + + monkeypatch.setattr( + reconstruction_dialog_module, "estimate_geometry_steps", estimate + ) + + dialog._apply_geometry_steps([0], 15.0) + + assert captured == {"frame": "hkl", "percentile": 15.0} + grid = dialog._grids()[0] + np.testing.assert_allclose(grid.step, (0.1, 0.2, 0.3)) + assert grid.to_spec().shape == (10, 10, 10) + assert "uncompressed" in dialog.grid_table.item(0, 14).text() + + dialog.close() + dialog._test_parent.close() + + +def test_live_summary_resolves_active_mask_instead_of_showing_null_asset( + tmp_path, +): + """The live overview distinguishes active inputs from frozen assets.""" + dialog = _dialog(tmp_path) + mask = np.zeros((5, 7), dtype=bool) + mask[1, 2] = True + mask[3, 4] = True + dialog.orgui.get_detector_mask = lambda shape: mask + config = SimpleNamespace( + corrections=CorrectionState(use_mask=True) + ) + + summary = dialog._live_correction_summary(config, mask.shape) + + assert "mask_asset" not in summary + assert "background_asset" not in summary + assert "background_variance_asset" not in summary + assert summary["active_inputs"]["mask"] == { + "status": "active", + "masked_pixels": 2, + "job_asset_path": "/mask", + } + assert summary["active_inputs"]["background"] == { + "status": "disabled" + } + + dialog.close() + dialog._test_parent.close() diff --git a/orgui/backend/beamline/id31_tools.py b/orgui/backend/beamline/id31_tools.py index 4cb4d62..4b1b8cd 100644 --- a/orgui/backend/beamline/id31_tools.py +++ b/orgui/backend/beamline/id31_tools.py @@ -135,6 +135,9 @@ def _scanTitle(h5file, name): # currently only set up for th scans in TOMO session, cbf fileformat class Fastscan(Scan): def __init__(self, fastscan_specfile, scanno): + self.fastscan_specfile = os.path.abspath(fastscan_specfile) + self.scanno = int(scanno) + self._scan_reference_args = [self.fastscan_specfile, self.scanno] id31_fastscan_spec = specfile.Specfile(fastscan_specfile) scan = id31_fastscan_spec[scanno - 1] filename_line = scan.header("C next_image_file")[0] @@ -170,10 +173,13 @@ def __init__(self, fastscan_specfile, scanno): ) ) scan = id31_fastscan_spec[scanno] - self.th = np.mean( - [scan.datacol("th_UpPos"), scan.datacol("th_DownPos")], axis=0 - ) + th_up = scan.datacol("th_UpPos") + th_down = scan.datacol("th_DownPos") + self.th = np.mean([th_up, th_down], axis=0) self.omega = -1 * self.th + self.omega_bounds_rad = np.deg2rad( + np.column_stack((-th_up, -th_down)) + ) self.exposure_time = scan.datacol("TrigTime") commands = scan.command().split() self.nopoints = int(commands[-4]) @@ -184,6 +190,7 @@ def __init__(self, fastscan_specfile, scanno): def set_th_offset(self, offset): self.th += offset self.omega = -1 * self.th + self.omega_bounds_rad -= np.deg2rad(offset) def set_image_folder(self, path_to_folder): # self.filenames = [None]*len(self.th) @@ -398,6 +405,14 @@ def __init__( data_1["measurement"]["th_trig"][: self.nopoints] + data_1["measurement"]["th_delta"][: self.nopoints] / 2 ) + th_start = data_1["measurement"]["th_trig"][: self.nopoints] + th_stop = ( + th_start + + data_1["measurement"]["th_delta"][: self.nopoints] + ) + self.omega_bounds_rad = np.deg2rad( + np.column_stack((-th_start, -th_stop)) + ) self.axisname = "th" self.mu = self.positioners["mu"] @@ -424,6 +439,21 @@ def __init__( data_1["measurement"]["linai_trig"][: self.nopoints] + data_1["measurement"]["linai_delta"][: self.nopoints] / 2 ) + linai_start = data_1["measurement"]["linai_trig"][ + : self.nopoints + ] + linai_stop = ( + linai_start + + data_1["measurement"]["linai_delta"][: self.nopoints] + ) + self.alpha_bounds_rad = np.deg2rad( + np.column_stack( + ( + lintomu.linai_to_mu(linai_start), + lintomu.linai_to_mu(linai_stop), + ) + ) + ) self.mu = lintomu.linai_to_mu(self.linai) self.axisname = "mu" self.th = self.positioners["th"] @@ -793,6 +823,14 @@ def __init__( data_1["measurement"]["th_trig"][: self.nopoints] + data_1["measurement"]["th_delta"][: self.nopoints] / 2 ) + th_start = data_1["measurement"]["th_trig"][: self.nopoints] + th_stop = ( + th_start + + data_1["measurement"]["th_delta"][: self.nopoints] + ) + self.omega_bounds_rad = np.deg2rad( + np.column_stack((-th_start, -th_stop)) + ) self.axisname = "th" self.mu = self.positioners["mu"] @@ -829,6 +867,21 @@ def __init__( data_1["measurement"]["linai_trig"][: self.nopoints] + data_1["measurement"]["linai_delta"][: self.nopoints] / 2 ) + linai_start = data_1["measurement"]["linai_trig"][ + : self.nopoints + ] + linai_stop = ( + linai_start + + data_1["measurement"]["linai_delta"][: self.nopoints] + ) + self.alpha_bounds_rad = np.deg2rad( + np.column_stack( + ( + lintomu.linai_to_mu(linai_start), + lintomu.linai_to_mu(linai_stop), + ) + ) + ) self.mu = lintomu.linai_to_mu(self.linai) self.axisname = "mu" self.th = self.positioners["th"] @@ -1054,6 +1107,8 @@ def get_p3_img(self, img): class BlissScan(Fastscan): def __init__(self, hdffilepath, scanname): hdffilepath = os.path.abspath(hdffilepath) + self.hdffilepath_orNode = hdffilepath + self._scan_reference_args = [hdffilepath, scanname] filepath, filename = os.path.split(hdffilepath) self.filepath = filepath # print(filepath,filename) diff --git a/orgui/backend/scans.py b/orgui/backend/scans.py index c746890..0c39d07 100644 --- a/orgui/backend/scans.py +++ b/orgui/backend/scans.py @@ -31,6 +31,11 @@ import numpy as np from abc import ABC, abstractmethod +from dataclasses import dataclass +import importlib +import os +from pathlib import Path +import runpy class Scan(ABC): @@ -146,17 +151,75 @@ def get_raw_img(self, i): """ raise NotImplementedError() + def exposure_angle_bounds(self, config, fallback="stationary"): + """Return exposure bounds in Vlieg order and radians. + + :param ConfigData config: + Active orGUI configuration providing fixed ``mu``, ``chi``, and + ``phi`` values. + :param str fallback: + ``"stationary"`` or the explicitly requested ``"midpoint"``. + :returns: + Array shaped ``(frames, 2, 4)`` ordered + ``alpha, omega, chi, phi``. + """ + return scan_exposure_angle_bounds(self, config, fallback=fallback) + class h5_Image: - def __init__(self, data): + def __init__(self, data, variance=None, processing_provenance=None): """Only the image data is required as numpy array. motors and counters do not need to be populated. """ self.img = data + self.variance = variance + self.processing_provenance = dict(processing_provenance or {}) self.motors = dict() self.counters = dict() +def load_scan_backend_file(filename, class_name=None): + """Load one scan backend class from an existing Python backend file. + + :param filename: + Backend Python file understood by orGUI's backend loader. + :param str class_name: + Optional qualified class name required by a serialized scan reference. + :returns: + ``(namespace_name, scan_class)``. + :raises ValueError: + If no unique matching :class:`Scan` implementation is present. + """ + filename = str(Path(filename).absolute()) + namespace = runpy.run_path(filename) + candidates = [] + seen = set() + for name, value in namespace.items(): + try: + is_scan = issubclass(value, Scan) and value is not Scan + except TypeError: + continue + if not is_scan or id(value) in seen: + continue + if class_name is not None and value.__qualname__ != class_name: + continue + seen.add(id(value)) + candidates.append((name, value)) + qualifier = f" named {class_name!r}" if class_name is not None else "" + if not candidates: + raise ValueError( + f"Found no Scan class{qualifier} in backend file {filename}" + ) + if len(candidates) > 1: + raise ValueError( + f"Found more than one Scan class{qualifier} in backend file " + f"{filename}. Only one is permitted" + ) + name, scan_class = candidates[0] + scan_class._orgui_backend_file = filename + return name, scan_class + + class SimulationScan(Scan): def __init__(self, detshape, axismin, axismax, points, axis="th", fixed=0.0): self.shape = detshape @@ -188,3 +251,325 @@ def set_raw_img(self, i, data): # for intensity simulation in the future. def parse_h5_node(cls, node): # unused pass + + +def _frame_values(values, count, name): + values = np.asarray(values, dtype=np.float64) + if values.ndim == 0 or values.size == 1: + return np.full(count, float(values.reshape(-1)[0])) + if values.ndim == 1 and values.size == count: + return values + raise ValueError(f"{name} must be scalar or contain one value per frame") + + +def _midpoint_bounds(values): + values = np.unwrap(np.asarray(values, dtype=np.float64)) + if values.size == 1: + return np.repeat(values[:, None], 2, axis=1) + edges = np.empty(values.size + 1, dtype=np.float64) + edges[1:-1] = 0.5 * (values[:-1] + values[1:]) + edges[0] = values[0] - 0.5 * (values[1] - values[0]) + edges[-1] = values[-1] + 0.5 * (values[-1] - values[-2]) + return np.column_stack((edges[:-1], edges[1:])) + + +def scan_exposure_angle_bounds(scan, config, fallback="stationary"): + """Return centralized exposure bounds for any loaded orGUI scan. + + Scan motor positions are degrees. Fixed diffractometer values from + :class:`ConfigData` and the returned bounds are radians. + """ + if fallback not in {"stationary", "midpoint"}: + raise ValueError("fallback must be 'stationary' or 'midpoint'") + if hasattr(scan, "subscans"): + bounds = np.concatenate( + [ + scan_exposure_angle_bounds(child, config, fallback=fallback) + for child in scan.subscans + ] + ) + indices = getattr(scan, "indices", None) + if indices is not None: + bounds = bounds[np.asarray(indices, dtype=np.int64)] + return np.ascontiguousarray(bounds) + count = len(scan) + if hasattr(scan, "mu"): + alpha = np.deg2rad(_frame_values(scan.mu, count, "alpha")) + else: + alpha = _frame_values(config.mu, count, "alpha") + omega = np.deg2rad( + _frame_values(getattr(scan, "omega", 0.0), count, "omega") + ) + centers = np.column_stack( + ( + alpha, + omega, + np.full(count, config.chi), + np.full(count, config.phi), + ) + ) + bounds = np.repeat(centers[:, None, :], 2, axis=1) + explicit = getattr(scan, "exposure_bounds_rad", None) + if explicit is not None: + explicit = np.asarray(explicit, dtype=np.float64) + if explicit.shape != bounds.shape: + raise ValueError("scan exposure_bounds_rad has an invalid shape") + return np.ascontiguousarray(explicit) + for axis, name in enumerate(("alpha", "omega", "chi", "phi")): + explicit_axis = getattr(scan, f"{name}_bounds_rad", None) + if explicit_axis is None: + continue + explicit_axis = np.asarray(explicit_axis, dtype=np.float64) + if explicit_axis.shape != (count, 2): + raise ValueError(f"scan {name}_bounds_rad has an invalid shape") + bounds[:, :, axis] = explicit_axis + if fallback == "midpoint": + angle_names = ("alpha", "omega", "chi", "phi") + for axis in range(4): + explicit_axis = getattr( + scan, f"{angle_names[axis]}_bounds_rad", None + ) + if explicit_axis is None: + bounds[:, :, axis] = _midpoint_bounds(centers[:, axis]) + return np.ascontiguousarray(bounds) + + +@dataclass(frozen=True) +class ScanReference: + """Serializable reference capable of reopening an orGUI scan.""" + + kind: str + module: str + class_name: str + parameters: dict + source_fingerprints: tuple[dict, ...] = () + + @staticmethod + def _fingerprint(path): + path = Path(path).absolute() + stat = path.stat() + return { + "path": str(path), + "size": int(stat.st_size), + "mtime_ns": int(stat.st_mtime_ns), + } + + @classmethod + def from_scan(cls, scan): + """Create a central reference for a currently loaded scan.""" + scan_class = type(scan) + module = scan_class.__module__ + class_name = scan_class.__qualname__ + backend_file = getattr(scan_class, "_orgui_backend_file", None) + if backend_file is None and module in {"", "__main__"}: + initializer = getattr(scan_class, "__init__", None) + backend_file = getattr(initializer, "__globals__", {}).get( + "__file__" + ) + if backend_file is not None: + backend_file = str(Path(backend_file).absolute()) + if not Path(backend_file).is_file(): + backend_file = None + if isinstance(scan, SimulationScan): + fixed = scan.mu if scan.axisname == "th" else scan.th + return cls( + "simulation", + module, + class_name, + { + "detshape": list(scan.shape), + "axismin": float(scan.axis[0]), + "axismax": float(scan.axis[-1]), + "points": len(scan), + "axis": scan.axisname, + "fixed": float(np.asarray(fixed).reshape(-1)[0]), + }, + ) + if module.endswith("interlacedScanLoader"): + return cls( + "interlaced", + module, + class_name, + { + "scans": [ + reference.to_dict() + for reference in map(cls.from_scan, scan.subscans) + ], + "sort": bool(scan.sort), + "axis": scan.axisname, + }, + ) + if module.endswith("universalScanLoader"): + paths = [ + str(Path(directory, filename).absolute()) + for directory, filenames in [scan.inpath] + for filename in filenames + ] + return cls( + "manual_images", + module, + class_name, + { + "filename": str(Path(scan.filename).absolute()), + "axis": scan.axisname, + "axismin": float(np.asarray(scan.axis)[0]), + "axismax": float(np.asarray(scan.axis)[-1]), + "fixed": float( + np.asarray( + scan.mu if scan.axisname == "th" else scan.th + ).reshape(-1)[0] + ), + }, + tuple(cls._fingerprint(path) for path in paths), + ) + source = getattr(scan, "hdffilepath_orNode", None) + if source is not None and not isinstance(source, str | os.PathLike): + source = getattr(source, "local_filename", None) + if source is None: + source = getattr(scan, "fastscan_specfile", None) + if source is None: + source = getattr(scan, "filename", None) + if source is None: + filenames = getattr(scan, "filenames", None) + if filenames: + source = filenames[0] + if source is None: + raise ValueError( + f"{module}.{class_name} does not expose a reloadable source" + ) + source = str(Path(source).absolute()) + scan_number = getattr(scan, "scanno", None) + if scan_number is None: + scan_number = getattr(scan, "scanno1", None) + if ( + isinstance(scan_number, str) + and scan_number.endswith(".1") + and hasattr(scan, "scanno2") + ): + scan_number = scan_number[:-2] + if scan_number is None: + scan_name = getattr(scan, "scanname", None) + if isinstance(scan_name, str): + suffix = scan_name.rsplit("_", 1)[-1] + try: + scan_number = int(suffix.split(".", 1)[0]) + except ValueError: + scan_number = None + parameters = {"source": source, "scan_number": scan_number} + if backend_file is not None: + parameters["backend_file"] = backend_file + constructor_args = getattr(scan, "_scan_reference_args", None) + if constructor_args is not None: + parameters["constructor_args"] = list(constructor_args) + if hasattr(scan, "offsetindex") and ( + int(scan.offsetindex) > 0 + or len(scan) < int(getattr(scan, "scandatapoints", len(scan))) + ): + parameters["slice"] = [ + int(scan.offsetindex), + int(scan.offsetindex) + len(scan), + ] + fingerprints = [cls._fingerprint(source)] + for filename in getattr(scan, "filenames", ()): + candidate = Path(filename).absolute() + if candidate.is_file() and str(candidate) != source: + fingerprints.append(cls._fingerprint(candidate)) + if backend_file is not None and backend_file != source: + fingerprints.append(cls._fingerprint(backend_file)) + return cls( + "backend_file" if backend_file is not None else "backend", + module, + class_name, + parameters, + tuple(fingerprints), + ) + + def to_dict(self): + """Return the JSON-compatible scan reference.""" + return { + "kind": self.kind, + "module": self.module, + "class_name": self.class_name, + "parameters": self.parameters, + "source_fingerprints": list(self.source_fingerprints), + } + + @classmethod + def from_dict(cls, values): + """Build a scan reference from JSON-compatible values.""" + values = dict(values) + values["source_fingerprints"] = tuple( + values.get("source_fingerprints", ()) + ) + return cls(**values) + + def verify(self): + """Verify that referenced source files have not changed.""" + for expected in self.source_fingerprints: + actual = self._fingerprint(expected["path"]) + if actual != expected: + raise RuntimeError( + f"Scan source changed after job preparation: " + f"{expected['path']}" + ) + + def open(self): + """Reopen this scan through the existing central backend class.""" + self.verify() + if self.kind == "backend_file": + _, scan_class = load_scan_backend_file( + self.parameters["backend_file"], self.class_name + ) + elif self.module == "": + # Jobs created before backend-file provenance was recorded can + # still run in the GUI process where that backend remains loaded. + from . import backends + + candidates = { + id(value): value + for value in backends.fscans.values() + if value.__qualname__ == self.class_name + } + if len(candidates) != 1: + raise RuntimeError( + "This job references a transient custom backend but does " + "not record its Python file. Load the same backend in " + "orGUI and retry, or prepare a new reconstruction job." + ) + scan_class = next(iter(candidates.values())) + else: + scan_class = importlib.import_module(self.module) + for component in self.class_name.split("."): + scan_class = getattr(scan_class, component) + if self.kind == "simulation": + return scan_class(**self.parameters) + if self.kind == "interlaced": + scans = [ + ScanReference.from_dict(value).open() + for value in self.parameters["scans"] + ] + return scan_class( + scans, self.parameters["sort"], self.parameters["axis"] + ) + if self.kind == "manual_images": + scan = scan_class(self.parameters["filename"]) + scan.set_axis( + self.parameters["axismin"], + self.parameters["axismax"], + self.parameters["axis"], + self.parameters["fixed"], + ) + return scan + constructor_args = self.parameters.get("constructor_args") + if constructor_args is not None: + scan = scan_class(*constructor_args) + else: + number = self.parameters.get("scan_number") + scan = ( + scan_class(self.parameters["source"], number) + if number is not None + else scan_class(self.parameters["source"]) + ) + if "slice" in self.parameters: + scan = scan.slice(*self.parameters["slice"]) + return scan diff --git a/orgui/backend/universalScanLoader.py b/orgui/backend/universalScanLoader.py index b156914..98f836a 100644 --- a/orgui/backend/universalScanLoader.py +++ b/orgui/backend/universalScanLoader.py @@ -101,10 +101,16 @@ def find_files(self): def set_axis(self, axismin, axismax, axis, fixedAxisValue): self.axis = np.linspace(axismin, axismax, self.nopoints) self.axisname = axis - if axis == "th": - self.th = self.axis - self.omega = -1 * self.th - self.mu = fixedAxisValue + if axis == "th": + self.th = self.axis + self.omega = -1 * self.th + self.mu = fixedAxisValue + elif axis == "mu": + self.mu = self.axis + self.th = fixedAxisValue + self.omega = -1 * self.th + else: + raise ValueError(f"{axis} is not an implemented scan axis.") def __len__(self): return self.nopoints diff --git a/orgui/datautils/xrayutils/AGENTS.md b/orgui/datautils/xrayutils/AGENTS.md index 1fc53f6..e57db94 100644 --- a/orgui/datautils/xrayutils/AGENTS.md +++ b/orgui/datautils/xrayutils/AGENTS.md @@ -14,11 +14,22 @@ This directory contains the highest-risk scientific code: - `CTRfilm.py`: film and epitaxy objects built on top of unit cells. - `CTRutil.py`, `CTRplotutil.py`: CTR parsing, data containers, plotting, and ANAROD-style import/export paths. +- `CTRdistributions.py`, `CTRopt.py`, `CTRoptics.py`, `CTRresolution.py`, + `CTRstacking.py`, `CTRsymmetry.py`: additional CTR model, optics, + resolution, and symmetry helpers. +- `reconstruction.py`: out-of-core reciprocal-space reconstruction (voxel + binning/reduction). `orgui/reconstruction_job.py`, `reconstruction_cli.py`, + and `reconstruction_cluster.py` (outside this directory, no nested + `AGENTS.md` of their own) build on it and follow the same conventions as + this file. +- `cpp/`: native C++ kernels (`CTRcalc_cpp.cpp`, + `reciprocal_reconstruction_cpp.cpp`) backing performance-critical CTR and + reconstruction paths. - `_CTRcalc_accel.py`: optional numba-accelerated kernels. - `element_data.py`, `unitcells/`: scattering-factor data and bundled reference structure files. -- `test/`: regression tests for lattice math, detector conversions, and CTR - calculations. +- `test/`: regression tests for lattice math, detector conversions, CTR + calculations, and reconstruction. Use the repository root instructions together with this file. @@ -75,6 +86,12 @@ rewrite stable code just to restate units. - Detector calibration should preserve round-trip behavior for pixel-to-angle and angle-to-pixel functions. - If adding an exported quantity, document its unit at the point of creation. +- `cpp/`: keep any Python fallback numerically consistent with the C++ + kernel, and rebuild the extension before trusting local test results after + touching this directory. +- `reconstruction_cluster.py` distributes work across processes/nodes; treat + partial-failure handling and result ordering as correctness, not just + performance. Within the CTR stack, the typical dependency direction is: @@ -104,6 +121,7 @@ Run the narrowest relevant regression test first: - `pytest orgui/datautils/xrayutils/test/test_CTRresolution.py` - `pytest orgui/datautils/xrayutils/test/test_CTRoptical_profile.py` - `pytest orgui/datautils/xrayutils/test/test_scattering_factor_cache.py` +- `pytest orgui/datautils/xrayutils/test/test_reconstruction*.py` Use `ruff check orgui/datautils/xrayutils` for local lint checks. If a change touches config-facing unit conversions, also inspect `examples/config_minimal` diff --git a/orgui/datautils/xrayutils/cpp/reciprocal_reconstruction_cpp.cpp b/orgui/datautils/xrayutils/cpp/reciprocal_reconstruction_cpp.cpp new file mode 100644 index 0000000..35ba889 --- /dev/null +++ b/orgui/datautils/xrayutils/cpp/reciprocal_reconstruction_cpp.cpp @@ -0,0 +1,1800 @@ +// Copyright (c) 2026 Timo Fuchs +// SPDX-License-Identifier: MIT + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include +#include +#include + +namespace py = pybind11; + +using FloatArray = py::array_t; +using BoolArray = py::array_t; +using Int64Array = py::array_t; +using UInt64Array = py::array_t; +using ContiguousUInt64Array = py::array_t< + std::uint64_t, + py::array::c_style | py::array::forcecast +>; + +std::string xxh128_hex(const XXH128_hash_t hash) { + XXH128_canonical_t canonical; + XXH128_canonicalFromHash(&canonical, hash); + constexpr char digits[] = "0123456789abcdef"; + std::string result(32, '0'); + for (std::size_t index = 0; index < sizeof(canonical.digest); ++index) { + const unsigned char value = canonical.digest[index]; + result[2 * index] = digits[value >> 4]; + result[2 * index + 1] = digits[value & 0x0f]; + } + return result; +} + +py::buffer_info contiguous_buffer(const py::buffer &buffer) { + py::buffer_info info = buffer.request(); + py::ssize_t expected_stride = info.itemsize; + for (py::ssize_t axis = info.ndim - 1; axis >= 0; --axis) { + if (info.shape[axis] > 1 && info.strides[axis] != expected_stride) { + throw py::value_error("XXH3 input buffer must be C-contiguous"); + } + expected_stride *= info.shape[axis]; + } + return info; +} + +std::string xxh3_128_buffer(const py::buffer &buffer) { + const py::buffer_info info = contiguous_buffer(buffer); + const std::size_t size = static_cast( + info.size * info.itemsize + ); + XXH128_hash_t hash; + { + py::gil_scoped_release release; + hash = XXH3_128bits(info.ptr, size); + } + return xxh128_hex(hash); +} + +py::dict merge_sorted_batches( + const ContiguousUInt64Array &left_chunk, + const ContiguousUInt64Array &left_local, + const FloatArray &left_intensity, + const FloatArray &left_variance, + const FloatArray &left_weight, + const ContiguousUInt64Array &left_contributors, + const ContiguousUInt64Array &right_chunk, + const ContiguousUInt64Array &right_local, + const FloatArray &right_intensity, + const FloatArray &right_variance, + const FloatArray &right_weight, + const ContiguousUInt64Array &right_contributors +) { + const auto lc = left_chunk.request(); + const auto ll = left_local.request(); + const auto li = left_intensity.request(); + const auto lv = left_variance.request(); + const auto lw = left_weight.request(); + const auto ln = left_contributors.request(); + const auto rc = right_chunk.request(); + const auto rl = right_local.request(); + const auto ri = right_intensity.request(); + const auto rv = right_variance.request(); + const auto rw = right_weight.request(); + const auto rn = right_contributors.request(); + const auto validate = [](const py::buffer_info &info, const char *name) { + if (info.ndim != 1) { + throw py::value_error(std::string(name) + " must be one-dimensional"); + } + }; + validate(lc, "left_chunk_id"); + validate(ll, "left_local_voxel_id"); + validate(li, "left_weighted_intensity"); + validate(lv, "left_weighted_variance"); + validate(lw, "left_weight"); + validate(ln, "left_contributors"); + validate(rc, "right_chunk_id"); + validate(rl, "right_local_voxel_id"); + validate(ri, "right_weighted_intensity"); + validate(rv, "right_weighted_variance"); + validate(rw, "right_weight"); + validate(rn, "right_contributors"); + const py::ssize_t left_size = lc.size; + const py::ssize_t right_size = rc.size; + if ( + ll.size != left_size || li.size != left_size || lv.size != left_size + || lw.size != left_size || ln.size != left_size + ) { + throw py::value_error("Left batch columns must have equal lengths"); + } + if ( + rl.size != right_size || ri.size != right_size || rv.size != right_size + || rw.size != right_size || rn.size != right_size + ) { + throw py::value_error("Right batch columns must have equal lengths"); + } + const auto *lc_data = static_cast(lc.ptr); + const auto *ll_data = static_cast(ll.ptr); + const auto *li_data = static_cast(li.ptr); + const auto *lv_data = static_cast(lv.ptr); + const auto *lw_data = static_cast(lw.ptr); + const auto *ln_data = static_cast(ln.ptr); + const auto *rc_data = static_cast(rc.ptr); + const auto *rl_data = static_cast(rl.ptr); + const auto *ri_data = static_cast(ri.ptr); + const auto *rv_data = static_cast(rv.ptr); + const auto *rw_data = static_cast(rw.ptr); + const auto *rn_data = static_cast(rn.ptr); + const auto less = []( + const std::uint64_t chunk_a, + const std::uint64_t local_a, + const std::uint64_t chunk_b, + const std::uint64_t local_b + ) { + return chunk_a < chunk_b + || (chunk_a == chunk_b && local_a < local_b); + }; + py::ssize_t output_size = 0; + { + py::gil_scoped_release release; + py::ssize_t left = 0; + py::ssize_t right = 0; + while (left < left_size && right < right_size) { + if ( + lc_data[left] == rc_data[right] + && ll_data[left] == rl_data[right] + ) { + ++left; + ++right; + } else if (less( + lc_data[left], + ll_data[left], + rc_data[right], + rl_data[right] + )) { + ++left; + } else { + ++right; + } + ++output_size; + } + output_size += left_size - left + right_size - right; + } + ContiguousUInt64Array output_chunk(output_size); + ContiguousUInt64Array output_local(output_size); + FloatArray output_intensity(output_size); + FloatArray output_variance(output_size); + FloatArray output_weight(output_size); + ContiguousUInt64Array output_contributors(output_size); + auto *oc_data = static_cast(output_chunk.request().ptr); + auto *ol_data = static_cast(output_local.request().ptr); + auto *oi_data = static_cast(output_intensity.request().ptr); + auto *ov_data = static_cast(output_variance.request().ptr); + auto *ow_data = static_cast(output_weight.request().ptr); + auto *on_data = static_cast( + output_contributors.request().ptr + ); + { + py::gil_scoped_release release; + py::ssize_t left = 0; + py::ssize_t right = 0; + py::ssize_t output = 0; + while (left < left_size || right < right_size) { + const bool take_left = right >= right_size || ( + left < left_size && less( + lc_data[left], + ll_data[left], + rc_data[right], + rl_data[right] + ) + ); + const bool equal = ( + left < left_size && right < right_size + && lc_data[left] == rc_data[right] + && ll_data[left] == rl_data[right] + ); + if (equal) { + oc_data[output] = lc_data[left]; + ol_data[output] = ll_data[left]; + oi_data[output] = li_data[left] + ri_data[right]; + ov_data[output] = lv_data[left] + rv_data[right]; + ow_data[output] = lw_data[left] + rw_data[right]; + on_data[output] = ln_data[left] + rn_data[right]; + ++left; + ++right; + } else if (take_left) { + oc_data[output] = lc_data[left]; + ol_data[output] = ll_data[left]; + oi_data[output] = li_data[left]; + ov_data[output] = lv_data[left]; + ow_data[output] = lw_data[left]; + on_data[output] = ln_data[left]; + ++left; + } else { + oc_data[output] = rc_data[right]; + ol_data[output] = rl_data[right]; + oi_data[output] = ri_data[right]; + ov_data[output] = rv_data[right]; + ow_data[output] = rw_data[right]; + on_data[output] = rn_data[right]; + ++right; + } + ++output; + } + } + py::dict result; + result["chunk_id"] = std::move(output_chunk); + result["local_voxel_id"] = std::move(output_local); + result["weighted_intensity"] = std::move(output_intensity); + result["weighted_variance"] = std::move(output_variance); + result["weight"] = std::move(output_weight); + result["contributors"] = std::move(output_contributors); + return result; +} + +struct Vec3 { + double x; + double y; + double z; +}; + +struct Mat3 { + std::array value{}; +}; + +struct PixelRays { + Vec3 base; + Vec3 du; + Vec3 dv; + Vec3 duv; +}; + +struct FrameRotation { + double sin_alpha; + double cos_alpha; + double sin_omega; + double cos_omega; + double sin_chi; + double cos_chi; + double sin_phi; + double cos_phi; +}; + +struct CoordinateTransform { + Mat3 matrix; + Vec3 offset; +}; + +Vec3 apply(const Mat3 &matrix, const Vec3 &vector) { + return { + matrix.value[0] * vector.x + matrix.value[1] * vector.y + + matrix.value[2] * vector.z, + matrix.value[3] * vector.x + matrix.value[4] * vector.y + + matrix.value[5] * vector.z, + matrix.value[6] * vector.x + matrix.value[7] * vector.y + + matrix.value[8] * vector.z, + }; +} + +Vec3 rotate_x(const Vec3 &vector, const double angle) { + const double sine = std::sin(angle); + const double cosine = std::cos(angle); + return { + vector.x, + cosine * vector.y - sine * vector.z, + sine * vector.y + cosine * vector.z, + }; +} + +Vec3 rotate_y(const Vec3 &vector, const double angle) { + const double sine = std::sin(angle); + const double cosine = std::cos(angle); + return { + cosine * vector.x + sine * vector.z, + vector.y, + -sine * vector.x + cosine * vector.z, + }; +} + +Vec3 rotate_z(const Vec3 &vector, const double angle) { + const double sine = std::sin(angle); + const double cosine = std::cos(angle); + return { + cosine * vector.x - sine * vector.y, + sine * vector.x + cosine * vector.y, + vector.z, + }; +} + +Vec3 rotate_x_sc( + const Vec3 &vector, + const double sine, + const double cosine +) { + return { + vector.x, + cosine * vector.y - sine * vector.z, + sine * vector.y + cosine * vector.z, + }; +} + +Vec3 rotate_y_sc( + const Vec3 &vector, + const double sine, + const double cosine +) { + return { + cosine * vector.x + sine * vector.z, + vector.y, + -sine * vector.x + cosine * vector.z, + }; +} + +Vec3 rotate_z_sc( + const Vec3 &vector, + const double sine, + const double cosine +) { + return { + cosine * vector.x - sine * vector.y, + sine * vector.x + cosine * vector.y, + vector.z, + }; +} + +enum class CoordinateFrame { + Lab, + Alpha, + Omega, + Chi, + Phi, + Crystal, + Hkl, +}; + +CoordinateFrame parse_frame(const std::string &frame) { + if (frame == "lab" || frame == "q_lab") { + return CoordinateFrame::Lab; + } + if (frame == "alpha" || frame == "q_alpha") { + return CoordinateFrame::Alpha; + } + if (frame == "omega" || frame == "q_omega") { + return CoordinateFrame::Omega; + } + if (frame == "chi" || frame == "q_chi") { + return CoordinateFrame::Chi; + } + if (frame == "phi" || frame == "q_phi") { + return CoordinateFrame::Phi; + } + if (frame == "crystal" || frame == "q_crystal") { + return CoordinateFrame::Crystal; + } + if (frame == "hkl") { + return CoordinateFrame::Hkl; + } + throw py::value_error( + "frame must be one of lab, alpha, omega, chi, phi, crystal, or hkl" + ); +} + +Mat3 matrix_from_array(const FloatArray &array, const char *name) { + const py::buffer_info info = array.request(); + if (info.ndim != 2 || info.shape[0] != 3 || info.shape[1] != 3) { + throw py::value_error(std::string(name) + " must have shape (3, 3)"); + } + const auto *data = static_cast(info.ptr); + Mat3 result; + std::copy(data, data + 9, result.value.begin()); + return result; +} + +std::array triple_from_array( + const FloatArray &array, + const char *name, + const bool strictly_positive +) { + const py::buffer_info info = array.request(); + if (info.ndim != 1 || info.shape[0] != 3) { + throw py::value_error(std::string(name) + " must have shape (3,)"); + } + const auto *data = static_cast(info.ptr); + std::array result{data[0], data[1], data[2]}; + for (const double value : result) { + if (!std::isfinite(value) || (strictly_positive && value <= 0.0)) { + throw py::value_error(std::string(name) + " contains an invalid value"); + } + } + return result; +} + +std::array shape_from_array( + const Int64Array &array, + const char *name +) { + const py::buffer_info info = array.request(); + if (info.ndim != 1 || info.shape[0] != 3) { + throw py::value_error(std::string(name) + " must have shape (3,)"); + } + const auto *data = static_cast(info.ptr); + std::array result{data[0], data[1], data[2]}; + for (const std::int64_t value : result) { + if (value <= 0) { + throw py::value_error(std::string(name) + " values must be positive"); + } + } + return result; +} + +struct Grid { + std::array minimum; + std::array step; + std::array shape; + std::array chunk_shape; + std::array chunk_grid; +}; + +struct RecordKey { + std::uint64_t chunk; + std::uint64_t local; +}; + +bool operator<(const RecordKey &left, const RecordKey &right) { + return left.chunk < right.chunk + || (left.chunk == right.chunk && left.local < right.local); +} + +bool operator==(const RecordKey &left, const RecordKey &right) { + return left.chunk == right.chunk && left.local == right.local; +} + +struct Record { + RecordKey key; + double weighted_intensity; + double weighted_variance; + double weight; + std::uint64_t contributors; +}; + +struct VoxelWeight { + std::uint64_t voxel; + double weight; +}; + +struct CachedVoxel { + std::uint64_t voxel = 0; + std::uint64_t generation = 0; + bool valid = false; +}; + +struct LatticeVoxel { + std::uint64_t voxel = 0; + bool valid = false; +}; + +struct PixelCoordinateCache { + std::size_t side = 0; + std::uint64_t generation = 0; + std::vector values; + + void begin_pixel(const int max_depth) { + // A dense dyadic cache is small through depth 3 (17^3 entries). + // Higher depths retain the uncached path to keep memory bounded. + if (max_depth > 3) { + side = 0; + return; + } + const std::size_t required_side = + (static_cast(1) << (max_depth + 1)) + 1; + if (side != required_side) { + side = required_side; + values.assign(side * side * side, CachedVoxel{}); + } + ++generation; + if (generation == 0) { + for (auto &value : values) { + value.generation = 0; + } + generation = 1; + } + } +}; + +bool operator<(const VoxelWeight &left, const VoxelWeight &right) { + return left.voxel < right.voxel; +} + +struct Cell { + double u0; + double u1; + double v0; + double v1; + double t0; + double t1; + double weight; + int depth; +}; + +struct BlockProfile { + std::uint64_t pixels_seen = 0; + std::uint64_t valid_pixels = 0; + std::uint64_t coordinate_evaluations = 0; + std::uint64_t voxel_weights = 0; + std::uint64_t maximum_weights_per_pixel = 0; + std::uint64_t unreduced_records = 0; + std::uint64_t reduced_records = 0; + std::uint64_t mapping_nanoseconds = 0; + std::uint64_t reduction_nanoseconds = 0; +}; + +class ReconstructionKernel { +public: + ReconstructionKernel( + const FloatArray &minimum, + const FloatArray &step, + const Int64Array &shape, + const Int64Array &chunk_shape, + const std::string &frame, + const double wavevector, + const FloatArray &ub_inverse, + const FloatArray &u_inverse, + const int max_depth = 2, + int threads = 1, + const std::size_t work_block_pixels = 4096, + const std::size_t memory_budget_bytes = 512ULL * 1024ULL * 1024ULL + ) + : frame_(parse_frame(frame)), + wavevector_(wavevector), + ub_inverse_(matrix_from_array(ub_inverse, "ub_inverse")), + u_inverse_(matrix_from_array(u_inverse, "u_inverse")), + max_depth_(max_depth), + threads_(threads), + work_block_pixels_(work_block_pixels), + memory_budget_bytes_(memory_budget_bytes) { + grid_.minimum = triple_from_array(minimum, "minimum", false); + grid_.step = triple_from_array(step, "step", true); + grid_.shape = shape_from_array(shape, "shape"); + grid_.chunk_shape = shape_from_array(chunk_shape, "chunk_shape"); + for (int axis = 0; axis < 3; ++axis) { + grid_.chunk_grid[axis] = static_cast( + (grid_.shape[axis] + grid_.chunk_shape[axis] - 1) + / grid_.chunk_shape[axis] + ); + } + if (!std::isfinite(wavevector_) || wavevector_ <= 0.0) { + throw py::value_error("wavevector must be finite and positive"); + } + if (max_depth_ < 0 || max_depth_ > 8) { + throw py::value_error("max_depth must be between 0 and 8"); + } + if (threads_ < 1) { + threads_ = 1; + } + if (work_block_pixels_ < 1) { + throw py::value_error("work_block_pixels must be positive"); + } + if (memory_budget_bytes_ < 1024 * 1024) { + throw py::value_error("memory_budget_bytes must be at least 1 MiB"); + } + } + + py::dict accumulate( + const FloatArray &intensity, + const FloatArray &variance, + const BoolArray &mask, + const FloatArray &corner_rays, + const FloatArray &angles_start, + const FloatArray &angles_end, + const bool profile = false + ) const { + const auto total_started = std::chrono::steady_clock::now(); + const py::buffer_info intensity_info = intensity.request(); + const py::buffer_info variance_info = variance.request(); + const py::buffer_info mask_info = mask.request(); + const py::buffer_info ray_info = corner_rays.request(); + const py::buffer_info start_info = angles_start.request(); + const py::buffer_info end_info = angles_end.request(); + validate_inputs( + intensity_info, + variance_info, + mask_info, + ray_info, + start_info, + end_info + ); + + const auto *intensity_data = static_cast(intensity_info.ptr); + const auto *variance_data = static_cast(variance_info.ptr); + const auto *mask_data = static_cast(mask_info.ptr); + const auto *ray_data = static_cast(ray_info.ptr); + const auto *start_data = static_cast(start_info.ptr); + const auto *end_data = static_cast(end_info.ptr); + bool stationary = true; + for (int index = 0; index < 4; ++index) { + stationary = stationary && start_data[index] == end_data[index]; + } + const std::vector rotations = + frame_rotations(start_data, end_data); + std::vector transforms; + transforms.reserve(rotations.size()); + for (const FrameRotation &rotation : rotations) { + transforms.push_back(coordinate_transform(rotation)); + } + const std::size_t rows = static_cast(intensity_info.shape[0]); + const std::size_t cols = static_cast(intensity_info.shape[1]); + const std::size_t pixels = rows * cols; + std::size_t worst_leaves = 1; + const std::size_t subdivision_children = stationary ? 4 : 8; + for (int depth = 0; depth < max_depth_; ++depth) { + worst_leaves *= subdivision_children; + } + const std::size_t estimated_bytes_per_pixel = + 128 + 2 * worst_leaves * sizeof(Record); + if ( + pixels > 0 + && estimated_bytes_per_pixel + > memory_budget_bytes_ / pixels + ) { + throw py::value_error( + "Detector tile exceeds the native memory budget at the configured " + "adaptive depth; use a smaller detector tile or increase " + "memory_budget_bytes" + ); + } + const std::size_t block_size = bounded_block_size(); + const std::size_t blocks = (pixels + block_size - 1) / block_size; + std::vector> block_results(blocks); + std::vector block_profiles( + profile ? blocks : 0 + ); + std::atomic next_block{0}; + + const auto blocks_started = std::chrono::steady_clock::now(); + { + py::gil_scoped_release release; + const int worker_count = static_cast( + std::min(static_cast(threads_), blocks) + ); + std::vector workers; + workers.reserve(static_cast(worker_count)); + for (int worker = 0; worker < worker_count; ++worker) { + workers.emplace_back([&, this]() { + while (true) { + const std::size_t block = next_block.fetch_add(1); + if (block >= blocks) { + break; + } + const std::size_t begin = block * block_size; + const std::size_t end = std::min(begin + block_size, pixels); + block_results[block] = accumulate_block( + begin, + end, + rows, + cols, + intensity_data, + variance_data, + mask_data, + ray_data, + transforms, + stationary, + profile ? &block_profiles[block] : nullptr + ); + } + }); + } + for (auto &worker : workers) { + worker.join(); + } + } + const auto blocks_finished = std::chrono::steady_clock::now(); + + std::vector records; + std::size_t record_count = 0; + const auto concatenate_started = std::chrono::steady_clock::now(); + { + py::gil_scoped_release release; + for (const auto &block : block_results) { + record_count += block.size(); + } + records.reserve(record_count); + for (auto &block : block_results) { + records.insert( + records.end(), + std::make_move_iterator(block.begin()), + std::make_move_iterator(block.end()) + ); + } + } + const auto concatenate_finished = std::chrono::steady_clock::now(); + const auto final_reduce_started = concatenate_finished; + { + py::gil_scoped_release release; + reduce_records(records); + } + const auto final_reduce_finished = std::chrono::steady_clock::now(); + const auto conversion_started = final_reduce_finished; + py::dict result = records_to_python(records); + const auto conversion_finished = std::chrono::steady_clock::now(); + if (profile) { + BlockProfile combined; + for (const BlockProfile &block : block_profiles) { + combined.pixels_seen += block.pixels_seen; + combined.valid_pixels += block.valid_pixels; + combined.coordinate_evaluations += block.coordinate_evaluations; + combined.voxel_weights += block.voxel_weights; + combined.maximum_weights_per_pixel = std::max( + combined.maximum_weights_per_pixel, + block.maximum_weights_per_pixel + ); + combined.unreduced_records += block.unreduced_records; + combined.reduced_records += block.reduced_records; + combined.mapping_nanoseconds += block.mapping_nanoseconds; + combined.reduction_nanoseconds += block.reduction_nanoseconds; + } + const auto seconds = [](const auto start, const auto stop) { + return std::chrono::duration(stop - start).count(); + }; + py::dict details; + details["pixels_seen"] = combined.pixels_seen; + details["valid_pixels"] = combined.valid_pixels; + details["coordinate_evaluations"] = + combined.coordinate_evaluations; + details["voxel_weights"] = combined.voxel_weights; + details["maximum_weights_per_pixel"] = + combined.maximum_weights_per_pixel; + details["unreduced_block_records"] = + combined.unreduced_records; + details["reduced_block_records"] = combined.reduced_records; + details["concatenated_records"] = record_count; + details["final_records"] = records.size(); + details["block_mapping_cpu_seconds"] = + static_cast(combined.mapping_nanoseconds) * 1.0e-9; + details["block_reduction_cpu_seconds"] = + static_cast(combined.reduction_nanoseconds) * 1.0e-9; + details["block_wall_seconds"] = + seconds(blocks_started, blocks_finished); + details["concatenate_seconds"] = + seconds(concatenate_started, concatenate_finished); + details["final_reduce_seconds"] = + seconds(final_reduce_started, final_reduce_finished); + details["python_conversion_seconds"] = + seconds(conversion_started, conversion_finished); + details["total_seconds"] = + seconds(total_started, conversion_finished); + result["_profile"] = std::move(details); + } + return result; + } + + py::dict configuration() const { + py::dict result; + result["max_depth"] = max_depth_; + result["threads"] = threads_; + result["work_block_pixels"] = work_block_pixels_; + result["memory_budget_bytes"] = memory_budget_bytes_; + return result; + } + + py::array_t coordinate( + const FloatArray &corner_rays, + const FloatArray &angles_start, + const FloatArray &angles_end, + const std::size_t row, + const std::size_t column, + const double u = 0.5, + const double v = 0.5, + const double t = 0.5 + ) const { + const py::buffer_info ray_info = corner_rays.request(); + const py::buffer_info start_info = angles_start.request(); + const py::buffer_info end_info = angles_end.request(); + if (ray_info.ndim != 3 || ray_info.shape[2] != 3) { + throw py::value_error("corner_rays must have shape (rows, columns, 3)"); + } + if ( + ray_info.shape[0] < 2 + || ray_info.shape[1] < 2 + || row + 1 >= static_cast(ray_info.shape[0]) + || column + 1 >= static_cast(ray_info.shape[1]) + ) { + throw py::value_error("Requested pixel is outside corner_rays"); + } + if (start_info.ndim != 1 || start_info.shape[0] != 4) { + throw py::value_error("angles_start must have shape (4,)"); + } + if (end_info.ndim != 1 || end_info.shape[0] != 4) { + throw py::value_error("angles_end must have shape (4,)"); + } + for (const double value : {u, v, t}) { + if (value < 0.0 || value > 1.0 || !std::isfinite(value)) { + throw py::value_error("u, v, and t must be finite values in [0, 1]"); + } + } + const PixelRays prepared_rays = pixel_rays( + row, + column, + static_cast(ray_info.shape[1] - 1), + static_cast(ray_info.ptr) + ); + const FrameRotation rotation = frame_rotation( + t, + static_cast(start_info.ptr), + static_cast(end_info.ptr) + ); + const Vec3 result = coordinate_at( + prepared_rays, + u, + v, + rotation + ); + py::array_t output(3); + auto *data = static_cast(output.request().ptr); + data[0] = result.x; + data[1] = result.y; + data[2] = result.z; + return output; + } + +private: + Grid grid_; + CoordinateFrame frame_; + double wavevector_; + Mat3 ub_inverse_; + Mat3 u_inverse_; + int max_depth_; + int threads_; + std::size_t work_block_pixels_; + std::size_t memory_budget_bytes_; + + std::size_t bounded_block_size() const { + // A record is deliberately overestimated to retain headroom for the + // adaptive cell stack and per-pixel temporary weights. + constexpr std::size_t estimated_bytes_per_pixel = 512; + const std::size_t per_worker = memory_budget_bytes_ + / static_cast(std::max(threads_, 1)); + return std::max( + 1, + std::min(work_block_pixels_, per_worker / estimated_bytes_per_pixel) + ); + } + + static void validate_inputs( + const py::buffer_info &intensity, + const py::buffer_info &variance, + const py::buffer_info &mask, + const py::buffer_info &rays, + const py::buffer_info &start, + const py::buffer_info &end + ) { + if (intensity.ndim != 2) { + throw py::value_error("intensity must be a two-dimensional array"); + } + if ( + variance.ndim != 2 + || variance.shape[0] != intensity.shape[0] + || variance.shape[1] != intensity.shape[1] + ) { + throw py::value_error("variance must have the same shape as intensity"); + } + if ( + mask.ndim != 2 + || mask.shape[0] != intensity.shape[0] + || mask.shape[1] != intensity.shape[1] + ) { + throw py::value_error("mask must have the same shape as intensity"); + } + if ( + rays.ndim != 3 + || rays.shape[0] != intensity.shape[0] + 1 + || rays.shape[1] != intensity.shape[1] + 1 + || rays.shape[2] != 3 + ) { + throw py::value_error( + "corner_rays must have shape (rows + 1, columns + 1, 3)" + ); + } + if (start.ndim != 1 || start.shape[0] != 4) { + throw py::value_error( + "angles_start must contain alpha, omega, chi, and phi" + ); + } + if (end.ndim != 1 || end.shape[0] != 4) { + throw py::value_error( + "angles_end must contain alpha, omega, chi, and phi" + ); + } + } + + PixelRays pixel_rays( + const std::size_t row, + const std::size_t column, + const std::size_t columns, + const double *rays + ) const { + const std::size_t stride = columns + 1; + const auto value = [rays, stride](const std::size_t r, const std::size_t c) { + const std::size_t offset = (r * stride + c) * 3; + return Vec3{rays[offset], rays[offset + 1], rays[offset + 2]}; + }; + const Vec3 r00 = value(row, column); + const Vec3 r10 = value(row + 1, column); + const Vec3 r01 = value(row, column + 1); + const Vec3 r11 = value(row + 1, column + 1); + return { + r00, + { + r10.x - r00.x, + r10.y - r00.y, + r10.z - r00.z, + }, + { + r01.x - r00.x, + r01.y - r00.y, + r01.z - r00.z, + }, + { + r11.x - r10.x - r01.x + r00.x, + r11.y - r10.y - r01.y + r00.y, + r11.z - r10.z - r01.z + r00.z, + }, + }; + } + + Vec3 ray_at( + const PixelRays &rays, + const double u, + const double v + ) const { + const double uv = u * v; + Vec3 ray{ + rays.base.x + u * rays.du.x + v * rays.dv.x + uv * rays.duv.x, + rays.base.y + u * rays.du.y + v * rays.dv.y + uv * rays.duv.y, + rays.base.z + u * rays.du.z + v * rays.dv.z + uv * rays.duv.z, + }; + const double norm = std::sqrt(ray.x * ray.x + ray.y * ray.y + ray.z * ray.z); + if (norm > 0.0) { + ray.x /= norm; + ray.y /= norm; + ray.z /= norm; + } + return ray; + } + + FrameRotation frame_rotation( + const double t, + const double *angles_start, + const double *angles_end + ) const { + std::array angles{}; + for (int index = 0; index < 4; ++index) { + angles[index] = angles_start[index] + + t * (angles_end[index] - angles_start[index]); + } + return { + std::sin(-angles[0]), + std::cos(-angles[0]), + std::sin(angles[1]), + std::cos(angles[1]), + std::sin(-angles[2]), + std::cos(-angles[2]), + std::sin(-angles[3]), + std::cos(-angles[3]), + }; + } + + std::vector frame_rotations( + const double *angles_start, + const double *angles_end + ) const { + const std::size_t side = + (static_cast(1) << (max_depth_ + 1)) + 1; + std::vector rotations; + rotations.reserve(side); + for (std::size_t index = 0; index < side; ++index) { + rotations.push_back( + frame_rotation( + static_cast(index) + / static_cast(side - 1), + angles_start, + angles_end + ) + ); + } + return rotations; + } + + Vec3 apply_frame_rotation( + Vec3 current, + const FrameRotation &rotation + ) const { + if (frame_ == CoordinateFrame::Lab) { + return current; + } + current = rotate_x_sc( + current, + rotation.sin_alpha, + rotation.cos_alpha + ); + if (frame_ == CoordinateFrame::Alpha) { + return current; + } + current = rotate_z_sc( + current, + rotation.sin_omega, + rotation.cos_omega + ); + if (frame_ == CoordinateFrame::Omega) { + return current; + } + current = rotate_y_sc( + current, + rotation.sin_chi, + rotation.cos_chi + ); + if (frame_ == CoordinateFrame::Chi) { + return current; + } + current = rotate_x_sc( + current, + rotation.sin_phi, + rotation.cos_phi + ); + if (frame_ == CoordinateFrame::Phi) { + return current; + } + if (frame_ == CoordinateFrame::Crystal) { + return apply(u_inverse_, current); + } + return apply(ub_inverse_, current); + } + + Vec3 coordinate_at( + const PixelRays &rays, + const double u, + const double v, + const FrameRotation &rotation + ) const { + const Vec3 ray = ray_at(rays, u, v); + Vec3 current{ + wavevector_ * ray.x, + wavevector_ * (ray.y - 1.0), + wavevector_ * ray.z, + }; + return apply_frame_rotation(current, rotation); + } + + CoordinateTransform coordinate_transform( + const FrameRotation &rotation + ) const { + const Vec3 x_axis = apply_frame_rotation( + {wavevector_, 0.0, 0.0}, rotation + ); + const Vec3 y_axis = apply_frame_rotation( + {0.0, wavevector_, 0.0}, rotation + ); + const Vec3 z_axis = apply_frame_rotation( + {0.0, 0.0, wavevector_}, rotation + ); + return { + { + { + x_axis.x, y_axis.x, z_axis.x, + x_axis.y, y_axis.y, z_axis.y, + x_axis.z, y_axis.z, z_axis.z, + } + }, + apply_frame_rotation({0.0, -wavevector_, 0.0}, rotation), + }; + } + + Vec3 coordinate_from_unit_ray( + const Vec3 &ray, + const CoordinateTransform &transform + ) const { + const Vec3 rotated = apply(transform.matrix, ray); + return { + rotated.x + transform.offset.x, + rotated.y + transform.offset.y, + rotated.z + transform.offset.z, + }; + } + + Vec3 coordinate_at( + const PixelRays &rays, + const double u, + const double v, + const CoordinateTransform &transform + ) const { + return coordinate_from_unit_ray(ray_at(rays, u, v), transform); + } + + bool voxel_id(const Vec3 &coordinate, std::uint64_t &voxel) const { + std::array index{}; + const std::array values{ + coordinate.x, + coordinate.y, + coordinate.z, + }; + for (int axis = 0; axis < 3; ++axis) { + if (!std::isfinite(values[axis])) { + return false; + } + index[axis] = static_cast( + std::floor((values[axis] - grid_.minimum[axis]) / grid_.step[axis]) + ); + if (index[axis] < 0 || index[axis] >= grid_.shape[axis]) { + return false; + } + } + voxel = ( + static_cast(index[0]) + * static_cast(grid_.shape[1]) + + static_cast(index[1]) + ) * static_cast(grid_.shape[2]) + + static_cast(index[2]); + return true; + } + + RecordKey record_key(const std::uint64_t voxel) const { + std::uint64_t remaining = voxel; + const std::uint64_t shape_z = + static_cast(grid_.shape[2]); + const std::uint64_t shape_y = + static_cast(grid_.shape[1]); + const std::uint64_t index_z = remaining % shape_z; + remaining /= shape_z; + const std::uint64_t index_y = remaining % shape_y; + const std::uint64_t index_x = remaining / shape_y; + const std::uint64_t chunk_x = static_cast( + index_x / static_cast(grid_.chunk_shape[0]) + ); + const std::uint64_t chunk_y = static_cast( + index_y / static_cast(grid_.chunk_shape[1]) + ); + const std::uint64_t chunk_z = static_cast( + index_z / static_cast(grid_.chunk_shape[2]) + ); + RecordKey key{}; + key.chunk = ( + chunk_x * grid_.chunk_grid[1] + chunk_y + ) * grid_.chunk_grid[2] + chunk_z; + const std::uint64_t local_x = static_cast( + index_x % static_cast(grid_.chunk_shape[0]) + ); + const std::uint64_t local_y = static_cast( + index_y % static_cast(grid_.chunk_shape[1]) + ); + const std::uint64_t local_z = static_cast( + index_z % static_cast(grid_.chunk_shape[2]) + ); + key.local = ( + local_x * static_cast(grid_.chunk_shape[1]) + local_y + ) * static_cast(grid_.chunk_shape[2]) + local_z; + return key; + } + + bool cached_voxel_key( + const PixelRays &rays, + const double u, + const double v, + const double t, + const std::vector &transforms, + PixelCoordinateCache &cache, + std::uint64_t &voxel, + BlockProfile *profile + ) const { + const double rotation_scale = + static_cast(transforms.size() - 1); + const std::size_t it = static_cast( + std::llround(t * rotation_scale) + ); + if (cache.side == 0) { + if (profile != nullptr) { + ++profile->coordinate_evaluations; + } + return voxel_id( + coordinate_at( + rays, + u, + v, + transforms[it] + ), + voxel + ); + } + const double scale = static_cast(cache.side - 1); + const std::size_t iu = static_cast(std::llround(u * scale)); + const std::size_t iv = static_cast(std::llround(v * scale)); + CachedVoxel &cached = cache.values[ + (iu * cache.side + iv) * cache.side + it + ]; + if (cached.generation != cache.generation) { + if (profile != nullptr) { + ++profile->coordinate_evaluations; + } + cached.valid = voxel_id( + coordinate_at( + rays, + u, + v, + transforms[it] + ), + cached.voxel + ); + cached.generation = cache.generation; + } + voxel = cached.voxel; + return cached.valid; + } + + LatticeVoxel stationary_lattice_voxel( + const PixelRays &rays, + const std::size_t u_index, + const std::size_t v_index, + const CoordinateTransform &transform, + BlockProfile *profile + ) const { + LatticeVoxel result; + if (profile != nullptr) { + ++profile->coordinate_evaluations; + } + if ( + (u_index == 0 || u_index == 8) + && (v_index == 0 || v_index == 8) + ) { + Vec3 ray = rays.base; + if (u_index == 8) { + ray.x += rays.du.x; + ray.y += rays.du.y; + ray.z += rays.du.z; + } + if (v_index == 8) { + ray.x += rays.dv.x; + ray.y += rays.dv.y; + ray.z += rays.dv.z; + } + if (u_index == 8 && v_index == 8) { + ray.x += rays.duv.x; + ray.y += rays.duv.y; + ray.z += rays.duv.z; + } + result.valid = voxel_id( + coordinate_from_unit_ray(ray, transform), + result.voxel + ); + return result; + } + result.valid = voxel_id( + coordinate_at( + rays, + static_cast(u_index) * 0.125, + static_cast(v_index) * 0.125, + transform + ), + result.voxel + ); + return result; + } + + static bool same_lattice_voxel( + const LatticeVoxel &first, + const LatticeVoxel &second, + const LatticeVoxel &third, + const LatticeVoxel &fourth + ) { + return first.valid + && second.valid + && third.valid + && fourth.valid + && first.voxel == second.voxel + && first.voxel == third.voxel + && first.voxel == fourth.voxel; + } + + void split_pixel_stationary_depth2( + const PixelRays &rays, + const CoordinateTransform &transform, + std::vector &weights, + BlockProfile *profile + ) const { + // Depth two uses exact eighth-pixel dyadics. Classify the root first, + // then materialize only the lattice points needed by unresolved + // children. This removes recursive stack/cache traffic without + // evaluating points that a voxel-equality test has already accepted. + std::array corners; + std::uint32_t evaluated = 0; + const auto corner_index = [](const std::size_t u, const std::size_t v) { + return (u / 2) * 5 + v / 2; + }; + const auto evaluate_corner = [&](const std::size_t u, const std::size_t v) { + const std::size_t index = corner_index(u, v); + const std::uint32_t bit = + std::uint32_t{1} << static_cast(index); + if ((evaluated & bit) == 0) { + corners[index] = + stationary_lattice_voxel(rays, u, v, transform, profile); + evaluated |= bit; + } + }; + evaluate_corner(0, 0); + evaluate_corner(8, 0); + evaluate_corner(0, 8); + evaluate_corner(8, 8); + if ( + same_lattice_voxel( + corners[corner_index(0, 0)], + corners[corner_index(8, 0)], + corners[corner_index(0, 8)], + corners[corner_index(8, 8)] + ) + ) { + weights.push_back({corners[corner_index(0, 0)].voxel, 1.0}); + return; + } + + // Only the 3x3 depth-one lattice is needed to classify the four + // children. The five additional quarter-pixel nodes inside a child + // are evaluated below only if that child actually needs splitting. + for (std::size_t u = 0; u <= 8; u += 4) { + for (std::size_t v = 0; v <= 8; v += 4) { + evaluate_corner(u, v); + } + } + + for (int child = 0; child < 4; ++child) { + const std::size_t u0 = (child & 1) != 0 ? 4 : 0; + const std::size_t v0 = (child & 2) != 0 ? 4 : 0; + const std::size_t u1 = u0 + 4; + const std::size_t v1 = v0 + 4; + if ( + same_lattice_voxel( + corners[corner_index(u0, v0)], + corners[corner_index(u1, v0)], + corners[corner_index(u0, v1)], + corners[corner_index(u1, v1)] + ) + ) { + weights.push_back({ + corners[corner_index(u0, v0)].voxel, + 0.25, + }); + continue; + } + // The four child corners are already available. These are the + // five new nodes shared by its four grandchildren. + evaluate_corner(u0 + 2, v0); + evaluate_corner(u0, v0 + 2); + evaluate_corner(u0 + 2, v0 + 2); + evaluate_corner(u1, v0 + 2); + evaluate_corner(u0 + 2, v1); + for (int grandchild = 0; grandchild < 4; ++grandchild) { + const std::size_t child_u0 = + u0 + ((grandchild & 1) != 0 ? 2 : 0); + const std::size_t child_v0 = + v0 + ((grandchild & 2) != 0 ? 2 : 0); + const std::size_t child_u1 = child_u0 + 2; + const std::size_t child_v1 = child_v0 + 2; + const LatticeVoxel &first = + corners[corner_index(child_u0, child_v0)]; + if ( + same_lattice_voxel( + first, + corners[corner_index(child_u1, child_v0)], + corners[corner_index(child_u0, child_v1)], + corners[corner_index(child_u1, child_v1)] + ) + ) { + weights.push_back({first.voxel, 0.0625}); + continue; + } + const LatticeVoxel centroid = stationary_lattice_voxel( + rays, + child_u0 + 1, + child_v0 + 1, + transform, + profile + ); + if (centroid.valid) { + weights.push_back({centroid.voxel, 0.0625}); + } + } + } + } + + void split_pixel( + const PixelRays &rays, + const std::vector &transforms, + const bool stationary, + std::vector &weights, + PixelCoordinateCache &cache, + std::vector &stack, + BlockProfile *profile + ) const { + cache.begin_pixel(max_depth_); + stack.clear(); + stack.push_back({ + 0.0, + 1.0, + 0.0, + 1.0, + stationary ? 0.5 : 0.0, + stationary ? 0.5 : 1.0, + 1.0, + 0, + }); + while (!stack.empty()) { + const Cell cell = stack.back(); + stack.pop_back(); + std::uint64_t first_voxel = 0; + bool first_valid = false; + bool all_same = true; + const int corner_count = stationary ? 4 : 8; + for (int corner = 0; corner < corner_count; ++corner) { + const double u = (corner & 1) != 0 ? cell.u1 : cell.u0; + const double v = (corner & 2) != 0 ? cell.v1 : cell.v0; + const double t = stationary + ? 0.5 + : ((corner & 4) != 0 ? cell.t1 : cell.t0); + std::uint64_t voxel = 0; + const bool valid = cached_voxel_key( + rays, + u, + v, + t, + transforms, + cache, + voxel, + profile + ); + if (corner == 0) { + first_valid = valid; + first_voxel = voxel; + } else if ( + valid != first_valid + || (valid && voxel != first_voxel) + ) { + all_same = false; + } + } + if (all_same && first_valid) { + weights.push_back({first_voxel, cell.weight}); + continue; + } + if (cell.depth >= max_depth_) { + std::uint64_t voxel = 0; + const bool valid = cached_voxel_key( + rays, + 0.5 * (cell.u0 + cell.u1), + 0.5 * (cell.v0 + cell.v1), + 0.5 * (cell.t0 + cell.t1), + transforms, + cache, + voxel, + profile + ); + if (valid) { + weights.push_back({voxel, cell.weight}); + } + continue; + } + const double um = 0.5 * (cell.u0 + cell.u1); + const double vm = 0.5 * (cell.v0 + cell.v1); + const double tm = 0.5 * (cell.t0 + cell.t1); + const int child_count = stationary ? 4 : 8; + const double child_weight = + cell.weight / static_cast(child_count); + for (int child = child_count - 1; child >= 0; --child) { + stack.push_back({ + (child & 1) != 0 ? um : cell.u0, + (child & 1) != 0 ? cell.u1 : um, + (child & 2) != 0 ? vm : cell.v0, + (child & 2) != 0 ? cell.v1 : vm, + stationary + ? 0.5 + : ((child & 4) != 0 ? tm : cell.t0), + stationary + ? 0.5 + : ((child & 4) != 0 ? cell.t1 : tm), + child_weight, + cell.depth + 1, + }); + } + } + } + + std::vector accumulate_block( + const std::size_t begin, + const std::size_t end, + const std::size_t rows, + const std::size_t columns, + const double *intensity, + const double *variance, + const bool *mask, + const double *rays, + const std::vector &transforms, + const bool stationary, + BlockProfile *profile + ) const { + (void)rows; + std::vector records; + std::vector weights; + PixelCoordinateCache coordinate_cache; + std::vector stack; + const std::size_t reserve_leaves = static_cast(1) + << std::min((stationary ? 2 : 3) * max_depth_, 9); + records.reserve(2 * (end - begin)); + weights.reserve(reserve_leaves); + stack.reserve(reserve_leaves); + const auto mapping_started = std::chrono::steady_clock::now(); + for (std::size_t flat = begin; flat < end; ++flat) { + if (profile != nullptr) { + ++profile->pixels_seen; + } + if ( + mask[flat] + || !std::isfinite(intensity[flat]) + || !std::isfinite(variance[flat]) + || variance[flat] < 0.0 + ) { + continue; + } + if (profile != nullptr) { + ++profile->valid_pixels; + } + const std::size_t row = flat / columns; + const std::size_t column = flat % columns; + const PixelRays prepared_rays = pixel_rays( + row, + column, + columns, + rays + ); + weights.clear(); + if (max_depth_ == 0) { + if (profile != nullptr) { + ++profile->coordinate_evaluations; + } + std::uint64_t voxel = 0; + if ( + voxel_id( + coordinate_at( + prepared_rays, + 0.5, + 0.5, + transforms[transforms.size() / 2] + ), + voxel + ) + ) { + weights.push_back({voxel, 1.0}); + } + } else if (stationary && max_depth_ == 2) { + split_pixel_stationary_depth2( + prepared_rays, + transforms[transforms.size() / 2], + weights, + profile + ); + } else { + split_pixel( + prepared_rays, + transforms, + stationary, + weights, + coordinate_cache, + stack, + profile + ); + } + if (profile != nullptr) { + profile->voxel_weights += weights.size(); + profile->maximum_weights_per_pixel = std::max( + profile->maximum_weights_per_pixel, + static_cast(weights.size()) + ); + } + std::sort(weights.begin(), weights.end()); + std::size_t offset = 0; + while (offset < weights.size()) { + const std::uint64_t voxel = weights[offset].voxel; + double weight = 0.0; + do { + weight += weights[offset].weight; + ++offset; + } while ( + offset < weights.size() + && weights[offset].voxel == voxel + ); + const RecordKey key = record_key(voxel); + records.push_back({ + key, + weight * intensity[flat], + weight * weight * variance[flat], + weight, + 1, + }); + } + } + const auto mapping_finished = std::chrono::steady_clock::now(); + if (profile != nullptr) { + profile->unreduced_records = records.size(); + profile->mapping_nanoseconds = static_cast( + std::chrono::duration_cast( + mapping_finished - mapping_started + ).count() + ); + } + const auto reduction_started = mapping_finished; + reduce_records(records); + if (profile != nullptr) { + profile->reduced_records = records.size(); + profile->reduction_nanoseconds = static_cast( + std::chrono::duration_cast( + std::chrono::steady_clock::now() - reduction_started + ).count() + ); + } + return records; + } + + static void reduce_records(std::vector &records) { + std::stable_sort( + records.begin(), + records.end(), + [](const Record &left, const Record &right) { + return left.key < right.key; + } + ); + std::size_t write = 0; + std::size_t read = 0; + while (read < records.size()) { + Record combined = records[read++]; + while (read < records.size() && records[read].key == combined.key) { + combined.weighted_intensity += records[read].weighted_intensity; + combined.weighted_variance += records[read].weighted_variance; + combined.weight += records[read].weight; + combined.contributors += records[read].contributors; + ++read; + } + records[write++] = combined; + } + records.resize(write); + } + + static py::dict records_to_python(const std::vector &records) { + const py::ssize_t size = static_cast(records.size()); + UInt64Array chunk_id(size); + UInt64Array local_voxel_id(size); + FloatArray weighted_intensity(size); + FloatArray weighted_variance(size); + FloatArray weight(size); + UInt64Array contributors(size); + auto *chunk_data = static_cast(chunk_id.request().ptr); + auto *local_data = static_cast(local_voxel_id.request().ptr); + auto *intensity_data = static_cast(weighted_intensity.request().ptr); + auto *variance_data = static_cast(weighted_variance.request().ptr); + auto *weight_data = static_cast(weight.request().ptr); + auto *contributors_data = static_cast( + contributors.request().ptr + ); + { + py::gil_scoped_release release; + for (py::ssize_t index = 0; index < size; ++index) { + const Record &record = records[static_cast(index)]; + chunk_data[index] = record.key.chunk; + local_data[index] = record.key.local; + intensity_data[index] = record.weighted_intensity; + variance_data[index] = record.weighted_variance; + weight_data[index] = record.weight; + contributors_data[index] = record.contributors; + } + } + py::dict result; + result["chunk_id"] = std::move(chunk_id); + result["local_voxel_id"] = std::move(local_voxel_id); + result["weighted_intensity"] = std::move(weighted_intensity); + result["weighted_variance"] = std::move(weighted_variance); + result["weight"] = std::move(weight); + result["contributors"] = std::move(contributors); + return result; + } +}; + +PYBIND11_MODULE(_reciprocal_reconstruction_cpp, module) { + module.doc() = + "Native reciprocal-space coordinate conversion and footprint accumulation."; + module.def( + "xxh3_128", + &xxh3_128_buffer, + py::arg("buffer"), + "Return the canonical XXH3-128 digest of a C-contiguous buffer." + ); + module.def( + "merge_sorted_batches", + &merge_sorted_batches, + py::arg("left_chunk_id"), + py::arg("left_local_voxel_id"), + py::arg("left_weighted_intensity"), + py::arg("left_weighted_variance"), + py::arg("left_weight"), + py::arg("left_contributors"), + py::arg("right_chunk_id"), + py::arg("right_local_voxel_id"), + py::arg("right_weighted_intensity"), + py::arg("right_weighted_variance"), + py::arg("right_weight"), + py::arg("right_contributors"), + "Linearly merge two sorted, already-reduced record batches." + ); + py::class_(module, "ReconstructionKernel") + .def( + py::init< + const FloatArray &, + const FloatArray &, + const Int64Array &, + const Int64Array &, + const std::string &, + double, + const FloatArray &, + const FloatArray &, + int, + int, + std::size_t, + std::size_t + >(), + py::arg("minimum"), + py::arg("step"), + py::arg("shape"), + py::arg("chunk_shape"), + py::arg("frame"), + py::arg("wavevector"), + py::arg("ub_inverse"), + py::arg("u_inverse"), + py::arg("max_depth") = 2, + py::arg("threads") = 1, + py::arg("work_block_pixels") = 4096, + py::arg("memory_budget_bytes") = 512ULL * 1024ULL * 1024ULL + ) + .def( + "accumulate", + &ReconstructionKernel::accumulate, + py::arg("intensity"), + py::arg("variance"), + py::arg("mask"), + py::arg("corner_rays"), + py::arg("angles_start"), + py::arg("angles_end"), + py::arg("profile") = false + ) + .def( + "coordinate", + &ReconstructionKernel::coordinate, + py::arg("corner_rays"), + py::arg("angles_start"), + py::arg("angles_end"), + py::arg("row"), + py::arg("column"), + py::arg("u") = 0.5, + py::arg("v") = 0.5, + py::arg("t") = 0.5 + ) + .def_property_readonly( + "configuration", + &ReconstructionKernel::configuration + ); +} diff --git a/orgui/datautils/xrayutils/reconstruction.py b/orgui/datautils/xrayutils/reconstruction.py new file mode 100644 index 0000000..5a8b4fe --- /dev/null +++ b/orgui/datautils/xrayutils/reconstruction.py @@ -0,0 +1,1760 @@ +"""Out-of-core reciprocal-space reconstruction. + +The numerical hot path is implemented by +``_reciprocal_reconstruction_cpp``. This module owns the scientific Python +boundary, uncertainty propagation, immutable Parquet task products, and final +NeXus/HDF5 serialization. + +Momentum-transfer grids use ``Angstrom^-1``. HKL grids use reciprocal lattice +units. All diffractometer angles accepted here are in radians. +""" + +from __future__ import annotations + +from collections.abc import Callable, Iterable, Mapping +from concurrent.futures import ThreadPoolExecutor, as_completed +from dataclasses import asdict, dataclass, field +from hashlib import sha256 +import importlib +import json +import math +import os +from pathlib import Path +from queue import Queue +from typing import Any + +import h5py +import numpy as np + + +_PARTIAL_COLUMNS = ( + "chunk_id", + "local_voxel_id", + "weighted_intensity", + "weighted_variance", + "weight", + "contributors", +) +_Q_FRAMES = {"lab", "alpha", "omega", "chi", "phi", "crystal"} +_ALL_FRAMES = _Q_FRAMES | {"hkl"} +_MIN_REDUCER_WORKER_MEMORY = 64 * 1024**2 + + +def _triple(values, name, dtype=float): + result = tuple(dtype(value) for value in values) + if len(result) != 3: + raise ValueError(f"{name} must contain exactly three values") + return result + + +@dataclass(frozen=True) +class _GridSpec: + """Describe a regular reciprocal-space grid. + + :param minimum: + Lower voxel edges in r.l.u. for ``hkl`` or ``Angstrom^-1`` for Q. + :param maximum: + Exclusive upper grid bounds in the same units as ``minimum``. + :param step: + Voxel widths in the same units as ``minimum``. + :param frame: + ``hkl`` or one of ``lab``, ``alpha``, ``omega``, ``chi``, ``phi``, + and ``crystal``. + :param chunk_shape: + HDF5 chunk shape in voxels. + :param name: + Optional HDF5 group name. + """ + + minimum: tuple[float, float, float] + maximum: tuple[float, float, float] + step: tuple[float, float, float] + frame: str + chunk_shape: tuple[int, int, int] = (64, 64, 64) + name: str | None = None + + def __post_init__(self): + object.__setattr__(self, "minimum", _triple(self.minimum, "minimum")) + object.__setattr__(self, "maximum", _triple(self.maximum, "maximum")) + object.__setattr__(self, "step", _triple(self.step, "step")) + object.__setattr__( + self, "chunk_shape", _triple(self.chunk_shape, "chunk_shape", int) + ) + frame = self.frame.removeprefix("q_").lower() + if frame not in _ALL_FRAMES: + raise ValueError(f"Unsupported reciprocal-space frame: {self.frame}") + object.__setattr__(self, "frame", frame) + if any(not math.isfinite(value) for value in self.minimum + self.maximum): + raise ValueError("Grid bounds must be finite") + if any(value <= 0 or not math.isfinite(value) for value in self.step): + raise ValueError("Grid steps must be finite and positive") + if any(upper <= lower for lower, upper in zip(self.minimum, self.maximum)): + raise ValueError("Each grid maximum must exceed its minimum") + if any(value <= 0 for value in self.chunk_shape): + raise ValueError("Chunk dimensions must be positive") + effective_chunk = tuple( + min(size, chunk) + for size, chunk in zip(self.shape, self.chunk_shape) + ) + chunk_bytes = math.prod(effective_chunk) * np.dtype(np.float64).itemsize + if chunk_bytes >= 2**32: + raise ValueError( + "The effective HDF5 chunk must be smaller than 4 GiB; " + f"{effective_chunk} requires {chunk_bytes / 1024**3:.2f} GiB " + "per float64 dataset" + ) + if self.name is not None and not self.name: + raise ValueError("Grid name cannot be empty") + + @property + def shape(self) -> tuple[int, int, int]: + """Grid shape in voxels.""" + return tuple( + int(math.ceil((upper - lower) / width)) + for lower, upper, width in zip(self.minimum, self.maximum, self.step) + ) + + @property + def effective_maximum(self) -> tuple[float, float, float]: + """Upper edge implied by the integer grid shape.""" + return tuple( + lower + size * width + for lower, size, width in zip(self.minimum, self.shape, self.step) + ) + + @property + def grid_name(self) -> str: + """Stable HDF5-safe name for this grid.""" + raw = self.name or ("hkl" if self.frame == "hkl" else f"q_{self.frame}") + return "".join(char if char.isalnum() or char in "_-" else "_" for char in raw) + + +@dataclass(frozen=True) +class _ReconstructionSpec: + """Configuration shared by mapping, reduction, and finalization.""" + + grids: tuple[_GridSpec, ...] + max_depth: int = 2 + threads: int = 1 + work_block_pixels: int = 4096 + memory_budget_bytes: int = 512 * 1024 * 1024 + partition_chunk_span: int = 256 + compression: str = "bitshuffle-lz4" + infer_angle_bounds: bool = False + metadata: Mapping[str, Any] = field(default_factory=dict) + + def __post_init__(self): + grids = tuple( + grid if isinstance(grid, _GridSpec) else _GridSpec(**grid) + for grid in self.grids + ) + object.__setattr__(self, "grids", grids) + if not grids: + raise ValueError("At least one output grid is required") + names = [grid.grid_name for grid in grids] + if len(names) != len(set(names)): + raise ValueError("Grid names must be unique") + if not 0 <= self.max_depth <= 8: + raise ValueError("max_depth must be between 0 and 8") + if self.threads < 1: + raise ValueError("threads must be positive") + if self.work_block_pixels < 1: + raise ValueError("work_block_pixels must be positive") + if self.memory_budget_bytes < 1024 * 1024: + raise ValueError("memory_budget_bytes must be at least 1 MiB") + if self.partition_chunk_span < 1: + raise ValueError("partition_chunk_span must be positive") + + def to_dict(self) -> dict[str, Any]: + """Return a JSON-serializable representation.""" + result = asdict(self) + result["metadata"] = dict(self.metadata) + return result + + @classmethod + def from_dict(cls, values: Mapping[str, Any]): + """Create a specification from JSON-compatible values.""" + data = dict(values) + data["grids"] = tuple(_GridSpec(**grid) for grid in data["grids"]) + return cls(**data) + + @property + def digest(self) -> str: + """SHA-256 digest identifying scientifically equivalent settings.""" + encoded = json.dumps( + self.to_dict(), sort_keys=True, separators=(",", ":"), default=str + ).encode("utf-8") + return sha256(encoded).hexdigest() + + +@dataclass +class _PartitionFile: + """One immutable Parquet object produced by a mapping task.""" + + grid_name: str + bucket: int + uri: str + rows: int + checksum: str + size_bytes: int | None = None + + +@dataclass +class _ChunkFile: + """One reduced Parquet shard belonging to one output spatial chunk.""" + + grid_name: str + chunk_id: int + shard_start: int + shard_stop: int + uri: str + rows: int + checksum: str + size_bytes: int | None = None + + +@dataclass +class _TaskManifest: + """Serializable, restart-safe mapping or reduction task manifest.""" + + kind: str + task_id: str + spec_hash: str + status: str + spec: dict[str, Any] + frame_range: tuple[int, int] | None = None + detector_tile: tuple[int, int, int, int] | None = None + partitions: list[_PartitionFile] = field(default_factory=list) + chunks: list[_ChunkFile] = field(default_factory=list) + source_tasks: list[str] = field(default_factory=list) + metadata: dict[str, Any] = field(default_factory=dict) + + def to_dict(self) -> dict[str, Any]: + """Return a JSON-compatible manifest.""" + result = asdict(self) + return result + + @classmethod + def from_dict(cls, values: Mapping[str, Any]): + """Deserialize a task manifest.""" + data = dict(values) + data["partitions"] = [ + value if isinstance(value, _PartitionFile) else _PartitionFile(**value) + for value in data.get("partitions", []) + ] + data["chunks"] = [ + value if isinstance(value, _ChunkFile) else _ChunkFile(**value) + for value in data.get("chunks", []) + ] + if data.get("frame_range") is not None: + data["frame_range"] = tuple(data["frame_range"]) + if data.get("detector_tile") is not None: + data["detector_tile"] = tuple(data["detector_tile"]) + return cls(**data) + + +def _native_module(): + try: + return importlib.import_module( + "orgui.datautils.xrayutils._reciprocal_reconstruction_cpp" + ) + except ImportError as exc: + raise RuntimeError( + "The native reciprocal-space reconstruction extension is unavailable. " + "Install orGUI from a wheel containing its C++ extensions or rebuild it " + "with Meson." + ) from exc + + +def _xxh3_128(array) -> str: + """Return the XXH3-128 fingerprint of one contiguous array or buffer.""" + value = np.ascontiguousarray(array) + return _native_module().xxh3_128(value) + + +def _detector_corner_rays(detector, detector_tile) -> np.ndarray: + """Calculate calibrated outgoing unit rays at detector-pixel corners. + + :param DetectorCalibration.Detector2D_SXRD detector: + Calibrated detector geometry. + :param detector_tile: + ``(row_start, row_stop, column_start, column_stop)`` using exclusive + stop indices. + :returns: + C-contiguous array shaped ``(rows + 1, columns + 1, 3)`` in the + laboratory frame. + """ + row_start, row_stop, column_start, column_stop = map(int, detector_tile) + if row_start < 0 or column_start < 0: + raise ValueError("Detector tile starts must be non-negative") + if row_stop <= row_start or column_stop <= column_start: + raise ValueError("Detector tile stops must exceed starts") + rows = np.arange(row_start, row_stop + 1, dtype=np.float64) - 0.5 + columns = np.arange(column_start, column_stop + 1, dtype=np.float64) - 0.5 + row_grid, column_grid = np.meshgrid(rows, columns, indexing="ij") + gamma_p, delta_p = detector.primBeamPoints(row_grid, column_grid) + cosine_gamma = np.cos(gamma_p) + rays = np.empty((*row_grid.shape, 3), dtype=np.float64) + rays[..., 0] = np.sin(delta_p) * cosine_gamma + rays[..., 1] = np.cos(delta_p) * cosine_gamma + rays[..., 2] = np.sin(gamma_p) + rays /= np.linalg.norm(rays, axis=-1, keepdims=True) + return np.ascontiguousarray(rays) + + +def _kernel_for_grid( + spec, + grid, + ub_calculator, + *, + threads=None, + memory_budget_bytes=None, +): + ub = np.asarray(ub_calculator.getUB(), dtype=np.float64) + u = np.asarray(ub_calculator.getU(), dtype=np.float64) + return _native_module().ReconstructionKernel( + np.asarray(grid.minimum, dtype=np.float64), + np.asarray(grid.step, dtype=np.float64), + np.asarray(grid.shape, dtype=np.int64), + np.asarray(grid.chunk_shape, dtype=np.int64), + grid.frame, + float(ub_calculator.getK()), + np.ascontiguousarray(np.linalg.inv(ub)), + np.ascontiguousarray(np.linalg.inv(u)), + spec.max_depth, + spec.threads if threads is None else threads, + spec.work_block_pixels, + ( + spec.memory_budget_bytes + if memory_budget_bytes is None + else memory_budget_bytes + ), + ) + + +def _empty_batch() -> dict[str, np.ndarray]: + return { + "chunk_id": np.empty(0, dtype=np.uint64), + "local_voxel_id": np.empty(0, dtype=np.uint64), + "weighted_intensity": np.empty(0, dtype=np.float64), + "weighted_variance": np.empty(0, dtype=np.float64), + "weight": np.empty(0, dtype=np.float64), + "contributors": np.empty(0, dtype=np.uint64), + } + + +def _reduce_batches(batches: Iterable[Mapping[str, np.ndarray]]): + levels = [] + for batch in batches: + if not np.asarray(batch["chunk_id"]).size: + continue + level = 0 + while level < len(levels) and levels[level] is not None: + batch = _merge_sorted_batches(levels[level], batch) + levels[level] = None + level += 1 + if level == len(levels): + levels.append(batch) + else: + levels[level] = batch + result = _empty_batch() + for batch in reversed(levels): + if batch is not None: + result = _merge_sorted_batches(result, batch) + return result + + +def _merge_sorted_batches(left, right): + """Linearly merge two sorted, already-reduced native record batches.""" + if not left["chunk_id"].size: + return right + if not right["chunk_id"].size: + return left + arguments = [ + np.ascontiguousarray(batch[name]) + for batch in (left, right) + for name in _PARTIAL_COLUMNS + ] + return _native_module().merge_sorted_batches(*arguments) + + +def _require_pyarrow(): + try: + import pyarrow as pa + import pyarrow.fs as pafs + import pyarrow.parquet as pq + except ImportError as exc: + raise RuntimeError( + "Parquet reconstruction storage requires the optional 'pyarrow' " + "dependency." + ) from exc + return pa, pafs, pq + + +def _filesystem_path(uri): + _, pafs, _ = _require_pyarrow() + text = os.fspath(uri) + if "://" not in text: + path = str(Path(text).absolute()) + return pafs.LocalFileSystem(), path + return pafs.FileSystem.from_uri(text) + + +def _write_parquet(batch, uri, metadata): + pa, _, pq = _require_pyarrow() + filesystem, path = _filesystem_path(uri) + parent = str(Path(path).parent).replace("\\", "/") + try: + filesystem.create_dir(parent, recursive=True) + except (NotImplementedError, OSError): + pass + table = pa.table({name: batch[name] for name in _PARTIAL_COLUMNS}) + encoded_metadata = { + str(key).encode("utf-8"): json.dumps(value, sort_keys=True).encode("utf-8") + for key, value in metadata.items() + } + table = table.replace_schema_metadata(encoded_metadata) + pq.write_table(table, path, filesystem=filesystem, compression="zstd") + + +def _read_parquet(uri): + _, _, pq = _require_pyarrow() + filesystem, path = _filesystem_path(uri) + table = pq.read_table(path, filesystem=filesystem, columns=list(_PARTIAL_COLUMNS)) + return {name: table[name].to_numpy() for name in _PARTIAL_COLUMNS} + + +class _ParquetRangeReader: + """Stream monotonically increasing key ranges from sorted Parquet.""" + + def __init__(self, uri, *, batch_size=131072, use_threads=True): + _, _, pq = _require_pyarrow() + filesystem, path = _filesystem_path(uri) + self.parquet = pq.ParquetFile(path, filesystem=filesystem) + names = self.parquet.schema_arrow.names + self.chunk_column = names.index("chunk_id") + self.local_column = names.index("local_voxel_id") + self.batch_size = int(batch_size) + self.use_threads = bool(use_threads) + self.iterator = None + self.batch = None + self.position = 0 + self.previous_key = None + + def _start(self, chunk_id, local_start): + row_groups = [] + for index in range(self.parquet.metadata.num_row_groups): + metadata = self.parquet.metadata.row_group(index) + chunk_statistics = metadata.column( + self.chunk_column + ).statistics + if ( + chunk_statistics is not None + and chunk_statistics.has_min_max + and int(chunk_statistics.max) < chunk_id + ): + continue + if ( + chunk_statistics is not None + and chunk_statistics.has_min_max + and int(chunk_statistics.max) == chunk_id + ): + local_statistics = metadata.column( + self.local_column + ).statistics + if ( + local_statistics is not None + and local_statistics.has_min_max + and int(local_statistics.max) < local_start + ): + continue + row_groups.append(index) + self.iterator = self.parquet.iter_batches( + batch_size=self.batch_size, + row_groups=row_groups, + columns=list(_PARTIAL_COLUMNS), + use_threads=self.use_threads, + ) + + def _advance(self): + try: + record_batch = next(self.iterator) + except StopIteration: + self.batch = None + return False + self.batch = { + name: record_batch.column(name).to_numpy(zero_copy_only=False) + for name in _PARTIAL_COLUMNS + } + self.position = 0 + return True + + def read(self, chunk_id, local_start, local_stop): + """Read one exact range without revisiting earlier Parquet data.""" + key = (int(chunk_id), int(local_start)) + if self.previous_key is not None and key < self.previous_key: + raise ValueError("Parquet ranges must be requested in sorted order") + self.previous_key = key + if self.iterator is None: + self._start(chunk_id, local_start) + parts = [] + while self.batch is not None or self._advance(): + chunks = self.batch["chunk_id"] + position = self.position + if position >= chunks.size: + self.batch = None + continue + position += int( + np.searchsorted( + chunks[position:], + np.uint64(chunk_id), + side="left", + ) + ) + if position >= chunks.size: + self.batch = None + continue + current_chunk = int(chunks[position]) + if current_chunk > chunk_id: + self.position = position + break + chunk_stop = int( + np.searchsorted( + chunks, + np.uint64(chunk_id), + side="right", + sorter=None, + ) + ) + local = self.batch["local_voxel_id"] + start = position + int( + np.searchsorted( + local[position:chunk_stop], + np.uint64(local_start), + side="left", + ) + ) + stop = position + int( + np.searchsorted( + local[position:chunk_stop], + np.uint64(local_stop), + side="left", + ) + ) + if stop > start: + parts.append( + { + name: values[start:stop] + for name, values in self.batch.items() + } + ) + self.position = stop + if stop < chunk_stop: + break + self.position = chunk_stop + if chunk_stop < chunks.size: + break + self.batch = None + if not parts: + return _empty_batch() + if len(parts) == 1: + return parts[0] + return { + name: np.concatenate([part[name] for part in parts]) + for name in _PARTIAL_COLUMNS + } + + +def _uri_checksum_and_size(uri): + filesystem, path = _filesystem_path(uri) + digest = sha256() + size = 0 + with filesystem.open_input_file(path) as stream: + while True: + block = stream.read(8 * 1024 * 1024) + if not block: + break + digest.update(block) + size += len(block) + return digest.hexdigest(), size + + +def _uri_checksum(uri): + return _uri_checksum_and_size(uri)[0] + + +def _uri_size(uri): + filesystem, path = _filesystem_path(uri) + size = int(filesystem.get_file_info(path).size) + if size < 0: + raise OSError(f"Could not determine scratch-file size: {uri}") + return size + + +def _mark_scratch_verified(file, verification_cache): + if verification_cache is not None: + verification_cache.add((file.uri, file.checksum)) + + +def _verify_scratch_file(file, verification_cache=None): + key = (file.uri, file.checksum) + if verification_cache is not None and key in verification_cache: + return True + if file.size_bytes is not None and _uri_size(file.uri) != file.size_bytes: + raise OSError(f"Scratch-file size mismatch: {file.uri}") + if _uri_checksum(file.uri) != file.checksum: + raise OSError(f"Checksum mismatch for {file.uri}") + if verification_cache is not None: + verification_cache.add(key) + return True + + +def _join_uri(base, *parts): + base = os.fspath(base).rstrip("/\\") + separator = "/" if "://" in base else os.sep + return separator.join((base, *(str(part).strip("/\\") for part in parts))) + + +def _jsonable(value): + if isinstance(value, Mapping): + return {str(key): _jsonable(item) for key, item in value.items()} + if isinstance(value, list | tuple): + return [_jsonable(item) for item in value] + if isinstance(value, np.ndarray): + return value.tolist() + if isinstance(value, np.generic): + return value.item() + if isinstance(value, str | int | float | bool) or value is None: + return value + return repr(value) + + +def _write_manifest(manifest: _TaskManifest, uri): + """Write a task manifest atomically on a local filesystem.""" + path = Path(uri) + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_name(path.name + ".tmp") + temporary.write_text( + json.dumps(manifest.to_dict(), indent=2, sort_keys=True), + encoding="utf-8", + ) + temporary.replace(path) + return str(path) + + +def _read_manifest(value) -> _TaskManifest: + """Read a manifest path or normalize an existing manifest.""" + if isinstance(value, _TaskManifest): + return value + if isinstance(value, Mapping): + return _TaskManifest.from_dict(value) + return _TaskManifest.from_dict( + json.loads(Path(value).read_text(encoding="utf-8")) + ) + + +def _map_frame_range( + spec: _ReconstructionSpec, + scan, + detector, + ub_calculator, + frame_range, + detector_tiles, + angle_bounds_rad, + output_uri, + *, + correction_pipeline: Callable, + job_digest: str, + image_payloads: Mapping[int, object] | None = None, + corner_rays: Mapping[tuple[int, int, int, int], np.ndarray] | None = None, + corner_rays_fingerprints: Mapping[ + tuple[int, int, int, int], str + ] + | None = None, + verification_cache: set[tuple[str, str]] | None = None, + kernel_threads: int | None = None, + kernel_memory_budget_bytes: int | None = None, + accumulation_budget_bytes: int = 64 * 1024 * 1024, + image_progress: Callable[[int, int, int], None] | None = None, +) -> _TaskManifest: + """Map a streaming frame range and write byte-bounded Parquet partials. + + ``angle_bounds_rad`` must have shape ``(frames, 2, 4)`` in the order + ``alpha, omega, chi, phi``. The central correction object may provide a + ``correct_frame`` method, which is called once before detector tiling. + This preserves full-detector pixel-repair neighborhoods and avoids + repeating static corrections for every tile. The tile callback remains + available for correction objects without that method. + + Every image is loaded once, then all detector tiles are processed. Reduced + native records are retained until ``accumulation_budget_bytes`` would be + exceeded, at which point they are reduced across images and written as one + deterministic Parquet segment. + + ``image_payloads``, ``corner_rays``, ``corner_rays_fingerprints``, + ``verification_cache``, ``kernel_threads``, + ``kernel_memory_budget_bytes``, and ``accumulation_budget_bytes`` are local + execution controls. They do not alter scientific values. + + ``image_progress`` is called after each completed image with the frame + index, retained-record bytes, and number of flushed segments. + """ + start, stop = map(int, frame_range) + if start < 0 or stop <= start or stop > len(scan): + raise ValueError("Invalid frame range") + detector_tiles = tuple( + tuple(map(int, detector_tile)) for detector_tile in detector_tiles + ) + if not detector_tiles: + raise ValueError("At least one detector tile is required") + detector_rows, detector_columns = detector.detector.shape + for row_start, row_stop, column_start, column_stop in detector_tiles: + if ( + row_start < 0 + or column_start < 0 + or row_stop > detector_rows + or column_stop > detector_columns + or row_stop <= row_start + or column_stop <= column_start + ): + raise ValueError("Invalid detector tile") + bounds = np.asarray(angle_bounds_rad, dtype=np.float64) + if bounds.shape != (stop - start, 2, 4): + raise ValueError("angle_bounds_rad must have shape (frame_count, 2, 4)") + if not np.all(np.isfinite(bounds)): + raise ValueError("angle_bounds_rad contains non-finite values") + if not callable(correction_pipeline): + raise TypeError("correction_pipeline must be the central job correction") + if not job_digest: + raise ValueError("job_digest is required") + if accumulation_budget_bytes < 1: + raise ValueError("accumulation_budget_bytes must be positive") + ray_arrays = {} + ray_fingerprints = {} + for detector_tile in detector_tiles: + row_start, row_stop, column_start, column_stop = detector_tile + rays = ( + _detector_corner_rays(detector, detector_tile) + if corner_rays is None or detector_tile not in corner_rays + else np.ascontiguousarray( + corner_rays[detector_tile], dtype=np.float64 + ) + ) + expected_ray_shape = ( + row_stop - row_start + 1, + column_stop - column_start + 1, + 3, + ) + if rays.shape != expected_ray_shape: + raise ValueError( + f"corner_rays must have shape {expected_ray_shape}" + ) + ray_arrays[detector_tile] = rays + fingerprint = ( + None + if corner_rays_fingerprints is None + else corner_rays_fingerprints.get(detector_tile) + ) + if fingerprint is None: + fingerprint = _xxh3_128(rays) + elif not isinstance(fingerprint, str) or len(fingerprint) != 32: + raise ValueError( + "Corner-ray fingerprints must be 32-character " + "XXH3-128 digests" + ) + ray_fingerprints[detector_tile] = fingerprint + ub = np.ascontiguousarray(ub_calculator.getUB(), dtype=np.float64) + u = np.ascontiguousarray(ub_calculator.getU(), dtype=np.float64) + angle_bounds_fingerprint = _xxh3_128(bounds) + ub_fingerprint = _xxh3_128(ub) + u_fingerprint = _xxh3_128(u) + wavevector = float(ub_calculator.getK()) + detector_config = ( + _jsonable(detector.get_config()) if hasattr(detector, "get_config") else None + ) + base_context = { + "job_sha256": job_digest, + "ub": ub.tolist(), + "u": u.tolist(), + "wavevector_A^-1": wavevector, + "detector": detector_config, + } + base_context_hash = sha256( + json.dumps(base_context, sort_keys=True, separators=(",", ":")).encode() + ).hexdigest() + scientific_components = { + "algorithm": "xxh3-128-components-v1", + "job_sha256": job_digest, + "angle_bounds_xxh3_128": angle_bounds_fingerprint, + "corner_rays_xxh3_128": [ + { + "detector_tile": list(detector_tile), + "digest": ray_fingerprints[detector_tile], + } + for detector_tile in detector_tiles + ], + "ub_xxh3_128": ub_fingerprint, + "u_xxh3_128": u_fingerprint, + "wavevector_A^-1": wavevector, + } + scientific_context_hash = sha256( + json.dumps( + scientific_components, sort_keys=True, separators=(",", ":") + ).encode("utf-8") + ).hexdigest() + scientific_context = { + **base_context, + "base_context_sha256": base_context_hash, + **scientific_components, + "scientific_context_sha256": scientific_context_hash, + "detector_tiles": [list(detector_tile) for detector_tile in detector_tiles], + } + task_seed = { + "spec": spec.digest, + "frame_range": [start, stop], + "detector_tiles": [ + list(detector_tile) for detector_tile in detector_tiles + ], + "scientific_context": scientific_context_hash, + } + task_id = sha256( + json.dumps(task_seed, sort_keys=True, separators=(",", ":")).encode() + ).hexdigest()[:24] + kernels = { + grid.grid_name: _kernel_for_grid( + spec, + grid, + ub_calculator, + threads=kernel_threads, + memory_budget_bytes=kernel_memory_budget_bytes, + ) + for grid in spec.grids + } + grid_batches: dict[str, list[Mapping[str, np.ndarray]]] = { + grid.grid_name: [] for grid in spec.grids + } + accumulated_bytes = 0 + segment_index = 0 + partitions = [] + + def flush(): + nonlocal accumulated_bytes, segment_index + if not any(grid_batches.values()): + return + for grid in spec.grids: + batch = _reduce_batches(grid_batches[grid.grid_name]) + if batch["chunk_id"].size == 0: + continue + buckets = batch["chunk_id"] // spec.partition_chunk_span + for bucket in np.unique(buckets): + selected = buckets == bucket + subset = { + name: values[selected] for name, values in batch.items() + } + uri = _join_uri( + output_uri, + f"grid={grid.grid_name}", + f"bucket={int(bucket):08d}", + ( + f"task={task_id}-" + f"segment={segment_index:06d}.parquet" + ), + ) + _write_parquet( + subset, + uri, + { + "spec_hash": spec.digest, + "task_id": task_id, + "segment": segment_index, + "grid_name": grid.grid_name, + "bucket": int(bucket), + }, + ) + checksum, size_bytes = _uri_checksum_and_size(uri) + partition = _PartitionFile( + grid_name=grid.grid_name, + bucket=int(bucket), + uri=uri, + rows=int(subset["chunk_id"].size), + checksum=checksum, + size_bytes=size_bytes, + ) + _mark_scratch_verified(partition, verification_cache) + partitions.append(partition) + for batches in grid_batches.values(): + batches.clear() + accumulated_bytes = 0 + segment_index += 1 + + for offset, frame_index in enumerate(range(start, stop)): + image_payload = ( + scan.get_raw_img(frame_index) + if image_payloads is None + else image_payloads[frame_index] + ) + image = np.asarray(image_payload.img) + frame_correction = getattr( + correction_pipeline, "correct_frame", None + ) + corrected_frame = ( + frame_correction(image_payload, image, frame_index) + if callable(frame_correction) + else None + ) + for detector_tile in detector_tiles: + row_start, row_stop, column_start, column_stop = detector_tile + selection = np.s_[ + row_start:row_stop, column_start:column_stop + ] + if corrected_frame is None: + intensity, variance, mask = correction_pipeline( + image_payload, + image[selection], + frame_index, + detector_tile, + ) + else: + intensity, variance, mask = ( + values[selection] for values in corrected_frame + ) + intensity = np.ascontiguousarray(intensity, dtype=np.float64) + variance = np.ascontiguousarray(variance, dtype=np.float64) + mask = np.ascontiguousarray(mask, dtype=bool) + for grid in spec.grids: + batch = kernels[grid.grid_name].accumulate( + intensity, + variance, + mask, + ray_arrays[detector_tile], + np.ascontiguousarray(bounds[offset, 0]), + np.ascontiguousarray(bounds[offset, 1]), + ) + batch_bytes = sum( + np.asarray(values).nbytes for values in batch.values() + ) + if ( + accumulated_bytes + and accumulated_bytes + batch_bytes + > accumulation_budget_bytes + ): + flush() + grid_batches[grid.grid_name].append(batch) + accumulated_bytes += batch_bytes + if accumulated_bytes >= accumulation_budget_bytes: + flush() + if image_progress is not None: + image_progress( + frame_index, + accumulated_bytes, + segment_index, + ) + flush() + scientific_context["accumulation_segments"] = segment_index + return _TaskManifest( + kind="map", + task_id=task_id, + spec_hash=spec.digest, + status="complete", + spec=spec.to_dict(), + frame_range=(start, stop), + detector_tile=None, + partitions=partitions, + metadata=scientific_context, + ) + + +def _map_partition( + spec: _ReconstructionSpec, + scan, + detector, + ub_calculator, + frame_range, + detector_tile, + angle_bounds_rad, + output_uri, + *, + correction_pipeline: Callable, + job_digest: str, + image_payloads: Mapping[int, object] | None = None, + corner_rays: np.ndarray | None = None, + corner_rays_fingerprint: str | None = None, + verification_cache: set[tuple[str, str]] | None = None, + kernel_threads: int | None = None, + kernel_memory_budget_bytes: int | None = None, + accumulation_budget_bytes: int = 64 * 1024 * 1024, + image_progress: Callable[[int, int, int], None] | None = None, +) -> _TaskManifest: + """Map one detector tile through the streaming range implementation.""" + tile = tuple(map(int, detector_tile)) + result = _map_frame_range( + spec, + scan, + detector, + ub_calculator, + frame_range, + (tile,), + angle_bounds_rad, + output_uri, + correction_pipeline=correction_pipeline, + job_digest=job_digest, + image_payloads=image_payloads, + corner_rays=None if corner_rays is None else {tile: corner_rays}, + corner_rays_fingerprints=( + None + if corner_rays_fingerprint is None + else {tile: corner_rays_fingerprint} + ), + verification_cache=verification_cache, + kernel_threads=kernel_threads, + kernel_memory_budget_bytes=kernel_memory_budget_bytes, + accumulation_budget_bytes=accumulation_budget_bytes, + image_progress=image_progress, + ) + result.detector_tile = tile + return result + + +def _reduce_partition( + manifest_set, + output_uri, + *, + verification_cache: set[tuple[str, str]] | None = None, + memory_budget_bytes: int = 1024**3, + workers: int = 1, + progress: Callable[[int, int, str], None] | None = None, + checkpoint_root=None, +) -> _TaskManifest: + """Externally reduce mapping partitions into bounded chunk shards. + + Contiguous shard ranges run concurrently with private forward-only + Parquet readers. The total memory budget is divided between workers. + """ + memory_budget_bytes = max(1, int(memory_budget_bytes)) + requested_workers = max(1, int(workers)) + worker_limit = max( + 1, + memory_budget_bytes // _MIN_REDUCER_WORKER_MEMORY, + ) + configured_workers = min(requested_workers, worker_limit) + worker_memory_budget = max( + 1, + memory_budget_bytes // configured_workers, + ) + manifests = [_read_manifest(value) for value in manifest_set] + if not manifests: + raise ValueError("At least one mapping manifest is required") + if any(manifest.kind != "map" for manifest in manifests): + raise ValueError("reduce_partition accepts mapping manifests only") + spec_hashes = {manifest.spec_hash for manifest in manifests} + if len(spec_hashes) != 1: + raise ValueError("All mapping manifests must use the same specification") + spec = _ReconstructionSpec.from_dict(manifests[0].spec) + base_contexts = { + manifest.metadata.get("base_context_sha256") for manifest in manifests + } + if len(base_contexts) != 1: + raise ValueError( + "Mapping manifests contain different geometry, inputs, or correction " + "identities" + ) + grouped: dict[tuple[str, int], list[_PartitionFile]] = {} + for manifest in manifests: + if manifest.status != "complete": + raise ValueError(f"Mapping task {manifest.task_id} is not complete") + for partition in manifest.partitions: + grouped.setdefault((partition.grid_name, partition.bucket), []).append( + partition + ) + source_tasks = sorted(manifest.task_id for manifest in manifests) + reduce_id = sha256( + (spec.digest + "\0" + "\0".join(source_tasks)).encode("utf-8") + ).hexdigest()[:24] + base_metadata = { + **{ + key: value + for key, value in manifests[0].metadata.items() + if key + not in { + "angle_bounds_sha256", + "corner_rays_sha256", + "angle_bounds_xxh3_128", + "corner_rays_xxh3_128", + "ub_xxh3_128", + "u_xxh3_128", + "scientific_context_sha256", + } + }, + "source_scientific_contexts": { + manifest.task_id: manifest.metadata.get( + "scientific_context_sha256" + ) + for manifest in manifests + }, + } + grid_by_name = {grid.grid_name: grid for grid in spec.grids} + local_span = max( + 256 * 1024, + min( + 4 * 1024 * 1024, + worker_memory_budget // (8 * 48), + ), + ) + plans = [] + for (grid_name, bucket), partitions in sorted(grouped.items()): + grid = grid_by_name[grid_name] + chunk_grid = tuple( + math.ceil(size / chunk) + for size, chunk in zip(grid.shape, grid.chunk_shape) + ) + total_chunks = math.prod(chunk_grid) + chunk_start = bucket * spec.partition_chunk_span + chunk_stop = min( + total_chunks, + (bucket + 1) * spec.partition_chunk_span, + ) + for chunk_id in range(chunk_start, chunk_stop): + coordinates = _chunk_coordinates(chunk_id, grid) + local_shape = tuple( + min(chunk, size - coordinate * chunk) + for coordinate, chunk, size in zip( + coordinates, + grid.chunk_shape, + grid.shape, + ) + ) + local_stop = ( + (local_shape[0] - 1) + * grid.chunk_shape[1] + * grid.chunk_shape[2] + + (local_shape[1] - 1) * grid.chunk_shape[2] + + local_shape[2] + ) + for shard_start in range(0, local_stop, local_span): + plans.append( + ( + grid_name, + bucket, + chunk_id, + shard_start, + min(shard_start + local_span, local_stop), + tuple(sorted(partitions, key=lambda item: item.uri)), + ) + ) + + checkpoint_path = ( + None + if checkpoint_root is None + else Path(checkpoint_root) / f"{reduce_id}.json" + ) + chunks = [] + completed_shards = set() + if checkpoint_path is not None and checkpoint_path.exists(): + try: + checkpoint = _read_manifest(checkpoint_path) + if ( + checkpoint.kind != "reduce" + or checkpoint.task_id != reduce_id + or checkpoint.spec_hash != spec.digest + or checkpoint.source_tasks != source_tasks + ): + raise ValueError("Reduction checkpoint identity mismatch") + for chunk in checkpoint.chunks: + if _verify_scratch_file(chunk, verification_cache): + chunks.append(chunk) + completed_shards.update( + checkpoint.metadata.get("completed_shards", ()) + ) + except (OSError, ValueError, RuntimeError, TypeError): + chunks = [] + completed_shards.clear() + + def shard_key(grid_name, chunk_id, shard_start, shard_stop): + return ( + f"{grid_name}:{chunk_id}:" + f"{shard_start}:{shard_stop}" + ) + + def checkpoint(status): + ordered_chunks = sorted( + chunks, + key=lambda item: ( + item.grid_name, + item.chunk_id, + item.shard_start, + item.shard_stop, + item.uri, + ), + ) + manifest = _TaskManifest( + kind="reduce", + task_id=reduce_id, + spec_hash=spec.digest, + status=status, + spec=spec.to_dict(), + chunks=ordered_chunks, + source_tasks=source_tasks, + metadata={ + **base_metadata, + "local_voxel_shard_span": local_span, + "reducer_worker_capacity": configured_workers, + "reducer_workers_used": actual_workers, + "reducer_memory_bytes_per_worker": worker_memory_budget, + "completed_shards": sorted(completed_shards), + }, + ) + if checkpoint_path is not None: + _write_manifest(manifest, checkpoint_path) + return manifest + + existing = { + shard_key( + chunk.grid_name, + chunk.chunk_id, + chunk.shard_start, + chunk.shard_stop, + ) + for chunk in chunks + } + checkpoint_interval = 64 + total_plans = len(plans) + completed = 0 + active_plans = [] + for plan in plans: + grid_name, _, chunk_id, shard_start, shard_stop, _ = plan + key = shard_key( + grid_name, + chunk_id, + shard_start, + shard_stop, + ) + if key in existing or key in completed_shards: + completed += 1 + if progress is not None: + progress( + completed, + total_plans, + f"Reusing reduced shard {completed}/{total_plans}", + ) + continue + active_plans.append(plan) + + actual_workers = min(configured_workers, len(active_plans)) + worker_memory_budget = max( + 1, + memory_budget_bytes // max(1, actual_workers), + ) + + active_partitions = { + (partition.uri, partition.checksum): partition + for plan in active_plans + for partition in plan[5] + } + partitions_to_verify = [ + partition + for key, partition in active_partitions.items() + if verification_cache is None or key not in verification_cache + ] + if partitions_to_verify: + verification_workers = min( + actual_workers, + len(partitions_to_verify), + ) + with ThreadPoolExecutor( + max_workers=verification_workers + ) as executor: + futures = { + executor.submit( + _verify_scratch_file, + partition, + None, + ): partition + for partition in partitions_to_verify + } + for verified, future in enumerate( + as_completed(futures), + start=1, + ): + partition = futures[future] + future.result() + _mark_scratch_verified(partition, verification_cache) + if progress is not None: + progress( + completed, + total_plans, + ( + "Verified mapping partition " + f"{verified}/{len(partitions_to_verify)} " + f"with {verification_workers} workers" + ), + ) + + result_queue = Queue() + + def reduce_plan_batch(plan_batch): + readers = {} + active_group = None + try: + for ( + grid_name, + bucket, + chunk_id, + shard_start, + shard_stop, + partitions, + ) in plan_batch: + group_key = (grid_name, bucket) + if group_key != active_group: + readers.clear() + active_group = group_key + for partition in partitions: + if partition.uri in readers: + continue + batch_rows = max( + 4096, + min( + 131072, + worker_memory_budget + // (max(1, len(partitions)) * 48 * 4), + ), + ) + readers[partition.uri] = _ParquetRangeReader( + partition.uri, + batch_size=batch_rows, + use_threads=actual_workers == 1, + ) + levels = [] + for partition in partitions: + batch = readers[partition.uri].read( + chunk_id, + shard_start, + shard_stop, + ) + if not batch["chunk_id"].size: + continue + level = 0 + while level < len(levels) and levels[level] is not None: + batch = _merge_sorted_batches(levels[level], batch) + levels[level] = None + level += 1 + if level == len(levels): + levels.append(batch) + else: + levels[level] = batch + reduced = _empty_batch() + for batch in reversed(levels): + if batch is not None: + reduced = _merge_sorted_batches(reduced, batch) + chunk = None + if reduced["chunk_id"].size: + uri = _join_uri( + output_uri, + f"grid={grid_name}", + ( + f"chunk={int(chunk_id):016d}-" + f"shard={shard_start:016d}-{shard_stop:016d}.parquet" + ), + ) + _write_parquet( + reduced, + uri, + { + "spec_hash": spec.digest, + "reduce_id": reduce_id, + "grid_name": grid_name, + "chunk_id": int(chunk_id), + "shard_start": shard_start, + "shard_stop": shard_stop, + "source_bucket": bucket, + }, + ) + checksum, size_bytes = _uri_checksum_and_size(uri) + chunk = _ChunkFile( + grid_name=grid_name, + chunk_id=int(chunk_id), + shard_start=shard_start, + shard_stop=shard_stop, + uri=uri, + rows=int(reduced["chunk_id"].size), + checksum=checksum, + size_bytes=size_bytes, + ) + result_queue.put( + ( + shard_key( + grid_name, + chunk_id, + shard_start, + shard_stop, + ), + chunk, + grid_name, + chunk_id, + shard_start, + shard_stop, + ) + ) + finally: + result_queue.put(None) + + if active_plans: + plans_per_worker, extra = divmod( + len(active_plans), + actual_workers, + ) + plan_batches = [] + start = 0 + for worker_index in range(actual_workers): + stop = start + plans_per_worker + (worker_index < extra) + plan_batches.append(active_plans[start:stop]) + start = stop + + with ThreadPoolExecutor(max_workers=actual_workers) as executor: + futures = [ + executor.submit(reduce_plan_batch, batch) + for batch in plan_batches + ] + finished_workers = 0 + while finished_workers < len(futures): + outcome = result_queue.get() + if outcome is None: + finished_workers += 1 + continue + ( + key, + chunk, + grid_name, + chunk_id, + shard_start, + shard_stop, + ) = outcome + if chunk is not None: + _mark_scratch_verified(chunk, verification_cache) + chunks.append(chunk) + completed_shards.add(key) + completed += 1 + if ( + checkpoint_path is not None + and ( + completed % checkpoint_interval == 0 + or completed == total_plans + ) + ): + checkpoint("running") + if progress is not None: + progress( + completed, + total_plans, + ( + f"Reduced shard {completed}/{total_plans} " + f"with {actual_workers} workers: " + f"{grid_name} chunk {chunk_id}, " + f"{shard_start}:{shard_stop}" + ), + ) + for future in futures: + future.result() + return checkpoint("complete") + + +def _compression_kwargs(name): + if name.startswith("database:"): + from ...app.database import FILTERS + + filter_name = name.partition(":")[2] + if filter_name not in FILTERS: + raise ValueError(f"Unknown orGUI database compression: {filter_name}") + selected = FILTERS[filter_name] + if selected is None: + return {} + if isinstance(selected, Mapping): + return dict(selected) + return {"compression": selected} + normalized = name.lower() + if normalized in {"none", "raw"}: + return {} + if normalized == "bitshuffle-lz4": + try: + import hdf5plugin + + return {"compression": hdf5plugin.Bitshuffle(cname="lz4")} + except ImportError: + return {"compression": "gzip", "compression_opts": 4, "shuffle": True} + if normalized == "lzf": + return {"compression": "lzf", "shuffle": True} + if normalized == "gzip": + return {"compression": "gzip", "compression_opts": 4, "shuffle": True} + raise ValueError(f"Unsupported HDF5 compression: {name}") + + +def _chunk_coordinates(chunk_id, grid): + chunk_grid = tuple( + math.ceil(size / chunk) for size, chunk in zip(grid.shape, grid.chunk_shape) + ) + chunk_x, remainder = divmod(chunk_id, chunk_grid[1] * chunk_grid[2]) + chunk_y, chunk_z = divmod(remainder, chunk_grid[2]) + return chunk_x, chunk_y, chunk_z + + +def _write_chunk( + group, + grid, + chunk_files, + verification_cache=None, +): + chunk_files = sorted(chunk_files, key=lambda item: item.shard_start) + if not chunk_files: + return + chunk_id = chunk_files[0].chunk_id + if any(chunk.chunk_id != chunk_id for chunk in chunk_files): + raise ValueError("Reduced shard group contains multiple chunks") + coordinates = _chunk_coordinates(chunk_id, grid) + starts = tuple( + coordinate * chunk + for coordinate, chunk in zip(coordinates, grid.chunk_shape) + ) + stops = tuple( + min(start + chunk, size) + for start, chunk, size in zip(starts, grid.chunk_shape, grid.shape) + ) + local_shape = tuple(stop - start for start, stop in zip(starts, stops)) + weighted_intensity = np.zeros(local_shape, dtype=np.float64) + weighted_variance = np.zeros(local_shape, dtype=np.float64) + weight = np.zeros(local_shape, dtype=np.float64) + contributors = np.zeros(local_shape, dtype=np.uint64) + for chunk_file in chunk_files: + _verify_scratch_file(chunk_file, verification_cache) + batch = _read_parquet(chunk_file.uri) + if batch["chunk_id"].size and np.any( + batch["chunk_id"] != np.uint64(chunk_id) + ): + raise ValueError(f"{chunk_file.uri} contains more than one chunk") + local = batch["local_voxel_id"].astype(np.uint64, copy=False) + if local.size and ( + np.any(local < np.uint64(chunk_file.shard_start)) + or np.any(local >= np.uint64(chunk_file.shard_stop)) + ): + raise ValueError( + f"{chunk_file.uri} contains records outside its shard" + ) + local_x, remainder = np.divmod( + local, np.uint64(grid.chunk_shape[1] * grid.chunk_shape[2]) + ) + local_y, local_z = np.divmod( + remainder, + np.uint64(grid.chunk_shape[2]), + ) + valid = ( + (local_x < local_shape[0]) + & (local_y < local_shape[1]) + & (local_z < local_shape[2]) + ) + if not np.all(valid): + raise ValueError( + f"{chunk_file.uri} contains an invalid local voxel ID" + ) + index = ( + local_x.astype(int), + local_y.astype(int), + local_z.astype(int), + ) + weighted_intensity[index] = batch["weighted_intensity"] + weighted_variance[index] = batch["weighted_variance"] + weight[index] = batch["weight"] + contributors[index] = batch["contributors"] + populated = weight > 0 + weighted_intensity[~populated] = np.nan + weighted_variance[~populated] = np.nan + np.divide( + weighted_intensity, + weight, + out=weighted_intensity, + where=populated, + ) + np.divide( + weighted_variance, + weight, + out=weighted_variance, + where=populated, + ) + np.divide( + weighted_variance, + weight, + out=weighted_variance, + where=populated, + ) + selection = tuple(slice(start, stop) for start, stop in zip(starts, stops)) + group["intensity"][selection] = weighted_intensity + group["variance"][selection] = weighted_variance + group["weight"][selection] = weight + group["contributors"][selection] = contributors + + +def _finalize_reconstruction( + manifest_set, + output_path, + *, + provenance: Mapping[str, Any] | None = None, + config=None, + chunk_progress: Callable | None = None, + verification_cache: set[tuple[str, str]] | None = None, +): + """Create a conventional HDF5 reconstruction from reduced chunk files.""" + manifests = [_read_manifest(value) for value in manifest_set] + if not manifests: + raise ValueError("At least one reduction manifest is required") + if any(manifest.kind != "reduce" for manifest in manifests): + raise ValueError("finalize_reconstruction accepts reduction manifests only") + spec_hashes = {manifest.spec_hash for manifest in manifests} + if len(spec_hashes) != 1: + raise ValueError("Reduction manifests use different specifications") + spec = _ReconstructionSpec.from_dict(manifests[0].spec) + grid_by_name = {grid.grid_name: grid for grid in spec.grids} + chunk_files = {} + for manifest in manifests: + if manifest.status != "complete": + raise ValueError(f"Reduction task {manifest.task_id} is not complete") + for chunk in manifest.chunks: + key = (chunk.grid_name, chunk.chunk_id) + chunk_files.setdefault(key, []).append(chunk) + for key, shards in chunk_files.items(): + ordered = sorted(shards, key=lambda item: item.shard_start) + for previous, current in zip(ordered, ordered[1:]): + if current.shard_start < previous.shard_stop: + raise ValueError(f"Overlapping reduced shards for chunk {key}") + output_path = Path(output_path).absolute() + output_path.parent.mkdir(parents=True, exist_ok=True) + temporary_path = output_path.with_name(output_path.name + ".tmp") + compression = _compression_kwargs(spec.compression) + with h5py.File(temporary_path, "w") as h5file: + entry = h5file.create_group("entry") + entry.attrs["NX_class"] = "NXentry" + process = entry.create_group("reconstruction") + process.attrs["NX_class"] = "NXprocess" + process.attrs["program"] = "orGUI" + process.attrs["spec_sha256"] = spec.digest + process.create_dataset( + "configuration_json", + data=json.dumps(spec.to_dict(), sort_keys=True), + ) + if provenance: + process.create_dataset( + "provenance_json", + data=json.dumps(dict(provenance), sort_keys=True, default=str), + ) + process.create_dataset( + "scientific_context_json", + data=json.dumps(manifests[0].metadata, sort_keys=True, default=str), + ) + process.create_dataset( + "marginal_variance_warning", + data=( + "Adaptive footprint splitting and pixel repair can introduce " + "cross-voxel or spatial covariance. Variance datasets contain " + "marginal variances only." + ), + ) + if config is not None: + from ...app.config_data import ConfigHandler + + ConfigHandler(create_dataset_args=compression).write_scan_config( + entry, config, source="reciprocal_space_reconstruction" + ) + results = process.create_group("results") + groups = {} + for grid in spec.grids: + group = results.create_group(grid.grid_name) + group.attrs["NX_class"] = "NXdata" + group.attrs["signal"] = "intensity" + group.attrs["axes"] = np.asarray( + ["h", "k", "l"] if grid.frame == "hkl" else ["qx", "qy", "qz"], + dtype=h5py.string_dtype(), + ) + group.attrs["coordinate_frame"] = grid.frame + group.attrs["units"] = "r.l.u." if grid.frame == "hkl" else "Angstrom^-1" + axis_names = ("h", "k", "l") if grid.frame == "hkl" else ("qx", "qy", "qz") + for axis, axis_name in enumerate(axis_names): + values = grid.minimum[axis] + ( + np.arange(grid.shape[axis], dtype=np.float64) + 0.5 + ) * grid.step[axis] + dataset = group.create_dataset(axis_name, data=values) + dataset.attrs["units"] = ( + "r.l.u." if grid.frame == "hkl" else "Angstrom^-1" + ) + dataset.attrs["long_name"] = f"{axis_name} voxel center" + hdf5_chunks = tuple( + min(size, chunk) + for size, chunk in zip(grid.shape, grid.chunk_shape) + ) + group.create_dataset( + "intensity", + shape=grid.shape, + dtype=np.float64, + chunks=hdf5_chunks, + fillvalue=np.nan, + **compression, + ) + group.create_dataset( + "variance", + shape=grid.shape, + dtype=np.float64, + chunks=hdf5_chunks, + fillvalue=np.nan, + **compression, + ) + group.create_dataset( + "weight", + shape=grid.shape, + dtype=np.float64, + chunks=hdf5_chunks, + fillvalue=0.0, + **compression, + ) + group.create_dataset( + "contributors", + shape=grid.shape, + dtype=np.uint64, + chunks=hdf5_chunks, + fillvalue=0, + **compression, + ) + groups[grid.grid_name] = group + for written, ((grid_name, _), chunk_shards) in enumerate( + sorted(chunk_files.items()), start=1 + ): + if grid_name not in grid_by_name: + raise ValueError(f"Unknown grid in chunk manifest: {grid_name}") + _write_chunk( + groups[grid_name], + grid_by_name[grid_name], + chunk_shards, + verification_cache, + ) + if chunk_progress is not None: + chunk_progress(written, len(chunk_files)) + h5file.flush() + temporary_path.replace(output_path) + with h5py.File(output_path, "r") as h5file: + results = h5file["entry/reconstruction/results"] + for grid in spec.grids: + group = results[grid.grid_name] + for dataset in ("intensity", "variance", "weight", "contributors"): + if group[dataset].shape != grid.shape: + raise OSError( + f"Final reconstruction dataset has invalid shape: " + f"{grid.grid_name}/{dataset}" + ) + digest = sha256() + with output_path.open("rb") as stream: + while block := stream.read(8 * 1024 * 1024): + digest.update(block) + result = { + "path": str(output_path), + "sha256": digest.hexdigest(), + "bytes": output_path.stat().st_size, + "chunks_written": len(chunk_files), + } + return result + + +__all__ = [] diff --git a/orgui/datautils/xrayutils/test/test_reconstruction.py b/orgui/datautils/xrayutils/test/test_reconstruction.py new file mode 100644 index 0000000..8a15c5b --- /dev/null +++ b/orgui/datautils/xrayutils/test/test_reconstruction.py @@ -0,0 +1,908 @@ +"""Regression tests for native reciprocal-space reconstruction.""" + +from __future__ import annotations + +import json +import threading + +import h5py +import numpy as np +import pytest + +from orgui.backend.scans import h5_Image +from orgui.app.config_data import CorrectionState +from orgui.app.mask_config import repair_intensity_variance +from orgui.datautils.xrayutils import HKLVlieg +import orgui.datautils.xrayutils.reconstruction as reconstruction +from orgui.reconstruction_job import _correction_pipeline, derive_grid +from orgui.datautils.xrayutils.reconstruction import ( + _GridSpec, + _ReconstructionSpec, + _finalize_reconstruction, + _map_frame_range, + _map_partition, + _merge_sorted_batches, + _read_parquet, + _reduce_batches, + _reduce_partition, + _write_manifest, +) + + +native = pytest.importorskip( + "orgui.datautils.xrayutils._reciprocal_reconstruction_cpp" +) + + +def test_native_merge_sorted_batches(): + """Native merging preserves order and reduces matching voxel keys.""" + + def batch(local, intensity, contributors): + size = len(local) + return { + "chunk_id": np.zeros(size, dtype=np.uint64), + "local_voxel_id": np.asarray(local, dtype=np.uint64), + "weighted_intensity": np.asarray(intensity, dtype=np.float64), + "weighted_variance": np.ones(size, dtype=np.float64), + "weight": np.ones(size, dtype=np.float64), + "contributors": np.asarray(contributors, dtype=np.uint64), + } + + merged = _merge_sorted_batches( + batch([1, 3], [10.0, 30.0], [1, 2]), + batch([2, 3], [20.0, 4.0], [3, 4]), + ) + + np.testing.assert_array_equal(merged["local_voxel_id"], [1, 2, 3]) + np.testing.assert_array_equal( + merged["weighted_intensity"], [10.0, 20.0, 34.0] + ) + np.testing.assert_array_equal(merged["weighted_variance"], [1.0, 1.0, 2.0]) + np.testing.assert_array_equal(merged["weight"], [1.0, 1.0, 2.0]) + np.testing.assert_array_equal(merged["contributors"], [1, 3, 6]) + + +def test_grid_validates_effective_hdf5_chunk_bytes(): + """Requested chunks are clamped to grid shape before HDF5 validation.""" + valid = _GridSpec( + minimum=(0.0, 0.0, 0.0), + maximum=(514.0, 514.0, 514.0), + step=(1.0, 1.0, 1.0), + frame="hkl", + chunk_shape=(1024, 1024, 1024), + ) + assert valid.shape == (514, 514, 514) + + with pytest.raises(ValueError, match="smaller than 4 GiB"): + _GridSpec( + minimum=(0.0, 0.0, 0.0), + maximum=(1024.0, 1024.0, 1024.0), + step=(1.0, 1.0, 1.0), + frame="hkl", + chunk_shape=(1024, 1024, 1024), + ) + + +def _kernel(frame="lab", *, threads=1, max_depth=2, ub=None): + if ub is None: + ub = np.eye(3) + return native.ReconstructionKernel( + np.array([-20.0, -20.0, -20.0]), + np.array([0.01, 0.01, 0.01]), + np.array([4000, 4000, 4000], dtype=np.int64), + np.array([64, 64, 64], dtype=np.int64), + frame, + 1.0, + np.linalg.inv(ub), + np.eye(3), + max_depth, + threads, + 16, + 1024 * 1024, + ) + + +def _constant_rays(rows, columns): + rays = np.zeros((rows + 1, columns + 1, 3), dtype=np.float64) + rays[..., 1] = 1.0 + return rays + + +def test_native_average_variance_for_two_pixels(): + kernel = _kernel() + result = kernel.accumulate( + np.array([[10.0, 20.0]]), + np.array([[10.0, 20.0]]), + np.zeros((1, 2), dtype=bool), + _constant_rays(1, 2), + np.zeros(4), + np.zeros(4), + ) + + assert result["chunk_id"].size == 1 + assert result["weighted_intensity"][0] == 30.0 + assert result["weighted_variance"][0] == 30.0 + assert result["weight"][0] == 2.0 + assert result["contributors"][0] == 2 + assert result["weighted_intensity"][0] / result["weight"][0] == 15.0 + assert result["weighted_variance"][0] / result["weight"][0] ** 2 == 7.5 + + +def test_arbitrary_rectangular_tiles_match_full_detector(): + """Shared corner rays make arbitrary tile boundaries numerically exact.""" + rows, columns = 3, 5 + yy, xx = np.meshgrid( + np.arange(rows + 1), + np.arange(columns + 1), + indexing="ij", + ) + rays = np.stack( + ( + (xx - columns / 2) * 0.02, + np.ones_like(xx), + (yy - rows / 2) * 0.02, + ), + axis=-1, + ) + rays /= np.linalg.norm(rays, axis=-1, keepdims=True) + intensity = np.arange(1, rows * columns + 1, dtype=np.float64).reshape( + rows, columns + ) + variance = intensity + 0.5 + mask = np.zeros((rows, columns), dtype=bool) + kernel = _kernel(max_depth=2) + full = kernel.accumulate( + intensity, + variance, + mask, + rays, + np.zeros(4), + np.zeros(4), + ) + tiles = ( + (0, 1, 0, 2), + (0, 1, 2, 5), + (1, 3, 0, 2), + (1, 3, 2, 5), + ) + tiled = _reduce_batches( + kernel.accumulate( + intensity[row_start:row_stop, column_start:column_stop], + variance[row_start:row_stop, column_start:column_stop], + mask[row_start:row_stop, column_start:column_stop], + rays[ + row_start : row_stop + 1, + column_start : column_stop + 1, + ], + np.zeros(4), + np.zeros(4), + ) + for row_start, row_stop, column_start, column_stop in tiles + ) + + for name in ("chunk_id", "local_voxel_id", "contributors"): + np.testing.assert_array_equal(tiled[name], full[name]) + for name in ( + "weighted_intensity", + "weighted_variance", + "weight", + ): + np.testing.assert_allclose(tiled[name], full[name], rtol=2e-15) + assert tiled["weight"].sum() == pytest.approx(rows * columns) + + +def test_native_xxh3_128_matches_known_vector(): + assert ( + native.xxh3_128(np.empty(0, dtype=np.uint8)) + == "99aa06d3014798d86001c324468d497f" + ) + + +def test_shared_correction_pipeline_propagates_factor_uncertainty(): + scan = _FakeScan([10.0]) + scan.exposure_time = np.array([2.0]) + scan.exposure_time_variance = np.array([0.25]) + config = type( + "Config", + (), + { + "corrections": CorrectionState( + use_background=True, normalize_exposure=True + ), + "detector": object(), + }, + )() + provenance = {} + correct = _correction_pipeline( + config, + scan, + { + "background": np.array([[2.0]]), + "background_variance": np.array([[1.0]]), + }, + provenance, + ) + + intensity, variance, mask = correct( + h5_Image(np.array([[10.0]])), + np.array([[10.0]]), + 0, + (0, 1, 0, 1), + ) + + assert intensity[0, 0] == 4.0 + assert variance[0, 0] == 3.75 + assert not mask[0, 0] + assert provenance["factor_uncertainty"]["exposure"] == "propagated" + + +def test_shared_correction_repairs_across_detector_tile_boundaries(): + scan = _FakeScan([0.0]) + config = type( + "Config", + (), + { + "corrections": CorrectionState( + use_mask=True, + repair_masked_pixels=True, + normalize_exposure=False, + ), + "detector": object(), + }, + )() + image = np.arange(49, dtype=np.float64).reshape(7, 7) + mask = np.zeros((7, 7), dtype=bool) + mask[3, 3] = True + payload = h5_Image(image) + correct = _correction_pipeline( + config, + scan, + { + "mask": mask, + "repair": { + "max_component_pixels": 4, + "max_span": 3, + "radius": 2, + "min_valid_neighbors": 6, + "row_gaps": np.empty((0, 2), dtype=np.int32), + "column_gaps": np.empty((0, 2), dtype=np.int32), + }, + }, + {}, + ) + + full = correct.correct_frame(payload, image, 0) + left = correct(payload, image[:, :4], 0, (0, 7, 0, 4)) + right = correct(payload, image[:, 4:], 0, (0, 7, 4, 7)) + + for index in range(3): + combined = np.concatenate((left[index], right[index]), axis=1) + np.testing.assert_array_equal(combined, full[index]) + assert not full[2][3, 3] + + +def test_pixel_repair_propagates_interpolation_weights(): + intensity = np.arange(9, dtype=np.float64).reshape(3, 3) + variance = np.ones((3, 3), dtype=np.float64) + mask = np.zeros((3, 3), dtype=bool) + mask[1, 1] = True + + _, repaired_variance, remaining, repaired = repair_intensity_variance( + intensity, + variance, + mask, + max_component_pixels=2, + max_span=2, + radius=2, + min_valid_neighbors=4, + ) + + assert repaired[1, 1] + assert not remaining[1, 1] + assert repaired_variance[1, 1] == 0.5 + + +def test_native_coordinate_frames_match_vlieg(): + lattice = HKLVlieg.Lattice([3.9, 4.1, 5.2], [90.0, 95.0, 110.0]) + calculator = HKLVlieg.UBCalculator(lattice, 70.0) + u = HKLVlieg.Rotation.from_euler("xyz", [0.17, -0.08, 0.11]).as_matrix() + calculator.setU(u) + position = np.deg2rad([0.3, 7.0, 3.0, -4.0, 2.0, 5.0]) + primary = HKLVlieg.primBeamAngles(position) + gamma_p = primary[2] + delta_p = primary[1] + ray = np.array( + [ + np.sin(delta_p) * np.cos(gamma_p), + np.cos(delta_p) * np.cos(gamma_p), + np.sin(gamma_p), + ] + ) + rays = np.broadcast_to(ray, (2, 2, 3)).copy() + sample_angles = position[[0, 3, 4, 5]] + q_alpha = np.asarray( + HKLVlieg.VliegAngles(calculator).QAlpha(*position[:3]) + ).reshape(3) + matrices = HKLVlieg.createVliegMatrices(position) + q_omega = matrices[3].T @ q_alpha + q_chi = matrices[4].T @ q_omega + q_phi = matrices[5].T @ q_chi + q_lab = matrices[0] @ q_alpha + expected = { + "lab": q_lab, + "alpha": q_alpha, + "omega": q_omega, + "chi": q_chi, + "phi": q_phi, + "crystal": np.linalg.inv(calculator.getU()) @ q_phi, + "hkl": np.linalg.inv(calculator.getUB()) @ q_phi, + } + + for frame, coordinates in expected.items(): + kernel = native.ReconstructionKernel( + np.full(3, -100.0), + np.ones(3), + np.full(3, 200, dtype=np.int64), + np.full(3, 16, dtype=np.int64), + frame, + calculator.getK(), + np.linalg.inv(calculator.getUB()), + np.linalg.inv(calculator.getU()), + ) + actual = kernel.coordinate( + rays, sample_angles, sample_angles, 0, 0 + ) + assert np.allclose(actual, coordinates, rtol=1e-13, atol=1e-13), frame + + +def test_native_coordinate_interpolates_detector_corner_rays(): + rays = np.array( + [ + [[-0.2, 0.97, -0.1], [0.3, 0.93, -0.15]], + [[-0.25, 0.94, 0.2], [0.35, 0.9, 0.25]], + ], + dtype=np.float64, + ) + rays /= np.linalg.norm(rays, axis=2, keepdims=True) + u = 0.37 + v = 0.61 + interpolated = ( + (1.0 - u) * (1.0 - v) * rays[0, 0] + + u * (1.0 - v) * rays[1, 0] + + (1.0 - u) * v * rays[0, 1] + + u * v * rays[1, 1] + ) + interpolated /= np.linalg.norm(interpolated) + expected = interpolated - np.array([0.0, 1.0, 0.0]) + + actual = _kernel(frame="lab").coordinate( + rays, + np.zeros(4), + np.array([0.1, -0.2, 0.3, -0.4]), + 0, + 0, + u, + v, + 0.73, + ) + + assert np.allclose(actual, expected, rtol=1e-14, atol=1e-14) + + +def test_native_results_do_not_depend_on_thread_count(): + rng = np.random.default_rng(12) + intensity = rng.random((8, 8)) + variance = intensity.copy() + mask = np.zeros_like(intensity, dtype=bool) + rays = _constant_rays(8, 8) + serial = _kernel(threads=1).accumulate( + intensity, variance, mask, rays, np.zeros(4), np.zeros(4) + ) + threaded = _kernel(threads=4).accumulate( + intensity, variance, mask, rays, np.zeros(4), np.zeros(4) + ) + + for name in serial: + assert np.array_equal(serial[name], threaded[name]), name + + +def test_center_only_profiles_one_coordinate_per_valid_pixel(): + intensity = np.ones((2, 3), dtype=np.float64) + variance = np.ones_like(intensity) + mask = np.zeros_like(intensity, dtype=bool) + mask[0, 1] = True + + result = _kernel(max_depth=0).accumulate( + intensity, + variance, + mask, + _constant_rays(2, 3), + np.zeros(4), + np.zeros(4), + profile=True, + ) + + profile = result.pop("_profile") + assert profile["valid_pixels"] == 5 + assert profile["coordinate_evaluations"] == 5 + assert profile["maximum_weights_per_pixel"] == 1 + + +def test_footprint_split_conserves_weight_and_pixel_variance(): + rays = np.empty((2, 2, 3), dtype=np.float64) + for column, x_value in enumerate((-0.2, 0.2)): + rays[:, column, 0] = x_value + rays[:, column, 1] = np.sqrt(1.0 - x_value**2) + rays[:, column, 2] = 0.0 + result = _kernel(max_depth=2).accumulate( + np.array([[10.0]]), + np.array([[10.0]]), + np.array([[False]]), + rays, + np.zeros(4), + np.zeros(4), + ) + + assert result["weight"].size > 1 + assert np.sum(result["weight"]) == 1.0 + assert np.sum(result["weighted_intensity"]) == 10.0 + assert np.allclose( + result["weighted_variance"], result["weight"] ** 2 * 10.0 + ) + assert np.all(result["contributors"] == 1) + + +class _FakeDetector: + detector = type("Detector", (), {"shape": (1, 1)})() + + def __init__(self): + self.ray_calculations = 0 + + def primBeamPoints(self, rows, columns): + self.ray_calculations += 1 + return np.zeros_like(rows), np.zeros_like(columns) + + +class _FakeUB: + def getUB(self): + return np.eye(3) + + def getU(self): + return np.eye(3) + + def getK(self): + return 1.0 + + +class _FakeScan: + shape = (1, 1) + + def __init__(self, values): + self.values = values + self.image_loads = 0 + + def __len__(self): + return len(self.values) + + def get_raw_img(self, index): + self.image_loads += 1 + return h5_Image(np.array([[self.values[index]]], dtype=np.float64)) + + def exposure_angle_bounds(self, config, fallback="stationary"): + return np.zeros((len(self), 2, 4), dtype=np.float64) + + +def test_derive_grid_uses_native_corner_ray_interface(): + config = type( + "Config", + (), + { + "detector": _FakeDetector(), + "ub_calculator": _FakeUB(), + }, + )() + scan = _FakeScan([1.0, 2.0]) + + for frame in ("lab", "alpha", "omega", "chi", "phi", "crystal", "hkl"): + grid = derive_grid(config, scan, frame=frame) + assert grid.frame == frame + assert np.all(np.isfinite(grid.minimum)) + assert np.all(np.isfinite(grid.maximum)) + assert np.all(np.asarray(grid.maximum) > np.asarray(grid.minimum)) + + +def test_map_partition_accepts_local_image_and_ray_caches(tmp_path): + pytest.importorskip("pyarrow") + grid = _GridSpec( + minimum=(-1.0, -1.0, -1.0), + maximum=(1.0, 1.0, 1.0), + step=(1.0, 1.0, 1.0), + frame="lab", + chunk_shape=(2, 2, 2), + ) + spec = _ReconstructionSpec(grids=(grid,), max_depth=0) + scan = _FakeScan([10.0]) + detector = _FakeDetector() + payload = h5_Image(np.array([[10.0]], dtype=np.float64)) + + _map_partition( + spec, + scan, + detector, + _FakeUB(), + (0, 1), + (0, 1, 0, 1), + np.zeros((1, 2, 4), dtype=np.float64), + tmp_path / "map", + correction_pipeline=lambda payload, raw, frame, tile: ( + raw.astype(np.float64), + np.maximum(raw, 0).astype(np.float64), + np.zeros(raw.shape, dtype=bool), + ), + job_digest="cache-test", + image_payloads={0: payload}, + corner_rays=_constant_rays(1, 1), + ) + + assert scan.image_loads == 0 + assert detector.ray_calculations == 0 + + +def test_map_frame_range_corrects_once_before_tiling(tmp_path): + pytest.importorskip("pyarrow") + + class Detector: + detector = type("PyfaiDetector", (), {"shape": (1, 2)})() + + def primBeamPoints(self, rows, columns): + return np.zeros_like(rows), np.zeros_like(columns) + + class Scan: + def __len__(self): + return 1 + + def get_raw_img(self, index): + return h5_Image(np.array([[10.0, 20.0]])) + + calls = {"frame": 0, "tile": 0} + + def correction(payload, raw, frame, tile): + calls["tile"] += 1 + raise AssertionError("tile correction must not be used") + + def correct_frame(payload, raw, frame): + calls["frame"] += 1 + return ( + raw.astype(np.float64), + np.maximum(raw, 0.0).astype(np.float64), + np.zeros(raw.shape, dtype=bool), + ) + + correction.correct_frame = correct_frame + grid = _GridSpec( + minimum=(-1.0, -1.0, -1.0), + maximum=(1.0, 1.0, 1.0), + step=(1.0, 1.0, 1.0), + frame="lab", + chunk_shape=(2, 2, 2), + ) + spec = _ReconstructionSpec(grids=(grid,), max_depth=0) + tiles = ((0, 1, 0, 1), (0, 1, 1, 2)) + + _map_frame_range( + spec, + Scan(), + Detector(), + _FakeUB(), + (0, 1), + tiles, + np.zeros((1, 2, 4), dtype=np.float64), + tmp_path / "map", + correction_pipeline=correction, + job_digest="full-frame-correction", + corner_rays={ + tiles[0]: _constant_rays(1, 1), + tiles[1]: _constant_rays(1, 1), + }, + ) + + assert calls == {"frame": 1, "tile": 0} + + +def test_streaming_accumulator_flushes_by_retained_bytes(tmp_path): + """A larger byte ceiling combines images before Parquet serialization.""" + pytest.importorskip("pyarrow") + grid = _GridSpec( + minimum=(-1.0, -1.0, -1.0), + maximum=(1.0, 1.0, 1.0), + step=(1.0, 1.0, 1.0), + frame="lab", + chunk_shape=(2, 2, 2), + ) + spec = _ReconstructionSpec(grids=(grid,), max_depth=0) + bounds = np.zeros((3, 2, 4), dtype=np.float64) + tile = (0, 1, 0, 1) + + def correction(payload, raw, frame, detector_tile): + return ( + raw.astype(np.float64), + np.maximum(raw, 0).astype(np.float64), + np.zeros(raw.shape, dtype=bool), + ) + + small_scan = _FakeScan([10.0, 20.0, 30.0]) + image_updates = [] + small = _map_frame_range( + spec, + small_scan, + _FakeDetector(), + _FakeUB(), + (0, 3), + (tile,), + bounds, + tmp_path / "small", + correction_pipeline=correction, + job_digest="small-accumulator", + corner_rays={tile: _constant_rays(1, 1)}, + accumulation_budget_bytes=1, + image_progress=lambda frame, retained, segments: image_updates.append( + (frame, retained, segments) + ), + ) + large_scan = _FakeScan([10.0, 20.0, 30.0]) + large = _map_frame_range( + spec, + large_scan, + _FakeDetector(), + _FakeUB(), + (0, 3), + (tile,), + bounds, + tmp_path / "large", + correction_pipeline=correction, + job_digest="large-accumulator", + corner_rays={tile: _constant_rays(1, 1)}, + accumulation_budget_bytes=1024**2, + ) + + assert small_scan.image_loads == 3 + assert large_scan.image_loads == 3 + assert [update[0] for update in image_updates] == [0, 1, 2] + assert small.metadata["accumulation_segments"] == 3 + assert large.metadata["accumulation_segments"] == 1 + assert len(small.partitions) == 3 + assert len(large.partitions) == 1 + small_batch = _reduce_batches( + _read_parquet(partition.uri) for partition in small.partitions + ) + large_batch = _reduce_batches( + _read_parquet(partition.uri) for partition in large.partitions + ) + for name in small_batch: + np.testing.assert_array_equal(small_batch[name], large_batch[name]) + + +def test_parquet_reduce_and_hdf5_finalize(tmp_path, monkeypatch): + pytest.importorskip("pyarrow") + checksum_calls = [] + original_checksum = reconstruction._uri_checksum_and_size + + def counted_checksum(uri): + checksum_calls.append(str(uri)) + return original_checksum(uri) + + monkeypatch.setattr( + reconstruction, "_uri_checksum_and_size", counted_checksum + ) + verification_cache = set() + grid = _GridSpec( + minimum=(-1.0, -1.0, -1.0), + maximum=(1.0, 1.0, 1.0), + step=(1.0, 1.0, 1.0), + frame="lab", + chunk_shape=(2, 2, 2), + ) + spec = _ReconstructionSpec( + grids=(grid,), + max_depth=1, + threads=2, + compression="gzip", + ) + scan = _FakeScan([10.0, 20.0]) + bounds = np.zeros((1, 2, 4), dtype=np.float64) + map_manifests = [] + for frame in range(2): + manifest = _map_partition( + spec, + scan, + _FakeDetector(), + _FakeUB(), + (frame, frame + 1), + (0, 1, 0, 1), + bounds, + tmp_path / "map", + correction_pipeline=lambda payload, raw, frame, tile: ( + raw.astype(np.float64), + np.maximum(raw, 0).astype(np.float64), + np.zeros(raw.shape, dtype=bool), + ), + job_digest="test-job", + verification_cache=verification_cache, + ) + path = tmp_path / f"map-{frame}.json" + _write_manifest(manifest, path) + map_manifests.append(path) + + reduced = _reduce_partition( + map_manifests, + tmp_path / "reduced", + verification_cache=verification_cache, + ) + reduced_path = tmp_path / "reduced.json" + _write_manifest(reduced, reduced_path) + output = tmp_path / "result.h5" + result = _finalize_reconstruction( + [reduced_path], + output, + verification_cache=verification_cache, + ) + + assert result["chunks_written"] == 1 + assert len(checksum_calls) == 3 + assert all( + partition.size_bytes + for path in map_manifests + for partition in reconstruction._read_manifest(path).partitions + ) + assert all(chunk.size_bytes for chunk in reduced.chunks) + with h5py.File(output, "r") as h5file: + group = h5file["entry/reconstruction/results/q_lab"] + assert group.attrs["NX_class"] == "NXdata" + assert group["intensity"][1, 1, 1] == 15.0 + assert group["variance"][1, 1, 1] == 7.5 + assert group["weight"][1, 1, 1] == 2.0 + assert group["contributors"][1, 1, 1] == 2 + assert np.isnan(group["intensity"][0, 0, 0]) + saved = json.loads(h5file["entry/reconstruction/configuration_json"][()]) + assert saved["grids"][0]["frame"] == "lab" + + +def test_reduce_shards_large_chunk_and_reuses_checkpoint(tmp_path, monkeypatch): + """Large chunks reduce by bounded local ranges and resume by shard.""" + pytest.importorskip("pyarrow") + grid = _GridSpec( + minimum=(0.0, 0.0, 0.0), + maximum=(1.0, 1.0, 300000.0), + step=(1.0, 1.0, 1.0), + frame="lab", + chunk_shape=(1, 1, 300000), + ) + spec = _ReconstructionSpec( + grids=(grid,), + max_depth=0, + threads=1, + compression="gzip", + ) + manifest_paths = [] + for task_index, intensities in enumerate( + ([2.0, 4.0, 6.0], [3.0, 5.0, 7.0]) + ): + batch = { + "chunk_id": np.zeros(3, dtype=np.uint64), + "local_voxel_id": np.array( + [1, 260000, 299999], dtype=np.uint64 + ), + "weighted_intensity": np.asarray( + intensities, dtype=np.float64 + ), + "weighted_variance": np.ones(3, dtype=np.float64), + "weight": np.ones(3, dtype=np.float64), + "contributors": np.ones(3, dtype=np.uint64), + } + partition_path = tmp_path / f"map-{task_index}.parquet" + reconstruction._write_parquet(batch, partition_path, {}) + checksum, size_bytes = reconstruction._uri_checksum_and_size( + partition_path + ) + manifest = reconstruction._TaskManifest( + kind="map", + task_id=f"map-{task_index}", + spec_hash=spec.digest, + status="complete", + spec=spec.to_dict(), + partitions=[ + reconstruction._PartitionFile( + grid_name=grid.grid_name, + bucket=0, + uri=str(partition_path), + rows=3, + checksum=checksum, + size_bytes=size_bytes, + ) + ], + metadata={"base_context_sha256": "same-context"}, + ) + manifest_path = tmp_path / f"map-{task_index}.json" + _write_manifest(manifest, manifest_path) + manifest_paths.append(manifest_path) + + progress = [] + verification_cache = set() + reducer_threads = set() + reducer_lock = threading.Lock() + reducer_barrier = threading.Barrier(2) + original_read = reconstruction._ParquetRangeReader.read + + def observed_read(reader, *args): + thread_id = threading.get_ident() + with reducer_lock: + first_read = thread_id not in reducer_threads + reducer_threads.add(thread_id) + if first_read: + reducer_barrier.wait(timeout=5) + return original_read(reader, *args) + + monkeypatch.setattr( + reconstruction._ParquetRangeReader, + "read", + observed_read, + ) + reduced = _reduce_partition( + manifest_paths, + tmp_path / "reduced", + verification_cache=verification_cache, + memory_budget_bytes=128 * 1024**2, + workers=2, + progress=lambda current, total, message: progress.append( + (current, total, message) + ), + checkpoint_root=tmp_path / "manifests", + ) + + assert reduced.status == "complete" + assert len(reduced.chunks) == 2 + assert {chunk.shard_start for chunk in reduced.chunks} == {0, 262144} + assert reduced.metadata["reducer_worker_capacity"] == 2 + assert reduced.metadata["reducer_workers_used"] == 2 + assert ( + reduced.metadata["reducer_memory_bytes_per_worker"] + == 64 * 1024**2 + ) + assert len(reducer_threads) == 2 + assert progress[-1][:2] == (2, 2) + + def unexpected_read(*args, **kwargs): + raise AssertionError("completed shards must not read source Parquet") + + monkeypatch.setattr( + reconstruction._ParquetRangeReader, "read", unexpected_read + ) + resumed = _reduce_partition( + manifest_paths, + tmp_path / "reduced", + verification_cache=verification_cache, + memory_budget_bytes=128 * 1024**2, + workers=2, + checkpoint_root=tmp_path / "manifests", + ) + assert resumed.status == "complete" + assert len(resumed.chunks) == 2 + + reduced_path = tmp_path / "reduced.json" + _write_manifest(resumed, reduced_path) + output = tmp_path / "sharded-result.h5" + _finalize_reconstruction( + [reduced_path], + output, + verification_cache=verification_cache, + ) + with h5py.File(output, "r") as h5file: + group = h5file["entry/reconstruction/results/q_lab"] + np.testing.assert_array_equal( + group["intensity"][0, 0, [1, 260000, 299999]], + [2.5, 4.5, 6.5], + ) + np.testing.assert_array_equal( + group["contributors"][0, 0, [1, 260000, 299999]], + [2, 2, 2], + ) diff --git a/orgui/datautils/xrayutils/test/test_reconstruction_angles.py b/orgui/datautils/xrayutils/test/test_reconstruction_angles.py new file mode 100644 index 0000000..9a2a65c --- /dev/null +++ b/orgui/datautils/xrayutils/test/test_reconstruction_angles.py @@ -0,0 +1,165 @@ +"""Tests for reconstruction exposure-angle preparation.""" + +import numpy as np + +from types import SimpleNamespace + +from orgui.backend import backends +from orgui.backend.interlacedScanLoader import InterlacedScan +from orgui.backend.scans import ( + ScanReference, + SimulationScan, + load_scan_backend_file, + scan_exposure_angle_bounds, +) + + +def test_stationary_exposure_bounds_broadcast_fixed_angles(): + """Stationary exposures have identical start and end positions.""" + scan = SimulationScan((1, 1), 0.0, 10.0, 2, axis="mu", fixed=3.0) + bounds = scan_exposure_angle_bounds( + scan, + SimpleNamespace(mu=0.0, chi=0.4, phi=0.5), + ) + + assert bounds.shape == (2, 2, 4) + np.testing.assert_allclose(bounds[:, 0], bounds[:, 1]) + np.testing.assert_allclose( + bounds[:, 0], + [ + [0.0, np.deg2rad(-3.0), 0.4, 0.5], + [np.deg2rad(10.0), np.deg2rad(-3.0), 0.4, 0.5], + ], + ) + + +def test_swept_exposure_bounds_use_adjacent_midpoints(): + """Continuous sweeps use midpoint edges and extrapolated end edges.""" + scan = SimulationScan((1, 1), 0.0, -20.0, 3, axis="th", fixed=0.1) + bounds = scan_exposure_angle_bounds( + scan, + SimpleNamespace(mu=0.1, chi=0.3, phi=0.4), + fallback="midpoint", + ) + + np.testing.assert_allclose( + bounds[:, :, 1], + np.deg2rad([[-5.0, 5.0], [5.0, 15.0], [15.0, 25.0]]), + ) + np.testing.assert_allclose(bounds[:, :, 0], np.deg2rad(0.1)) + + +def test_exact_bounds_override_fallback_and_survive_interlacing(): + config = SimpleNamespace(mu=0.0, chi=0.0, phi=0.0) + first = SimulationScan((1, 1), 0.0, 1.0, 2) + second = SimulationScan((1, 1), 2.0, 3.0, 2) + first.omega_bounds_rad = np.array([[0.0, 0.01], [0.01, 0.02]]) + second.omega_bounds_rad = np.array([[0.02, 0.03], [0.03, 0.04]]) + scan = InterlacedScan([first, second], False, "th") + + bounds = scan_exposure_angle_bounds(scan, config, fallback="midpoint") + + np.testing.assert_allclose( + bounds[:, :, 1], + np.concatenate( + [first.omega_bounds_rad, second.omega_bounds_rad] + ), + ) + + +def test_scan_reference_reopens_simulated_and_interlaced_scans(): + first = SimulationScan((2, 3), 0.0, 1.0, 2) + second = SimulationScan((2, 3), 2.0, 3.0, 2) + scan = InterlacedScan([first, second], False, "th") + + reopened = ScanReference.from_dict( + ScanReference.from_scan(scan).to_dict() + ).open() + + assert isinstance(reopened, InterlacedScan) + assert len(reopened) == 4 + np.testing.assert_allclose(reopened.axis, scan.axis) + + +def test_scan_reference_reopens_run_path_backend(tmp_path): + backend_file = tmp_path / "custom_backend.py" + source_file = tmp_path / "data.h5" + source_file.touch() + backend_file.write_text( + "\n".join( + [ + "import numpy as np", + "from orgui.backend.scans import Scan, h5_Image", + "", + "class CustomScan(Scan):", + " def __init__(self, source, number=None):", + " self.hdffilepath_orNode = source", + " self.scanno = number", + " self.axisname = 'th'", + " self.axis = np.array([0.0])", + " self.th = self.axis", + " self.omega = -self.axis", + " self.mu = 0.0", + " def __len__(self):", + " return 1", + " def get_raw_img(self, index):", + " return h5_Image(np.ones((1, 1)))", + " @classmethod", + " def parse_h5_node(cls, node):", + " return {}", + ] + ), + encoding="utf-8", + ) + _, scan_class = load_scan_backend_file(backend_file) + scan = scan_class(str(source_file), 7) + + reference = ScanReference.from_scan(scan) + reopened = ScanReference.from_dict(reference.to_dict()).open() + + assert reference.kind == "backend_file" + assert reference.parameters["backend_file"] == str(backend_file.absolute()) + assert type(reopened).__qualname__ == "CustomScan" + assert reopened.scanno == 7 + + +def test_old_run_path_reference_uses_loaded_backend(tmp_path, monkeypatch): + backend_file = tmp_path / "legacy_backend.py" + source_file = tmp_path / "data.h5" + source_file.touch() + backend_file.write_text( + "\n".join( + [ + "import numpy as np", + "from orgui.backend.scans import Scan, h5_Image", + "", + "class LegacyScan(Scan):", + " def __init__(self, source):", + " self.hdffilepath_orNode = source", + " self.axisname = 'th'", + " self.axis = np.array([0.0])", + " self.th = self.axis", + " self.omega = -self.axis", + " self.mu = 0.0", + " def __len__(self):", + " return 1", + " def get_raw_img(self, index):", + " return h5_Image(np.ones((1, 1)))", + " @classmethod", + " def parse_h5_node(cls, node):", + " return {}", + ] + ), + encoding="utf-8", + ) + _, scan_class = load_scan_backend_file(backend_file) + monkeypatch.setitem(backends.fscans, "legacy-test", scan_class) + reference = ScanReference( + kind="backend", + module="", + class_name="LegacyScan", + parameters={"source": str(source_file), "scan_number": None}, + source_fingerprints=(ScanReference._fingerprint(source_file),), + ) + + assert type(reference.open()).__qualname__ == "LegacyScan" diff --git a/orgui/datautils/xrayutils/test/test_reconstruction_cli.py b/orgui/datautils/xrayutils/test/test_reconstruction_cli.py new file mode 100644 index 0000000..cb54961 --- /dev/null +++ b/orgui/datautils/xrayutils/test/test_reconstruction_cli.py @@ -0,0 +1,111 @@ +"""Tests for reciprocal-space reconstruction command dispatch.""" + +import sys +from unittest import mock + +from orgui import main as orgui_main +from orgui import reconstruction_cli + + +def test_reconstruction_main_accepts_forwarded_arguments(): + """The reconstruction entry point accepts an explicit argument vector.""" + with mock.patch.object( + reconstruction_cli, "build_parser", wraps=reconstruction_cli.build_parser + ): + try: + reconstruction_cli.main(["--help"], prog="orGUI-rsmap") + except SystemExit as error: + assert error.code == 0 + + +def test_orgui_dispatches_rsmap_without_starting_gui(): + """``orGUI rsmap`` forwards remaining arguments to the reconstruction CLI.""" + with ( + mock.patch.object(sys, "argv", ["orGUI", "rsmap", "run", "job.json"]), + mock.patch( + "orgui.reconstruction_cli.main", return_value=17 + ) as reconstruction_main, + ): + result = orgui_main.main() + + assert result == 17 + reconstruction_main.assert_called_once_with( + ["run", "job.json"], prog="orGUI rsmap" + ) + + +def test_cli_exposes_prepared_job_and_cluster_commands(): + """Cluster stages consume prepared jobs without a specification CLI.""" + parser = reconstruction_cli.build_parser() + help_text = parser.format_help() + for command in ( + "run", + "resume", + "status", + "cluster-map", + "cluster-finalize", + "cluster-scripts", + ): + assert command in help_text + choices = parser._subparsers._group_actions[0].choices + for removed in ("preview", "map", "reduce", "finalize"): + assert removed not in choices + + +def test_cli_reports_progress_on_stderr(capsys): + """Long-running CLI jobs report progress without corrupting JSON stdout.""" + + def fake_run(path, *, progress): + progress(0, 4, "Mapping images 0/2") + progress(1, 4, "Mapping images 1/2") + progress(2, 4, "Reducing 1 mapping task") + progress(4, 4, "Complete") + return {"status": "complete", "output_path": path} + + with mock.patch.object(reconstruction_cli, "run_job", fake_run): + reconstruction_cli.main(["run", "job.json"]) + + captured = capsys.readouterr() + assert "Mapping images 0/2" in captured.err + assert "Reducing 1 mapping task" in captured.err + assert "Complete" in captured.err + assert '"status": "complete"' in captured.out + + +def test_cluster_map_cli_passes_array_resources(capsys): + """Array task IDs and scheduler resources reach the cluster map helper.""" + captured = {} + + def fake_map(path, task_index, *, cpus, memory_bytes, progress): + captured.update( + path=path, + task_index=task_index, + cpus=cpus, + memory_bytes=memory_bytes, + ) + progress(1, 1, "Complete") + return {"status": "complete"} + + with mock.patch.object( + reconstruction_cli, "run_cluster_map_task", fake_map + ): + reconstruction_cli.main( + [ + "cluster-map", + "job.json", + "--task-index", + "7", + "--cpus", + "4", + "--memory-gib", + "16", + ] + ) + + assert captured == { + "path": "job.json", + "task_index": 7, + "cpus": 4, + "memory_bytes": 16 * 1024**3, + } + assert '"status": "complete"' in capsys.readouterr().out diff --git a/orgui/datautils/xrayutils/test/test_reconstruction_cluster.py b/orgui/datautils/xrayutils/test/test_reconstruction_cluster.py new file mode 100644 index 0000000..35fb4de --- /dev/null +++ b/orgui/datautils/xrayutils/test/test_reconstruction_cluster.py @@ -0,0 +1,85 @@ +"""Tests for reciprocal-space cluster execution helpers.""" + +from types import SimpleNamespace + +import pytest + +from orgui.reconstruction_cluster import ( + ClusterSettings, + generate_cluster_scripts, +) + + +@pytest.mark.parametrize("scheduler", ["sge", "slurm"]) +def test_cluster_scripts_create_map_array_and_dependent_finalizer( + tmp_path, scheduler +): + """Generated scripts allocate separate map and reduction resources.""" + settings = ClusterSettings( + scheduler=scheduler, + job_name="sample-rsmap", + working_directory=str(tmp_path / "work"), + python_executable="/cluster/env/bin/python", + environment_setup="module load hdf5\nsource activate orgui", + queue="long", + account="beamtime", + array_cpus=4, + array_memory_gib=16, + array_walltime="12:00:00", + array_concurrency=8, + reduce_cpus=24, + reduce_memory_gib=96, + reduce_walltime="08:00:00", + sge_parallel_environment="smp", + sge_memory_resource="h_vmem", + ) + job = SimpleNamespace( + expected_map_tasks=17, + cluster_settings=settings.to_dict(), + ) + job_path = tmp_path / "prepared job.json" + + result = generate_cluster_scripts( + job_path, + job, + output_directory=tmp_path / "scripts", + ) + + map_text = (tmp_path / "scripts" / f"sample-rsmap-map.{scheduler}").read_text( + encoding="utf-8" + ) + finalize_text = ( + tmp_path / "scripts" / f"sample-rsmap-finalize.{scheduler}" + ).read_text(encoding="utf-8") + submit_text = ( + tmp_path / "scripts" / "sample-rsmap-submit.sh" + ).read_text(encoding="utf-8") + assert result["array_tasks"] == 17 + assert "cluster-map" in map_text + assert "cluster-finalize" in finalize_text + assert "--cpus 4" in map_text + assert "--cpus 24" in finalize_text + assert "module load hdf5" in map_text + assert "prepared job.json'" in map_text + if scheduler == "sge": + assert "#$ -t 1-17" in map_text + assert "#$ -tc 8" in map_text + assert "#$ -pe smp 4" in map_text + assert "#$ -l h_vmem=4G" in map_text + assert "$((SGE_TASK_ID - 1))" in map_text + assert "-hold_jid" in submit_text + else: + assert "#SBATCH --array=0-16%8" in map_text + assert "#SBATCH --cpus-per-task=4" in map_text + assert "#SBATCH --mem=16G" in map_text + assert "${SLURM_ARRAY_TASK_ID}" in map_text + assert "--dependency=\"afterok:${map_job_id}\"" in submit_text + + +def test_cluster_settings_reject_scheduler_mismatched_directives(): + """Extra directives must use the selected scheduler's syntax.""" + with pytest.raises(ValueError, match="#\\$"): + ClusterSettings( + scheduler="sge", + extra_array_directives="#SBATCH --constraint=fast", + ) diff --git a/orgui/datautils/xrayutils/test/test_reconstruction_geometry_steps.py b/orgui/datautils/xrayutils/test/test_reconstruction_geometry_steps.py new file mode 100644 index 0000000..2f886b4 --- /dev/null +++ b/orgui/datautils/xrayutils/test/test_reconstruction_geometry_steps.py @@ -0,0 +1,78 @@ +"""Tests for geometry-matched reciprocal-space grid sampling.""" + +from types import SimpleNamespace + +import numpy as np + +import orgui.datautils.xrayutils.reconstruction as reconstruction_module +from orgui.reconstruction_job import estimate_geometry_steps + + +class _Detector: + detector = SimpleNamespace(shape=(4, 5)) + + @staticmethod + def primBeamPoints(rows, columns): + return np.zeros_like(rows), np.zeros_like(columns) + + +class _UBCalculator: + @staticmethod + def getK(): + return 1.0 + + @staticmethod + def getUB(): + return np.eye(3) + + @staticmethod + def getU(): + return np.eye(3) + + +class _Scan: + def __len__(self): + return 2 + + @staticmethod + def exposure_angle_bounds(config, fallback): + bounds = np.zeros((2, 2, 4), dtype=np.float64) + bounds[1, :, 0] = 1.0 + return bounds + + +class _Kernel: + def __init__(self, *args): + pass + + @staticmethod + def coordinate(rays, start, end, row, column, u, v, t): + angle = start[0] + t * (end[0] - start[0]) + return np.asarray((2.0 * u, 3.0 * v, 4.0 * angle)) + + +def test_local_jacobian_propagates_uniform_pixel_and_scan_widths(monkeypatch): + """Jacobian columns propagate to one-sigma axis projections.""" + monkeypatch.setattr( + reconstruction_module, + "_native_module", + lambda: SimpleNamespace(ReconstructionKernel=_Kernel), + ) + config = SimpleNamespace( + detector=_Detector(), + ub_calculator=_UBCalculator(), + corrections=SimpleNamespace(excluded_frames=()), + ) + + steps = estimate_geometry_steps( + config, + _Scan(), + percentile=37.0, + detector_samples=2, + frame_samples=2, + ) + + np.testing.assert_allclose( + steps, + np.asarray((2.0, 3.0, 4.0)) / np.sqrt(12.0), + ) diff --git a/orgui/datautils/xrayutils/test/test_reconstruction_job.py b/orgui/datautils/xrayutils/test/test_reconstruction_job.py new file mode 100644 index 0000000..68656ce --- /dev/null +++ b/orgui/datautils/xrayutils/test/test_reconstruction_job.py @@ -0,0 +1,243 @@ +"""End-to-end tests for the centralized reconstruction job.""" + +from concurrent.futures import ThreadPoolExecutor +from hashlib import sha256 +import json +from pathlib import Path +import threading +import time + +import h5py +import numpy as np +import pytest + +from orgui.app.config_data import ConfigData +from orgui.app.database import config_data_to_json +from orgui.backend.scans import ScanReference, SimulationScan +from orgui.datautils.xrayutils import CTRcalc, DetectorCalibration, HKLVlieg +from orgui.reconstruction_job import ( + ReconstructionGrid, + ReconstructionJob, + job_status, + reconstruction_execution_settings, + run_cluster_finalize, + run_cluster_map_task, + run_job, + write_job, +) + + +pytest.importorskip("pyarrow") +pytest.importorskip( + "orgui.datautils.xrayutils._reciprocal_reconstruction_cpp" +) + + +def _config(): + cell = CTRcalc.UnitCell([3.0, 3.0, 3.0], [90.0, 90.0, 90.0]) + cell.addAtom("Pt", [0.0, 0.0, 0.0], 0.1, 0.1, 1.0) + ub = HKLVlieg.UBCalculator(cell, 70.0) + ub.defaultU() + detector = DetectorCalibration.Detector2D_SXRD() + detector.setFit2D( + 729.0, 0.5, 0.5, pixelX=172.0, pixelY=172.0 + ) + detector.set_wavelength(ub.getLambda() * 1e-10) + detector.detector.shape = (2, 2) + detector.detector.max_shape = (2, 2) + return ConfigData(detector, cell, ub) + + +def test_central_job_runs_resumes_and_cleans_verified_scratch( + tmp_path, monkeypatch +): + scan = SimulationScan((2, 2), 0.0, 1.0, 2) + assets = tmp_path / "scratch" / "job-assets.nxs" + assets.parent.mkdir() + with h5py.File(assets, "w") as h5file: + h5file.attrs["orgui_job_assets"] = 1 + assets_digest = sha256(assets.read_bytes()).hexdigest() + grid = ReconstructionGrid( + minimum=(-20.0, -20.0, -20.0), + maximum=(20.0, 20.0, 20.0), + step=(20.0, 20.0, 20.0), + frame="lab", + chunk_shape=(2, 2, 2), + ) + scan_reference = ScanReference.from_scan(scan).to_dict() + source_digest = sha256( + json.dumps( + scan_reference, sort_keys=True, separators=(",", ":") + ).encode() + ).hexdigest() + job = ReconstructionJob( + config=config_data_to_json(_config()), + scan_reference=scan_reference, + grids=[grid.__dict__], + scratch_path=str(assets.parent), + output_path=str(tmp_path / "result.h5"), + compression="Raw", + assets_path=str(assets), + assets_sha256=assets_digest, + source_fingerprint_sha256=source_digest, + threads_per_image=1, + accumulation_budget_bytes=None, + runtime_threads=2, + runtime_memory_bytes=64 * 1024 * 1024, + tile_shape=(1, 1), + frame_batch=1, + expected_map_tasks=2, + ) + job_path = tmp_path / "job.json" + write_job(job, job_path) + + active_loads = 0 + maximum_active_loads = 0 + load_lock = threading.Lock() + original_get_raw_img = SimulationScan.get_raw_img + + def observed_get_raw_img(self, index): + nonlocal active_loads, maximum_active_loads + with load_lock: + active_loads += 1 + maximum_active_loads = max(maximum_active_loads, active_loads) + try: + time.sleep(0.02) + return original_get_raw_img(self, index) + finally: + with load_lock: + active_loads -= 1 + + monkeypatch.setattr( + SimulationScan, "get_raw_img", observed_get_raw_img + ) + + settings = reconstruction_execution_settings( + job, scan=scan, config=job.config_data + ) + assert settings["thread_budget"] == 2 + assert settings["native_threads_per_image"] == 1 + assert settings["detector_tile_shape"] == (1, 1) + assert settings["map_tasks"] == 2 + layout = settings["parallel_layouts"][0] + assert layout["exposure"] == "stationary" + assert layout["concurrent_image_workers"] == 2 + assert layout["native_threads_per_image"] == 1 + assert layout["tiles_per_image"] == 4 + assert layout["memory_per_image_MiB"] == pytest.approx(32.0, abs=0.01) + assert layout["accumulation_MiB_per_worker"] == pytest.approx( + 10.33, abs=0.02 + ) + + job.threads_per_image = 2 + combined = reconstruction_execution_settings( + job, scan=scan, config=job.config_data + ) + assert combined["native_threads_per_image"] == 2 + assert combined["parallel_layouts"][0]["concurrent_image_workers"] == 1 + assert combined["parallel_layouts"][0]["native_threads_per_image"] == 2 + job.threads_per_image = 1 + + progress_updates = [] + first = run_job( + job_path, + progress=lambda value, maximum, message: progress_updates.append( + (value, maximum, message) + ), + ) + resumed = run_job(job_path) + + assert maximum_active_loads == 2 + assert any( + "Mapping images 1/2" in message + for _value, _maximum, message in progress_updates + ) + assert any( + "Mapping images 2/2" in message + for _value, _maximum, message in progress_updates + ) + assert first == resumed + assert first["status"] == "complete" + assert not assets.exists() + with h5py.File(first["output_path"], "r") as h5file: + assert "entry/configuration" in h5file + group = h5file["entry/reconstruction/results/q_lab"] + assert np.nanmax(group["intensity"]) == 10.0 + assert np.nanmax(group["contributors"]) == 8 + assert job_status(job_path)["output_sha256"] == first["output_sha256"] + + +def test_cluster_array_tasks_do_not_mutate_job_and_finalizer_collects_them( + tmp_path, +): + """Independent array tasks publish manifests before one finalizer writes.""" + scan = SimulationScan((2, 2), 0.0, 1.0, 2) + assets = tmp_path / "scratch" / "job-assets.nxs" + assets.parent.mkdir() + with h5py.File(assets, "w") as h5file: + h5file.attrs["orgui_job_assets"] = 1 + grid = ReconstructionGrid( + minimum=(-20.0, -20.0, -20.0), + maximum=(20.0, 20.0, 20.0), + step=(20.0, 20.0, 20.0), + frame="lab", + chunk_shape=(2, 2, 2), + ) + scan_reference = ScanReference.from_scan(scan).to_dict() + job = ReconstructionJob( + config=config_data_to_json(_config()), + scan_reference=scan_reference, + grids=[grid.__dict__], + scratch_path=str(assets.parent), + output_path=str(tmp_path / "cluster-result.h5"), + compression="Raw", + assets_path=str(assets), + assets_sha256=sha256(assets.read_bytes()).hexdigest(), + source_fingerprint_sha256=sha256( + json.dumps( + scan_reference, sort_keys=True, separators=(",", ":") + ).encode() + ).hexdigest(), + threads_per_image=1, + accumulation_budget_bytes=None, + runtime_threads=2, + runtime_memory_bytes=64 * 1024 * 1024, + tile_shape=(1, 1), + frame_batch=1, + expected_map_tasks=2, + ) + job_path = tmp_path / "cluster-job.json" + write_job(job, job_path) + frozen_job = job_path.read_bytes() + + with ThreadPoolExecutor(max_workers=2) as executor: + futures = [ + executor.submit( + run_cluster_map_task, + job_path, + task_index, + cpus=1, + memory_bytes=64 * 1024 * 1024, + ) + for task_index in range(2) + ] + first, second = (future.result() for future in futures) + assert job_path.read_bytes() == frozen_job + assert first["frame_range"] == [0, 1] + assert second["frame_range"] == [1, 2] + assert job_status(job_path)["map_tasks"] == { + "completed": 2, + "pending": 0, + "total": 2, + } + assert job_path.read_bytes() == frozen_job + + result = run_cluster_finalize( + job_path, + cpus=2, + memory_bytes=128 * 1024 * 1024, + ) + + assert result["status"] == "complete" + assert Path(result["output_path"]).exists() + assert job_status(job_path)["map_tasks"]["completed"] == 2 diff --git a/orgui/main.py b/orgui/main.py index e0084ab..8ec51cf 100644 --- a/orgui/main.py +++ b/orgui/main.py @@ -78,12 +78,19 @@ def writable_file(path: str): epilog = """The reflection list allows calculation of binned reciprocal space using the HESXRD backends of binoculars.""" -usage = "orGUI [options] configfile" +usage = """orGUI [options] [configfile] + orGUI rsmap ...""" defaultconfigfile = os.path.expanduser("~/orgui") def main(): + """Launch the GUI or dispatch an orGUI command.""" + if len(sys.argv) > 1 and sys.argv[1] == "rsmap": + from .reconstruction_cli import main as reconstruction_main + + return reconstruction_main(sys.argv[2:], prog="orGUI rsmap") + parser = ArgumentParser(usage=usage, description=description, epilog=epilog) parser.add_argument( "configfile", diff --git a/orgui/reconstruction_cli.py b/orgui/reconstruction_cli.py new file mode 100644 index 0000000..25175a4 --- /dev/null +++ b/orgui/reconstruction_cli.py @@ -0,0 +1,131 @@ +"""Command-line execution of UI-prepared reciprocal-space jobs.""" + +from __future__ import annotations + +import argparse +import json +import sys +import time + +from .reconstruction_cluster import generate_cluster_scripts +from .reconstruction_job import ( + job_status, + read_job, + run_cluster_finalize, + run_cluster_map_task, + run_job, +) + + +def _progress_operation(operation, *args, **kwargs): + last_report = 0.0 + + def update(value, maximum, message): + nonlocal last_report + now = time.monotonic() + if ( + value in {0, maximum} + or now - last_report >= 1.0 + or not message.startswith("Mapping images") + ): + percent = 100.0 * value / max(1, maximum) + print( + f"{percent:6.2f}% {message}", + file=sys.stderr, + flush=True, + ) + last_report = now + + result = operation(*args, progress=update, **kwargs) + print(json.dumps(result, indent=2, sort_keys=True)) + + +def _run(arguments): + _progress_operation(run_job, arguments.job) + + +def _cluster_map(arguments): + _progress_operation( + run_cluster_map_task, + arguments.job, + arguments.task_index, + cpus=arguments.cpus, + memory_bytes=int(arguments.memory_gib * 1024**3), + ) + + +def _cluster_finalize(arguments): + _progress_operation( + run_cluster_finalize, + arguments.job, + cpus=arguments.cpus, + memory_bytes=int(arguments.memory_gib * 1024**3), + ) + + +def _cluster_scripts(arguments): + job = read_job(arguments.job) + result = generate_cluster_scripts( + arguments.job, + job, + output_directory=arguments.output_directory, + ) + print(json.dumps(result, indent=2, sort_keys=True)) + + +def _status(arguments): + print(json.dumps(job_status(arguments.job), indent=2, sort_keys=True)) + + +def build_parser(prog=None): + """Build the reciprocal-space job command parser.""" + parser = argparse.ArgumentParser( + prog=prog, + description="Execute UI-prepared orGUI reciprocal-space jobs", + ) + commands = parser.add_subparsers(dest="command", required=True) + for name, help_text in ( + ("run", "Run a prepared job"), + ("resume", "Verify and resume an interrupted job"), + ): + command = commands.add_parser(name, help=help_text) + command.add_argument("job", help="UI-prepared reconstruction job JSON") + command.set_defaults(handler=_run) + status = commands.add_parser("status", help="Inspect a prepared job") + status.add_argument("job") + status.set_defaults(handler=_status) + cluster_map = commands.add_parser( + "cluster-map", + help="Execute one prepared map-array task", + ) + cluster_map.add_argument("job") + cluster_map.add_argument("--task-index", type=int, required=True) + cluster_map.add_argument("--cpus", type=int, required=True) + cluster_map.add_argument("--memory-gib", type=float, required=True) + cluster_map.set_defaults(handler=_cluster_map) + cluster_finalize = commands.add_parser( + "cluster-finalize", + help="Reduce and finalize completed cluster map tasks", + ) + cluster_finalize.add_argument("job") + cluster_finalize.add_argument("--cpus", type=int, required=True) + cluster_finalize.add_argument("--memory-gib", type=float, required=True) + cluster_finalize.set_defaults(handler=_cluster_finalize) + scripts = commands.add_parser( + "cluster-scripts", + help="Generate scheduler scripts from a prepared job", + ) + scripts.add_argument("job") + scripts.add_argument("--output-directory") + scripts.set_defaults(handler=_cluster_scripts) + return parser + + +def main(argv=None, prog=None): + """Run the reciprocal-space job command-line interface.""" + arguments = build_parser(prog=prog).parse_args(argv) + arguments.handler(arguments) + + +if __name__ == "__main__": + main() diff --git a/orgui/reconstruction_cluster.py b/orgui/reconstruction_cluster.py new file mode 100644 index 0000000..0598ee9 --- /dev/null +++ b/orgui/reconstruction_cluster.py @@ -0,0 +1,322 @@ +"""Batch-script generation for UI-prepared reconstruction jobs.""" + +from __future__ import annotations + +from dataclasses import asdict, dataclass +import math +from pathlib import Path +import re +import shlex + + +_SCHEDULERS = {"sge", "slurm"} +_JOB_NAME = re.compile(r"^[A-Za-z0-9_.-]+$") + + +@dataclass(frozen=True) +class ClusterSettings: + """Describe scheduler resources for reconstruction map and reduce jobs. + + These settings control execution only. Scientific configuration remains + frozen in the reconstruction job descriptor. + """ + + scheduler: str = "sge" + job_name: str = "orgui-rsmap" + script_directory: str = "" + working_directory: str = "" + python_executable: str = "python" + environment_setup: str = "" + queue: str = "" + account: str = "" + array_cpus: int = 4 + array_memory_gib: float = 16.0 + array_walltime: str = "24:00:00" + array_concurrency: int = 0 + reduce_cpus: int = 24 + reduce_memory_gib: float = 64.0 + reduce_walltime: str = "24:00:00" + sge_parallel_environment: str = "smp" + sge_memory_resource: str = "h_vmem" + extra_array_directives: str = "" + extra_reduce_directives: str = "" + + def __post_init__(self): + scheduler = self.scheduler.lower() + object.__setattr__(self, "scheduler", scheduler) + if scheduler not in _SCHEDULERS: + raise ValueError("Scheduler must be 'sge' or 'slurm'") + if not self.job_name or not _JOB_NAME.fullmatch(self.job_name): + raise ValueError( + "Cluster job name may contain letters, numbers, '.', '_', " + "and '-' only" + ) + if not self.python_executable.strip(): + raise ValueError("Python executable cannot be empty") + for name in ("array_cpus", "reduce_cpus"): + if int(getattr(self, name)) < 1: + raise ValueError(f"{name} must be at least one") + for name in ("array_memory_gib", "reduce_memory_gib"): + if float(getattr(self, name)) <= 0: + raise ValueError(f"{name} must be positive") + if int(self.array_concurrency) < 0: + raise ValueError("array_concurrency cannot be negative") + for name in ("array_walltime", "reduce_walltime"): + if not str(getattr(self, name)).strip(): + raise ValueError(f"{name} cannot be empty") + if scheduler == "sge": + if not self.sge_parallel_environment.strip(): + raise ValueError("SGE parallel environment cannot be empty") + if not self.sge_memory_resource.strip(): + raise ValueError("SGE memory resource cannot be empty") + self._validate_directives( + self.extra_array_directives, "extra_array_directives" + ) + self._validate_directives( + self.extra_reduce_directives, "extra_reduce_directives" + ) + + def _validate_directives(self, value, name): + prefix = "#$" if self.scheduler == "sge" else "#SBATCH" + invalid = [ + line + for line in str(value).splitlines() + if line.strip() and not line.lstrip().startswith(prefix) + ] + if invalid: + raise ValueError( + f"Every non-empty {name} line must start with {prefix}" + ) + + def to_dict(self): + """Return JSON-compatible cluster settings.""" + return asdict(self) + + @classmethod + def from_dict(cls, values): + """Build settings from a reconstruction job dictionary.""" + return cls(**dict(values or {})) + + +def _quote(value): + return shlex.quote(str(value).replace("\\", "/")) + + +def _script_preamble(settings, working_directory): + lines = [ + "set -euo pipefail", + f"cd {_quote(working_directory)}", + ] + if settings.environment_setup.strip(): + lines.extend(settings.environment_setup.rstrip().splitlines()) + return lines + + +def _python_command(settings, arguments): + executable = _quote(settings.python_executable) + return ( + f"{executable} -m orgui.reconstruction_cli " + + " ".join(_quote(value) for value in arguments) + ) + + +def _sge_directives(settings, *, array_tasks=None): + is_array = array_tasks is not None + cpus = settings.array_cpus if is_array else settings.reduce_cpus + memory = ( + settings.array_memory_gib + if is_array + else settings.reduce_memory_gib + ) + walltime = ( + settings.array_walltime + if is_array + else settings.reduce_walltime + ) + suffix = "map" if is_array else "finalize" + lines = [ + f"#$ -N {settings.job_name}-{suffix}", + "#$ -cwd", + f"#$ -pe {settings.sge_parallel_environment} {cpus}", + ( + f"#$ -l {settings.sge_memory_resource}=" + f"{math.ceil(memory / cpus)}G" + ), + f"#$ -l h_rt={walltime}", + ] + if is_array: + lines.append(f"#$ -t 1-{array_tasks}") + if settings.array_concurrency: + lines.append(f"#$ -tc {settings.array_concurrency}") + if settings.queue: + lines.append(f"#$ -q {settings.queue}") + if settings.account: + lines.append(f"#$ -P {settings.account}") + extra = ( + settings.extra_array_directives + if is_array + else settings.extra_reduce_directives + ) + lines.extend(line for line in extra.splitlines() if line.strip()) + return lines + + +def _slurm_directives(settings, *, array_tasks=None): + is_array = array_tasks is not None + cpus = settings.array_cpus if is_array else settings.reduce_cpus + memory = ( + settings.array_memory_gib + if is_array + else settings.reduce_memory_gib + ) + walltime = ( + settings.array_walltime + if is_array + else settings.reduce_walltime + ) + suffix = "map" if is_array else "finalize" + lines = [ + f"#SBATCH --job-name={settings.job_name}-{suffix}", + f"#SBATCH --cpus-per-task={cpus}", + f"#SBATCH --mem={math.ceil(memory)}G", + f"#SBATCH --time={walltime}", + ] + if is_array: + array = f"0-{array_tasks - 1}" + if settings.array_concurrency: + array += f"%{settings.array_concurrency}" + lines.append(f"#SBATCH --array={array}") + if settings.queue: + lines.append(f"#SBATCH --partition={settings.queue}") + if settings.account: + lines.append(f"#SBATCH --account={settings.account}") + extra = ( + settings.extra_array_directives + if is_array + else settings.extra_reduce_directives + ) + lines.extend(line for line in extra.splitlines() if line.strip()) + return lines + + +def generate_cluster_scripts(job_path, job, output_directory=None): + """Generate map-array, finalizer, and dependency-submission scripts. + + :param job_path: + Prepared reconstruction job JSON. + :param ReconstructionJob job: + Decoded prepared job. + :param output_directory: + Optional directory replacing the frozen cluster script directory. + :returns: + Paths keyed by ``map``, ``finalize``, and ``submit``. + :rtype: dict + """ + if job.expected_map_tasks < 1: + raise ValueError("Prepared job contains no mapping tasks") + settings = ClusterSettings.from_dict(job.cluster_settings) + job_path = Path(job_path).absolute() + directory = Path( + output_directory + or settings.script_directory + or job_path.parent / f"{job_path.stem}-cluster" + ).absolute() + directory.mkdir(parents=True, exist_ok=True) + working_directory = Path( + settings.working_directory or job_path.parent + ).absolute() + suffix = "sge" if settings.scheduler == "sge" else "slurm" + map_path = directory / f"{settings.job_name}-map.{suffix}" + finalize_path = directory / f"{settings.job_name}-finalize.{suffix}" + submit_path = directory / f"{settings.job_name}-submit.sh" + preamble = _script_preamble(settings, working_directory) + + if settings.scheduler == "sge": + map_index = '"$((SGE_TASK_ID - 1))"' + map_directives = _sge_directives( + settings, array_tasks=job.expected_map_tasks + ) + finalize_directives = _sge_directives(settings) + else: + map_index = '"${SLURM_ARRAY_TASK_ID}"' + map_directives = _slurm_directives( + settings, array_tasks=job.expected_map_tasks + ) + finalize_directives = _slurm_directives(settings) + + map_command = _python_command( + settings, + ( + "cluster-map", + job_path, + "--task-index", + "__ARRAY_INDEX__", + "--cpus", + settings.array_cpus, + "--memory-gib", + settings.array_memory_gib, + ), + ).replace(_quote("__ARRAY_INDEX__"), map_index) + finalize_command = _python_command( + settings, + ( + "cluster-finalize", + job_path, + "--cpus", + settings.reduce_cpus, + "--memory-gib", + settings.reduce_memory_gib, + ), + ) + map_text = "\n".join( + ["#!/usr/bin/env bash", *map_directives, "", *preamble, map_command, ""] + ) + finalize_text = "\n".join( + [ + "#!/usr/bin/env bash", + *finalize_directives, + "", + *preamble, + finalize_command, + "", + ] + ) + if settings.scheduler == "sge": + submit_lines = [ + "#!/usr/bin/env bash", + "set -euo pipefail", + f"map_job_id=$(qsub -terse {_quote(map_path)})", + 'map_job_id="${map_job_id%%.*}"', + ( + f'qsub -hold_jid "$map_job_id" ' + f"{_quote(finalize_path)}" + ), + 'echo "Submitted SGE map array ${map_job_id} and finalizer"', + "", + ] + else: + submit_lines = [ + "#!/usr/bin/env bash", + "set -euo pipefail", + f"map_job_id=$(sbatch --parsable {_quote(map_path)})", + 'map_job_id="${map_job_id%%;*}"', + ( + f'sbatch --dependency="afterok:${{map_job_id}}" ' + f"{_quote(finalize_path)}" + ), + 'echo "Submitted Slurm map array ${map_job_id} and finalizer"', + "", + ] + for path, text in ( + (map_path, map_text), + (finalize_path, finalize_text), + (submit_path, "\n".join(submit_lines)), + ): + path.write_text(text, encoding="utf-8", newline="\n") + return { + "scheduler": settings.scheduler, + "map": str(map_path), + "finalize": str(finalize_path), + "submit": str(submit_path), + "array_tasks": job.expected_map_tasks, + } diff --git a/orgui/reconstruction_job.py b/orgui/reconstruction_job.py new file mode 100644 index 0000000..cf36183 --- /dev/null +++ b/orgui/reconstruction_job.py @@ -0,0 +1,1875 @@ +"""Central orGUI job model for reciprocal-space reconstruction.""" + +from __future__ import annotations + +from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait +from dataclasses import asdict, dataclass, field +from hashlib import sha256 +import json +import math +from pathlib import Path +from queue import Empty, SimpleQueue +import shutil +from threading import Event +from typing import Any + +import h5py +import numpy as np + +from .app.config_data import ConfigData +from .app.database import FILTERS, config_data_from_json, config_data_to_json +from .app.mask_config import create_pixel_repair_plan +from .backend.scans import ScanReference +from .datautils.xrayutils.reconstruction import ( + _GridSpec, + _MIN_REDUCER_WORKER_MEMORY, + _ReconstructionSpec, + _TaskManifest, + _detector_corner_rays, + _finalize_reconstruction, + _map_frame_range, + _read_manifest, + _reduce_partition, + _verify_scratch_file, + _write_manifest, + _xxh3_128, +) + + +JOB_SCHEMA_VERSION = 3 +ACCURACY_DEPTHS = { + "center": 0, + "low": 1, + "balanced": 2, + "high": 3, + "very_high": 4, + "maximum": 5, +} +AUTO_MAX_FRAMES_PER_TASK = 64 +ACCUMULATION_TRANSIENT_FACTOR = 3 +AUTO_MAX_ACCUMULATION_BYTES = 2 * 1024**3 + + +def _build_metadata(): + from . import __version__ + + try: + from ._build_config import BUILD_CONFIG + except ImportError: + build_config = {"build_config": "unavailable-from-source-tree"} + else: + build_config = dict(BUILD_CONFIG) + return { + "orgui_version": __version__, + "numpy_version": np.__version__, + "h5py_version": h5py.__version__, + "native": build_config, + } + + +def _sha256_file(path): + digest = sha256() + with Path(path).open("rb") as stream: + while block := stream.read(8 * 1024 * 1024): + digest.update(block) + return digest.hexdigest() + + +def _atomic_json(path, values): + path = Path(path).absolute() + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_name(path.name + ".tmp") + temporary.write_text( + json.dumps(values, indent=2, sort_keys=True), encoding="utf-8" + ) + temporary.replace(path) + + +def _compression_name(value): + for name, available in FILTERS.items(): + try: + if value is available or value == available: + return name + except Exception: + continue + raise ValueError("The active database compression is not registered") + + +@dataclass(frozen=True) +class ReconstructionGrid: + """One user-selected reciprocal-space output grid.""" + + minimum: tuple[float, float, float] + maximum: tuple[float, float, float] + step: tuple[float, float, float] + frame: str = "hkl" + name: str | None = None + chunk_shape: tuple[int, int, int] = (64, 64, 64) + + def to_spec(self): + """Return the internal native/storage grid representation.""" + return _GridSpec(**asdict(self)) + + +@dataclass +class ReconstructionJob: + """Immutable experiment snapshot plus mutable deterministic task state.""" + + config: dict[str, Any] + scan_reference: dict[str, Any] + grids: list[dict[str, Any]] + scratch_path: str + output_path: str + compression: str + assets_path: str + assets_sha256: str + source_fingerprint_sha256: str + threads_per_image: int + accumulation_budget_bytes: int | None + build_metadata: dict[str, Any] = field(default_factory=dict) + runtime_threads: int = 1 + runtime_memory_bytes: int = 6_000 * 1024 * 1024 + angle_fallback: str = "stationary" + accuracy: str = "balanced" + advanced_depth: int | None = None + thread_override: int | None = None + memory_override_bytes: int | None = None + frame_batch: int | None = None + tile_shape: tuple[int, int] | None = None + work_block_pixels: int | None = None + partition_chunk_span: int | None = None + user_note: str = "" + status: str = "prepared" + expected_map_tasks: int = 0 + map_manifests: list[str] = field(default_factory=list) + reduction_manifest: str | None = None + output_sha256: str | None = None + correction_provenance: dict[str, Any] = field(default_factory=dict) + cleanup_errors: list[str] = field(default_factory=list) + cluster_settings: dict[str, Any] = field(default_factory=dict) + schema_version: int = JOB_SCHEMA_VERSION + + def to_dict(self): + """Return a JSON-compatible job descriptor.""" + return asdict(self) + + @classmethod + def from_dict(cls, values): + """Build a reconstruction job from JSON-compatible values.""" + values = dict(values) + if values.get("schema_version") != JOB_SCHEMA_VERSION: + raise ValueError("Unsupported reconstruction job schema") + if values.get("tile_shape") is not None: + values["tile_shape"] = tuple(values["tile_shape"]) + return cls(**values) + + @property + def digest(self): + """Return the immutable scientific job digest.""" + values = self.to_dict() + for key in ( + "status", + "map_manifests", + "reduction_manifest", + "output_sha256", + "correction_provenance", + "cleanup_errors", + ): + values.pop(key, None) + encoded = json.dumps( + values, sort_keys=True, separators=(",", ":") + ).encode() + return sha256(encoded).hexdigest() + + @property + def config_data(self): + """Return the central configuration snapshot.""" + return config_data_from_json(self.config) + + @property + def scan(self): + """Reopen and verify the referenced scan.""" + return ScanReference.from_dict(self.scan_reference).open() + + def internal_spec(self): + """Build the internal native/storage specification.""" + depth = ( + self.advanced_depth + if self.advanced_depth is not None + else ACCURACY_DEPTHS[self.accuracy] + ) + memory = self.memory_override_bytes or self.runtime_memory_bytes + threads = self.thread_override or self.runtime_threads + work_block = self.work_block_pixels or max( + 1024, min(65536, memory // max(threads * 256, 1)) + ) + chunk_span = self.partition_chunk_span or max( + 16, min(4096, memory // (64**3 * 32)) + ) + return _ReconstructionSpec( + grids=tuple(_GridSpec(**grid) for grid in self.grids), + max_depth=depth, + threads=threads, + work_block_pixels=work_block, + memory_budget_bytes=memory, + partition_chunk_span=chunk_span, + compression=f"database:{self.compression}", + infer_angle_bounds=self.angle_fallback == "midpoint", + ) + + +def read_job(path): + """Read a reconstruction job JSON file.""" + return ReconstructionJob.from_dict( + json.loads(Path(path).read_text(encoding="utf-8")) + ) + + +def write_job(job, path): + """Atomically write a reconstruction job JSON file.""" + _atomic_json(path, job.to_dict()) + return str(Path(path).absolute()) + + +def _snapshot_assets(gui, config, assets_path): + assets_path = Path(assets_path).absolute() + assets_path.parent.mkdir(parents=True, exist_ok=True) + temporary = assets_path.with_name(assets_path.name + ".tmp") + correction = config.corrections + shape = tuple(gui.ubcalc.detectorCal.detector.shape) + mask = gui.get_detector_mask(shape) if correction.use_mask else None + if correction.use_mask and mask is None: + raise ValueError( + "Pixel-mask correction is enabled, but no active mask matches " + f"the detector image shape {shape}. Load a matching mask or " + "disable pixel-mask correction before preparing the job." + ) + background = ( + getattr(gui, "background_image", None) + if correction.use_background + else None + ) + background_variance = getattr(gui, "background_variance", None) + with h5py.File(temporary, "w") as h5file: + h5file.attrs["NX_class"] = "NXroot" + h5file.attrs["orgui_job_assets"] = JOB_SCHEMA_VERSION + if mask is not None: + h5file.create_dataset("mask", data=np.asarray(mask, dtype=np.bool_)) + correction.mask_asset = "/mask" + if background is not None: + h5file.create_dataset( + "background", data=np.asarray(background, dtype=np.float64) + ) + correction.background_asset = "/background" + if background_variance is not None: + h5file.create_dataset( + "background_variance", + data=np.asarray(background_variance, dtype=np.float64), + ) + correction.background_variance_asset = "/background_variance" + if correction.repair_masked_pixels: + enabled, settings, rows, columns = gui._repair_config_for_image(shape) + correction.repair_masked_pixels = bool(enabled) + if enabled: + repair = h5file.create_group("repair") + repair.attrs["max_component_pixels"] = settings.max_component_pixels + repair.attrs["max_span"] = settings.max_span + repair.attrs["radius"] = settings.radius + repair.attrs["min_valid_neighbors"] = settings.min_valid_neighbors + repair.create_dataset("row_gaps", data=rows) + repair.create_dataset("column_gaps", data=columns) + temporary.replace(assets_path) + return _sha256_file(assets_path) + + +def derive_grid(config, scan, *, frame="hkl", name=None): + """Estimate grid coverage and spacing from current experiment geometry.""" + detector = config.detector + rows, columns = detector.detector.shape + detector_rows = np.array([-0.5, -0.5, rows - 0.5, rows - 0.5]) + detector_columns = np.array( + [-0.5, columns - 0.5, -0.5, columns - 0.5] + ) + gamma, delta = detector.primBeamPoints( + detector_rows, detector_columns + ) + cosine_gamma = np.cos(gamma) + edge_rays = np.column_stack( + ( + np.sin(delta) * cosine_gamma, + np.cos(delta) * cosine_gamma, + np.sin(gamma), + ) + ) + corner_rays = np.ascontiguousarray(edge_rays.reshape(2, 2, 3)) + bounds = scan.exposure_angle_bounds(config, fallback="stationary") + from .datautils.xrayutils.reconstruction import _native_module + + ub = config.ub_calculator + kernel = _native_module().ReconstructionKernel( + np.full(3, -1e12), + np.ones(3), + np.full(3, 2_000_000_000, dtype=np.int64), + np.ones(3, dtype=np.int64), + frame, + float(ub.getK()), + np.ascontiguousarray(np.linalg.inv(ub.getUB())), + np.ascontiguousarray(np.linalg.inv(ub.getU())), + 0, + 1, + 1, + 1024 * 1024, + ) + coordinates = [] + for angles_start, angles_end in bounds: + angles_start = np.ascontiguousarray(angles_start, dtype=np.float64) + angles_end = np.ascontiguousarray(angles_end, dtype=np.float64) + for u in (0.0, 1.0): + for v in (0.0, 1.0): + for t in (0.0, 1.0): + coordinates.append( + kernel.coordinate( + corner_rays, + angles_start, + angles_end, + 0, + 0, + u, + v, + t, + ) + ) + coordinates = np.asarray(coordinates) + minimum = np.min(coordinates, axis=0) + maximum = np.max(coordinates, axis=0) + extent = np.maximum(maximum - minimum, 1e-6) + samples = max(64, min(512, max(rows, columns, len(scan)))) + step = extent / samples + minimum -= step + maximum += step + return ReconstructionGrid( + tuple(minimum), + tuple(maximum), + tuple(step), + frame=frame, + name=name, + ) + + +def estimate_geometry_steps( + config, + scan, + *, + frame="hkl", + percentile=10.0, + detector_samples=5, + frame_samples=32, +): + """Estimate axis steps from local detector/scan geometry Jacobians. + + Detector-pixel rows and columns are treated as unit-width uniform + variables. The third Jacobian direction is the exposure sweep, or the + adjacent-frame center displacement for stationary exposures. The returned + values are the requested percentile of the local one-sigma axis + projections in r.l.u. for ``hkl`` or ``Angstrom^-1`` for Q frames. + + :param config: + Central :class:`ConfigData` experiment snapshot. + :param scan: + Active scan backend providing exposure angle bounds in radians. + :param str frame: + Output coordinate frame. + :param float percentile: + Robust percentile of sampled local axis resolutions. + :param int detector_samples: + Number of representative pixels sampled along each detector axis. + :param int frame_samples: + Maximum number of representative scan frames. + :returns: + Three geometry-matched axis steps. + :rtype: tuple[float, float, float] + :raises ValueError: + If the percentile or sampling counts are invalid, or the scan does not + span sufficient reciprocal-space dimensions. + """ + percentile = float(percentile) + if not 0.0 < percentile <= 100.0: + raise ValueError("Resolution percentile must be above 0 and at most 100") + detector_samples = int(detector_samples) + frame_samples = int(frame_samples) + if detector_samples < 2 or frame_samples < 1: + raise ValueError( + "Use at least two detector samples and one frame sample" + ) + + detector = config.detector + rows, columns = detector.detector.shape + bounds = np.asarray( + scan.exposure_angle_bounds(config, fallback="stationary"), + dtype=np.float64, + ) + if bounds.shape != (len(scan), 2, 4): + raise ValueError( + "Exposure angle bounds must have shape (frames, 2, 4)" + ) + excluded = set(config.corrections.excluded_frames) + included = np.asarray( + [index for index in range(len(scan)) if index not in excluded], + dtype=np.int64, + ) + if included.size == 0: + raise ValueError("No included frames are available for resolution sampling") + + detector_rows = np.unique( + np.rint( + np.linspace(0, rows - 1, min(detector_samples, rows)) + ).astype(np.int64) + ) + detector_columns = np.unique( + np.rint( + np.linspace(0, columns - 1, min(detector_samples, columns)) + ).astype(np.int64) + ) + sampled_positions = np.unique( + np.rint( + np.linspace(0, included.size - 1, min(frame_samples, included.size)) + ).astype(np.int64) + ) + sampled_frames = included[sampled_positions] + + from .datautils.xrayutils.reconstruction import _native_module + + ub = config.ub_calculator + kernel = _native_module().ReconstructionKernel( + np.full(3, -1e12), + np.ones(3), + np.full(3, 2_000_000_000, dtype=np.int64), + np.ones(3, dtype=np.int64), + frame, + float(ub.getK()), + np.ascontiguousarray(np.linalg.inv(ub.getUB())), + np.ascontiguousarray(np.linalg.inv(ub.getU())), + 0, + 1, + 1, + 1024 * 1024, + ) + + def pixel_rays(row, column): + ray_rows = np.asarray([row - 0.5, row + 0.5]) + ray_columns = np.asarray([column - 0.5, column + 0.5]) + row_grid, column_grid = np.meshgrid( + ray_rows, ray_columns, indexing="ij" + ) + gamma, delta = detector.primBeamPoints(row_grid, column_grid) + cosine_gamma = np.cos(gamma) + rays = np.empty((2, 2, 3), dtype=np.float64) + rays[..., 0] = np.sin(delta) * cosine_gamma + rays[..., 1] = np.cos(delta) * cosine_gamma + rays[..., 2] = np.sin(gamma) + rays /= np.linalg.norm(rays, axis=-1, keepdims=True) + return np.ascontiguousarray(rays) + + def coordinate(rays, frame_index, u, v, t): + return kernel.coordinate( + rays, + np.ascontiguousarray(bounds[frame_index, 0]), + np.ascontiguousarray(bounds[frame_index, 1]), + 0, + 0, + u, + v, + t, + ) + + local_sigmas = [] + scan_spans = [] + for frame_index in sampled_frames: + included_position = int(np.searchsorted(included, frame_index)) + neighbor_index = None + if included.size > 1: + neighbor_position = ( + included_position + 1 + if included_position + 1 < included.size + else included_position - 1 + ) + neighbor_index = int(included[neighbor_position]) + for row in detector_rows: + for column in detector_columns: + rays = pixel_rays(int(row), int(column)) + row_vector = ( + coordinate(rays, frame_index, 1.0, 0.5, 0.5) + - coordinate(rays, frame_index, 0.0, 0.5, 0.5) + ) + column_vector = ( + coordinate(rays, frame_index, 0.5, 1.0, 0.5) + - coordinate(rays, frame_index, 0.5, 0.0, 0.5) + ) + scan_vector = ( + coordinate(rays, frame_index, 0.5, 0.5, 1.0) + - coordinate(rays, frame_index, 0.5, 0.5, 0.0) + ) + if ( + np.linalg.norm(scan_vector) + <= np.finfo(np.float64).eps + and neighbor_index is not None + ): + scan_vector = ( + coordinate(rays, neighbor_index, 0.5, 0.5, 0.5) + - coordinate(rays, frame_index, 0.5, 0.5, 0.5) + ) + scan_spans.append(np.linalg.norm(scan_vector)) + jacobian = np.column_stack( + (row_vector, column_vector, scan_vector) + ) + local_sigmas.append( + np.sqrt(np.sum(jacobian * jacobian, axis=1) / 12.0) + ) + + local_sigmas = np.asarray(local_sigmas, dtype=np.float64) + if not scan_spans or max(scan_spans) <= np.finfo(np.float64).eps: + raise ValueError( + "The scan has no exposure sweep or adjacent-frame displacement; " + "a three-dimensional geometry-matched step cannot be estimated" + ) + steps = np.percentile(local_sigmas, percentile, axis=0) + if np.any(~np.isfinite(steps)) or np.any(steps <= 0): + raise ValueError( + "The sampled geometry does not span all three output axes; " + "geometry-matched 3-D steps cannot be estimated" + ) + return tuple(float(value) for value in steps) + + +def prepare_job( + gui, + job_path, + *, + grids, + scratch_path, + output_path, + accuracy="balanced", + advanced_depth=None, + compression_override=None, + angle_fallback="stationary", + user_note="", + thread_override=None, + memory_override_bytes=None, + frame_batch=None, + tile_shape=None, + work_block_pixels=None, + partition_chunk_span=None, + threads_per_image=4, + accumulation_budget_bytes=None, + cluster_settings=None, +): + """Freeze current orGUI state into an immutable reconstruction job.""" + if gui.fscan is None: + raise RuntimeError("Load a scan before preparing reconstruction") + if accuracy not in ACCURACY_DEPTHS: + raise ValueError(f"Unknown accuracy preset: {accuracy}") + if angle_fallback not in {"stationary", "midpoint"}: + raise ValueError("Angle fallback must be stationary or midpoint") + if advanced_depth is not None and not 0 <= advanced_depth <= 8: + raise ValueError("Advanced split depth must be between 0 and 8") + if compression_override is not None and compression_override not in FILTERS: + raise ValueError( + f"Unknown HDF5 compression override: {compression_override}" + ) + if int(threads_per_image) < 1: + raise ValueError("Threads per image must be at least one") + if ( + accumulation_budget_bytes is not None + and int(accumulation_budget_bytes) < 1024**2 + ): + raise ValueError( + "Per-worker accumulation budget must be at least 1 MiB" + ) + grid_values = [ + asdict(grid) if isinstance(grid, ReconstructionGrid) else dict(grid) + for grid in grids + ] + if not grid_values: + raise ValueError("At least one reconstruction grid is required") + for values in grid_values: + _GridSpec(**values) + config = ConfigData.from_gui(gui) + for monitor in config.corrections.monitor_corrections: + if not hasattr(gui.fscan, monitor): + raise ValueError( + f"Active scan has no monitor counter named {monitor!r}" + ) + reference = ScanReference.from_scan(gui.fscan) + scratch = Path(scratch_path).absolute() + for component in (scratch, *scratch.parents): + if component.exists() and not component.is_dir(): + raise ValueError( + "Scratch directory cannot be created because a path " + f"component is a file: {component}" + ) + scratch.mkdir(parents=True, exist_ok=True) + assets = scratch / "job-assets.nxs" + assets_sha256 = _snapshot_assets(gui, config, assets) + fingerprint = sha256( + json.dumps( + reference.to_dict(), sort_keys=True, separators=(",", ":") + ).encode() + ).hexdigest() + job = ReconstructionJob( + config=config_data_to_json(config), + scan_reference=reference.to_dict(), + grids=grid_values, + scratch_path=str(scratch), + output_path=str(Path(output_path).absolute()), + compression=( + compression_override + if compression_override is not None + else _compression_name(gui.database.compression) + ), + assets_path=str(assets), + assets_sha256=assets_sha256, + source_fingerprint_sha256=fingerprint, + threads_per_image=int(threads_per_image), + accumulation_budget_bytes=( + None + if accumulation_budget_bytes is None + else int(accumulation_budget_bytes) + ), + build_metadata=_build_metadata(), + runtime_threads=max(1, int(gui.numberthreads)), + runtime_memory_bytes=max(1, int(gui.maxMemory * 1024 * 1024)), + accuracy=accuracy, + advanced_depth=advanced_depth, + angle_fallback=angle_fallback, + user_note=user_note, + thread_override=thread_override, + memory_override_bytes=memory_override_bytes, + frame_batch=frame_batch, + tile_shape=tile_shape, + work_block_pixels=work_block_pixels, + partition_chunk_span=partition_chunk_span, + cluster_settings=dict(cluster_settings or {}), + ) + ranges, tiles = _execution_layout(job, gui.fscan, config) + job.expected_map_tasks = len(ranges) + write_job(job, job_path) + return job + + +def verify_job(job): + """Verify immutable job assets and source fingerprints.""" + ScanReference.from_dict(job.scan_reference).verify() + reference_digest = sha256( + json.dumps( + job.scan_reference, sort_keys=True, separators=(",", ":") + ).encode() + ).hexdigest() + if reference_digest != job.source_fingerprint_sha256: + raise RuntimeError("Reconstruction scan-reference digest mismatch") + if _sha256_file(job.assets_path) != job.assets_sha256: + raise RuntimeError("Reconstruction job assets checksum mismatch") + return True + + +def _load_assets(job): + result = {} + with h5py.File(job.assets_path, "r") as h5file: + for name in ("mask", "background", "background_variance"): + if name in h5file: + result[name] = h5file[name][()] + if "repair" in h5file: + repair = h5file["repair"] + result["repair"] = { + "max_component_pixels": int( + repair.attrs["max_component_pixels"] + ), + "max_span": int(repair.attrs["max_span"]), + "radius": int(repair.attrs["radius"]), + "min_valid_neighbors": int( + repair.attrs["min_valid_neighbors"] + ), + "row_gaps": repair["row_gaps"][()], + "column_gaps": repair["column_gaps"][()], + } + return result + + +def _correction_pipeline(config, scan, assets, provenance): + correction = config.corrections + detector = config.detector + background = ( + np.asarray(assets["background"], dtype=np.float64) + if correction.use_background and "background" in assets + else None + ) + background_variance = ( + np.asarray(assets["background_variance"], dtype=np.float64) + if background is not None and "background_variance" in assets + else None + ) + static_mask = ( + np.ascontiguousarray(assets["mask"], dtype=bool) + if correction.use_mask and "mask" in assets + else None + ) + static_factor = None + if correction.use_solid_angle: + static_factor = 1.0 / np.asarray( + detector.solidAngleArray(), dtype=np.float64 + ) + provenance.setdefault("factor_uncertainty", {})[ + "solid_angle" + ] = "deterministic-no-uncertainty" + if correction.use_polarization: + polarization = np.asarray( + detector.polarization( + factor=detector._polFactor, axis_offset=detector._polAxis + ), + dtype=np.float64, + ) + if static_factor is None: + static_factor = 1.0 / polarization + else: + static_factor /= polarization + provenance.setdefault("factor_uncertainty", {})[ + "polarization" + ] = "deterministic-no-uncertainty" + if static_factor is not None: + static_factor = np.ascontiguousarray(static_factor, dtype=np.float64) + static_factor_squared = np.square(static_factor) + else: + static_factor_squared = None + + repair_plan = None + if ( + correction.repair_masked_pixels + and static_mask is not None + and "repair" in assets + ): + repair_plan = create_pixel_repair_plan( + static_mask, **dict(assets["repair"]) + ) + provenance["repair_plan"] = repair_plan.configuration() + provenance["repair_covariance"] = "marginal-variance-only" + + exposure = ( + np.asarray(scan.exposure_time, dtype=np.float64) + if correction.normalize_exposure and hasattr(scan, "exposure_time") + else None + ) + exposure_variance = ( + np.asarray(scan.exposure_time_variance, dtype=np.float64) + if exposure is not None and hasattr(scan, "exposure_time_variance") + else None + ) + monitor_values = { + name: np.asarray(getattr(scan, name), dtype=np.float64) + for name in correction.monitor_corrections + } + monitor_variances = { + name: ( + np.asarray(getattr(scan, f"{name}_variance"), dtype=np.float64) + if hasattr(scan, f"{name}_variance") + else None + ) + for name in correction.monitor_corrections + } + + def frame_value(values, frame_index): + return ( + float(values) + if values.ndim == 0 + else float(values[frame_index]) + ) + + def apply_factor(intensity, variance, factor, factor_variance, name): + if factor_variance is not None: + variance *= factor**2 + variance += intensity**2 * factor_variance + intensity *= factor + provenance.setdefault("factor_uncertainty", {})[ + name + ] = "propagated" + else: + intensity *= factor + variance *= factor**2 + provenance.setdefault("factor_uncertainty", {})[ + name + ] = "deterministic-no-uncertainty" + + def correct_frame(payload, raw, frame_index): + source_variance = getattr(payload, "variance", None) + image_provenance = getattr(payload, "processing_provenance", None) + if image_provenance: + provenance.setdefault("image_processing", {}).update( + dict(image_provenance) + ) + if source_variance is None: + variance = np.maximum(raw, 0.0) + provenance["variance_fallback"] = "clipped-current-image-poisson" + else: + variance = np.asarray(source_variance, dtype=np.float64).copy() + intensity = np.asarray(raw, dtype=np.float64).copy() + if background is not None: + intensity -= background + if background_variance is not None: + variance += background_variance + else: + provenance["background_variance"] = "deterministic" + if static_mask is None: + mask = np.zeros(intensity.shape, dtype=bool) + else: + mask = static_mask.copy() + if repair_plan is not None: + mask, _ = repair_plan.apply_inplace(intensity, variance) + if static_factor is not None: + intensity *= static_factor + variance *= static_factor_squared + if exposure is not None: + value = frame_value(exposure, frame_index) + if value <= 0 or not math.isfinite(value): + raise ValueError("Exposure time must be finite and positive") + factor = 1.0 / value + factor_variance = None + if exposure_variance is not None: + value_variance = frame_value( + exposure_variance, frame_index + ) + factor_variance = value_variance / value**4 + apply_factor( + intensity, + variance, + factor, + factor_variance, + "exposure", + ) + elif correction.normalize_exposure: + provenance["exposure_normalization"] = "unavailable" + for name, values in monitor_values.items(): + value = frame_value(values, frame_index) + if value == 0 or not math.isfinite(value): + raise ValueError(f"Monitor {name} must be finite and nonzero") + monitor_variance = monitor_variances[name] + factor_variance = None + if monitor_variance is not None: + value_variance = frame_value(monitor_variance, frame_index) + factor_variance = value_variance / value**4 + apply_factor( + intensity, + variance, + 1.0 / value, + factor_variance, + f"monitor:{name}", + ) + mask |= ~np.isfinite(intensity) | ~np.isfinite(variance) + return intensity, variance, mask + + def correct(payload, raw, frame_index, tile): + """Compatibility call shape used by internal focused tests.""" + row_start, row_stop, column_start, column_stop = tile + full_raw = np.asarray(getattr(payload, "img", raw)) + intensity, variance, mask = correct_frame( + payload, full_raw, frame_index + ) + selection = np.s_[ + row_start:row_stop, column_start:column_stop + ] + return ( + intensity[selection], + variance[selection], + mask[selection], + ) + + correct.correct_frame = correct_frame + correct.repair_plan = repair_plan + return correct + + +def _included_ranges(count, excluded, batch_size): + included = [index for index in range(count) if index not in excluded] + ranges = [] + start = None + previous = None + for index in included: + if ( + start is None + or index != previous + 1 + or index - start >= batch_size + ): + if start is not None: + ranges.append((start, previous + 1)) + start = index + previous = index + if start is not None: + ranges.append((start, previous + 1)) + return ranges + + +def _execution_layout(job, scan, config): + rows, columns = config.detector.detector.shape + effective_memory = ( + job.memory_override_bytes or job.runtime_memory_bytes + ) + depth = ( + job.advanced_depth + if job.advanced_depth is not None + else ACCURACY_DEPTHS[job.accuracy] + ) + worst_leaves = 8**depth + estimated_native_bytes_per_pixel = 128 + 2 * worst_leaves * 40 + tile_pixels = max( + 1, + min( + 1024**2, + effective_memory + // max( + (len(job.grids) + 2) * estimated_native_bytes_per_pixel, + 1, + ), + ), + ) + tile_side = max(1, int(math.sqrt(tile_pixels))) + tile_rows, tile_columns = job.tile_shape or ( + min(tile_side, rows), + min(tile_side, columns), + ) + included_count = max( + 1, + len(scan) + - len( + { + frame + for frame in config.corrections.excluded_frames + if 0 <= frame < len(scan) + } + ), + ) + if job.frame_batch is None: + spec = job.internal_spec() + native_threads = max( + 1, min(job.threads_per_image, spec.threads) + ) + image_workers = max(1, spec.threads // native_threads) + target_tasks = max(1, image_workers * 4) + frame_batch = min( + AUTO_MAX_FRAMES_PER_TASK, + max(1, math.ceil(included_count / target_tasks)), + ) + else: + frame_batch = job.frame_batch + ranges = _included_ranges( + len(scan), set(config.corrections.excluded_frames), frame_batch + ) + tiles = [ + ( + row, + min(row + tile_rows, rows), + column, + min(column + tile_columns, columns), + ) + for row in range(0, rows, tile_rows) + for column in range(0, columns, tile_columns) + ] + return ranges, tiles + + +def _frame_parallelism( + spec, + tiles, + memory_bytes, + *, + stationary, + frames_per_task=1, + threads_per_image=1, + accumulation_budget_bytes=None, +): + """Derive bounded parallelism across images and within each image. + + Each frame worker processes that frame's detector tiles sequentially. The + configured CPU budget is divided between concurrent image workers and + native threads used by each image. The worker count is also bounded by a + depth-aware working-set estimate. Native blocks are reduced before task + records are retained, so the unattainable bound in which every footprint + leaf survives for the complete detector must not be multiplied by the + number of image workers. + + :param _ReconstructionSpec spec: + Frozen reconstruction compute settings. + :param iterable tiles: + Detector tiles as ``(row_start, row_stop, column_start, + column_stop)`` tuples. + :param int memory_bytes: + Total memory budget in bytes. + :param bool stationary: + Whether exposure start and end angles are identical. + :param int frames_per_task: + Number of images streamed by one resumable task. Images are not + retained together. + :param int threads_per_image: + Requested native reconstruction threads for each concurrent image. + :param accumulation_budget_bytes: + Optional requested retained-record bytes per worker. + :returns: + ``(image_workers, kernel_threads, per_worker_memory_bytes, + accumulation_budget_bytes)``. + :rtype: tuple[int, int, int, int] + """ + tiles = list(tiles) + if not tiles: + minimum = 1024**2 + return 1, 1, max(minimum, memory_bytes), minimum + tile_pixels = [ + (row_stop - row_start) * (column_stop - column_start) + for row_start, row_stop, column_start, column_stop in tiles + ] + detector_pixels = sum(tile_pixels) + children = 4 if stationary else 8 + depth_scale = 1 << max(0, spec.max_depth - 2) + sweep_scale = 2 if children == 8 else 1 + bytes_per_pixel = 128 * depth_scale * sweep_scale + image_memory = detector_pixels * bytes_per_pixel + worker_memory = max( + 1024**2, + image_memory, + ) + kernel_threads = max( + 1, min(int(threads_per_image), max(1, int(spec.threads))) + ) + cpu_workers = max(1, int(spec.threads) // kernel_threads) + minimum_accumulation = 1024**2 + if accumulation_budget_bytes is None: + memory_workers = max(1, memory_bytes // worker_memory) + else: + requested = max( + minimum_accumulation, int(accumulation_budget_bytes) + ) + required = ( + worker_memory + + ACCUMULATION_TRANSIENT_FACTOR * requested + ) + memory_workers = max(1, memory_bytes // required) + image_workers = max( + 1, + min(cpu_workers, memory_workers), + ) + safe_accumulation = max( + minimum_accumulation, + ( + memory_bytes // image_workers - worker_memory + ) + // ACCUMULATION_TRANSIENT_FACTOR, + ) + if accumulation_budget_bytes is None: + accumulation = min( + AUTO_MAX_ACCUMULATION_BYTES, safe_accumulation + ) + else: + accumulation = min( + max(minimum_accumulation, int(accumulation_budget_bytes)), + safe_accumulation, + ) + per_worker_memory = ( + worker_memory + + ACCUMULATION_TRANSIENT_FACTOR * accumulation + ) + return ( + image_workers, + kernel_threads, + per_worker_memory, + accumulation, + ) + + +def reconstruction_execution_settings(job, scan=None, config=None): + """Return the effective map-task and native execution layout. + + Automatic values are derived with the same functions used by + :func:`run_job`; unset advanced job fields therefore remain automatic. + + :param ReconstructionJob job: + Prepared reconstruction job. + :param scan: + Optional already-open scan. The job scan reference is opened when + omitted. + :param ConfigData config: + Optional decoded configuration snapshot. + :returns: + JSON-compatible detected execution settings. + :rtype: dict + """ + scan = job.scan if scan is None else scan + config = job.config_data if config is None else config + spec = job.internal_spec() + ranges, tiles = _execution_layout(job, scan, config) + bounds = scan.exposure_angle_bounds( + config, fallback=job.angle_fallback + ) + ray_cache_bytes = sum( + (tile[1] - tile[0] + 1) + * (tile[3] - tile[2] + 1) + * 3 + * np.dtype(np.float64).itemsize + for tile in tiles + ) + cache_detector_rays = ( + ray_cache_bytes <= spec.memory_budget_bytes // 8 + ) + scheduler_memory = max( + 1024**2, + spec.memory_budget_bytes + - (ray_cache_bytes if cache_detector_rays else 0), + ) + layouts = {} + for start, stop in ranges: + task_bounds = bounds[start:stop] + stationary = bool( + np.array_equal(task_bounds[:, 0], task_bounds[:, 1]) + ) + ( + image_workers, + kernel_threads, + per_worker_memory, + accumulation_budget, + ) = ( + _frame_parallelism( + spec, + tiles, + scheduler_memory, + stationary=stationary, + frames_per_task=stop - start, + threads_per_image=job.threads_per_image, + accumulation_budget_bytes=job.accumulation_budget_bytes, + ) + ) + key = ( + "stationary" if stationary else "swept", + image_workers, + kernel_threads, + per_worker_memory, + accumulation_budget, + ) + layouts[key] = { + "exposure": key[0], + "concurrent_image_workers": image_workers, + "native_threads_per_image": kernel_threads, + "tiles_per_image": len(tiles), + "memory_per_image_MiB": per_worker_memory / 1024**2, + "accumulation_MiB_per_worker": ( + accumulation_budget / 1024**2 + ), + } + if tiles: + tile_shape = ( + max(row_stop - row_start for row_start, row_stop, _, _ in tiles), + max( + column_stop - column_start + for _, _, column_start, column_stop in tiles + ), + ) + else: + tile_shape = (0, 0) + return { + "thread_budget": spec.threads, + "maximum_parallel_reducer_workers": min( + spec.threads, + max( + 1, + spec.memory_budget_bytes + // _MIN_REDUCER_WORKER_MEMORY, + ), + ), + "native_threads_per_image": max( + 1, min(job.threads_per_image, spec.threads) + ), + "memory_budget_MiB": spec.memory_budget_bytes / 1024**2, + "accumulation_budget_MiB_per_worker": min( + ( + layout["accumulation_MiB_per_worker"] + for layout in layouts.values() + ), + default=1.0, + ), + "frames_per_task": max( + (stop - start for start, stop in ranges), default=0 + ), + "detector_tile_shape": tile_shape, + "native_work_block_pixels": spec.work_block_pixels, + "parquet_chunk_span": spec.partition_chunk_span, + "frame_tasks": len(ranges), + "detector_tiles": len(tiles), + "map_tasks": len(ranges), + "detector_ray_cache_MiB": ( + ray_cache_bytes / 1024**2 if cache_detector_rays else 0.0 + ), + "parallel_layouts": list(layouts.values()), + } + + +def job_status(path): + """Return verified completion state for a job JSON.""" + job = read_job(path) + map_manifests = list(job.map_manifests) + if job.status != "complete": + verify_job(job) + scan = job.scan + config = job.config_data + ranges, tiles = _execution_layout(job, scan, config) + discovered = _discover_map_manifests( + job, + ranges, + tiles, + job.internal_spec(), + verify_partitions=False, + ) + map_manifests = list(discovered.values()) + result = { + "status": job.status, + "job_sha256": job.digest, + "map_tasks": { + "completed": len(map_manifests), + "pending": max( + 0, job.expected_map_tasks - len(map_manifests) + ), + "total": job.expected_map_tasks, + }, + "reduction_manifest": job.reduction_manifest, + "output_path": job.output_path, + "output_sha256": job.output_sha256, + "grids": [ + { + "name": _GridSpec(**grid).grid_name, + "frame": grid["frame"], + "shape": _GridSpec(**grid).shape, + } + for grid in job.grids + ], + "cleanup_errors": job.cleanup_errors, + } + return result + + +def _valid_map_manifest( + path, + frame_range, + detector_tiles, + *, + spec_hash, + job_digest, + verification_cache=None, + verify_partitions=True, +): + try: + manifest = _read_manifest(path) + return ( + manifest.kind == "map" + and manifest.status == "complete" + and manifest.spec_hash == spec_hash + and manifest.metadata.get("job_sha256") == job_digest + and tuple(manifest.frame_range) == tuple(frame_range) + and manifest.detector_tile is None + and manifest.metadata.get("detector_tiles") + == [list(tile) for tile in detector_tiles] + and ( + not verify_partitions + or all( + _verify_scratch_file(partition, verification_cache) + for partition in manifest.partitions + ) + ) + ) + except (OSError, ValueError, RuntimeError): + return False + + +def _discover_map_manifests( + job, + ranges, + tiles, + spec, + *, + verification_cache=None, + only_ranges=None, + verify_partitions=True, +): + manifest_root = Path(job.scratch_path) / "manifests" + candidates = { + str(Path(value).absolute()) for value in job.map_manifests + } + if manifest_root.exists(): + candidates.update( + str(path.absolute()) for path in manifest_root.glob("*.json") + ) + expected = { + tuple(frame_range) + for frame_range in ( + ranges if only_ranges is None else only_ranges + ) + } + existing = {} + for manifest_path in sorted(candidates): + try: + manifest = _read_manifest(manifest_path) + frame_range = tuple(manifest.frame_range) + if ( + frame_range in expected + and _valid_map_manifest( + manifest_path, + frame_range, + tiles, + spec_hash=spec.digest, + job_digest=job.digest, + verification_cache=verification_cache, + verify_partitions=verify_partitions, + ) + ): + existing[frame_range] = manifest_path + except (OSError, TypeError, ValueError, RuntimeError): + continue + return existing + + +def _base_provenance(job, config): + return { + **job.correction_provenance, + "job_sha256": job.digest, + "user_note": job.user_note, + "source_fingerprint_sha256": job.source_fingerprint_sha256, + "correction_state": config.corrections.to_dict(), + "cross_voxel_covariance": "marginal variances only", + "native_build": job.build_metadata, + } + + +def _merge_provenance(target, source): + for key, value in source.items(): + if isinstance(value, dict) and isinstance(target.get(key), dict): + _merge_provenance(target[key], value) + else: + target[key] = value + + +def run_cluster_map_task( + path, + task_index, + *, + cpus=1, + memory_bytes=1024**3, + progress=None, +): + """Execute one deterministic cluster map-array task. + + The task writes only its immutable map manifest and Parquet partitions; + it never mutates the shared reconstruction job JSON. + """ + path = Path(path).absolute() + job = read_job(path) + verify_job(job) + scan = job.scan + config = job.config_data + assets = _load_assets(job) + spec = job.internal_spec() + ranges, tiles = _execution_layout(job, scan, config) + task_index = int(task_index) + if task_index < 0 or task_index >= len(ranges): + raise IndexError( + f"Map task index {task_index} is outside 0..{len(ranges) - 1}" + ) + frame_range = ranges[task_index] + cpus = max(1, int(cpus)) + memory_bytes = max(1024**2, int(memory_bytes)) + verification_cache = set() + existing = _discover_map_manifests( + job, + ranges, + tiles, + spec, + verification_cache=verification_cache, + only_ranges=(frame_range,), + ) + if tuple(frame_range) in existing: + return { + "status": "complete", + "reused": True, + "task_index": task_index, + "frame_range": list(frame_range), + "manifest": existing[tuple(frame_range)], + } + bounds = scan.exposure_angle_bounds( + config, fallback=job.angle_fallback + ) + task_bounds = bounds[frame_range[0] : frame_range[1]] + stationary = bool( + np.array_equal(task_bounds[:, 0], task_bounds[:, 1]) + ) + execution_spec = _ReconstructionSpec.from_dict( + { + **spec.to_dict(), + "threads": cpus, + "memory_budget_bytes": memory_bytes, + } + ) + _, kernel_threads, _, accumulation_budget = _frame_parallelism( + execution_spec, + tiles, + memory_bytes, + stationary=stationary, + frames_per_task=frame_range[1] - frame_range[0], + threads_per_image=cpus, + accumulation_budget_bytes=job.accumulation_budget_bytes, + ) + provenance = _base_provenance(job, config) + correct = _correction_pipeline(config, scan, assets, provenance) + ray_cache = {} + for tile in tiles: + rays = _detector_corner_rays(config.detector, tile) + ray_cache[tuple(tile)] = (rays, _xxh3_128(rays)) + + completed = 0 + + def image_progress(frame_index, retained_bytes, segments): + nonlocal completed + completed += 1 + if progress is not None: + progress( + completed, + frame_range[1] - frame_range[0], + ( + f"Cluster map task {task_index}; frame {frame_index}; " + f"{retained_bytes / 1024**2:.0f} MiB retained; " + f"{segments} segments flushed" + ), + ) + + manifest = _map_frame_range( + spec, + scan, + config.detector, + config.ub_calculator, + frame_range, + tiles, + task_bounds, + Path(job.scratch_path) / "map", + correction_pipeline=correct, + job_digest=job.digest, + corner_rays={ + tile: values[0] for tile, values in ray_cache.items() + }, + corner_rays_fingerprints={ + tile: values[1] for tile, values in ray_cache.items() + }, + verification_cache=verification_cache, + kernel_threads=kernel_threads, + kernel_memory_budget_bytes=memory_bytes, + accumulation_budget_bytes=accumulation_budget, + image_progress=image_progress, + ) + manifest.metadata["correction_provenance"] = provenance + manifest_root = Path(job.scratch_path) / "manifests" + manifest_root.mkdir(parents=True, exist_ok=True) + manifest_path = manifest_root / f"{manifest.task_id}.json" + _write_manifest(manifest, manifest_path) + return { + "status": "complete", + "reused": False, + "task_index": task_index, + "frame_range": list(frame_range), + "manifest": str(manifest_path), + "partitions": len(manifest.partitions), + } + + +def run_cluster_finalize( + path, + *, + cpus=1, + memory_bytes=1024**3, + progress=None, +): + """Verify all cluster map tasks, then reduce and finalize their job.""" + path = Path(path).absolute() + job = read_job(path) + if job.status == "complete": + return job_status(path) + verify_job(job) + scan = job.scan + config = job.config_data + spec = job.internal_spec() + ranges, tiles = _execution_layout(job, scan, config) + verification_cache = set() + existing = _discover_map_manifests( + job, + ranges, + tiles, + spec, + verification_cache=verification_cache, + ) + missing = [ + index + for index, frame_range in enumerate(ranges) + if tuple(frame_range) not in existing + ] + if missing: + preview = ", ".join(map(str, missing[:20])) + suffix = "..." if len(missing) > 20 else "" + raise RuntimeError( + f"Cannot finalize: {len(missing)} map array tasks are missing " + f"({preview}{suffix})" + ) + job.map_manifests = [ + existing[tuple(frame_range)] for frame_range in ranges + ] + provenance = _base_provenance(job, config) + for manifest_path in job.map_manifests: + manifest = _read_manifest(manifest_path) + _merge_provenance( + provenance, + manifest.metadata.get("correction_provenance", {}), + ) + job.correction_provenance = provenance + write_job(job, path) + return run_job( + path, + progress=progress, + execution_threads=max(1, int(cpus)), + execution_memory_bytes=max(1024**2, int(memory_bytes)), + ) + + +def run_job( + path, + *, + progress=None, + execution_threads=None, + execution_memory_bytes=None, +): + """Run or resume one prepared reconstruction job.""" + path = Path(path).absolute() + job = read_job(path) + if job.status == "complete": + return job_status(path) + verify_job(job) + scan = job.scan + config = job.config_data + assets = _load_assets(job) + spec = job.internal_spec() + bounds = scan.exposure_angle_bounds( + config, fallback=job.angle_fallback + ) + ranges, tiles = _execution_layout(job, scan, config) + manifest_root = Path(job.scratch_path) / "manifests" + manifest_root.mkdir(parents=True, exist_ok=True) + verification_cache = set() + existing = _discover_map_manifests( + job, + ranges, + tiles, + spec, + verification_cache=verification_cache, + ) + job.map_manifests = list(existing.values()) + provenance = _base_provenance(job, config) + correct = _correction_pipeline(config, scan, assets, provenance) + total_tasks = len(ranges) + completed_tasks = sum( + tuple(frame_range) in existing + for frame_range in ranges + ) + total_images = sum(stop - start for start, stop in ranges) + completed_images = sum( + stop - start + for start, stop in ranges + if (start, stop) in existing + ) + if progress is not None: + progress( + completed_images, + total_images + 2, + ( + f"Mapping images {completed_images}/{total_images} " + f"({completed_tasks}/{total_tasks} tasks committed)" + ), + ) + effective_memory = ( + max(1024**2, int(execution_memory_bytes)) + if execution_memory_bytes is not None + else job.memory_override_bytes or job.runtime_memory_bytes + ) + reducer_threads = ( + max(1, int(execution_threads)) + if execution_threads is not None + else spec.threads + ) + ray_cache_bytes = sum( + (tile[1] - tile[0] + 1) + * (tile[3] - tile[2] + 1) + * 3 + * np.dtype(np.float64).itemsize + for tile in tiles + ) + cache_detector_rays = ray_cache_bytes <= effective_memory // 8 + ray_cache = {} + if cache_detector_rays: + for tile in tiles: + rays = _detector_corner_rays(config.detector, tile) + ray_cache[tuple(tile)] = (rays, _xxh3_128(rays)) + + pending_frames = [ + frame_range + for frame_range in ranges + if tuple(frame_range) not in existing + ] + + worker_limits = [] + kernel_thread_limits = [] + accumulation_limits = [] + scheduler_memory = max( + 1024**2, + effective_memory - (ray_cache_bytes if cache_detector_rays else 0), + ) + for frame_range in pending_frames: + task_bounds = bounds[frame_range[0] : frame_range[1]] + stationary = bool( + np.array_equal(task_bounds[:, 0], task_bounds[:, 1]) + ) + ( + image_limit, + native_threads, + _worker_memory, + accumulation_limit, + ) = _frame_parallelism( + spec, + tiles, + scheduler_memory, + stationary=stationary, + frames_per_task=frame_range[1] - frame_range[0], + threads_per_image=job.threads_per_image, + accumulation_budget_bytes=job.accumulation_budget_bytes, + ) + worker_limits.append(image_limit) + kernel_thread_limits.append(native_threads) + accumulation_limits.append(accumulation_limit) + image_workers = ( + min(len(pending_frames), min(worker_limits)) + if worker_limits + else 1 + ) + kernel_threads = ( + min(kernel_thread_limits) if kernel_thread_limits else 1 + ) + accumulation_budget = ( + min(accumulation_limits) if accumulation_limits else 1024**2 + ) + # The native budget is a per-call guard. The scheduler independently + # bounds the aggregate worker working set above. + kernel_memory_budget = effective_memory + progress_events = SimpleQueue() + cancellation = Event() + mapped_images = completed_images + parallel_mapping = image_workers > 1 + + def publish_image_progress(frame_index, retained_bytes, segments): + if cancellation.is_set(): + raise RuntimeError("Reconstruction mapping cancelled") + event = (frame_index, retained_bytes, segments) + if parallel_mapping: + progress_events.put(event) + else: + report_image_progress(event) + + def report_image_progress(event): + nonlocal mapped_images + frame_index, retained_bytes, segments = event + mapped_images += 1 + if progress is not None: + progress( + mapped_images, + total_images + 2, + ( + f"Mapping images {mapped_images}/{total_images}; " + f"frame {frame_index}; " + f"{completed_tasks}/{total_tasks} tasks committed; " + f"{retained_bytes / 1024**2:.0f} MiB retained; " + f"{segments} segments flushed" + ), + ) + + def drain_progress_events(): + while True: + try: + event = progress_events.get_nowait() + except Empty: + return + report_image_progress(event) + + def map_frame(frame_range): + task_bounds = bounds[frame_range[0] : frame_range[1]] + return frame_range, _map_frame_range( + spec, + scan, + config.detector, + config.ub_calculator, + frame_range, + tiles, + task_bounds, + Path(job.scratch_path) / "map", + correction_pipeline=correct, + job_digest=job.digest, + corner_rays={ + tile: values[0] for tile, values in ray_cache.items() + }, + corner_rays_fingerprints={ + tile: values[1] for tile, values in ray_cache.items() + }, + verification_cache=verification_cache, + kernel_threads=kernel_threads, + kernel_memory_budget_bytes=kernel_memory_budget, + accumulation_budget_bytes=accumulation_budget, + image_progress=publish_image_progress, + ) + + executor = ( + None + if image_workers == 1 + else ThreadPoolExecutor( + max_workers=image_workers, + thread_name_prefix="orgui-rsmap-image", + ) + ) + try: + for wave_start in range(0, len(pending_frames), image_workers): + wave = pending_frames[ + wave_start : wave_start + image_workers + ] + if executor is None: + mapped_frames = [map_frame(frame_task) for frame_task in wave] + else: + future_set = { + executor.submit(map_frame, frame_task) + for frame_task in wave + } + mapped_frames = [] + while future_set: + done, future_set = wait( + future_set, + timeout=0.1, + return_when=FIRST_COMPLETED, + ) + drain_progress_events() + mapped_frames.extend( + future.result() for future in done + ) + drain_progress_events() + mapped_frames.sort(key=lambda item: item[0]) + # This bounded worker wave has finished, so correction provenance + # is stable while the coordinator updates resumable state. + for _frame_range, manifest in mapped_frames: + manifest.metadata["correction_provenance"] = provenance + manifest_path = ( + manifest_root / f"{manifest.task_id}.json" + ) + _write_manifest(manifest, manifest_path) + job.map_manifests.append(str(manifest_path)) + completed_tasks += 1 + job.correction_provenance = provenance + write_job(job, path) + if progress is not None: + progress( + mapped_images, + total_images + 2, + ( + f"Mapping images {mapped_images}/{total_images}; " + f"{completed_tasks}/{total_tasks} tasks committed " + f"({image_workers} image workers, " + f"{kernel_threads} native threads/image, " + f"{accumulation_budget / 1024**2:.0f} MiB " + "accumulation/worker)" + ), + ) + except BaseException: + cancellation.set() + raise + finally: + if executor is not None: + executor.shutdown(wait=True, cancel_futures=True) + if progress is not None: + progress( + total_images, + total_images + 2, + f"Reducing {len(job.map_manifests)} mapping tasks", + ) + if job.map_manifests: + reduced = _reduce_partition( + job.map_manifests, + Path(job.scratch_path) / "reduced", + verification_cache=verification_cache, + memory_budget_bytes=effective_memory, + workers=reducer_threads, + checkpoint_root=manifest_root, + progress=( + None + if progress is None + else lambda done, count, message: progress( + total_images, + total_images + 2, + f"Reducing shards {done}/{count}: {message}", + ) + ), + ) + else: + reduced = _TaskManifest( + kind="reduce", + task_id=sha256((spec.digest + job.digest).encode()).hexdigest()[:24], + spec_hash=spec.digest, + status="complete", + spec=spec.to_dict(), + metadata={"job_sha256": job.digest, "empty_input": True}, + ) + reduced_path = manifest_root / f"{reduced.task_id}.json" + _write_manifest(reduced, reduced_path) + job.reduction_manifest = str(reduced_path) + job.status = "finalizing" + write_job(job, path) + if progress is not None: + progress( + total_images + 1, + total_images + 2, + "Finalizing HDF5", + ) + result = _finalize_reconstruction( + [reduced_path], + job.output_path, + provenance={ + **provenance, + "scan_reference": job.scan_reference, + }, + config=config, + verification_cache=verification_cache, + chunk_progress=( + None + if progress is None + else lambda written, count: progress( + total_images + 1, + total_images + 2, + f"Finalizing HDF5 chunk {written}/{count}", + ) + ), + ) + job.output_sha256 = result["sha256"] + job.status = "complete" + for target in ( + Path(job.scratch_path) / "map", + Path(job.scratch_path) / "reduced", + manifest_root, + Path(job.assets_path), + ): + try: + if target.is_dir(): + shutil.rmtree(target) + elif target.exists(): + target.unlink() + except OSError as error: + job.cleanup_errors.append(f"{target}: {error}") + write_job(job, path) + if progress is not None: + progress(total_images + 2, total_images + 2, "Complete") + return job_status(path) diff --git a/pyproject.toml b/pyproject.toml index ff2379b..e95ec0b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -52,9 +52,10 @@ Citation = "https://doi.org/10.5281/zenodo.12592485" console = ['qtconsole'] speedup = ['numba'] extendedfilesupport = ['ase', 'hdf5plugin'] +reconstruction = ['pyarrow >= 15', 'hdf5plugin'] symmetry = ['pyxtal'] -full = ['qtconsole', 'ase', 'numba' , 'PyOpenGL', 'hdf5plugin', 'pyxtal', 'py3Dmol'] -full-devel = ['qtconsole', 'ase', 'numba' , 'PyOpenGL', 'hdf5plugin', 'pyxtal', 'py3Dmol', 'ruff', 'pyupgrade'] +full = ['qtconsole', 'ase', 'numba' , 'PyOpenGL', 'hdf5plugin', 'pyarrow >= 15', 'pyxtal', 'py3Dmol'] +full-devel = ['qtconsole', 'ase', 'numba' , 'PyOpenGL', 'hdf5plugin', 'pyarrow >= 15', 'pyxtal', 'py3Dmol', 'ruff', 'pyupgrade'] [build-system] requires = ["meson-python", "pybind11", "setuptools-scm>=8", "numpy"] @@ -68,6 +69,7 @@ namespaces = true [project.scripts] # orGUI='orgui.main:main' +orGUI-rsmap='orgui.reconstruction_cli:main' [tool.setuptools] include-package-data = false