From 96d3f6fc1e2e0b6556e3244de01a130c5930f81a Mon Sep 17 00:00:00 2001 From: "T. Fuchs" Date: Mon, 27 Jul 2026 01:27:32 -0400 Subject: [PATCH 1/6] feat(phys): add native out-of-core reciprocal-space reconstruction - add a C++17 reciprocal-space kernel for deterministic HKL and Q mapping with adaptive pixel and exposure-footprint subdivision - propagate pixel variances through weighted voxel accumulation while preserving source-pixel correlations and recording marginal-variance limits - add parallel image loading, correction, tiling, native mapping, and byte-bounded in-memory accumulation - implement immutable Parquet map partitions, bounded external reduction, resumable manifests, checksums, and one-write-per-chunk HDF5 finalization - freeze detector calibration, crystal, UB, scan references, corrections, exclusions, masks, backgrounds, and repair data from authoritative orGUI state - add exact and fallback exposure-angle bounds for supported scan backends and serializable references for regular, sliced, interlaced, manual, simulated, and external-backend scans - add a reusable native pixel-repair plan with interpolation variance propagation for constant detector masks - add the reciprocal-space reconstruction dialog with grid coverage, editable intervals, file-size estimates, geometry-matched Jacobian steps, HDF5 settings, performance controls, preview, run, and resume workflows - add shared global HDF5 chunk settings and an optional reconstruction compression override using the central filter registry - add orGUI rsmap and orGUI-rsmap run, resume, and status commands that consume UI-prepared job descriptors - write standalone NeXus-style HDF5 results with axes, intensity, variance, weight, contributors, configuration, build metadata, and provenance - register completed files as external orGUI results without copying large reconstructed arrays into the active database - add focused numerical, variance, geometry, correction, determinism, storage, CLI, dialog, resume, and native-kernel regression tests - add mapping, alignment, storage, and reducer benchmarks - document the complete reconstruction workflow, every user and job parameter, scientific equations, storage architecture, CLI usage, and practical tuning --- .../benchmark_reconstruction_alignment.py | 149 ++ .../benchmark_reconstruction_mapping.py | 392 ++++ benchmarks/benchmark_reconstruction_reduce.py | 139 ++ .../benchmark_reconstruction_storage.py | 171 ++ doc/source/index.rst | 1 + .../reciprocal_space_reconstruction.rst | 570 ++++++ meson.build | 8 + orgui/app/HDF5SettingsDialog.py | 147 ++ orgui/app/QScanSelector.py | 47 +- orgui/app/ReconstructionDialog.py | 1331 ++++++++++++ orgui/app/_roi_sum_accel.py | 10 +- orgui/app/config_data.py | 177 +- orgui/app/cpp/roi_sum_cpp.cpp | 393 ++++ orgui/app/database.py | 33 + orgui/app/mask_config.py | 178 ++ orgui/app/orGUI.py | 59 +- orgui/app/test/test_config_data.py | 97 +- orgui/app/test/test_mask_config.py | 68 +- orgui/app/test/test_reconstruction_dialog.py | 416 ++++ orgui/backend/beamline/id31_tools.py | 61 +- orgui/backend/scans.py | 387 +++- orgui/backend/universalScanLoader.py | 14 +- .../cpp/reciprocal_reconstruction_cpp.cpp | 1800 +++++++++++++++++ orgui/datautils/xrayutils/reconstruction.py | 1612 +++++++++++++++ .../xrayutils/test/test_reconstruction.py | 879 ++++++++ .../test/test_reconstruction_angles.py | 165 ++ .../xrayutils/test/test_reconstruction_cli.py | 64 + .../test_reconstruction_geometry_steps.py | 78 + .../xrayutils/test/test_reconstruction_job.py | 163 ++ orgui/main.py | 9 +- orgui/reconstruction_cli.py | 67 + orgui/reconstruction_job.py | 1595 +++++++++++++++ pyproject.toml | 6 +- 33 files changed, 11235 insertions(+), 51 deletions(-) create mode 100644 benchmarks/benchmark_reconstruction_alignment.py create mode 100644 benchmarks/benchmark_reconstruction_mapping.py create mode 100644 benchmarks/benchmark_reconstruction_reduce.py create mode 100644 benchmarks/benchmark_reconstruction_storage.py create mode 100644 doc/source/reciprocal_space_reconstruction.rst create mode 100644 orgui/app/HDF5SettingsDialog.py create mode 100644 orgui/app/ReconstructionDialog.py create mode 100644 orgui/app/test/test_reconstruction_dialog.py create mode 100644 orgui/datautils/xrayutils/cpp/reciprocal_reconstruction_cpp.cpp create mode 100644 orgui/datautils/xrayutils/reconstruction.py create mode 100644 orgui/datautils/xrayutils/test/test_reconstruction.py create mode 100644 orgui/datautils/xrayutils/test/test_reconstruction_angles.py create mode 100644 orgui/datautils/xrayutils/test/test_reconstruction_cli.py create mode 100644 orgui/datautils/xrayutils/test/test_reconstruction_geometry_steps.py create mode 100644 orgui/datautils/xrayutils/test/test_reconstruction_job.py create mode 100644 orgui/reconstruction_cli.py create mode 100644 orgui/reconstruction_job.py 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..d682e83 --- /dev/null +++ b/doc/source/reciprocal_space_reconstruction.rst @@ -0,0 +1,570 @@ +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. + +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. +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. + +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. + +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. + +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/meson.build b/meson.build index 80e95da..29349e7 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/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 379e322..41e2f16 100644 --- a/orgui/app/QScanSelector.py +++ b/orgui/app/QScanSelector.py @@ -51,7 +51,6 @@ from . import qutils from .QReflectionSelector import QReflectionAnglesDialog from .QHKLDialog import HKLDialog -import runpy from contextlib import contextmanager @@ -725,11 +724,25 @@ def __init__(self, parentmainwindow, parent=None): self.useLorentzBox.setEnabled(False) self.useSolidAngleBox = qt.QCheckBox("Solid angle correction") self.usePolarizationBox = qt.QCheckBox("Polarization correction") + self.normalizeExposureBox = qt.QCheckBox("Normalize by exposure time") + self.normalizeExposureBox.setChecked(True) + self.monitorCorrectionsEdit = qt.QLineEdit() + self.monitorCorrectionsEdit.setPlaceholderText( + "Optional scan counters, comma-separated" + ) + self.monitorCorrectionsEdit.setToolTip( + "Counters are applied as divisive monitor normalizations." + ) optionsGroupLayout.addWidget(self.useMaskBox, 0, 0) optionsGroupLayout.addWidget(self.useLorentzBox, 1, 0) optionsGroupLayout.addWidget(self.useSolidAngleBox, 0, 1) optionsGroupLayout.addWidget(self.usePolarizationBox, 1, 1) + optionsGroupLayout.addWidget(self.normalizeExposureBox, 2, 0, 1, 2) + optionsGroupLayout.addWidget( + qt.QLabel("Monitor corrections:"), 3, 0 + ) + optionsGroupLayout.addWidget(self.monitorCorrectionsEdit, 3, 1) optionsGroup.setLayout(optionsGroupLayout) @@ -847,6 +860,12 @@ def set_integration_options(self, ddict): self.useSolidAngleBox.setChecked(ddict[key]) elif key == "polarization": self.usePolarizationBox.setChecked(ddict[key]) + elif key == "normalizeExposure": + self.normalizeExposureBox.setChecked(ddict[key]) + elif key == "monitorCorrections": + self.monitorCorrectionsEdit.setText( + ", ".join(map(str, ddict[key])) + ) elif key == "advanced": self.roioptions.set_parameters(ddict[key]) @@ -855,6 +874,12 @@ def get_integration_options(self): ddict["mask"] = self.useMaskBox.isChecked() ddict["solidAngle"] = self.useSolidAngleBox.isChecked() ddict["polarization"] = self.usePolarizationBox.isChecked() + ddict["normalizeExposure"] = self.normalizeExposureBox.isChecked() + ddict["monitorCorrections"] = tuple( + value.strip() + for value in self.monitorCorrectionsEdit.text().split(",") + if value.strip() + ) ddict["advanced"] = self.roioptions.get_parameters() return ddict @@ -1059,25 +1084,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..7483830 --- /dev/null +++ b/orgui/app/ReconstructionDialog.py @@ -0,0 +1,1331 @@ +"""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_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 + 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._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), + ("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.", + "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) + return widget + + 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 + + @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")) + + 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) + ) + 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 + ), + ) + 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 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.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..b912f03 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( + options.get("normalizeExposure", True) + ), + monitor_corrections=tuple( + options.get("monitorCorrections", ()) + ), + 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,46 @@ 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, + "normalizeExposure": self.corrections.normalize_exposure, + "monitorCorrections": ( + 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 +513,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 +538,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 +559,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 3d02683..2612cdd 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 @@ -558,6 +571,26 @@ def add_nxdict(self, nxentry, update_mode="add", h5path="/"): self.hdf5model.synchronizeH5pyObject(self.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 1497146..bc6ccf3 100644 --- a/orgui/app/orGUI.py +++ b/orgui/app/orGUI.py @@ -66,6 +66,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 @@ -158,6 +160,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() @@ -409,7 +413,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) @@ -515,6 +519,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") @@ -554,6 +560,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.""" @@ -2472,23 +2501,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..beec0ef 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,89 @@ 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, + "normalizeExposure": True, + "monitorCorrections": (), + } + ), + maskManager=SimpleNamespace( + settings=SimpleNamespace(pixel_repair=repair) + ), + excludedImagesDialog=SimpleNamespace( + getData=lambda: np.empty(0, dtype=np.int64) + ), + ) + + captured = ConfigData.from_gui(gui) + + assert captured.corrections.repair_masked_pixels is True + assert captured.corrections.use_mask is True + + +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..f92a894 --- /dev/null +++ b/orgui/app/test/test_reconstruction_dialog.py @@ -0,0 +1,416 @@ +"""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.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) >= 5 + 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", + "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.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_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 4a50cb5..154c124 100644 --- a/orgui/backend/beamline/id31_tools.py +++ b/orgui/backend/beamline/id31_tools.py @@ -123,6 +123,9 @@ def __init__(self, filename): # 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] @@ -158,10 +161,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]) @@ -172,6 +178,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) @@ -363,6 +370,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"] @@ -389,6 +404,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"] @@ -758,6 +788,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"] @@ -794,6 +832,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"] @@ -1019,6 +1072,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 1c0d715..725b6af 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): @@ -108,17 +113,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 @@ -150,3 +213,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/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..f1ef94d --- /dev/null +++ b/orgui/datautils/xrayutils/reconstruction.py @@ -0,0 +1,1612 @@ +"""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 dataclasses import asdict, dataclass, field +from hashlib import sha256 +import importlib +import json +import math +import os +from pathlib import Path +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"} + + +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): + _, _, 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.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), + ) + + 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, + progress: Callable[[int, int, str], None] | None = None, + checkpoint_root=None, +) -> _TaskManifest: + """Externally reduce mapping partitions into bounded chunk shards.""" + 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, + int(memory_budget_bytes) // (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): + manifest = _TaskManifest( + kind="reduce", + task_id=reduce_id, + spec_hash=spec.digest, + status=status, + spec=spec.to_dict(), + chunks=chunks, + source_tasks=source_tasks, + metadata={ + **base_metadata, + "local_voxel_shard_span": local_span, + "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 + } + readers = {} + active_group = None + checkpoint_interval = 64 + total_plans = len(plans) + completed = 0 + for ( + grid_name, + bucket, + chunk_id, + shard_start, + shard_stop, + partitions, + ) in plans: + key = shard_key( + grid_name, + chunk_id, + shard_start, + shard_stop, + ) + group_key = (grid_name, bucket) + if group_key != active_group: + readers.clear() + active_group = group_key + 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 + for partition in partitions: + _verify_scratch_file(partition, verification_cache) + if partition.uri not in readers: + batch_rows = max( + 4096, + min( + 131072, + int(memory_budget_bytes) + // (max(1, len(partitions)) * 48 * 4), + ), + ) + readers[partition.uri] = _ParquetRangeReader( + partition.uri, + batch_size=batch_rows, + ) + 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) + 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, + ) + _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"{grid_name} chunk {chunk_id}, " + f"{shard_start}:{shard_stop}" + ), + ) + 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..70ae86f --- /dev/null +++ b/orgui/datautils/xrayutils/test/test_reconstruction.py @@ -0,0 +1,879 @@ +"""Regression tests for native reciprocal-space reconstruction.""" + +from __future__ import annotations + +import json + +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() + reduced = _reduce_partition( + manifest_paths, + tmp_path / "reduced", + verification_cache=verification_cache, + memory_budget_bytes=1, + 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 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=1, + 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..3bbe640 --- /dev/null +++ b/orgui/datautils/xrayutils/test/test_reconstruction_cli.py @@ -0,0 +1,64 @@ +"""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_only_exposes_job_commands(): + """The CLI has no independent specification or stage interfaces.""" + parser = reconstruction_cli.build_parser() + help_text = parser.format_help() + for command in ("run", "resume", "status"): + assert command in help_text + for removed in ("preview", "map", "reduce", "finalize"): + assert removed not in help_text + + +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 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..2e9a58e --- /dev/null +++ b/orgui/datautils/xrayutils/test/test_reconstruction_job.py @@ -0,0 +1,163 @@ +"""End-to-end tests for the centralized reconstruction job.""" + +from hashlib import sha256 +import json +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_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"] diff --git a/orgui/main.py b/orgui/main.py index 1384527..8c65c69 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..b181199 --- /dev/null +++ b/orgui/reconstruction_cli.py @@ -0,0 +1,67 @@ +"""Command-line execution of UI-prepared reciprocal-space jobs.""" + +from __future__ import annotations + +import argparse +import json +import sys +import time + +from .reconstruction_job import job_status, run_job + + +def _run(arguments): + 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 = run_job(arguments.job, progress=update) + 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) + 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_job.py b/orgui/reconstruction_job.py new file mode 100644 index 0000000..6f5e189 --- /dev/null +++ b/orgui/reconstruction_job.py @@ -0,0 +1,1595 @@ +"""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, + _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) + 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, +): + """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, + ) + 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, + "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) + result = { + "status": job.status, + "job_sha256": job.digest, + "map_tasks": { + "completed": len(job.map_manifests), + "pending": max( + 0, job.expected_map_tasks - len(job.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, + } + if job.status != "complete": + verify_job(job) + return result + + +def _valid_map_manifest( + path, + frame_range, + detector_tiles, + *, + spec_hash, + job_digest, + verification_cache=None, +): + 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 all( + _verify_scratch_file(partition, verification_cache) + for partition in manifest.partitions + ) + ) + except (OSError, ValueError, RuntimeError): + return False + + +def run_job(path, *, progress=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 = {} + for manifest_path in job.map_manifests: + try: + manifest = _read_manifest(manifest_path) + key = tuple(manifest.frame_range) + if _valid_map_manifest( + manifest_path, + key, + tiles, + spec_hash=spec.digest, + job_digest=job.digest, + verification_cache=verification_cache, + ): + existing[key] = manifest_path + except (OSError, ValueError, RuntimeError): + continue + job.map_manifests = list(existing.values()) + provenance = { + **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, + } + 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 = ( + job.memory_override_bytes or job.runtime_memory_bytes + ) + 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_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, + 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 b75701d..ccd8d80 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -51,9 +51,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"] @@ -67,6 +68,7 @@ namespaces = true [project.scripts] # orGUI='orgui.main:main' +orGUI-rsmap='orgui.reconstruction_cli:main' [tool.setuptools] include-package-data = false From 912c7a01722c5a38ef06341c89f26ba996dd3142 Mon Sep 17 00:00:00 2001 From: "T. Fuchs" Date: Mon, 27 Jul 2026 19:05:56 -0400 Subject: [PATCH 2/6] perf: rsmap parallel reduction of submaps --- .../reciprocal_space_reconstruction.rst | 12 +- orgui/datautils/xrayutils/reconstruction.py | 360 ++++++++++++------ .../xrayutils/test/test_reconstruction.py | 33 +- orgui/reconstruction_job.py | 10 + 4 files changed, 306 insertions(+), 109 deletions(-) diff --git a/doc/source/reciprocal_space_reconstruction.rst b/doc/source/reciprocal_space_reconstruction.rst index d682e83..b866d85 100644 --- a/doc/source/reciprocal_space_reconstruction.rst +++ b/doc/source/reciprocal_space_reconstruction.rst @@ -392,7 +392,8 @@ The execution stages are: 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. + 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. @@ -401,6 +402,15 @@ 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 ----- diff --git a/orgui/datautils/xrayutils/reconstruction.py b/orgui/datautils/xrayutils/reconstruction.py index f1ef94d..5a8b4fe 100644 --- a/orgui/datautils/xrayutils/reconstruction.py +++ b/orgui/datautils/xrayutils/reconstruction.py @@ -12,6 +12,7 @@ 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 @@ -19,6 +20,7 @@ import math import os from pathlib import Path +from queue import Queue from typing import Any import h5py @@ -35,6 +37,7 @@ ) _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): @@ -424,7 +427,7 @@ def _read_parquet(uri): class _ParquetRangeReader: """Stream monotonically increasing key ranges from sorted Parquet.""" - def __init__(self, uri, *, batch_size=131072): + 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) @@ -432,6 +435,7 @@ def __init__(self, uri, *, batch_size=131072): 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 @@ -469,6 +473,7 @@ def _start(self, chunk_id, local_start): batch_size=self.batch_size, row_groups=row_groups, columns=list(_PARTIAL_COLUMNS), + use_threads=self.use_threads, ) def _advance(self): @@ -1016,10 +1021,26 @@ def _reduce_partition( *, 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.""" + """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") @@ -1076,7 +1097,7 @@ def _reduce_partition( 256 * 1024, min( 4 * 1024 * 1024, - int(memory_budget_bytes) // (8 * 48), + worker_memory_budget // (8 * 48), ), ) plans = [] @@ -1155,17 +1176,30 @@ def shard_key(grid_name, chunk_id, 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=chunks, + 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), }, ) @@ -1182,29 +1216,18 @@ def checkpoint(status): ) for chunk in chunks } - readers = {} - active_group = None checkpoint_interval = 64 total_plans = len(plans) completed = 0 - for ( - grid_name, - bucket, - chunk_id, - shard_start, - shard_stop, - partitions, - ) in plans: + 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, ) - group_key = (grid_name, bucket) - if group_key != active_group: - readers.clear() - active_group = group_key if key in existing or key in completed_shards: completed += 1 if progress is not None: @@ -1214,98 +1237,223 @@ def checkpoint(status): f"Reusing reduced shard {completed}/{total_plans}", ) continue - for partition in partitions: - _verify_scratch_file(partition, verification_cache) - if partition.uri not in readers: - batch_rows = max( - 4096, - min( - 131072, - int(memory_budget_bytes) - // (max(1, len(partitions)) * 48 * 4), - ), - ) - readers[partition.uri] = _ParquetRangeReader( - partition.uri, - batch_size=batch_rows, - ) - levels = [] - for partition in partitions: - batch = readers[partition.uri].read( + 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, - ) - 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) - 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, - ) - _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, + 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 ( - f"Reduced shard {completed}/{total_plans}: " - f"{grid_name} chunk {chunk_id}, " - f"{shard_start}:{shard_stop}" - ), - ) + 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") diff --git a/orgui/datautils/xrayutils/test/test_reconstruction.py b/orgui/datautils/xrayutils/test/test_reconstruction.py index 70ae86f..8a15c5b 100644 --- a/orgui/datautils/xrayutils/test/test_reconstruction.py +++ b/orgui/datautils/xrayutils/test/test_reconstruction.py @@ -3,6 +3,7 @@ from __future__ import annotations import json +import threading import h5py import numpy as np @@ -827,11 +828,31 @@ def test_reduce_shards_large_chunk_and_reuses_checkpoint(tmp_path, monkeypatch): 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=1, + memory_budget_bytes=128 * 1024**2, + workers=2, progress=lambda current, total, message: progress.append( (current, total, message) ), @@ -841,6 +862,13 @@ def test_reduce_shards_large_chunk_and_reuses_checkpoint(tmp_path, monkeypatch): 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): @@ -853,7 +881,8 @@ def unexpected_read(*args, **kwargs): manifest_paths, tmp_path / "reduced", verification_cache=verification_cache, - memory_budget_bytes=1, + memory_budget_bytes=128 * 1024**2, + workers=2, checkpoint_root=tmp_path / "manifests", ) assert resumed.status == "complete" diff --git a/orgui/reconstruction_job.py b/orgui/reconstruction_job.py index 6f5e189..b17fc09 100644 --- a/orgui/reconstruction_job.py +++ b/orgui/reconstruction_job.py @@ -22,6 +22,7 @@ from .backend.scans import ScanReference from .datautils.xrayutils.reconstruction import ( _GridSpec, + _MIN_REDUCER_WORKER_MEMORY, _ReconstructionSpec, _TaskManifest, _detector_corner_rays, @@ -1166,6 +1167,14 @@ def reconstruction_execution_settings(job, scan=None, config=None): 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) ), @@ -1524,6 +1533,7 @@ def map_frame(frame_range): Path(job.scratch_path) / "reduced", verification_cache=verification_cache, memory_budget_bytes=effective_memory, + workers=spec.threads, checkpoint_root=manifest_root, progress=( None From 9204490cd494416bef3ab043c64f946bef83742a Mon Sep 17 00:00:00 2001 From: "T. Fuchs" Date: Mon, 27 Jul 2026 20:59:26 -0400 Subject: [PATCH 3/6] feat: add SGE and Slurm batch arrays for rsmap reconstruction --- .../reciprocal_space_reconstruction.rst | 123 +++++++ orgui/app/ReconstructionDialog.py | 318 ++++++++++++++++ orgui/app/test/test_reconstruction_dialog.py | 50 ++- .../xrayutils/test/test_reconstruction_cli.py | 55 ++- .../test/test_reconstruction_cluster.py | 85 +++++ .../xrayutils/test/test_reconstruction_job.py | 80 +++++ orgui/reconstruction_cli.py | 70 +++- orgui/reconstruction_cluster.py | 322 +++++++++++++++++ orgui/reconstruction_job.py | 340 ++++++++++++++++-- 9 files changed, 1400 insertions(+), 43 deletions(-) create mode 100644 orgui/datautils/xrayutils/test/test_reconstruction_cluster.py create mode 100644 orgui/reconstruction_cluster.py diff --git a/doc/source/reciprocal_space_reconstruction.rst b/doc/source/reciprocal_space_reconstruction.rst index b866d85..446a28f 100644 --- a/doc/source/reciprocal_space_reconstruction.rst +++ b/doc/source/reciprocal_space_reconstruction.rst @@ -449,6 +449,125 @@ The direct alias is equivalent: ``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 ------------------------ @@ -542,6 +661,10 @@ inspection, scheduling, and 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 ------------------ diff --git a/orgui/app/ReconstructionDialog.py b/orgui/app/ReconstructionDialog.py index 7483830..04dc3c1 100644 --- a/orgui/app/ReconstructionDialog.py +++ b/orgui/app/ReconstructionDialog.py @@ -10,6 +10,10 @@ from silx.gui import qt from .. import logger_utils +from ..reconstruction_cluster import ( + ClusterSettings, + generate_cluster_scripts, +) from ..reconstruction_job import ( ACCURACY_DEPTHS, ReconstructionGrid, @@ -138,6 +142,7 @@ def __init__(self, orgui, parent=None): 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) @@ -155,6 +160,7 @@ def __init__(self, orgui, parent=None): ("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), ): @@ -165,6 +171,10 @@ def __init__(self, orgui, parent=None): "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] @@ -641,6 +651,277 @@ def _paths_tab(self): 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,) @@ -732,6 +1013,10 @@ def _set_default_paths(self, stem): 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) @@ -1078,6 +1363,18 @@ def _show_execution_settings(self, job, *, scan=None, config=None): 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): @@ -1130,6 +1427,7 @@ def _prepare(self): * 1024 * 1024 ), + cluster_settings=self._cluster_settings().to_dict(), ) self._show_execution_settings( job, @@ -1273,6 +1571,25 @@ def run_local(self): 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.""" @@ -1304,6 +1621,7 @@ def open_job(self): 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)) diff --git a/orgui/app/test/test_reconstruction_dialog.py b/orgui/app/test/test_reconstruction_dialog.py index f92a894..04faa3f 100644 --- a/orgui/app/test/test_reconstruction_dialog.py +++ b/orgui/app/test/test_reconstruction_dialog.py @@ -42,6 +42,7 @@ def test_no_scan_actions_are_reported_without_raising(tmp_path, caplog): dialog.preview() dialog.prepare() dialog.run_local() + dialog.create_cluster_scripts() dialog.resume() dialog.refresh_live_state() @@ -49,7 +50,7 @@ def test_no_scan_actions_are_reported_without_raising(tmp_path, caplog): records = [ record for record in caplog.records if hasattr(record, "show_dialog") ] - assert len(records) >= 5 + 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) @@ -191,6 +192,11 @@ def test_settings_are_grouped_and_have_tooltips(tmp_path): "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 @@ -234,6 +240,24 @@ def test_settings_are_grouped_and_have_tooltips(tmp_path): 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, @@ -248,6 +272,30 @@ def test_settings_are_grouped_and_have_tooltips(tmp_path): 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) diff --git a/orgui/datautils/xrayutils/test/test_reconstruction_cli.py b/orgui/datautils/xrayutils/test/test_reconstruction_cli.py index 3bbe640..cb54961 100644 --- a/orgui/datautils/xrayutils/test/test_reconstruction_cli.py +++ b/orgui/datautils/xrayutils/test/test_reconstruction_cli.py @@ -34,14 +34,22 @@ def test_orgui_dispatches_rsmap_without_starting_gui(): ) -def test_cli_only_exposes_job_commands(): - """The CLI has no independent specification or stage interfaces.""" +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"): + 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 help_text + assert removed not in choices def test_cli_reports_progress_on_stderr(capsys): @@ -62,3 +70,42 @@ def fake_run(path, *, progress): 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_job.py b/orgui/datautils/xrayutils/test/test_reconstruction_job.py index 2e9a58e..68656ce 100644 --- a/orgui/datautils/xrayutils/test/test_reconstruction_job.py +++ b/orgui/datautils/xrayutils/test/test_reconstruction_job.py @@ -1,7 +1,9 @@ """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 @@ -18,6 +20,8 @@ ReconstructionJob, job_status, reconstruction_execution_settings, + run_cluster_finalize, + run_cluster_map_task, run_job, write_job, ) @@ -161,3 +165,79 @@ def observed_get_raw_img(self, index): 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/reconstruction_cli.py b/orgui/reconstruction_cli.py index b181199..25175a4 100644 --- a/orgui/reconstruction_cli.py +++ b/orgui/reconstruction_cli.py @@ -7,10 +7,17 @@ import sys import time -from .reconstruction_job import job_status, run_job +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 _run(arguments): +def _progress_operation(operation, *args, **kwargs): last_report = 0.0 def update(value, maximum, message): @@ -29,7 +36,40 @@ def update(value, maximum, message): ) last_report = now - result = run_job(arguments.job, progress=update) + 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)) @@ -54,6 +94,30 @@ def build_parser(prog=None): 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 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 index b17fc09..cf36183 100644 --- a/orgui/reconstruction_job.py +++ b/orgui/reconstruction_job.py @@ -146,6 +146,7 @@ class ReconstructionJob: 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): @@ -560,6 +561,7 @@ def prepare_job( 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: @@ -646,6 +648,7 @@ def prepare_job( 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) @@ -1205,13 +1208,27 @@ def reconstruction_execution_settings(job, scan=None, config=None): 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(job.map_manifests), + "completed": len(map_manifests), "pending": max( - 0, job.expected_map_tasks - len(job.map_manifests) + 0, job.expected_map_tasks - len(map_manifests) ), "total": job.expected_map_tasks, }, @@ -1228,8 +1245,6 @@ def job_status(path): ], "cleanup_errors": job.cleanup_errors, } - if job.status != "complete": - verify_job(job) return result @@ -1241,6 +1256,7 @@ def _valid_map_manifest( spec_hash, job_digest, verification_cache=None, + verify_partitions=True, ): try: manifest = _read_manifest(path) @@ -1253,16 +1269,279 @@ def _valid_map_manifest( and manifest.detector_tile is None and manifest.metadata.get("detector_tiles") == [list(tile) for tile in detector_tiles] - and all( - _verify_scratch_file(partition, verification_cache) - for partition in manifest.partitions + and ( + not verify_partitions + or all( + _verify_scratch_file(partition, verification_cache) + for partition in manifest.partitions + ) ) ) except (OSError, ValueError, RuntimeError): return False -def run_job(path, *, progress=None): +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) @@ -1280,32 +1559,15 @@ def run_job(path, *, progress=None): manifest_root = Path(job.scratch_path) / "manifests" manifest_root.mkdir(parents=True, exist_ok=True) verification_cache = set() - existing = {} - for manifest_path in job.map_manifests: - try: - manifest = _read_manifest(manifest_path) - key = tuple(manifest.frame_range) - if _valid_map_manifest( - manifest_path, - key, - tiles, - spec_hash=spec.digest, - job_digest=job.digest, - verification_cache=verification_cache, - ): - existing[key] = manifest_path - except (OSError, ValueError, RuntimeError): - continue + existing = _discover_map_manifests( + job, + ranges, + tiles, + spec, + verification_cache=verification_cache, + ) job.map_manifests = list(existing.values()) - provenance = { - **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, - } + provenance = _base_provenance(job, config) correct = _correction_pipeline(config, scan, assets, provenance) total_tasks = len(ranges) completed_tasks = sum( @@ -1328,7 +1590,14 @@ def run_job(path, *, progress=None): ), ) effective_memory = ( - job.memory_override_bytes or job.runtime_memory_bytes + 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) @@ -1494,6 +1763,7 @@ def map_frame(frame_range): # 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" ) @@ -1533,7 +1803,7 @@ def map_frame(frame_range): Path(job.scratch_path) / "reduced", verification_cache=verification_cache, memory_budget_bytes=effective_memory, - workers=spec.threads, + workers=reducer_threads, checkpoint_root=manifest_root, progress=( None From 187376e5b79ee354bf9b199feb5a9a73a113b4a6 Mon Sep 17 00:00:00 2001 From: "T. Fuchs" Date: Mon, 27 Jul 2026 21:09:46 -0400 Subject: [PATCH 4/6] docs: add doc/AGENTS.md and an automated changelog/doc-sync workflow --- .github/workflows/docs-sync.yml | 101 +++++++++++++++++++++++ AGENTS.md | 2 + doc/AGENTS.md | 140 ++++++++++++++++++++++++++++++++ 3 files changed, 243 insertions(+) create mode 100644 .github/workflows/docs-sync.yml create mode 100644 doc/AGENTS.md diff --git a/.github/workflows/docs-sync.yml b/.github/workflows/docs-sync.yml new file mode 100644 index 0000000..78850a6 --- /dev/null +++ b/.github/workflows/docs-sync.yml @@ -0,0 +1,101 @@ +name: Docs Sync + +"on": + push: + branches: + - master + +permissions: + contents: write + id-token: write + +concurrency: + group: docs-sync + cancel-in-progress: false + +jobs: + doc-sync: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v6 + with: + fetch-depth: 0 + + - uses: anthropics/claude-code-action@v1 + with: + claude_code_oauth_token: ${{ secrets.CLAUDE_CODE_OAUTH_TOKEN }} + prompt: | + REPO: ${{ github.repository }} + COMMIT: ${{ github.sha }} + + You are running the orGUI automated doc-sync job, triggered by a push to + master. Read `doc/AGENTS.md` first — it is the authoritative reference for + what counts as user-facing, how to write CHANGELOG.md entries, how topic + rst pages are organized, and what never to hand-edit. Also skim the root + `AGENTS.md` instruction tree and follow whichever nested `AGENTS.md` + (`orgui/app/AGENTS.md`, `orgui/datautils/xrayutils/AGENTS.md`) covers the + files a given commit touches. + + ## Determine the commit range + + Find the latest release tag reachable from this commit: + `git describe --tags --abbrev=0 --match "v*" ${{ github.sha }}`. Review + every commit in `..${{ github.sha }}` (newest first — prioritize + the most recent commits since they're most likely undocumented), using + `git log --oneline` for the list and `git show ` to judge intent from + the actual diff, not just the message. + + ## Prepare the review branch + + All edits land on a dedicated `docs/auto-sync` branch, never on `master`, + and this job never opens a pull request: + + 1. `git fetch origin docs/auto-sync` (ignore failure if it doesn't exist + yet). + 2. If `origin/docs/auto-sync` exists: `git checkout -B docs/auto-sync + origin/docs/auto-sync`, then `git merge origin/master --no-edit` to + bring in the commit that triggered this run. If the merge conflicts, + stop and report the conflict instead of forcing a resolution. + 3. If it doesn't exist yet: `git checkout -b docs/auto-sync`. + + ## For each commit in range + + 1. Decide if it is user-facing per `doc/AGENTS.md`'s criteria (GUI/CLI + behavior, config fields, calculation results/scientific conventions, + public API contracts, packaging — not pure refactors/tests/CI/comment-only + changes). Use judgment from the diff, not the commit message wording or + the `cliff.toml` grouping heuristic. + 2. If user-facing: check whether `CHANGELOG.md`'s `## [Unreleased]` section + already covers it (read its current content first — do not duplicate). + If not, add a short prose bullet under a fitting existing or new theme + heading matching the existing style (see the `[1.5.0]` section as the + pattern). Do not write `###`-style generated-group headings by hand. + 3. If the change clearly affects an existing `doc/source/*.rst` topic page + (e.g. reconstruction changes -> `reciprocal_space_reconstruction.rst`, + CTR model changes -> `ctr_structure_factors.rst`), check if that page + already reflects it; if not, make a small, localized edit. Don't create + a new page unless nothing existing fits, and if you do, add it to the + toctree in `source/index.rst`. Skip the edit if you can't verify a + scientific convention/unit statement from the diff — check the + underlying code and its local `AGENTS.md` instead of guessing. + 4. Skip anything not clearly user-facing, and skip anything already + documented. + + ## Regenerate and commit + + After any `CHANGELOG.md` edits, regenerate the release notes page: + `python3 doc/generate_release_notes.py`. Never hand-edit + `doc/source/release_notes.rst`. + + If you made any changes, commit them with a Conventional Commit message, + `docs` type, e.g. `docs: sync changelog and topic pages for `. Then + push: try `git push origin HEAD:docs/auto-sync` first; if that's rejected + as non-fast-forward (it shouldn't be, since you just merged), use + `git push --force-with-lease origin HEAD:docs/auto-sync` — this branch is + fully owned by this job, so force-with-lease is expected there, but never + push to `master` and never use a bare `--force`. + + If there was nothing new to document, make no commit and do not push. + + claude_args: | + --allowedTools "Read,Write,Edit,Grep,Glob,Bash(git:*),Bash(python3:*)" diff --git a/AGENTS.md b/AGENTS.md index d18c166..f37a4cb 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -21,6 +21,8 @@ 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. +- `doc/AGENTS.md`: Sphinx documentation source, the changelog/release-notes + workflow, and what counts as a user-facing change worth documenting. Do not add another nested `AGENTS.md` unless it supplies critical local guidance that would otherwise distract unrelated work. diff --git a/doc/AGENTS.md b/doc/AGENTS.md new file mode 100644 index 0000000..7252d5b --- /dev/null +++ b/doc/AGENTS.md @@ -0,0 +1,140 @@ +# AGENTS.md + +## Scope + +This directory contains the Sphinx documentation source and the tooling that +keeps it in sync with releases: + +- `source/*.rst`: hand-maintained topic and reference pages (`geometry.rst`, + `image_integration.rst`, `reciprocal_space_reconstruction.rst`, + `ctr_structure_factors.rst`, `beamline_backends.rst`, ...), wired together by + the toctree in `source/index.rst`. +- `source/api.rst`: autodoc/autosummary entry point into the `orgui` package. +- `source/release_notes.rst`: **generated**. Do not hand-edit; see below. +- `source/release_procedure.rst`: the authoritative step-by-step release + checklist (changelog, tagging, publishing, verification). +- `source/developer.rst`: developer-facing toctree root. +- `source/conf.py`: Sphinx configuration (`furo` theme, `nbsphinx`, autodoc). +- `generate_release_notes.py`: extracts hand-written release notes from + `../CHANGELOG.md` into `source/release_notes.rst`. +- `build/`, `_build/`: build output, not source. + +Use the repository root `AGENTS.md` together with this file. + +## Two Documentation Tracks + +1. **Topic/reference docs** (`source/*.rst`): describe features, workflows, + geometry conventions, and APIs. Hand-maintained, evolve with the code they + document. +2. **Changelog / release notes** (`../CHANGELOG.md` -> + `generate_release_notes.py` -> `source/release_notes.rst`): a chronological + record of user-facing changes, curated at release time per + `source/release_procedure.rst`. + +A given code change may need updates on either track, both, or neither. +Internal refactors typically need neither. + +## Updating Topic Pages + +- Write plain reST matching the existing files. (The repository root + `AGENTS.md` describes Markedly Structured Text as a future direction for + *docstrings*; the `source/*.rst` files today are conventional Sphinx reST — + match what is already there rather than introducing a new dialect.) +- Before creating a new page, look for an existing one that already covers the + topic (e.g. reconstruction changes usually belong in + `reciprocal_space_reconstruction.rst`, CTR model changes in + `ctr_structure_factors.rst`). Prefer extending an existing page. +- If you do add a new page, add it to the toctree in `source/index.rst` — + pages not listed there are unreachable from the built docs. +- State units and coordinate conventions the same way the code does. Do not + invent a different convention for the doc page than the one in the + docstring or `orgui/datautils/xrayutils/AGENTS.md` / `orgui/app/AGENTS.md`. +- Keep edits localized to the feature that changed. Don't reflow or rewrite + unrelated prose on the same page. + +## Updating the Changelog + +`CHANGELOG.md` has an `## [Unreleased]` section at the top. Add short, +high-signal prose bullets there for user-facing changes, grouped under a +theme heading consistent with the existing style (e.g. "Scientific and +analysis additions:", "CLI and scripting highlights:", "Other user-visible +changes:") — see the existing `[1.5.0]` section for the pattern. + +This is **not** the same format as the auto-generated `### Added` / `### +Fixed` / `### Removed` groups that `git cliff` inserts at release time (see +`cliff.toml` and `source/release_procedure.rst` section 2.1) — those are +generated from conventional-commit messages only when a release is cut. Do +not hand-write `###`-style group headings into the Unreleased section. + +Before adding a bullet, check whether the change is already covered by an +existing Unreleased bullet — do not duplicate. + +Never hand-edit `source/release_notes.rst`. After changing +`CHANGELOG.md`, regenerate it: + +```powershell +python doc\generate_release_notes.py +``` + +## What Counts As User-Facing + +Document it (changelog and/or topic page): + +- GUI behavior, dialogs, CLI flags/behavior +- config file fields or file-format changes (see `examples/config_*`) +- calculation results, scientific conventions, or anything that changes + previously produced numbers +- public API functions/classes and their documented contracts +- installation/packaging changes + +Usually skip: + +- pure internal refactors with no observable behavior change +- test-only changes +- CI/build tooling changes +- comment- or docstring-only clarifications that don't change a documented + contract + +`cliff.toml`'s commit parsers are a coarse heuristic for release-time grouping +(e.g. it groups anything matching `test` under "Fixed") — use judgment about +user impact when deciding whether something needs a changelog/doc entry, not +the cliff group label. + +## Regenerating and Validating + +```powershell +python doc\generate_release_notes.py +python -m sphinx -b html doc\source doc\build\html +``` + +Treat new Sphinx build warnings (broken references, toctree issues) as +signal, not noise. + +## Avoid + +- Do not invent user-facing prose for a change you can't verify from its + commit message and diff. If the intent is unclear, say so rather than + guessing. +- Do not touch scientific unit/convention statements in doc pages without + checking the underlying code and its local `AGENTS.md`. +- Do not edit or reflow already-released `CHANGELOG.md` sections. +- Do not remove or rewrite existing hand-written notes when adding new ones. + +## Automated Doc-Sync + +Two automations walk commits since the last release tag and draft +`CHANGELOG.md` / topic-page updates for gaps they find. Both follow the rules +in this file and never touch `master` or hand-edit `source/release_notes.rst` +directly; this file is the durable reference they read for conventions. + +- **`.github/workflows/docs-sync.yml`** (durable): runs on every push to + `master` via the `anthropics/claude-code-action`. It force-updates a + `docs/auto-sync` branch with any missing changelog/doc entries and never + opens a pull request — review and merge that branch manually. Requires a + `CLAUDE_CODE_OAUTH_TOKEN` repository secret (generated via `claude + setup-token` from a Claude Pro/Max subscription, not a metered API key). +- A local, session-scoped `CronCreate` task (set up ad hoc in a Claude Code + session, not checked into the repo) may also run the same procedure against + a local sibling worktree for the same `docs/auto-sync` branch, without + pushing. It expires after 7 days or when the session ends — treat the + GitHub Actions workflow as the source of truth if the two ever disagree. From 7459b3832d1ab26280d4e0552b12ca8987308771 Mon Sep 17 00:00:00 2001 From: "T. Fuchs" Date: Mon, 27 Jul 2026 21:11:23 -0400 Subject: [PATCH 5/6] docs: add CLAUDE.md and close AGENTS.md coverage gaps Add a minimal CLAUDE.md that imports AGENTS.md so Claude Code picks up the existing guidance tree. Extend the AGENTS.md files to cover code that had grown outside their scope: the top-level reconstruction pipeline, orgui/backend/, and the xrayutils cpp/ native kernels, plus refresh stale file/test lists (reconstruction.py, ReconstructionDialog.py, test_reconstruction* suite). Co-Authored-By: Claude Sonnet 5 --- AGENTS.md | 7 +++++++ CLAUDE.md | 1 + orgui/app/AGENTS.md | 3 +++ orgui/datautils/xrayutils/AGENTS.md | 22 ++++++++++++++++++++-- 4 files changed, 31 insertions(+), 2 deletions(-) create mode 100644 CLAUDE.md 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/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/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/datautils/xrayutils/AGENTS.md b/orgui/datautils/xrayutils/AGENTS.md index 891d8c4..718095a 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: @@ -100,6 +117,7 @@ Run the narrowest relevant regression test first: - `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` Use `ruff check orgui/datautils/xrayutils` for local lint checks. If a change touches config-facing unit conversions, also inspect `examples/config_minimal` From 07c6d631804c837165542f298e0b7079eb547290 Mon Sep 17 00:00:00 2001 From: "T. Fuchs" Date: Tue, 4 Aug 2026 17:14:05 -0400 Subject: [PATCH 6/6] refactor(RS): Moved the normalization controls out of the shared integration panel --- CHANGELOG.md | 14 ++++++ .../reciprocal_space_reconstruction.rst | 11 +++++ doc/source/release_notes.rst | 7 +++ orgui/app/QScanSelector.py | 26 ---------- orgui/app/ReconstructionDialog.py | 49 +++++++++++++++++++ orgui/app/config_data.py | 10 ++-- orgui/app/test/test_config_data.py | 29 ++++++++++- 7 files changed, 112 insertions(+), 34 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 23ebb22..82218ac 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -17,6 +17,20 @@ Scientific and analysis additions: shared viewer. Covalent radii are interpreted consistently by both backends and can be adjusted with the dimensionless ``radius_scale`` parameter. +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/doc/source/reciprocal_space_reconstruction.rst b/doc/source/reciprocal_space_reconstruction.rst index 446a28f..c8db2ec 100644 --- a/doc/source/reciprocal_space_reconstruction.rst +++ b/doc/source/reciprocal_space_reconstruction.rst @@ -103,6 +103,17 @@ Corrections and Exclusions ``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 ---------------------- diff --git a/doc/source/release_notes.rst b/doc/source/release_notes.rst index 498dd3d..69b091e 100644 --- a/doc/source/release_notes.rst +++ b/doc/source/release_notes.rst @@ -10,6 +10,13 @@ Unreleased (2026-07-19) Scientific and analysis additions: - Added optional CTR intensity-resolution modeling with constant or gamma-dependent box and Gaussian functions. Calculations can convolve irregular existing L points or sample the crystal structure factor with deterministic quadrature. +- Added py3Dmol atom-sphere rendering for Jupyter notebooks. ``plot3d`` now selects py3Dmol automatically in a notebook, can be directed to either py3Dmol or Mayavi explicitly, and can incrementally add unit cells to a shared viewer. Covalent radii are interpreted consistently by both backends and can be adjusted with the dimensionless ``radius_scale`` parameter. + +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/orgui/app/QScanSelector.py b/orgui/app/QScanSelector.py index 41e2f16..f832044 100644 --- a/orgui/app/QScanSelector.py +++ b/orgui/app/QScanSelector.py @@ -724,25 +724,11 @@ def __init__(self, parentmainwindow, parent=None): self.useLorentzBox.setEnabled(False) self.useSolidAngleBox = qt.QCheckBox("Solid angle correction") self.usePolarizationBox = qt.QCheckBox("Polarization correction") - self.normalizeExposureBox = qt.QCheckBox("Normalize by exposure time") - self.normalizeExposureBox.setChecked(True) - self.monitorCorrectionsEdit = qt.QLineEdit() - self.monitorCorrectionsEdit.setPlaceholderText( - "Optional scan counters, comma-separated" - ) - self.monitorCorrectionsEdit.setToolTip( - "Counters are applied as divisive monitor normalizations." - ) optionsGroupLayout.addWidget(self.useMaskBox, 0, 0) optionsGroupLayout.addWidget(self.useLorentzBox, 1, 0) optionsGroupLayout.addWidget(self.useSolidAngleBox, 0, 1) optionsGroupLayout.addWidget(self.usePolarizationBox, 1, 1) - optionsGroupLayout.addWidget(self.normalizeExposureBox, 2, 0, 1, 2) - optionsGroupLayout.addWidget( - qt.QLabel("Monitor corrections:"), 3, 0 - ) - optionsGroupLayout.addWidget(self.monitorCorrectionsEdit, 3, 1) optionsGroup.setLayout(optionsGroupLayout) @@ -860,12 +846,6 @@ def set_integration_options(self, ddict): self.useSolidAngleBox.setChecked(ddict[key]) elif key == "polarization": self.usePolarizationBox.setChecked(ddict[key]) - elif key == "normalizeExposure": - self.normalizeExposureBox.setChecked(ddict[key]) - elif key == "monitorCorrections": - self.monitorCorrectionsEdit.setText( - ", ".join(map(str, ddict[key])) - ) elif key == "advanced": self.roioptions.set_parameters(ddict[key]) @@ -874,12 +854,6 @@ def get_integration_options(self): ddict["mask"] = self.useMaskBox.isChecked() ddict["solidAngle"] = self.useSolidAngleBox.isChecked() ddict["polarization"] = self.usePolarizationBox.isChecked() - ddict["normalizeExposure"] = self.normalizeExposureBox.isChecked() - ddict["monitorCorrections"] = tuple( - value.strip() - for value in self.monitorCorrectionsEdit.text().split(",") - if value.strip() - ) ddict["advanced"] = self.roioptions.get_parameters() return ddict diff --git a/orgui/app/ReconstructionDialog.py b/orgui/app/ReconstructionDialog.py index 04dc3c1..33575a3 100644 --- a/orgui/app/ReconstructionDialog.py +++ b/orgui/app/ReconstructionDialog.py @@ -135,6 +135,10 @@ def __init__(self, orgui, parent=None): 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) @@ -264,8 +268,53 @@ def _data_tab(self): "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) diff --git a/orgui/app/config_data.py b/orgui/app/config_data.py index b912f03..395f449 100644 --- a/orgui/app/config_data.py +++ b/orgui/app/config_data.py @@ -379,10 +379,10 @@ def from_gui(cls, gui): getattr(repair, "gap_size_px", 1) ), normalize_exposure=bool( - options.get("normalizeExposure", True) + getattr(gui, "reconstruction_normalize_exposure", True) ), monitor_corrections=tuple( - options.get("monitorCorrections", ()) + getattr(gui, "reconstruction_monitor_corrections", ()) ), excluded_frames=tuple( sorted( @@ -441,12 +441,10 @@ def apply_to_gui(self, gui): "mask": self.corrections.use_mask, "solidAngle": self.corrections.use_solid_angle, "polarization": self.corrections.use_polarization, - "normalizeExposure": self.corrections.normalize_exposure, - "monitorCorrections": ( - self.corrections.monitor_corrections - ), } ) + 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 diff --git a/orgui/app/test/test_config_data.py b/orgui/app/test/test_config_data.py index beec0ef..212e8a2 100644 --- a/orgui/app/test/test_config_data.py +++ b/orgui/app/test/test_config_data.py @@ -160,8 +160,6 @@ def test_enabled_pixel_repair_implies_mask_correction(): "mask": False, "solidAngle": False, "polarization": False, - "normalizeExposure": True, - "monitorCorrections": (), } ), maskManager=SimpleNamespace( @@ -170,12 +168,39 @@ def test_enabled_pixel_repair_implies_mask_correction(): 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):