diff --git a/pyCHX/Two_Time_Correlation_Function.py b/pyCHX/Two_Time_Correlation_Function.py index e4f9c45..4eb1c32 100644 --- a/pyCHX/Two_Time_Correlation_Function.py +++ b/pyCHX/Two_Time_Correlation_Function.py @@ -17,6 +17,12 @@ from tqdm import tqdm from pyCHX._optional import imshow +from pyCHX._performance import ( + diagonal_nanmean, + normalize_diagonal_means, + numba_thread_limit, + physical_core_count, +) from pyCHX.chx_libs import RUN_GUI from pyCHX.chx_libs import colors from pyCHX.chx_libs import colors as colors_array @@ -944,18 +950,28 @@ def get_one_time_from_two_time(g12, norms=None, nopr=None): """ m, n, noqs = g12.shape - if norms is None: - g2f12 = np.array([np.nanmean(g12.diagonal(i), axis=1) for i in range(m)]) - else: - g2f12 = np.zeros([m, noqs]) - for q in range(noqs): - yn = norms[:, q] - g2f12[:, q] = np.array( + if g12.size < 1_000_000: + if norms is None: + return np.array([np.nanmean(g12.diagonal(delay), axis=1) for delay in range(m)]) + output = np.empty((m, noqs), dtype=np.float64) + for roi_index in range(noqs): + roi_norm = norms[:, roi_index] + output[:, roi_index] = np.asarray( [ - np.nanmean(g12[:, :, q].diagonal(i)) / (np.average(yn[i:]) * np.average(yn[: m - i]) * nopr[q]) - for i in range(m) + np.nanmean(g12[:, :, roi_index].diagonal(delay)) + / (np.average(roi_norm[delay:]) * np.average(roi_norm[: m - delay]) * nopr[roi_index]) + for delay in range(m) ] ) + return output + with numba_thread_limit(min(physical_core_count(), max(1, m * noqs))): + g2f12 = diagonal_nanmean(np.asarray(g12)) + if norms is not None: + g2f12 = normalize_diagonal_means( + g2f12, + np.asarray(norms, dtype=np.float64), + np.asarray(nopr, dtype=np.float64), + ) return g2f12 diff --git a/pyCHX/_performance.py b/pyCHX/_performance.py new file mode 100644 index 0000000..7773a03 --- /dev/null +++ b/pyCHX/_performance.py @@ -0,0 +1,371 @@ +"""Private CPU, memory, and compiled kernels used by performance-sensitive paths.""" + +from __future__ import annotations + +import os +from contextlib import contextmanager + +import numpy as np +from numba import get_num_threads, njit, prange, set_num_threads + + +def affinity_cpu_ids(): + """Return the logical CPUs on which this process may run.""" + try: + return tuple(sorted(os.sched_getaffinity(0))) + except (AttributeError, OSError): + return tuple(range(os.cpu_count() or 1)) + + +def physical_core_count(cpu_ids=None): + """Count physical cores in the current CPU affinity mask. + + Linux exposes package/core identifiers in sysfs. Falling back to the + affinity-sized logical count is conservative on platforms without that + topology information and, importantly, never escapes a scheduler cpuset. + """ + if cpu_ids is None: + cpu_ids = affinity_cpu_ids() + cpu_ids = tuple(cpu_ids) + cores = set() + try: + for cpu in cpu_ids: + topology = f"/sys/devices/system/cpu/cpu{cpu}/topology" + with open(os.path.join(topology, "physical_package_id")) as stream: + package = int(stream.read()) + with open(os.path.join(topology, "core_id")) as stream: + core = int(stream.read()) + cores.add((package, core)) + except (OSError, ValueError): + return max(1, len(cpu_ids)) + return max(1, len(cores)) + + +def available_memory_bytes(): + """Best-effort available-memory estimate honoring common cgroup limits.""" + available = None + try: + available = os.sysconf("SC_AVPHYS_PAGES") * os.sysconf("SC_PAGE_SIZE") + except (AttributeError, OSError, ValueError): + pass + + for current_name, maximum_name in ( + ("/sys/fs/cgroup/memory.current", "/sys/fs/cgroup/memory.max"), + ("/sys/fs/cgroup/memory/memory.usage_in_bytes", "/sys/fs/cgroup/memory/memory.limit_in_bytes"), + ): + try: + with open(current_name) as stream: + current = int(stream.read().strip()) + with open(maximum_name) as stream: + maximum_text = stream.read().strip() + if maximum_text != "max": + cgroup_available = max(0, int(maximum_text) - current) + available = cgroup_available if available is None else min(available, cgroup_available) + except (OSError, ValueError): + continue + return max(1, available or 512 * 1024**2) + + +@contextmanager +def numba_thread_limit(limit): + """Temporarily constrain Numba's parallel worker pool.""" + previous = get_num_threads() + limit = max(1, min(int(limit), previous)) + set_num_threads(limit) + try: + yield + finally: + set_num_threads(previous) + + +@njit(cache=True, nogil=True) +def sparse_scatter_normalized( + positions, + values, + lookup, + output, + output_row, + source_frame, + norm_1d, + norm_2d, + norm_columns, + imgsum, + mean_int_sets, + qind, + normalization_flags, +): + """Scatter one sparse frame into a selected-pixel row with fused norms.""" + has_norm_1d = normalization_flags[0] + has_norm_2d = normalization_flags[1] + has_imgsum = normalization_flags[2] + has_mean = normalization_flags[3] + for source_index in range(positions.size): + position = positions[source_index] + if position < 0 or position >= lookup.size: + continue + destination = lookup[position] + if destination < 0: + continue + divisor = 1.0 + if has_mean: + divisor *= mean_int_sets[source_frame, qind[destination] - 1] + if has_imgsum: + divisor *= imgsum[source_frame] + if has_norm_2d: + divisor *= norm_2d[source_frame, norm_columns[destination]] + elif has_norm_1d: + divisor *= norm_1d[norm_columns[destination]] + output[output_row, destination] = values[source_index] / divisor + + +@njit(cache=True, nogil=True) +def sparse_add_image(positions, values, flattened_output): + for index in range(positions.size): + position = positions[index] + if 0 <= position < flattened_output.size: + flattened_output[position] += values[index] + + +@njit(cache=True, nogil=True) +def sparse_roi_sums(positions, values, roi_lookup, output_row): + for index in range(positions.size): + position = positions[index] + if 0 <= position < roi_lookup.size: + roi_index = roi_lookup[position] + if roi_index >= 0: + output_row[roi_index] += values[index] + + +@njit(cache=True, nogil=True) +def sparse_frame_sum(values): + total = 0.0 + for index in range(values.size): + total += values[index] + return total + + +@njit(cache=True, nogil=True) +def process_one_time_block( + frames, + buf, + correlation, + past_intensity, + future_intensity, + images_per_level, + track_level, + current, + bad_counts, + level_offsets, + intensity_buf, + correlation_all, + past_all, + future_all, + calculate_error, +): + """Update one ROI's multi-tau state for a contiguous frame block.""" + num_levels, num_bufs, pixel_count = buf.shape + for frame_index in range(frames.shape[0]): + current[0] = (1 + current[0]) % num_bufs + buffer_number = current[0] - 1 + for pixel in range(pixel_count): + buf[0, buffer_number, pixel] = frames[frame_index, pixel] + + level = 0 + processing = True + while processing: + images_per_level[level] += 1 + minimum_delay = num_bufs // 2 if level else 0 + future = buf[level, buffer_number] + future_bad = False + for pixel in range(pixel_count): + if np.isnan(future[pixel]): + future_bad = True + break + + future_sum = 0.0 + if not future_bad and not calculate_error: + for pixel in range(pixel_count): + future_sum += future[pixel] + intensity_buf[level, buffer_number] = future_sum + + stop_delay = min(images_per_level[level], num_bufs) + for delay in range(minimum_delay, stop_delay): + time_index = level * num_bufs // 2 + delay + delay_number = (buffer_number - delay) % num_bufs + past = buf[level, delay_number] + bad = future_bad + if not bad: + for pixel in range(pixel_count): + if np.isnan(past[pixel]): + bad = True + break + local_index = time_index - level_offsets[level] + normalize = images_per_level[level] - delay - bad_counts[level, local_index] + if bad: + bad_counts[level, local_index] += 1 + elif calculate_error: + for pixel in range(pixel_count): + product = past[pixel] * future[pixel] + correlation_all[time_index, pixel] += ( + product - correlation_all[time_index, pixel] + ) / normalize + past_all[time_index, pixel] += (past[pixel] - past_all[time_index, pixel]) / normalize + future_all[time_index, pixel] += ( + future[pixel] - future_all[time_index, pixel] + ) / normalize + else: + product_sum = 0.0 + past_sum = intensity_buf[level, delay_number] + for pixel in range(pixel_count): + product_sum += past[pixel] * future[pixel] + product_mean = product_sum / pixel_count + past_mean = past_sum / pixel_count + future_mean = future_sum / pixel_count + correlation[time_index] += (product_mean - correlation[time_index]) / normalize + past_intensity[time_index] += (past_mean - past_intensity[time_index]) / normalize + future_intensity[time_index] += (future_mean - future_intensity[time_index]) / normalize + + if level + 1 >= num_levels: + processing = False + else: + level += 1 + if not track_level[level]: + track_level[level] = True + processing = False + else: + previous = 1 + (current[level - 1] - 2) % num_bufs + current[level] = 1 + current[level] % num_bufs + buffer_number = current[level] - 1 + for pixel in range(pixel_count): + buf[level, buffer_number, pixel] = ( + buf[level - 1, previous - 1, pixel] + buf[level - 1, current[level - 1] - 1, pixel] + ) / 2.0 + track_level[level] = False + + +@njit(cache=True, nogil=True, error_model="numpy") +def mirror_and_normalize_two_time(matrix, row_norm, pixel_count): + """Normalize the computed upper triangle and mirror it in place.""" + frame_count = matrix.shape[0] + for row in range(frame_count): + for column in range(row, frame_count): + value = matrix[row, column] + value /= row_norm[column] + value /= row_norm[row] + value /= pixel_count + matrix[row, column] = value + matrix[column, row] = value + + +@njit(cache=True, nogil=True, parallel=True, error_model="numpy") +def mirror_and_normalize_two_time_parallel(matrix, row_norm, pixel_count): + """Parallel large-matrix variant of :func:`mirror_and_normalize_two_time`.""" + frame_count = matrix.shape[0] + for row in prange(frame_count): + for column in range(row, frame_count): + value = matrix[row, column] + value /= row_norm[column] + value /= row_norm[row] + value /= pixel_count + matrix[row, column] = value + matrix[column, row] = value + + +@njit(cache=True, nogil=True) +def mirror_two_time(matrix): + """Copy an upper-triangular symmetric BLAS result to its lower half.""" + frame_count = matrix.shape[0] + for row in range(frame_count): + for column in range(row + 1, frame_count): + matrix[column, row] = matrix[row, column] + + +@njit(cache=True, nogil=True, parallel=True) +def mirror_two_time_parallel(matrix): + """Parallel large-matrix variant of :func:`mirror_two_time`.""" + frame_count = matrix.shape[0] + for row in prange(frame_count): + for column in range(row + 1, frame_count): + matrix[column, row] = matrix[row, column] + + +@njit(cache=True, nogil=True, parallel=True, error_model="numpy") +def store_symmetric_two_time_batch( + upper_triangles, + row_norms, + pixel_counts, + pre_normalized, + output, + output_start, + batch_count, +): + """Write several upper-triangular ROI results into the public C-order output.""" + frame_count = output.shape[0] + for task in prange(frame_count * frame_count): + row = task // frame_count + column = task - row * frame_count + source_row = row + source_column = column + if row > column: + source_row = column + source_column = row + for batch_index in range(batch_count): + value = upper_triangles[source_row, source_column, batch_index] + if not pre_normalized[batch_index]: + value /= row_norms[source_column, batch_index] + value /= row_norms[source_row, batch_index] + value /= pixel_counts[batch_index] + output[row, column, output_start + batch_index] = value + + +@njit(cache=True, nogil=True, parallel=True) +def diagonal_nanmean(g12): + """Reduce upper diagonals of a C-order (time, time, ROI) array.""" + frame_count = g12.shape[0] + roi_count = g12.shape[2] + output = np.empty((frame_count, roi_count), dtype=np.float64) + for task in prange(frame_count * roi_count): + delay = task // roi_count + roi_index = task - delay * roi_count + total = 0.0 + count = 0 + for frame in range(frame_count - delay): + value = g12[frame, frame + delay, roi_index] + if not np.isnan(value): + total += value + count += 1 + output[delay, roi_index] = total / count if count else np.nan + return output + + +@njit(cache=True, nogil=True, parallel=True) +def normalize_diagonal_means(diagonal_means, norms, pixel_counts): + frame_count, roi_count = diagonal_means.shape + output = np.empty_like(diagonal_means) + prefixes = np.empty((frame_count + 1, roi_count), dtype=np.float64) + nan_prefixes = np.zeros((frame_count + 1, roi_count), dtype=np.int64) + prefixes[0, :] = 0.0 + for frame in range(frame_count): + for roi_index in range(roi_count): + value = norms[frame, roi_index] + if np.isnan(value): + prefixes[frame + 1, roi_index] = prefixes[frame, roi_index] + nan_prefixes[frame + 1, roi_index] = nan_prefixes[frame, roi_index] + 1 + else: + prefixes[frame + 1, roi_index] = prefixes[frame, roi_index] + value + nan_prefixes[frame + 1, roi_index] = nan_prefixes[frame, roi_index] + for task in prange(frame_count * roi_count): + delay = task // roi_count + roi_index = task - delay * roi_count + count = frame_count - delay + future_nans = nan_prefixes[frame_count, roi_index] - nan_prefixes[delay, roi_index] + past_nans = nan_prefixes[count, roi_index] + if future_nans or past_nans: + output[delay, roi_index] = np.nan + else: + future_mean = (prefixes[frame_count, roi_index] - prefixes[delay, roi_index]) / count + past_mean = prefixes[count, roi_index] / count + output[delay, roi_index] = diagonal_means[delay, roi_index] / ( + future_mean * past_mean * pixel_counts[roi_index] + ) + return output diff --git a/pyCHX/benchmarks/__init__.py b/pyCHX/benchmarks/__init__.py new file mode 100644 index 0000000..424a584 --- /dev/null +++ b/pyCHX/benchmarks/__init__.py @@ -0,0 +1 @@ +"""Opt-in pyCHX benchmark programs.""" diff --git a/pyCHX/benchmarks/benchmark_xpcs.py b/pyCHX/benchmarks/benchmark_xpcs.py new file mode 100644 index 0000000..20a6d84 --- /dev/null +++ b/pyCHX/benchmarks/benchmark_xpcs.py @@ -0,0 +1,196 @@ +"""Opt-in benchmark for the production XPCS compressed-data call path. + +This module never runs as part of the test suite. Invoke it with +``python -m pyCHX.benchmarks.benchmark_xpcs --help``. +""" + +from __future__ import annotations + +import argparse +import json +import multiprocessing +import os +import resource +import threading +import time +from pathlib import Path + +import numpy as np + +from pyCHX.chx_compress import Multifile, compress_eigerdata +from pyCHX.chx_correlationc import Get_Pixel_Arrayc, auto_two_Arrayc +from pyCHX.chx_correlationp import cal_g2p +from pyCHX.Two_Time_Correlation_Function import get_one_time_from_two_time + + +def _rss_bytes(): + # Linux reports ru_maxrss in KiB; macOS reports bytes. + value = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss + return int(value if os.uname().sysname == "Darwin" else value * 1024) + + +def _process_rss_and_threads(pid): + try: + fields = {} + for line in Path(f"/proc/{pid}/status").read_text().splitlines(): + if line.startswith(("VmRSS:", "Threads:")): + key, value, *_ = line.split() + fields[key.rstrip(":")] = int(value) + return fields.get("VmRSS", 0) * 1024, fields.get("Threads", 0) + except OSError: + return 0, 0 + + +def _measure(name, function, traversed_bytes): + before_rss = _rss_bytes() + peaks = {"rss": before_rss, "threads": threading.active_count(), "processes": 1} + stopped = threading.Event() + + def monitor(): + while not stopped.wait(0.02): + pids = [os.getpid(), *(child.pid for child in multiprocessing.active_children())] + snapshots = [_process_rss_and_threads(pid) for pid in pids] + peaks["rss"] = max(peaks["rss"], sum(item[0] for item in snapshots)) + peaks["threads"] = max(peaks["threads"], sum(item[1] for item in snapshots)) + peaks["processes"] = max(peaks["processes"], len(pids)) + + monitor_thread = threading.Thread(target=monitor, daemon=True) + monitor_thread.start() + started = time.perf_counter() + try: + result = function() + finally: + elapsed = time.perf_counter() - started + stopped.set() + monitor_thread.join() + return result, { + "name": name, + "seconds": elapsed, + "peak_rss_bytes": max(peaks["rss"], _rss_bytes()), + "rss_before_bytes": before_rss, + "peak_threads": peaks["threads"], + "peak_processes": peaks["processes"], + "bytes_traversed": traversed_bytes, + } + + +def main(argv=None): + parser = argparse.ArgumentParser() + parser.add_argument("--cmp", required=True, type=Path) + parser.add_argument("--roi", required=True, type=Path, help="NumPy .npy ROI-label mask") + parser.add_argument("--norm", type=Path, help="Optional NumPy .npy per-pixel normalization") + parser.add_argument("--imgsum", type=Path, help="Optional NumPy .npy frame-intensity normalization") + parser.add_argument("--bad-frames", type=Path, help="Optional NumPy .npy bad-frame indices") + parser.add_argument("--beg", type=int, default=0) + parser.add_argument("--end", required=True, type=int) + parser.add_argument("--num-buf", type=int, default=8) + parser.add_argument("--output", type=Path) + parser.add_argument("--eiger-master", type=Path, help="Also benchmark compression from this Eiger master") + parser.add_argument("--mask", type=Path, help="Detector mask required with --eiger-master") + parser.add_argument("--metadata", type=Path, help="Compression metadata JSON required with --eiger-master") + parser.add_argument("--compression-output", type=Path, help="Base name for temporary benchmark CMP output") + parser.add_argument("--images-per-file", type=int, default=100) + parser.add_argument("--compression-frames", type=int) + parser.add_argument("--num-sub", type=int, default=128) + parser.add_argument("--compression-bins", type=int, default=1) + parser.add_argument("--compression-bytes", type=int, choices=(2, 4, 8), default=4) + parser.add_argument("--bad-pixel-threshold", type=float, default=1e15) + parser.add_argument("--bad-pixel-low-threshold", type=float, default=0) + parser.add_argument("--hot-pixel-threshold", type=float, default=2**30) + parser.add_argument("--no-reverse", action="store_true") + parser.add_argument("--rot90", action="store_true") + args = parser.parse_args(argv) + + roi_mask = np.load(args.roi, allow_pickle=False) + norm = np.load(args.norm, allow_pickle=False) if args.norm else None + imgsum = np.load(args.imgsum, allow_pickle=False) if args.imgsum else None + bad_frames = np.load(args.bad_frames, allow_pickle=False) if args.bad_frames else [] + pixel_list = np.flatnonzero(roi_mask.ravel() > 0) + measurements = [] + + if args.eiger_master: + if not args.mask or not args.metadata or not args.compression_output: + parser.error("--eiger-master requires --mask, --metadata, and --compression-output") + detector_mask = np.load(args.mask, allow_pickle=False) + metadata = json.loads(args.metadata.read_text()) + metadata["pixel_mask"] = detector_mask.copy() + frame_count = args.compression_frames or args.end + raw_bytes = frame_count * detector_mask.size * np.dtype(np.uint32).itemsize + for pass_name in ("cold", "warm"): + destination = args.compression_output.with_name( + f"{args.compression_output.stem}-{pass_name}{args.compression_output.suffix or '.cmp'}" + ) + _, metric = _measure( + f"compress_eigerdata_{pass_name}", + lambda destination=destination: compress_eigerdata( + np.empty(frame_count, dtype=np.uint8), + detector_mask.copy(), + metadata.copy(), + os.fspath(destination), + force_compress=True, + para_compress=True, + num_sub=args.num_sub, + bins=args.compression_bins, + nobytes=args.compression_bytes, + bad_pixel_threshold=args.bad_pixel_threshold, + bad_pixel_low_threshold=args.bad_pixel_low_threshold, + hot_pixel_threshold=args.hot_pixel_threshold, + dtypes="uid", + with_pickle=False, + direct_load_data=True, + data_path=os.fspath(args.eiger_master), + images_per_file=args.images_per_file, + copy_rawdata=False, + reverse=not args.no_reverse, + rot90=args.rot90, + func_images_per_file=lambda _path: args.images_per_file, + ), + raw_bytes, + ) + measurements.append(metric) + + for pass_name in ("cold", "warm"): + with Multifile(os.fspath(args.cmp), args.beg, args.end) as compressed: + compressed._reset_io_counters() + data, metric = _measure( + f"Get_Pixel_Arrayc_{pass_name}", + lambda: Get_Pixel_Arrayc(compressed, pixel_list, norm=norm).get_data(), + 0, + ) + metric["bytes_traversed"] = compressed._bytes_traversed + measurements.append(metric) + with Multifile(os.fspath(args.cmp), args.beg, args.end) as compressed: + compressed._reset_io_counters() + (_, _), metric = _measure( + f"cal_g2p_{pass_name}", + lambda: cal_g2p( + compressed, + roi_mask, + bad_frame_list=bad_frames, + num_buf=args.num_buf, + imgsum=imgsum, + norm=norm, + ), + 0, + ) + metric["bytes_traversed"] = compressed._bytes_traversed + measurements.append(metric) + two_time, metric = _measure( + f"auto_two_Arrayc_{pass_name}", lambda: auto_two_Arrayc(data, roi_mask), data.nbytes + ) + measurements.append(metric) + _, metric = _measure( + f"get_one_time_from_two_time_{pass_name}", + lambda: get_one_time_from_two_time(two_time), + two_time.nbytes, + ) + measurements.append(metric) + + report = json.dumps(measurements, indent=2) + if args.output: + args.output.write_text(report + "\n") + print(report) + + +if __name__ == "__main__": + main() diff --git a/pyCHX/chx_compress.py b/pyCHX/chx_compress.py index f41eb3d..d671d3b 100644 --- a/pyCHX/chx_compress.py +++ b/pyCHX/chx_compress.py @@ -4,6 +4,7 @@ import shutil import struct import sys +import tempfile import time from multiprocessing import Pool, cpu_count @@ -15,6 +16,12 @@ from matplotlib.figure import Figure from tqdm import tqdm +from pyCHX._performance import ( + physical_core_count, + sparse_add_image, + sparse_frame_sum, + sparse_roi_sums, +) from pyCHX.chx_generic_functions import ( copy_data, create_time_slice, @@ -23,8 +30,6 @@ get_eigerImage_per_file, get_sid_filenames, load_data, - reverse_updown, - rot90_clockwise, ) from pyCHX.chx_handlers import EigerImages, db from pyCHX.chx_libs import RUN_GUI @@ -44,7 +49,58 @@ def _make_pool(task_count): """Create no more worker processes than either tasks or available CPUs.""" if task_count < 1: raise ValueError("at least one multiprocessing task is required") - return Pool(processes=min(task_count, cpu_count())) + return Pool(processes=min(task_count, _available_cpu_count())) + + +def _available_cpu_count(): + """Return affinity-constrained physical cores, avoiding SMT oversubscription.""" + detected = cpu_count() + try: + affinity = os.sched_getaffinity(0) + except (AttributeError, OSError): + return detected + return min(detected, physical_core_count(affinity)) + + +def _write_sparse_frame(stream, positions, values): + """Write one sparse CMP frame using the legacy native binary layout.""" + stream.write(np.asarray(len(positions), dtype=np.uint32).tobytes()) + if len(positions): + stream.write(np.asarray(positions, dtype=np.int32).tobytes()) + stream.write(np.ascontiguousarray(values).tobytes()) + + +def _publish_file(source, destination): + """Copy *source* beside *destination* and atomically publish it.""" + destination = os.path.abspath(destination) + destination_dir = os.path.dirname(destination) + descriptor, temporary = tempfile.mkstemp( + dir=destination_dir, + prefix=".%s." % os.path.basename(destination), + suffix=".tmp", + ) + os.close(descriptor) + try: + shutil.copyfile(source, temporary) + shutil.copymode(source, temporary) + os.replace(temporary, destination) + finally: + if os.path.exists(temporary): + os.remove(temporary) + + +def _staged_init_compress_eigerdata(images, mask, md, filename, new_path, **kwargs): + """Run serial compression locally, then publish its completed output.""" + staging_dir = tempfile.mkdtemp(prefix="pychx-compress-", dir=new_path) + staged_filename = os.path.join(staging_dir, os.path.basename(filename)) + try: + result = init_compress_eigerdata(images, mask, md, staged_filename, **kwargs) + _publish_file(staged_filename, filename) + if kwargs.get("with_pickle", True): + _publish_file(staged_filename + ".pkl", filename + ".pkl") + return result + finally: + shutil.rmtree(staging_dir, ignore_errors=True) def _collect_pool_results(pool, results, show_progress=False): @@ -76,6 +132,84 @@ def _frame_bin_edges(frame_count, bins): return np.column_stack((starts, stops)) +def _read_eiger_contiguous(images, start, stop): + """Read a contiguous Eiger range with one HDF5 slice per data file.""" + pieces = [] + while start < stop: + file_number = start // images.images_per_file + dataset = images._entry[f"data_{file_number + 1:06d}"] + local_start = start - file_number * images.images_per_file + local_stop = min(dataset.shape[0], stop - file_number * images.images_per_file) + pieces.append(dataset[local_start:local_stop]) + start += local_stop - local_start + if len(pieces) == 1: + return pieces[0] + return np.concatenate(pieces, axis=0) + + +def _iter_binned_images(images, start, stop, bins, reverse=False, rot90=False): + """Yield binned images, using bounded HDF5 block reads when available.""" + direct_eiger = all(hasattr(images, attribute) for attribute in ("_entry", "images_per_file")) + if not direct_eiger: + sliced = images[start:stop] + if bins == 1: + yield from sliced + return + for local_start, local_stop in _frame_bin_edges(stop - start, bins): + yield np.average(sliced[local_start:local_stop], axis=0) + return + + first_key = images.valid_keys[0] + frame_bytes = int(np.prod(images._entry[first_key].shape[1:])) * images._entry[first_key].dtype.itemsize + target_frames = max(1, (256 * 1024**2) // max(1, frame_bytes)) + block_frames = max(bins, (target_frames // bins) * bins) + for block_start in range(start, stop, block_frames): + block_stop = min(stop, block_start + block_frames) + block = _read_eiger_contiguous(images, block_start, block_stop) + if reverse: + block = block[:, ::-1, :] + if rot90: + block = np.rot90(block, axes=(1, 2)) + if bins == 1: + yield from block + continue + for local_start, local_stop in _frame_bin_edges(len(block), bins): + yield np.average(block[local_start:local_stop], axis=0) + + +def _publish_compressed_segments(staged_filename, destination, segment_count): + """Concatenate segment files directly into an atomic destination sibling.""" + sources = [staged_filename + "-header"] + [ + staged_filename + "_temp-%i.tmp" % index for index in range(segment_count) + ] + if not all(os.path.exists(source) for source in sources): + # Retains compatibility with callers/tests that replace the public + # combination and publication helpers. + combine_compressed(staged_filename, segment_count, del_old=True) + _publish_file(staged_filename, destination) + return + destination = os.path.abspath(destination) + descriptor, temporary = tempfile.mkstemp( + dir=os.path.dirname(destination), + prefix=".%s." % os.path.basename(destination), + suffix=".tmp", + ) + source_mode = os.stat(sources[0]).st_mode + try: + with os.fdopen(descriptor, "wb") as output: + for source_name in sources: + with open(source_name, "rb") as source: + shutil.copyfileobj(source, output, length=16 * 1024**2) + os.remove(source_name) + output.flush() + os.fsync(output.fileno()) + os.chmod(temporary, source_mode) + os.replace(temporary, destination) + finally: + if os.path.exists(temporary): + os.remove(temporary) + + def map_async(pool, fun, args): return pool.map_async(run_dill_encoded, (dill.dumps((fun, args)),)) @@ -181,11 +315,12 @@ def compress_eigerdata( new_path=new_path, ) else: - return init_compress_eigerdata( + return _staged_init_compress_eigerdata( images, mask, md, filename, + new_path, bad_pixel_threshold=bad_pixel_threshold, hot_pixel_threshold=hot_pixel_threshold, bad_pixel_low_threshold=bad_pixel_low_threshold, @@ -195,6 +330,8 @@ def compress_eigerdata( direct_load_data=direct_load_data, data_path=data_path, images_per_file=images_per_file, + reverse=reverse, + rot90=rot90, ) else: if not os.path.exists(filename): @@ -224,11 +361,12 @@ def compress_eigerdata( new_path=new_path, ) else: - return init_compress_eigerdata( + return _staged_init_compress_eigerdata( images, mask, md, filename, + new_path, bad_pixel_threshold=bad_pixel_threshold, hot_pixel_threshold=hot_pixel_threshold, bad_pixel_low_threshold=bad_pixel_low_threshold, @@ -238,6 +376,8 @@ def compress_eigerdata( direct_load_data=direct_load_data, data_path=data_path, images_per_file=images_per_file, + reverse=reverse, + rot90=rot90, ) else: print("Using already created compressed file with filename as :%s." % filename) @@ -295,16 +435,25 @@ def read_compressed_eigerdata( with Multifile(filename, beg, end) as FD: imgsum = np.zeros(FD.end - FD.beg, dtype=np.float64) avg_img = np.zeros([FD.md["ncols"], FD.md["nrows"]], dtype=np.float64) - imgsum, bad_frame_list_ = get_each_frame_intensityc( - FD, - sampling=1, - bad_pixel_threshold=bad_pixel_threshold, - bad_pixel_low_threshold=bad_pixel_low_threshold, - hot_pixel_threshold=hot_pixel_threshold, - plot_=False, - bad_frame_list=bad_frame_list, - ) - avg_img = get_avg_imgc(FD, beg=None, end=None, sampling=1, plot_=False, bad_frame_list=bad_frame_list_) + supplied_bad = set() if bad_frame_list is None else set(np.atleast_1d(bad_frame_list).tolist()) + detected_bad = [] + good_count = 0 + flattened_average = avg_img.ravel() + for output_index, frame_index in enumerate(range(FD.beg, FD.end)): + positions, values = FD._raw_frame_view(frame_index) + frame_sum = sparse_frame_sum(values) + imgsum[output_index] = frame_sum + is_bad = frame_sum > bad_pixel_threshold or frame_sum <= bad_pixel_low_threshold + if is_bad: + detected_bad.append(frame_index) + if not is_bad and frame_index not in supplied_bad: + sparse_add_image(positions, values, flattened_average) + good_count += 1 + bad_frame_list_ = np.unique(np.asarray([*supplied_bad, *detected_bad], dtype=np.int64)) + if good_count: + avg_img /= good_count + else: + avg_img.fill(np.nan) return mask, avg_img, imgsum, bad_frame_list_ @@ -334,6 +483,7 @@ def para_compress_eigerdata( ): data_path_ = data_path + raw_data_copied = False if dtypes == "uid": uid = md["uid"] # images if not direct_load_data: @@ -350,24 +500,29 @@ def para_compress_eigerdata( print("Copying...") copy_data(data_path, new_path) # print(data_path, new_path) - new_master_file = new_path + os.path.basename(data_path) + new_master_file = os.path.join(new_path, os.path.basename(data_path)) data_path_ = new_master_file images_ = EigerImages(new_master_file, images_per_file, md) + raw_data_copied = True # print(md) - if reverse: - images_ = reverse_updown(images_) # Why not np.flipud? - if rot90: - images_ = rot90_clockwise(images_) + try: + N = len(images_) + finally: + images_.close() - N = len(images_) + if not direct_load_data: + N = len(images_) else: N = len(images) if cpu_core_number == 0: - cpu_core_number = cpu_count() + cpu_core_number = _available_cpu_count() + else: + cpu_core_number = min(cpu_core_number, _available_cpu_count()) - N = len(_frame_bin_edges(N, bins)) + raw_image_count = N + N = len(_frame_bin_edges(raw_image_count, bins)) Nf = int(np.ceil(N / num_sub)) if Nf > cpu_core_number: print("The process number is larger than %s (current server's core threads)" % cpu_core_number) @@ -375,67 +530,71 @@ def para_compress_eigerdata( num_sub = int(np.ceil(N / cpu_core_number)) Nf = int(np.ceil(N / num_sub)) print("The sub compressed file number was changed from %s to %s" % (num_sub_old, num_sub)) - create_compress_header(md, filename + "-header", nobytes, bins, rot90=rot90) - # print( 'done for header here') - # print(data_path_, images_per_file) - results = para_segment_compress_eigerdata( - images=images, - mask=mask, - md=md, - filename=filename, - num_sub=num_sub, - bad_pixel_threshold=bad_pixel_threshold, - hot_pixel_threshold=hot_pixel_threshold, - bad_pixel_low_threshold=bad_pixel_low_threshold, - nobytes=nobytes, - bins=bins, - dtypes=dtypes, - num_max_para_process=num_max_para_process, - reverse=reverse, - rot90=rot90, - direct_load_data=direct_load_data, - data_path=data_path_, - images_per_file=images_per_file, - ) + staging_dir = tempfile.mkdtemp(prefix="pychx-compress-", dir=new_path) + staged_filename = os.path.join(staging_dir, os.path.basename(filename)) + try: + create_compress_header(md, staged_filename + "-header", nobytes, bins, rot90=rot90) + segment_results = _iter_parallel_segment_results( + images=images, + mask=mask, + md=md, + filename=staged_filename, + num_sub=num_sub, + bad_pixel_threshold=bad_pixel_threshold, + hot_pixel_threshold=hot_pixel_threshold, + bad_pixel_low_threshold=bad_pixel_low_threshold, + nobytes=nobytes, + bins=bins, + dtypes=dtypes, + num_max_para_process=num_max_para_process, + reverse=reverse, + rot90=rot90, + direct_load_data=direct_load_data, + data_path=data_path_, + images_per_file=images_per_file, + image_count=raw_image_count, + segment_count=Nf, + ) - res_ = [results[k].get() for k in list(sorted(results.keys()))] - imgsum = np.zeros(N) - bad_frame_list = np.zeros(N, dtype=bool) - good_count = 0 - for i in range(Nf): - mask_, avg_img_, imgsum_, bad_frame_list_ = res_[i] - imgsum[i * num_sub : (i + 1) * num_sub] = imgsum_ - bad_frame_list[i * num_sub : (i + 1) * num_sub] = bad_frame_list_ - segment_good_count = len(imgsum_) - np.count_nonzero(bad_frame_list_) - if i == 0: - mask = mask_ - avg_img = np.zeros_like(avg_img_, dtype=np.float64) + imgsum = np.zeros(N) + bad_frame_list = np.zeros(N, dtype=bool) + good_count = 0 + for i, segment_result in segment_results: + mask_, avg_img_, imgsum_, bad_frame_list_ = segment_result + imgsum[i * num_sub : (i + 1) * num_sub] = imgsum_ + bad_frame_list[i * num_sub : (i + 1) * num_sub] = bad_frame_list_ + segment_good_count = len(imgsum_) - np.count_nonzero(bad_frame_list_) + if i == 0: + mask = mask_ + avg_img = np.zeros_like(avg_img_, dtype=np.float64) + else: + mask *= mask_ + if segment_good_count and not np.any(np.isnan(avg_img_)): + avg_img += avg_img_ * segment_good_count + good_count += segment_good_count + + bad_frame_list = np.where(bad_frame_list)[0] + if good_count: + avg_img /= good_count else: - mask *= mask_ - if segment_good_count and not np.any(np.isnan(avg_img_)): - avg_img += avg_img_ * segment_good_count - good_count += segment_good_count - - bad_frame_list = np.where(bad_frame_list)[0] - if good_count: - avg_img /= good_count - else: - avg_img.fill(np.nan) + avg_img.fill(np.nan) - if len(bad_frame_list): - print("Bad frame list are: %s" % bad_frame_list) - else: - print("No bad frames are involved.") - print("Combining the seperated compressed files together...") - combine_compressed(filename, Nf, del_old=True) - del results - del res_ - if with_pickle: - with open(filename + ".pkl", "wb") as stream: - pkl.dump([mask, avg_img, imgsum, bad_frame_list], stream) - if copy_rawdata: - delete_data(data_path, new_path) - return mask, avg_img, imgsum, bad_frame_list + if len(bad_frame_list): + print("Bad frame list are: %s" % bad_frame_list) + else: + print("No bad frames are involved.") + print("Combining the seperated compressed files together...") + _publish_compressed_segments(staged_filename, filename, Nf) + if with_pickle: + staged_pickle = staged_filename + ".pkl" + with open(staged_pickle, "wb") as stream: + pkl.dump([mask, avg_img, imgsum, bad_frame_list], stream) + _publish_file(staged_pickle, filename + ".pkl") + return mask, avg_img, imgsum, bad_frame_list + finally: + shutil.rmtree(staging_dir, ignore_errors=True) + if raw_data_copied: + delete_data(data_path, new_path) def combine_compressed(filename, Nf, del_old=True): @@ -478,74 +637,273 @@ def para_segment_compress_eigerdata( parallelly compressed eiger data without header, this function is for parallel compress """ if dtypes == "uid": - uid = md["uid"] # images + uid = md["uid"] if not direct_load_data: detector = get_detector(db[uid]) - images_ = load_data(uid, detector, reverse=reverse, rot90=rot90) + image_count = len(load_data(uid, detector, reverse=reverse, rot90=rot90)) else: - images_ = EigerImages(data_path, images_per_file, md) - if reverse: - images_ = reverse_updown(images_) - if rot90: - images_ = rot90_clockwise(images_) - - N = len(images_) - + probe = EigerImages(data_path, images_per_file, md) + try: + image_count = len(probe) + finally: + probe.close() else: - N = len(images) + image_count = len(images) # N = int( np.ceil( N/ bins ) ) num_sub *= bins - Nf = int(np.ceil(N / num_sub)) + Nf = int(np.ceil(image_count / num_sub)) print("It will create %i temporary files for parallel compression." % Nf) if Nf > num_max_para_process: - N_runs = int(np.ceil(Nf / float(num_max_para_process))) - print("The parallel run number: %s is larger than num_max_para_process: %s" % (Nf, num_max_para_process)) - else: - N_runs = 1 + print("The segment count %s exceeds the concurrent worker limit %s" % (Nf, num_max_para_process)) + worker_count = min(Nf, num_max_para_process, _available_cpu_count()) + print("Pool processes: %s" % worker_count) + pool = Pool( + processes=worker_count, + initializer=_compression_worker_init, + initargs=( + images, + mask, + md, + bad_pixel_threshold, + hot_pixel_threshold, + bad_pixel_low_threshold, + nobytes, + bins, + dtypes, + reverse, + rot90, + direct_load_data, + data_path, + images_per_file, + ), + ) result = {} - # print( mask_filename )# + '*'* 10 + 'here' ) - for nr in range(N_runs): - if (nr + 1) * num_max_para_process > Nf: - inputs = range(num_max_para_process * nr, Nf) - else: - inputs = range(num_max_para_process * nr, num_max_para_process * (nr + 1)) - _ = [filename + "_temp-%i.tmp" % i for i in inputs] - # print( nr, inputs, ) - pool = _make_pool(len(inputs)) # , maxtasksperchild=1000 ) - print("Pool processes: %s" % len(inputs)) - for i in inputs: - if i * num_sub <= N: - result[i] = pool.apply_async( - segment_compress_eigerdata, - [ - images, - mask, - md, - filename + "_temp-%i.tmp" % i, - bad_pixel_threshold, - hot_pixel_threshold, - bad_pixel_low_threshold, - nobytes, - bins, - i * num_sub, - (i + 1) * num_sub, - dtypes, - reverse, - rot90, - direct_load_data, - data_path, - images_per_file, - ], - ) - + try: + for i in range(Nf): + start = i * num_sub + stop = min((i + 1) * num_sub, image_count) + result[i] = pool.apply_async( + _compression_worker_run, + ((filename + "_temp-%i.tmp" % i, start, stop),), + ) pool.close() + for async_result in result.values(): + async_result.wait() pool.join() + except BaseException: pool.terminate() + pool.join() + raise return result +_COMPRESSION_WORKER_CONTEXT = None + + +def _compression_worker_init( + images, + mask, + md, + bad_pixel_threshold, + hot_pixel_threshold, + bad_pixel_low_threshold, + nobytes, + bins, + dtypes, + reverse, + rot90, + direct_load_data, + data_path, + images_per_file, +): + """Initialize a persistent compression worker once.""" + global _COMPRESSION_WORKER_CONTEXT + apply_reverse = False + apply_rot90 = False + if dtypes == "uid": + if direct_load_data: + images = EigerImages(data_path, images_per_file, md) + apply_reverse = reverse + apply_rot90 = rot90 + else: + detector = get_detector(db[md["uid"]]) + images = load_data(md["uid"], detector, reverse=reverse, rot90=rot90) + _COMPRESSION_WORKER_CONTEXT = ( + images, + np.asarray(mask), + bad_pixel_threshold, + hot_pixel_threshold, + bad_pixel_low_threshold, + nobytes, + bins, + apply_reverse, + apply_rot90, + ) + + +def _compression_worker_run(descriptor): + filename, start, stop = descriptor + ( + images, + mask, + bad_pixel_threshold, + hot_pixel_threshold, + bad_pixel_low_threshold, + nobytes, + bins, + reverse, + rot90, + ) = _COMPRESSION_WORKER_CONTEXT + return _compress_segment( + images, + mask.copy(), + filename, + bad_pixel_threshold, + hot_pixel_threshold, + bad_pixel_low_threshold, + nobytes, + bins, + start, + stop, + reverse, + rot90, + ) + + +def _compression_worker_run_indexed(descriptor): + segment_index, filename, start, stop = descriptor + return segment_index, _compression_worker_run((filename, start, stop)) + + +def _iter_parallel_segment_results( + *, + images, + mask, + md, + filename, + num_sub, + bad_pixel_threshold, + hot_pixel_threshold, + bad_pixel_low_threshold, + nobytes, + bins, + dtypes, + reverse, + rot90, + num_max_para_process, + direct_load_data, + data_path, + images_per_file, + image_count, + segment_count, +): + """Yield completed segment reductions in order while workers stay alive.""" + raw_segment_size = num_sub * bins + descriptors = [ + ( + segment_index, + filename + "_temp-%i.tmp" % segment_index, + segment_index * raw_segment_size, + min((segment_index + 1) * raw_segment_size, image_count), + ) + for segment_index in range(segment_count) + ] + worker_count = min(segment_count, num_max_para_process, _available_cpu_count()) + print("Pool processes: %s" % worker_count) + pool = Pool( + processes=worker_count, + initializer=_compression_worker_init, + initargs=( + images, + mask, + md, + bad_pixel_threshold, + hot_pixel_threshold, + bad_pixel_low_threshold, + nobytes, + bins, + dtypes, + reverse, + rot90, + direct_load_data, + data_path, + images_per_file, + ), + ) + try: + iterator = pool.imap(_compression_worker_run_indexed, descriptors, chunksize=1) + pool.close() + yield from iterator + except BaseException: + pool.terminate() + raise + finally: + pool.join() + + +def _compress_segment( + images, + mask, + filename, + bad_pixel_threshold, + hot_pixel_threshold, + bad_pixel_low_threshold, + nobytes, + bins, + start, + stop, + reverse=False, + rot90=False, +): + """Compress one raw-frame range without constructing another reader.""" + if nobytes == 2: + dtype = np.int16 + elif nobytes == 4: + dtype = np.int32 + elif nobytes == 8: + dtype = np.float64 + else: + print("Wrong type of nobytes, only support 2 [np.int16] or 4 [np.int32]") + dtype = np.int32 + if bins != 1: + dtype = np.float64 + + output_count = len(_frame_bin_edges(stop - start, bins)) + imgsum = np.zeros(output_count) + avg_img = np.zeros(mask.shape, dtype=np.float64) + good_count = 0 + with open(filename, "wb") as stream: + for output_index, source_image in enumerate( + _iter_binned_images(images, start, stop, bins, reverse, rot90) + ): + image = np.asarray(source_image, dtype=dtype) + mask &= image < hot_pixel_threshold + flattened = image.ravel() + positions = np.flatnonzero((flattened > 0) & mask.ravel()) + values = flattened[positions] + imgsum[output_index] = values.sum() + if ( + len(positions) == 0 + or imgsum[output_index] > bad_pixel_threshold + or imgsum[output_index] <= bad_pixel_low_threshold + ): + _write_sparse_frame(stream, (), ()) + else: + avg_img.ravel()[positions] += values + good_count += 1 + _write_sparse_frame(stream, positions, values) + if good_count: + avg_img /= good_count + else: + avg_img.fill(np.nan) + bad_frames = (imgsum > bad_pixel_threshold) | (imgsum <= bad_pixel_low_threshold) + sys.stdout.write("#") + sys.stdout.flush() + return mask, avg_img, imgsum, bad_frames + + def segment_compress_eigerdata( images, mask, @@ -569,82 +927,41 @@ def segment_compress_eigerdata( Create a compressed eiger data without header, this function is for parallel compress for parallel compress don't pass any non-scalar parameters """ + owned_images = None + apply_reverse = False + apply_rot90 = False if dtypes == "uid": - uid = md["uid"] # images + uid = md["uid"] if not direct_load_data: detector = get_detector(db[uid]) - images = load_data(uid, detector, reverse=reverse, rot90=rot90)[N1:N2] - else: - images = EigerImages(data_path, images_per_file, md)[N1:N2] - if reverse: - images = reverse_updown(EigerImages(data_path, images_per_file, md))[N1:N2] - if rot90: - images = rot90_clockwise(images) - else: - images = images[N1:N2] - - Nimg_ = len(images) - M, N = images[0].shape - avg_img = np.zeros([M, N], dtype=np.float64) - _ = float(avg_img.size) - n = 0 - good_count = 0 - # frac = 0.0 - if nobytes == 2: - dtype = np.int16 - elif nobytes == 4: - dtype = np.int32 - elif nobytes == 8: - dtype = np.float64 - else: - print("Wrong type of nobytes, only support 2 [np.int16] or 4 [np.int32]") - dtype = np.int32 - - time_edge = _frame_bin_edges(Nimg_, bins) - Nimg = len(time_edge) - # print( time_edge, Nimg_, Nimg, bins, N1, N2 ) - imgsum = np.zeros(Nimg) - if bins != 1: - # print('The frames will be binned by %s'%bins) - dtype = np.float64 - - fp = open(filename, "wb") - for n in range(Nimg): - t1, t2 = time_edge[n] - if bins != 1: - img = np.array(np.average(images[t1:t2], axis=0), dtype=dtype) + source = load_data(uid, detector, reverse=reverse, rot90=rot90) else: - img = np.array(images[t1], dtype=dtype) - mask &= img < hot_pixel_threshold - p = np.where((np.ravel(img) > 0) * np.ravel(mask))[0] # don't use masked data - v = np.ravel(np.array(img, dtype=dtype))[p] - dlen = len(p) - imgsum[n] = v.sum() - if (dlen == 0) or (imgsum[n] > bad_pixel_threshold) or (imgsum[n] <= bad_pixel_low_threshold): - dlen = 0 - fp.write(struct.pack("@I", dlen)) - else: - np.ravel(avg_img)[p] += v - good_count += 1 - fp.write(struct.pack("@I", dlen)) - fp.write(struct.pack("@{}i".format(dlen), *p)) - if bins == 1: - fp.write(struct.pack("@{}{}".format(dlen, "ih"[nobytes == 2]), *v)) - else: - fp.write(struct.pack("@{}{}".format(dlen, "dd"[nobytes == 2]), *v)) # n +=1 - del p, v, img - fp.flush() - fp.close() - if good_count: - avg_img /= good_count + owned_images = EigerImages(data_path, images_per_file, md) + source = owned_images + apply_reverse = reverse + apply_rot90 = rot90 else: - avg_img.fill(np.nan) - bad_frame_list = (np.array(imgsum) > bad_pixel_threshold) | (np.array(imgsum) <= bad_pixel_low_threshold) - sys.stdout.write("#") - sys.stdout.flush() - # del images, mask, avg_img, imgsum, bad_frame_list - # print( 'Should release memory here') - return mask, avg_img, imgsum, bad_frame_list + source = images + start = 0 if N1 is None else N1 + stop = len(source) if N2 is None else min(N2, len(source)) + try: + return _compress_segment( + source, + mask, + filename, + bad_pixel_threshold, + hot_pixel_threshold, + bad_pixel_low_threshold, + nobytes, + bins, + start, + stop, + apply_reverse, + apply_rot90, + ) + finally: + if owned_images is not None: + owned_images.close() def create_compress_header(md, filename, nobytes=4, bins=1, rot90=False): @@ -834,8 +1151,17 @@ def init_compress_eigerdata( fp.write(Header) + owned_images = None + apply_reverse = False + apply_rot90 = False + if direct_load_data: + owned_images = EigerImages(data_path, images_per_file, md) + images = owned_images + apply_reverse = reverse + apply_rot90 = rot90 + Nimg_ = len(images) - avg_img = np.zeros_like(images[0], dtype=np.float64) + avg_img = np.zeros(mask.shape, dtype=np.float64) Nopix = float(avg_img.size) n = 0 good_count = 0 @@ -857,35 +1183,32 @@ def init_compress_eigerdata( if bins != 1: print("The frames will be binned by %s" % bins) - for n in tqdm(range(Nimg)): - t1, t2 = time_edge[n] - img = np.average(images[t1:t2], axis=0) - mask &= img < hot_pixel_threshold - p = np.where((np.ravel(img) > 0) & np.ravel(mask))[0] # don't use masked data - v = np.ravel(np.array(img, dtype=dtype))[p] - dlen = len(p) - imgsum[n] = v.sum() - if (imgsum[n] > bad_pixel_threshold) or (imgsum[n] <= bad_pixel_low_threshold): - # if imgsum[n] >=bad_pixel_threshold : - dlen = 0 - fp.write(struct.pack("@I", dlen)) - else: - np.ravel(avg_img)[p] += v - good_count += 1 - frac += dlen / Nopix - # s_fmt ='@I{}i{}{}'.format( dlen,dlen,'ih'[nobytes==2]) - fp.write(struct.pack("@I", dlen)) - fp.write(struct.pack("@{}i".format(dlen), *p)) - if bins == 1: - if nobytes != 8: - fp.write(struct.pack("@{}{}".format(dlen, "ih"[nobytes == 2]), *v)) - else: - fp.write(struct.pack("@{}{}".format(dlen, "dd"[nobytes == 2]), *v)) + try: + image_iterator = _iter_binned_images( + images, + 0, + Nimg_, + bins, + reverse=apply_reverse, + rot90=apply_rot90, + ) + for n, image in enumerate(tqdm(image_iterator, total=Nimg)): + mask &= image < hot_pixel_threshold + p = np.where((np.ravel(image) > 0) & np.ravel(mask))[0] # don't use masked data + v = np.ravel(np.array(image, dtype=dtype))[p] + dlen = len(p) + imgsum[n] = v.sum() + if (imgsum[n] > bad_pixel_threshold) or (imgsum[n] <= bad_pixel_low_threshold): + _write_sparse_frame(fp, (), ()) else: - fp.write(struct.pack("@{}{}".format(dlen, "dd"[nobytes == 2]), *v)) - # n +=1 - - fp.close() + np.ravel(avg_img)[p] += v + good_count += 1 + frac += dlen / Nopix + _write_sparse_frame(fp, p, v) + finally: + fp.close() + if owned_images is not None: + owned_images.close() if good_count: frac /= good_count else: @@ -968,6 +1291,9 @@ def __init__(self, filename, beg, end, reverse=False): self.filename = filename # br: bytes read br = self.FID.read(1024) + if len(br) != 1024: + self.FID.close() + raise ValueError("malformed compressed file: incomplete 1024-byte header") self.beg = beg self.end = end self.reverse = reverse @@ -989,12 +1315,18 @@ def __init__(self, filename, beg, end, reverse=False): "cols_end", ] - _ = struct.unpack("@16s", br[:16]) + version = struct.unpack("@16s", br[:16])[0] + if version != b"Version-COMP0001": + self.FID.close() + raise ValueError("unsupported compressed file header") md_temp = struct.unpack("@8d7I916x", br[16:]) self.md = dict(zip(ms_keys, md_temp)) self.imgread = 0 self.recno = 0 + self._mmap = None + self._frame_offsets = None + self._bytes_traversed = 0 if reverse: nrows = self.md["nrows"] @@ -1018,22 +1350,98 @@ def __init__(self, filename, beg, end, reverse=False): self.valtype = np.uint32 elif self.byts == 8: self.valtype = np.float64 + else: + self.FID.close() + raise ValueError("malformed compressed file: bytes per value must be 2, 4, or 8") # now convert pieces of these bytes to our data - self.dlen = np.fromfile(self.FID, dtype=np.int32, count=1)[0] + first_length = np.fromfile(self.FID, dtype=np.int32, count=1) + if first_length.size != 1: + self.FID.close() + raise ValueError("malformed compressed file: missing first frame") + self.dlen = first_length[0] + if self.dlen < 0: + self.FID.close() + raise ValueError("malformed compressed file: negative sparse-frame length") # now read first image # print "Opened file. Bytes per data is {0img.shape = (self.rows,self.cols)}".format(self.byts) def _readHeader(self): - self.dlen = np.fromfile(self.FID, dtype=np.int32, count=1)[0] + length = np.fromfile(self.FID, dtype=np.int32, count=1) + if length.size != 1: + raise ValueError("malformed compressed file: truncated frame header") + self.dlen = length[0] + if self.dlen < 0: + raise ValueError("malformed compressed file: negative sparse-frame length") def _readImageRaw(self): p = np.fromfile(self.FID, dtype=np.int32, count=self.dlen) v = np.fromfile(self.FID, dtype=self.valtype, count=self.dlen) + if p.size != self.dlen or v.size != self.dlen: + raise ValueError("malformed compressed file: truncated sparse-frame payload") self.imgread = 1 return (p, v) + def _ensure_index(self): + """Build and validate an in-memory frame-offset index on first use.""" + if self._frame_offsets is not None: + return self._frame_offsets + import mmap + + if self.FID.closed: + raise ValueError("I/O operation on closed compressed file") + mapped = mmap.mmap(self.FID.fileno(), 0, access=mmap.ACCESS_READ) + offsets = [] + position = 1024 + file_size = len(mapped) + try: + for frame in range(self.end): + if position + 4 > file_size: + raise ValueError("malformed compressed file: requested frame range extends past end of file") + length = struct.unpack_from("@i", mapped, position)[0] + if length < 0: + raise ValueError("malformed compressed file: negative sparse-frame length") + next_position = position + 4 + length * (4 + self.byts) + if next_position > file_size: + raise ValueError("malformed compressed file: truncated sparse-frame payload") + offsets.append(position) + position = next_position + except BaseException: + mapped.close() + raise + self._mmap = mapped + self._frame_offsets = np.asarray(offsets, dtype=np.int64) + self._bytes_traversed += 4 * len(offsets) + return self._frame_offsets + + def _raw_frame_view(self, n): + """Return read-only zero-copy position/value views for private consumers.""" + if n < self.beg or n >= self.end: + raise IndexError("Error, record out of range") + offsets = self._ensure_index() + offset = int(offsets[n]) + length = struct.unpack_from("@i", self._mmap, offset)[0] + positions = np.frombuffer(self._mmap, dtype=np.int32, count=length, offset=offset + 4) + values = np.frombuffer( + self._mmap, + dtype=self.valtype, + count=length, + offset=offset + 4 + length * np.dtype(np.int32).itemsize, + ) + positions.flags.writeable = False + values.flags.writeable = False + self._bytes_traversed += 4 + length * (np.dtype(np.int32).itemsize + self.byts) + return positions, values + + def _iter_raw_frames(self, indices): + self._ensure_index() + for index in indices: + yield index, self._raw_frame_view(index) + + def _reset_io_counters(self): + self._bytes_traversed = 0 + def _readImage(self): p, v = self._readImageRaw() img = np.zeros((self.md["ncols"], self.md["nrows"])) @@ -1088,7 +1496,38 @@ def rdrawframe(self, n): def close(self): """Close the compressed-data file.""" - self.FID.close() + if self._mmap is not None: + try: + self._mmap.close() + except BufferError: + # A private zero-copy view may briefly outlive this object. + pass + self._mmap = None + self._frame_offsets = None + if not self.FID.closed: + self.FID.close() + + def reopen(self): + """Reopen a closed compressed file and reset its sequential cursor.""" + if not self.FID.closed: + return self + replacement = type(self)(self.filename, self.beg, self.end, self.reverse) + self.__dict__.update(replacement.__dict__) + return self + + def __getstate__(self): + state = self.__dict__.copy() + state["_closed"] = self.FID.closed + state.pop("FID", None) + state.pop("_mmap", None) + state.pop("_frame_offsets", None) + return state + + def __setstate__(self, state): + replacement = type(self)(state["filename"], state["beg"], state["end"], state["reverse"]) + self.__dict__.update(replacement.__dict__) + if state.get("_closed", False): + self.close() def __enter__(self): return self @@ -1309,8 +1748,11 @@ def get_avg_imgc( tqdm(sample_indices, desc="Averaging %s images" % len(sample_indices)) if show_progress else sample_indices ) for index in indices: - p, v = FD.rdrawframe(index) - np.ravel(avg_img)[p] += v + if hasattr(FD, "_raw_frame_view"): + p, v = FD._raw_frame_view(index) + else: + p, v = FD.rdrawframe(index) + sparse_add_image(p, v, np.ravel(avg_img)) if sample_indices: avg_img /= len(sample_indices) @@ -1369,7 +1811,7 @@ def mean_intensityc(FD, labeled_array, sampling=1, index=None, multi_cor=False): """ qind, pixelist = roi.extract_label_indices(labeled_array) - sx, sy = (FD.rdframe(FD.beg)).shape + sx, sy = FD.md["ncols"], FD.md["nrows"] if labeled_array.shape != (sx, sy): raise ValueError( " `image` shape (%d, %d) in FD is not equal to the labeled_array shape (%d, %d)" @@ -1401,37 +1843,15 @@ def mean_intensityc(FD, labeled_array, sampling=1, index=None, multi_cor=False): sample_indices = range(FD.beg, FD.end, sampling) mean_intensity = np.zeros([len(sample_indices), len(index)]) - # fra_pix = np.zeros_like( pixelist, dtype=np.float64) - timg = np.zeros(FD.md["ncols"] * FD.md["nrows"], dtype=np.int32) - timg[pixelist] = np.arange(1, len(pixelist) + 1) - # maxqind = max(qind) + roi_lookup = np.full(FD.md["ncols"] * FD.md["nrows"], -1, dtype=np.int64) + roi_lookup[pixelist] = qind - 1 norm = np.bincount(qind, minlength=len(index) + 1)[1:] - n = 0 - # for i in tqdm(range( FD.beg , FD.end )): - if not multi_cor: - for i in tqdm(sample_indices, desc="Get ROI intensity of each frame"): - p, v = FD.rdrawframe(i) - w = np.where(timg[p])[0] - pxlist = timg[p[w]] - 1 - mean_intensity[n] = np.bincount(qind[pxlist], weights=v[w], minlength=len(index) + 1)[1:] - n += 1 - else: - ring_masks = [np.array(labeled_array == label, dtype=np.int64) for label in index] - inputs = range(len(ring_masks)) - go_through_FD(FD) - pool = _make_pool(len(inputs)) - print("Starting assign the tasks...") - results = {} - for i in tqdm(inputs): - results[i] = apply_async(pool, _get_mean_intensity_one_q, (FD, sampling, ring_masks[i])) - print("Starting running the tasks...") - res = _collect_pool_results(pool, results, show_progress=True) - # return res - for i in inputs: - mean_intensity[:, i] = res[i] - print("ROI mean_intensit calculation is DONE!") - del results - del res + for output_row, frame_index in enumerate(tqdm(sample_indices, desc="Get ROI intensity of each frame")): + if hasattr(FD, "_raw_frame_view"): + positions, values = FD._raw_frame_view(frame_index) + else: + positions, values = FD.rdrawframe(frame_index) + sparse_roi_sums(positions, values, roi_lookup, mean_intensity[output_row]) mean_intensity /= norm return mean_intensity, index @@ -1480,9 +1900,12 @@ def get_each_frame_intensityc( imgsum = np.zeros(len(sample_indices)) n = 0 for i in tqdm(sample_indices, desc="Get each frame intensity"): - p, v = FD.rdrawframe(i) - if len(p) > 0: - imgsum[n] = np.sum(v) + if hasattr(FD, "_raw_frame_view"): + _, v = FD._raw_frame_view(i) + else: + _, v = FD.rdrawframe(i) + if len(v) > 0: + imgsum[n] = sparse_frame_sum(v) n += 1 if plot_: diff --git a/pyCHX/chx_compress_analysis.py b/pyCHX/chx_compress_analysis.py index f88de62..793477d 100644 --- a/pyCHX/chx_compress_analysis.py +++ b/pyCHX/chx_compress_analysis.py @@ -9,6 +9,7 @@ from tqdm import tqdm from pyCHX._optional import imshow +from pyCHX._performance import sparse_scatter_normalized from pyCHX.chx_compress import Multifile as _Multifile from pyCHX.chx_compress import compress_eigerdata as _compress_eigerdata from pyCHX.chx_compress import get_avg_imgc @@ -128,24 +129,35 @@ def cal_waterfallc( # pre-allocate an array for performance # might be able to use list comprehension to make this faster - watf = np.zeros([int((FD.end - FD.beg) / sampling), len(qind)]) - - # fra_pix = np.zeros_like( pixelist, dtype=np.float64) - - timg = np.zeros(FD.md["ncols"] * FD.md["nrows"], dtype=np.int32) - timg[pixelist] = np.arange(1, len(pixelist) + 1) - - # maxqind = max(qind) - _ = np.bincount(qind)[1:] - n = 0 - # for i in tqdm(range( FD.beg , FD.end )): - for i in tqdm(range(FD.beg, FD.end, sampling), desc="Get waterfall for q index=%s" % qindex): - p, v = FD.rdrawframe(i) - w = np.where(timg[p])[0] - pxlist = timg[p[w]] - 1 - - watf[n][pxlist] = v[w] - n += 1 + watf = np.zeros([len(range(FD.beg, FD.end, sampling)), len(qind)]) + + lookup = np.full(FD.md["ncols"] * FD.md["nrows"], -1, dtype=np.int64) + lookup[pixelist] = np.arange(len(pixelist), dtype=np.int64) + dummy_float_1d = np.ones(1, dtype=np.float64) + dummy_float_2d = np.ones((1, 1), dtype=np.float64) + dummy_int = np.zeros(len(pixelist), dtype=np.int64) + flags = np.zeros(4, dtype=np.bool_) + frames = range(FD.beg, FD.end, sampling) + for output_row, frame_index in enumerate(tqdm(frames, desc="Get waterfall for q index=%s" % qindex)): + if hasattr(FD, "_raw_frame_view"): + positions, values = FD._raw_frame_view(frame_index) + else: + positions, values = FD.rdrawframe(frame_index) + sparse_scatter_normalized( + positions, + values, + lookup, + watf, + output_row, + 0, + dummy_float_1d, + dummy_float_2d, + dummy_int, + dummy_float_1d, + dummy_float_2d, + dummy_int, + flags, + ) if bin_waterfall: watf_ = watf.copy() diff --git a/pyCHX/chx_correlationc.py b/pyCHX/chx_correlationc.py index e055edc..19d7a74 100644 --- a/pyCHX/chx_correlationc.py +++ b/pyCHX/chx_correlationc.py @@ -8,14 +8,28 @@ import logging from collections import namedtuple +from concurrent.futures import ThreadPoolExecutor import matplotlib.pyplot as plt import numpy as np import skbeam.core.roi as roi +from scipy.linalg.blas import dsyrk from skbeam.core.roi import extract_label_indices from skbeam.core.utils import multi_tau_lags +from threadpoolctl import threadpool_limits from tqdm import tqdm +from pyCHX._performance import ( + available_memory_bytes, + mirror_and_normalize_two_time, + mirror_and_normalize_two_time_parallel, + mirror_two_time, + mirror_two_time_parallel, + numba_thread_limit, + physical_core_count, + sparse_scatter_normalized, + store_symmetric_two_time_batch, +) from pyCHX.chx_generic_functions import plot1D from pyCHX.chx_libs import markers @@ -35,6 +49,39 @@ def _one_time_process( buf_no, norm, lev_len, +): + """Run the one-time kernel without an external intensity cache.""" + return _one_time_process_cached( + buf, + G, + past_intensity_norm, + future_intensity_norm, + label_array, + num_bufs, + num_pixels, + img_per_level, + level, + buf_no, + norm, + lev_len, + None, + ) + + +def _one_time_process_cached( + buf, + G, + past_intensity_norm, + future_intensity_norm, + label_array, + num_bufs, + num_pixels, + img_per_level, + level, + buf_no, + norm, + lev_len, + intensity_buf, ): """Reference implementation of the inner loop of multi-tau one time correlation @@ -68,6 +115,8 @@ def _one_time_process( to track bad images lev_len : array length of each level + intensity_buf : array, optional + cached, unnormalized ROI sums for each ring-buffer slot Notes ----- .. math:: @@ -81,30 +130,40 @@ def _one_time_process( # in multi-tau correlation, the subsequent levels have half as many # buffers as the first i_min = num_bufs // 2 if level else 0 - # maxqind=G.shape[1] + level_offset = lev_len[:level].sum() + future_img = buf[level, buf_no] + future_has_nan = np.isnan(future_img).any() + if not future_has_nan: + future_binned = np.bincount(label_array, weights=future_img)[1:] + if intensity_buf is not None: + intensity_buf[level, buf_no] = future_binned + product = np.empty_like(future_img) for i in range(i_min, min(img_per_level[level], num_bufs)): # compute the index into the autocorrelation matrix t_index = int(level * num_bufs / 2 + i) delay_no = (buf_no - i) % num_bufs # get the images for correlating past_img = buf[level, delay_no] - future_img = buf[level, buf_no] # find the normalization that can work both for bad_images # and good_images - ind = int(t_index - lev_len[:level].sum()) + ind = int(t_index - level_offset) normalize = img_per_level[level] - i - norm[level + 1][ind] # take out the past_ing and future_img created using bad images # (bad images are converted to np.nan array) - if np.isnan(past_img).any() or np.isnan(future_img).any(): + if future_has_nan or np.isnan(past_img).any(): norm[level + 1][ind] += 1 else: - for w, arr in zip( - [past_img * future_img, past_img, future_img], [G, past_intensity_norm, future_intensity_norm] - ): - binned = np.bincount(label_array, weights=w)[1:] - # nonz = np.where(w)[0] - # binned = np.bincount(label_array[nonz], weights=w[nonz], minlength=maxqind+1 )[1:] - arr[t_index] += (binned / num_pixels - arr[t_index]) / normalize + np.multiply(past_img, future_img, out=product) + product_binned = np.bincount(label_array, weights=product)[1:] + if intensity_buf is None: + past_binned = np.bincount(label_array, weights=past_img)[1:] + else: + past_binned = intensity_buf[level, delay_no] + G[t_index] += (product_binned / num_pixels - G[t_index]) / normalize + past_intensity_norm[t_index] += (past_binned / num_pixels - past_intensity_norm[t_index]) / normalize + future_intensity_norm[t_index] += ( + future_binned / num_pixels - future_intensity_norm[t_index] + ) / normalize return None # modifies arguments in place! @@ -171,21 +230,24 @@ def _one_time_process_error( # in multi-tau correlation, the subsequent levels have half as many # buffers as the first i_min = num_bufs // 2 if level else 0 - # maxqind=G.shape[1] + level_offset = lev_len[:level].sum() + future_img = buf[level, buf_no] + future_has_nan = np.isnan(future_img).any() + if not future_has_nan: + product = np.empty_like(future_img) for i in range(i_min, min(img_per_level[level], num_bufs)): # compute the index into the autocorrelation matrix t_index = int(level * num_bufs / 2 + i) delay_no = (buf_no - i) % num_bufs # get the images for correlating past_img = buf[level, delay_no] - future_img = buf[level, buf_no] # find the normalization that can work both for bad_images # and good_images - ind = int(t_index - lev_len[:level].sum()) + ind = int(t_index - level_offset) normalize = img_per_level[level] - i - norm[level + 1][ind] # take out the past_ing and future_img created using bad images # (bad images are converted to np.nan array) - if np.isnan(past_img).any() or np.isnan(future_img).any(): + if future_has_nan or np.isnan(past_img).any(): norm[level + 1][ind] += 1 else: @@ -197,8 +259,9 @@ def _one_time_process_error( # #binned = np.bincount(label_array[nonz], weights=w[nonz], minlength=maxqind+1 )[1:] # arr[t_index] += ((binned / num_pixels - # arr[t_index]) / normalize) + np.multiply(past_img, future_img, out=product) for w, arr in zip( - [past_img * future_img, past_img, future_img], + [product, past_img, future_img], [ G_err, past_intensity_norm_err, @@ -536,7 +599,7 @@ def lazy_one_time( # create a shorthand reference to the results and state named tuple s = internal_state - qind, pixelist = roi.extract_label_indices(labels) + pixelist = s.pixel_list # iterate over the images to compute multi-tau correlation fra_pix = np.zeros_like(pixelist, dtype=np.float64) @@ -546,32 +609,36 @@ def lazy_one_time( if bad_frame_list is None: bad_frame_list = [] + bad_frames = set(bad_frame_list) + has_imgsum_norm = imgsum is not None + has_pixel_norm = norm is not None + pixel_norm_is_2d = has_pixel_norm and len(norm.shape) > 1 for i in tqdm(range(FD.beg, FD.end)): - if i in bad_frame_list: + if i in bad_frames: fra_pix[:] = np.nan else: p, v = FD.rdrawframe(i) - w = np.where(timg[p])[0] - pxlist = timg[p[w]] - 1 - - if imgsum is None: - if norm is None: - fra_pix[pxlist] = v[w] + mapped_pixels = timg[p] + selected = mapped_pixels != 0 + pxlist = mapped_pixels[selected] - 1 + values = v[selected] + + if not has_imgsum_norm: + if not has_pixel_norm: + fra_pix[pxlist] = values else: - S = norm.shape - if len(S) > 1: - fra_pix[pxlist] = v[w] / norm[i, pxlist] # -1.0 + if pixel_norm_is_2d: + fra_pix[pxlist] = values / norm[i, pxlist] # -1.0 else: - fra_pix[pxlist] = v[w] / norm[pxlist] # -1.0 + fra_pix[pxlist] = values / norm[pxlist] # -1.0 else: - if norm is None: - fra_pix[pxlist] = v[w] / imgsum[i] + if not has_pixel_norm: + fra_pix[pxlist] = values / imgsum[i] else: - S = norm.shape - if len(S) > 1: - fra_pix[pxlist] = v[w] / imgsum[i] / norm[i, pxlist] + if pixel_norm_is_2d: + fra_pix[pxlist] = values / imgsum[i] / norm[i, pxlist] else: - fra_pix[pxlist] = v[w] / imgsum[i] / norm[pxlist] + fra_pix[pxlist] = values / imgsum[i] / norm[pxlist] level = 0 # increment buffer s.cur[0] = (1 + s.cur[0]) % num_bufs @@ -634,9 +701,13 @@ def lazy_one_time( prev = 1 + (s.cur[level - 1] - 2) % num_bufs s.cur[level] = 1 + s.cur[level] % num_bufs - s.buf[level, s.cur[level] - 1] = ( - s.buf[level - 1, prev - 1] + s.buf[level - 1, s.cur[level - 1] - 1] - ) / 2 + level_buffer = s.buf[level, s.cur[level] - 1] + np.add( + s.buf[level - 1, prev - 1], + s.buf[level - 1, s.cur[level - 1] - 1], + out=level_buffer, + ) + level_buffer /= 2 # make the track_level zero once that level is processed s.track_level[level] = False @@ -1033,31 +1104,37 @@ def lazy_two_time( two_time_internal_state = _init_state_two_time(num_levels, num_bufs, labels, num_frames) # create a shorthand reference to the results and state named tuple s = two_time_internal_state - qind, pixelist = roi.extract_label_indices(labels) + pixelist = s.pixel_list # iterate over the images to compute multi-tau correlation fra_pix = np.zeros_like(pixelist, dtype=np.float64) timg = np.zeros(FD.md["ncols"] * FD.md["nrows"], dtype=np.int32) timg[pixelist] = np.arange(1, len(pixelist) + 1) if bad_frame_list is None: bad_frame_list = [] + bad_frames = set(bad_frame_list) + has_imgsum_norm = imgsum is not None + has_pixel_norm = norm is not None + intensity_buf = _create_intensity_buffer(s.buf, s.label_array, len(s.num_pixels)) for i in tqdm(range(FD.beg, FD.end)): - if i in bad_frame_list: + if i in bad_frames: fra_pix[:] = np.nan else: p, v = FD.rdrawframe(i) - w = np.where(timg[p])[0] - pxlist = timg[p[w]] - 1 - if imgsum is None: - if norm is None: - fra_pix[pxlist] = v[w] + mapped_pixels = timg[p] + selected = mapped_pixels != 0 + pxlist = mapped_pixels[selected] - 1 + values = v[selected] + if not has_imgsum_norm: + if not has_pixel_norm: + fra_pix[pxlist] = values else: - fra_pix[pxlist] = v[w] / norm[pxlist] # -1.0 + fra_pix[pxlist] = values / norm[pxlist] # -1.0 else: - if norm is None: - fra_pix[pxlist] = v[w] / imgsum[i] + if not has_pixel_norm: + fra_pix[pxlist] = values / imgsum[i] else: - fra_pix[pxlist] = v[w] / imgsum[i] / norm[pxlist] + fra_pix[pxlist] = values / imgsum[i] / norm[pxlist] level = 0 # increment buffer @@ -1068,7 +1145,7 @@ def lazy_two_time( # Put the ROI pixels into the ring buffer. s.buf[0, s.cur[0] - 1] = fra_pix fra_pix[:] = 0 - _two_time_process( + _two_time_process_cached( s.buf, s.g2, s.label_array, @@ -1079,6 +1156,7 @@ def lazy_two_time( s.current_img_time, level=0, buf_no=s.cur[0] - 1, + intensity_buf=intensity_buf, ) # time frame for each level s.time_ind[0].append(s.current_img_time) @@ -1095,9 +1173,13 @@ def lazy_two_time( prev = 1 + (s.cur[level - 1] - 2) % num_bufs s.cur[level] = 1 + s.cur[level] % num_bufs s.count_level[level] = 1 + s.count_level[level] - s.buf[level, s.cur[level] - 1] = ( - s.buf[level - 1, prev - 1] + s.buf[level - 1, s.cur[level - 1] - 1] - ) / 2 + level_buffer = s.buf[level, s.cur[level] - 1] + np.add( + s.buf[level - 1, prev - 1], + s.buf[level - 1, s.cur[level - 1] - 1], + out=level_buffer, + ) + level_buffer /= 2 t1_idx = (s.count_level[level] - 1) * 2 @@ -1110,7 +1192,7 @@ def lazy_two_time( # for multi-tau levels greater than one # Again, this is modifying things in place. See comment # on previous call above. - _two_time_process( + _two_time_process_cached( s.buf, s.g2, s.label_array, @@ -1121,6 +1203,7 @@ def lazy_two_time( current_img_time, level=level, buf_no=s.cur[level] - 1, + intensity_buf=intensity_buf, ) level += 1 @@ -1150,6 +1233,44 @@ def two_time_state_to_results(state): def _two_time_process( buf, g2, label_array, num_bufs, num_pixels, img_per_level, lag_steps, current_img_time, level, buf_no +): + """Run the two-time kernel without an external intensity cache.""" + return _two_time_process_cached( + buf, + g2, + label_array, + num_bufs, + num_pixels, + img_per_level, + lag_steps, + current_img_time, + level, + buf_no, + None, + ) + + +def _create_intensity_buffer(buf, label_array, num_rois): + """Create ROI-sum caches matching the current correlation buffers.""" + intensity_buf = np.empty((*buf.shape[:2], num_rois), dtype=np.float64) + for level in range(buf.shape[0]): + for buf_no in range(buf.shape[1]): + intensity_buf[level, buf_no] = np.bincount(label_array, weights=buf[level, buf_no])[1:] + return intensity_buf + + +def _two_time_process_cached( + buf, + g2, + label_array, + num_bufs, + num_pixels, + img_per_level, + lag_steps, + current_img_time, + level, + buf_no, + intensity_buf, ): """ Parameters @@ -1189,21 +1310,34 @@ def _two_time_process( else: i_min = num_bufs // 2 - for i in range(i_min, min(img_per_level[level], num_bufs)): + future_img = buf[level, buf_no] + future_binned = np.bincount(label_array, weights=future_img)[1:] + if intensity_buf is not None: + intensity_buf[level, buf_no] = future_binned + + i_max = min(img_per_level[level], num_bufs) + if i_min >= i_max: + return + + product = np.empty_like(future_img) + for i in range(i_min, i_max): t_index = level * num_bufs / 2 + i delay_no = (buf_no - i) % num_bufs past_img = buf[level, delay_no] - future_img = buf[level, buf_no] # print( np.sum( past_img ), np.sum( future_img )) # get the matrix of correlation function without normalizations - tmp_binned = np.bincount(label_array, weights=past_img * future_img)[1:] + np.multiply(past_img, future_img, out=product) + tmp_binned = np.bincount(label_array, weights=product)[1:] # get the matrix of past intensity normalizations - pi_binned = np.bincount(label_array, weights=past_img)[1:] + if intensity_buf is None: + pi_binned = np.bincount(label_array, weights=past_img)[1:] + else: + pi_binned = intensity_buf[level, delay_no] # get the matrix of future intensity normalizations - fi_binned = np.bincount(label_array, weights=future_img)[1:] + fi_binned = future_binned tind1 = current_img_time - 1 tind2 = current_img_time - lag_steps[int(t_index)] - 1 @@ -1593,63 +1727,63 @@ def get_data(self): Return: 2-D array, shape as (len(images), len(pixellist)) """ - data_array = np.zeros([self.length, len(self.pixelist)], dtype=np.float64) - # fra_pix = np.zeros_like( pixelist, dtype=np.float64) - timg = np.zeros(self.FD.md["ncols"] * self.FD.md["nrows"], dtype=np.int32) - timg[self.pixelist] = np.arange(1, len(self.pixelist) + 1) - - if self.mean_int_sets is not None: - # Mean_Int_Qind = np.array( self.qind.copy(), dtype=np.float) - Mean_Int_Qind = np.ones(len(self.qind), dtype=np.float64) - noqs = len(np.unique(self.qind)) - nopr = np.bincount(self.qind - 1) - noprs = np.concatenate([np.array([0]), np.cumsum(nopr)]) - qind_ = np.zeros_like(self.qind) - for j in range(noqs): - qind_[noprs[j] : noprs[j + 1]] = np.where(self.qind == j + 1)[0] - - n = 0 - for i in tqdm(range(self.beg, self.end)): - p, v = self.FD.rdrawframe(i) - w = np.where(timg[p])[0] - pxlist = timg[p[w]] - 1 - - if self.mean_int_sets is not None: # for normalization of each averaged ROI of each frame - for j in range(noqs): - # if i ==100: - # if j==0: - # print( self.mean_int_sets[i][j] ) - # print( qind_[ noprs[j]: noprs[j+1] ] ) - Mean_Int_Qind[qind_[noprs[j] : noprs[j + 1]]] = self.mean_int_sets[i][j] - norm_Mean_Int_Qind = Mean_Int_Qind[pxlist] # self.mean_int_set or Mean_Int_Qind[pxlist] - - # if i==100: - # print( i, Mean_Int_Qind[ self.qind== 11 ]) - - # print('Do norm_mean_int here') - # if i ==10: - # print( norm_Mean_Int_Qind ) - else: - norm_Mean_Int_Qind = 1.0 - if self.imgsum is not None: - norm_imgsum = self.imgsum[i] - else: - norm_imgsum = 1.0 - if self.norm is not None: - if len((self.norm).shape) > 1: - norm_avgimg_roi = self.norm[i][pxlist] - # print('here') + pixelist = np.asarray(self.pixelist, dtype=np.int64) + data_array = np.zeros([self.length, len(pixelist)], dtype=np.float64, order="C") + lookup = np.full(self.FD.md["ncols"] * self.FD.md["nrows"], -1, dtype=np.int64) + lookup[pixelist] = np.arange(len(pixelist), dtype=np.int64) + + has_mean_norm = self.mean_int_sets is not None + has_imgsum_norm = self.imgsum is not None + has_pixel_norm = self.norm is not None + pixel_norm_is_2d = has_pixel_norm and np.ndim(self.norm) > 1 + norm_1d = np.asarray( + self.norm if has_pixel_norm and not pixel_norm_is_2d else np.ones(1), dtype=np.float64 + ) + norm_2d = np.asarray(self.norm if pixel_norm_is_2d else np.ones((1, 1)), dtype=np.float64) + imgsum = np.asarray(self.imgsum if has_imgsum_norm else np.ones(1), dtype=np.float64) + mean_int_sets = np.asarray(self.mean_int_sets if has_mean_norm else np.ones((1, 1)), dtype=np.float64) + qind = np.asarray(self.qind if has_mean_norm else np.zeros(len(pixelist)), dtype=np.int64) + norm_columns = np.arange(len(pixelist), dtype=np.int64) + flags = np.asarray( + [has_pixel_norm and not pixel_norm_is_2d, pixel_norm_is_2d, has_imgsum_norm, has_mean_norm], + dtype=np.bool_, + ) + def fill_range(start, stop): + for frame_index in range(start, stop): + if hasattr(self.FD, "_raw_frame_view"): + positions, values = self.FD._raw_frame_view(frame_index) else: - norm_avgimg_roi = self.norm[pxlist] - else: - norm_avgimg_roi = 1.0 + positions, values = self.FD.rdrawframe(frame_index) + sparse_scatter_normalized( + positions, + values, + lookup, + data_array, + frame_index - self.beg, + frame_index, + norm_1d, + norm_2d, + norm_columns, + imgsum, + mean_int_sets, + qind, + flags, + ) - norms = norm_Mean_Int_Qind * norm_imgsum * norm_avgimg_roi - # if i==100: - # print(norm_Mean_Int_Qind[:100]) - data_array[n][pxlist] = v[w] / norms - n += 1 + if hasattr(self.FD, "_raw_frame_view") and self.length > 1: + self.FD._ensure_index() + worker_count = min(physical_core_count(), self.length) + chunk_size = max(1, (self.length + worker_count - 1) // worker_count) + with ThreadPoolExecutor(max_workers=worker_count) as executor: + futures = [ + executor.submit(fill_range, start, min(start + chunk_size, self.end)) + for start in range(self.beg, self.end, chunk_size) + ] + for future in futures: + future.result() + else: + fill_range(self.beg, self.end) return data_array @@ -1671,6 +1805,114 @@ def _select_two_time_rois(rois, index): return qind, selected_labels, pixel_counts +def _symmetric_two_time_product(data, row_norm, pixel_count): + """Compute one ROI matrix, pre-normalizing production-sized float inputs.""" + correlation, pre_normalized = _upper_two_time_product(data, row_norm, pixel_count) + use_parallel_finish = data.shape[0] >= 1_024 + if pre_normalized: + if use_parallel_finish: + mirror_two_time_parallel(correlation) + else: + mirror_two_time(correlation) + else: + if use_parallel_finish: + mirror_and_normalize_two_time_parallel(correlation, row_norm, pixel_count) + else: + mirror_and_normalize_two_time(correlation, row_norm, pixel_count) + return correlation + + +def _upper_two_time_product(data, row_norm, pixel_count, output=None): + """Compute one ROI's upper triangle, optionally into a reusable buffer.""" + pre_normalized = False + if np.issubdtype(data.dtype, np.floating): + if data.size >= 262_144 and np.all(np.isfinite(row_norm)) and np.all(row_norm != 0): + fortran_data = np.asarray(data, dtype=np.float64, order="F") + if np.all(row_norm == 1): + normalized = fortran_data + else: + normalized = fortran_data / row_norm[:, None] + kwargs = {} + if output is not None: + kwargs = {"c": output, "overwrite_c": 1} + correlation = dsyrk(1.0 / pixel_count, normalized, lower=0, trans=0, **kwargs) + pre_normalized = True + else: + kwargs = {} + if output is not None: + kwargs = {"c": output, "overwrite_c": 1} + correlation = dsyrk( + 1.0, + np.asarray(data, dtype=np.float64), + lower=0, + trans=0, + **kwargs, + ) + else: + correlation = np.dot(data, data.T).astype(np.float64) + + if output is not None and not np.shares_memory(correlation, output): + output[:, :] = correlation + correlation = output + return correlation, pre_normalized + + +def _two_time_batch_size(frame_count, roi_count): + """Choose a cache-friendly batch while bounding its square work buffers.""" + matrix_bytes = max(1, frame_count * frame_count * np.dtype(np.float64).itemsize) + memory_batch_count = max(1, int(available_memory_bytes() * 0.05) // matrix_bytes) + return min(roi_count, 16, memory_batch_count) + + +def _fill_two_time_output( + data_pixel, + roi_pixel_indices, + pixel_counts, + output, + norm=None, + use_data_mean=True, +): + """Calculate ROI matrices in bounded batches and write the public layout.""" + frame_count = data_pixel.shape[0] + roi_count = len(roi_pixel_indices) + batch_size = _two_time_batch_size(frame_count, roi_count) + upper_triangles = np.empty((frame_count, frame_count, batch_size), dtype=np.float64, order="F") + row_norms = np.empty((frame_count, batch_size), dtype=np.float64) + batch_pixel_counts = np.empty(batch_size, dtype=np.int64) + pre_normalized = np.empty(batch_size, dtype=np.bool_) + + for roi_index in tqdm(range(roi_count)): + batch_start = (roi_index // batch_size) * batch_size + batch_index = roi_index - batch_start + batch_count = min(batch_size, roi_count - batch_start) + pixel_indices = roi_pixel_indices[roi_index] + selected_data = data_pixel[:, pixel_indices] + if use_data_mean: + row_norm = np.average(selected_data, axis=1) + elif norm is None: + row_norm = np.ones(frame_count, dtype=np.float64) + else: + row_norm = np.average(norm[:, pixel_indices], axis=1) + row_norms[:, batch_index] = row_norm + batch_pixel_counts[batch_index] = pixel_counts[roi_index] + _, pre_normalized[batch_index] = _upper_two_time_product( + selected_data, + row_norm, + pixel_counts[roi_index], + output=upper_triangles[:, :, batch_index], + ) + if batch_index + 1 == batch_count: + store_symmetric_two_time_batch( + upper_triangles, + row_norms, + batch_pixel_counts, + pre_normalized, + output, + batch_start, + batch_count, + ) + + def auto_two_Arrayc(data_pixel, rois, index=None): """ Dec 16, 2015, Y.G.@CHX @@ -1694,6 +1936,7 @@ def auto_two_Arrayc(data_pixel, rois, index=None): """ qind, qlist, nopr = _select_two_time_rois(rois, index) + roi_pixel_indices = [np.flatnonzero(qind == label) for label in qlist] noframes = data_pixel.shape[0] # print( qlist ) try: @@ -1707,17 +1950,13 @@ def auto_two_Arrayc(data_pixel, rois, index=None): DO = False if DO: - for i, qi in enumerate(tqdm(qlist)): - # print (qi-1) - pixelist_qi = np.where(qind == qi)[0] - # print (pixelist_qi.shape, data_pixel[qi].shape) - data_pixel_qi = data_pixel[:, pixelist_qi] - sum1 = (np.average(data_pixel_qi, axis=1)).reshape(1, noframes) - sum2 = sum1.T - # print( qi, qlist, ) - # print( g12b[:,:,qi -1 ] ) - g12b[:, :, i] = np.dot(data_pixel_qi, data_pixel_qi.T) / sum1 / sum2 / nopr[i] - return g12b + core_count = physical_core_count() + with ( + threadpool_limits(limits=core_count, user_api="blas"), + numba_thread_limit(core_count), + ): + _fill_two_time_output(data_pixel, roi_pixel_indices, nopr, g12b) + return np.ascontiguousarray(g12b) def auto_two_Arrayc_ExplicitNorm(data_pixel, rois, norm=None, index=None): @@ -1758,18 +1997,21 @@ def auto_two_Arrayc_ExplicitNorm(data_pixel, rois, norm=None, index=None): """TO be done here """ DO = False if DO: - for i, qi in enumerate(tqdm(qlist)): - pixelist_qi = np.where(qind == qi)[0] - data_pixel_qi = data_pixel[:, pixelist_qi] - if norm is not None: - norm1 = norm[:, pixelist_qi] - sum1 = (np.average(norm1, axis=1)).reshape(1, noframes) - sum2 = sum1.T - else: - sum1 = 1 - sum2 = 1 - g12b[:, :, i] = np.dot(data_pixel_qi, data_pixel_qi.T) / sum1 / sum2 / nopr[i] - return g12b + roi_pixel_indices = [np.flatnonzero(qind == label) for label in qlist] + core_count = physical_core_count() + with ( + threadpool_limits(limits=core_count, user_api="blas"), + numba_thread_limit(core_count), + ): + _fill_two_time_output( + data_pixel, + roi_pixel_indices, + nopr, + g12b, + norm=norm, + use_data_mean=False, + ) + return np.ascontiguousarray(g12b) def two_time_norm(data_pixel, rois, index=None): diff --git a/pyCHX/chx_correlationp.py b/pyCHX/chx_correlationp.py index 11dfae0..7e2c928 100644 --- a/pyCHX/chx_correlationp.py +++ b/pyCHX/chx_correlationp.py @@ -7,17 +7,20 @@ from __future__ import absolute_import, division, print_function import logging +from concurrent.futures import ThreadPoolExecutor import numpy as np import skbeam.core.roi as roi -from skbeam.core.roi import extract_label_indices +from skbeam.core.utils import multi_tau_lags from tqdm import tqdm -from pyCHX.chx_compress import _collect_pool_results, _make_pool, apply_async, pass_FD -from pyCHX.chx_correlationc import _one_time_process as _one_time_processp +from pyCHX._performance import available_memory_bytes, process_one_time_block, sparse_scatter_normalized +from pyCHX.chx_compress import _available_cpu_count, _collect_pool_results, _make_pool, apply_async, pass_FD +from pyCHX.chx_correlationc import _create_intensity_buffer +from pyCHX.chx_correlationc import _one_time_process_cached as _one_time_processp_cached from pyCHX.chx_correlationc import _one_time_process_error as _one_time_process_errorp from pyCHX.chx_correlationc import _select_two_time_rois -from pyCHX.chx_correlationc import _two_time_process as _two_time_processp +from pyCHX.chx_correlationc import _two_time_process_cached as _two_time_processp_cached from pyCHX.chx_correlationc import _validate_and_transform_inputs logger = logging.getLogger(__name__) @@ -151,39 +154,44 @@ def lazy_two_timep( # create a shorthand reference to the results and state named tuple s = internal_state - qind, pixelist = roi.extract_label_indices(labels) + pixelist = s.pixel_list # iterate over the images to compute multi-tau correlation fra_pix = np.zeros_like(pixelist, dtype=np.float64) timg = np.zeros(FD.md["ncols"] * FD.md["nrows"], dtype=np.int32) timg[pixelist] = np.arange(1, len(pixelist) + 1) if bad_frame_list is None: bad_frame_list = [] + bad_frames = set(bad_frame_list) + has_imgsum_norm = imgsum is not None + has_pixel_norm = norm is not None + pixel_norm_is_2d = has_pixel_norm and len(norm.shape) > 1 + intensity_buf = _create_intensity_buffer(s.buf, s.label_array, len(s.num_pixels)) for i in range(FD.beg, FD.end): - if i in bad_frame_list: + if i in bad_frames: fra_pix[:] = np.nan else: p, v = FD.rdrawframe(i) - w = np.where(timg[p])[0] - pxlist = timg[p[w]] - 1 - if imgsum is None: - if norm is None: - fra_pix[pxlist] = v[w] + mapped_pixels = timg[p] + selected = mapped_pixels != 0 + pxlist = mapped_pixels[selected] - 1 + values = v[selected] + if not has_imgsum_norm: + if not has_pixel_norm: + fra_pix[pxlist] = values else: - S = norm.shape - if len(S) > 1: - fra_pix[pxlist] = v[w] / norm[i, pxlist] # -1.0 + if pixel_norm_is_2d: + fra_pix[pxlist] = values / norm[i, pxlist] # -1.0 else: - fra_pix[pxlist] = v[w] / norm[pxlist] # -1.0 + fra_pix[pxlist] = values / norm[pxlist] # -1.0 else: - if norm is None: - fra_pix[pxlist] = v[w] / imgsum[i] + if not has_pixel_norm: + fra_pix[pxlist] = values / imgsum[i] else: - S = norm.shape - if len(S) > 1: - fra_pix[pxlist] = v[w] / imgsum[i] / norm[i, pxlist] + if pixel_norm_is_2d: + fra_pix[pxlist] = values / imgsum[i] / norm[i, pxlist] else: - fra_pix[pxlist] = v[w] / imgsum[i] / norm[pxlist] + fra_pix[pxlist] = values / imgsum[i] / norm[pxlist] level = 0 # increment buffer s.cur[0] = (1 + s.cur[0]) % num_bufs @@ -194,7 +202,7 @@ def lazy_two_timep( # Put the ROI pixels into the ring buffer. s.buf[0, s.cur[0] - 1] = fra_pix fra_pix[:] = 0 - _two_time_processp( + _two_time_processp_cached( s.buf, s.g2, s.label_array, @@ -205,6 +213,7 @@ def lazy_two_timep( s.current_img_time, level=0, buf_no=s.cur[0] - 1, + intensity_buf=intensity_buf, ) # time frame for each level s.time_ind[0].append(s.current_img_time) @@ -221,9 +230,13 @@ def lazy_two_timep( prev = 1 + (s.cur[level - 1] - 2) % num_bufs s.cur[level] = 1 + s.cur[level] % num_bufs s.count_level[level] = 1 + s.count_level[level] - s.buf[level, s.cur[level] - 1] = ( - s.buf[level - 1, prev - 1] + s.buf[level - 1, s.cur[level - 1] - 1] - ) / 2 + level_buffer = s.buf[level, s.cur[level] - 1] + np.add( + s.buf[level - 1, prev - 1], + s.buf[level - 1, s.cur[level - 1] - 1], + out=level_buffer, + ) + level_buffer /= 2 t1_idx = (s.count_level[level] - 1) * 2 current_img_time = ((s.time_ind[level - 1])[t1_idx] + (s.time_ind[level - 1])[t1_idx + 1]) / 2.0 # time frame for each level @@ -234,7 +247,7 @@ def lazy_two_timep( # for multi-tau levels greater than one # Again, this is modifying things in place. See comment # on previous call above. - _two_time_processp( + _two_time_processp_cached( s.buf, s.g2, s.label_array, @@ -245,6 +258,7 @@ def lazy_two_timep( current_img_time, level=level, buf_no=s.cur[level] - 1, + intensity_buf=intensity_buf, ) level += 1 @@ -277,19 +291,12 @@ def cal_c12p(FD, ring_mask, bad_frame_list=None, good_start=0, num_buf=8, num_le roi_labels = np.unique(ring_mask) roi_labels = roi_labels[roi_labels > 0] ring_masks = [np.array(ring_mask == label, dtype=np.int64) for label in roi_labels] - qind, pixelist = roi.extract_label_indices(ring_mask) + qind, _ = roi.extract_label_indices(ring_mask) if norm is not None: - S = norm.shape - if len(S) > 1: - norms = [ - norm[:, np.isin(pixelist, extract_label_indices(np.array(ring_mask == i, dtype=np.int64))[1])] - for i in roi_labels - ] + if len(norm.shape) > 1: + norms = [norm[:, qind == label] for label in roi_labels] else: - norms = [ - norm[np.isin(pixelist, extract_label_indices(np.array(ring_mask == i, dtype=np.int64))[1])] - for i in roi_labels - ] + norms = [norm[qind == label] for label in roi_labels] inputs = range(len(ring_masks)) pool = _make_pool(len(inputs)) internal_state = None @@ -408,6 +415,8 @@ def __init__(self, num_levels, num_bufs, labels, cal_error=False): self.past_intensity_all = np.zeros_like(self.G_all) # matrix for normalizing G into g2 self.future_intensity_all = np.zeros_like(self.G_all) + else: + self.intensity_buf = np.zeros((num_levels, num_bufs, num_rois), dtype=np.float64) def __getstate__(self): """This is called before pickling.""" @@ -434,40 +443,45 @@ def lazy_one_timep( internal_state = _internal_statep(num_levels, num_bufs, labels, cal_error) # create a shorthand reference to the results and state named tuple s = internal_state - qind, pixelist = roi.extract_label_indices(labels) + pixelist = s.pixel_list # iterate over the images to compute multi-tau correlation fra_pix = np.zeros_like(pixelist, dtype=np.float64) timg = np.zeros(FD.md["ncols"] * FD.md["nrows"], dtype=np.int32) timg[pixelist] = np.arange(1, len(pixelist) + 1) if bad_frame_list is None: bad_frame_list = [] + intensity_buf = getattr(s, "intensity_buf", None) + bad_frames = set(bad_frame_list) + has_imgsum_norm = imgsum is not None + has_pixel_norm = norm is not None + pixel_norm_is_2d = has_pixel_norm and len(norm.shape) > 1 # for i in tqdm(range( FD.beg , FD.end )): for i in range(FD.beg, FD.end): - if i in bad_frame_list: + if i in bad_frames: fra_pix[:] = np.nan else: p, v = FD.rdrawframe(i) - w = np.where(timg[p])[0] - pxlist = timg[p[w]] - 1 - if imgsum is None: - if norm is None: + mapped_pixels = timg[p] + selected = mapped_pixels != 0 + pxlist = mapped_pixels[selected] - 1 + values = v[selected] + if not has_imgsum_norm: + if not has_pixel_norm: # print ('here') - fra_pix[pxlist] = v[w] + fra_pix[pxlist] = values else: - S = norm.shape - if len(S) > 1: - fra_pix[pxlist] = v[w] / norm[i, pxlist] # -1.0 + if pixel_norm_is_2d: + fra_pix[pxlist] = values / norm[i, pxlist] # -1.0 else: - fra_pix[pxlist] = v[w] / norm[pxlist] # -1.0 + fra_pix[pxlist] = values / norm[pxlist] # -1.0 else: - if norm is None: - fra_pix[pxlist] = v[w] / imgsum[i] + if not has_pixel_norm: + fra_pix[pxlist] = values / imgsum[i] else: - S = norm.shape - if len(S) > 1: - fra_pix[pxlist] = v[w] / imgsum[i] / norm[i, pxlist] + if pixel_norm_is_2d: + fra_pix[pxlist] = values / imgsum[i] / norm[i, pxlist] else: - fra_pix[pxlist] = v[w] / imgsum[i] / norm[pxlist] + fra_pix[pxlist] = values / imgsum[i] / norm[pxlist] level = 0 # increment buffer @@ -503,7 +517,7 @@ def lazy_one_timep( s.future_intensity_all, ) else: - _one_time_processp( + _one_time_processp_cached( s.buf, s.G, s.past_intensity, @@ -516,6 +530,7 @@ def lazy_one_timep( buf_no, s.norm, s.lev_len, + intensity_buf, ) # print (s.G) @@ -531,9 +546,13 @@ def lazy_one_timep( prev = 1 + (s.cur[level - 1] - 2) % num_bufs s.cur[level] = 1 + s.cur[level] % num_bufs - s.buf[level, s.cur[level] - 1] = ( - s.buf[level - 1, prev - 1] + s.buf[level - 1, s.cur[level - 1] - 1] - ) / 2 + level_buffer = s.buf[level, s.cur[level] - 1] + np.add( + s.buf[level - 1, prev - 1], + s.buf[level - 1, s.cur[level - 1] - 1], + out=level_buffer, + ) + level_buffer /= 2 # make the track_level zero once that level is processed s.track_level[level] = False @@ -561,7 +580,7 @@ def lazy_one_timep( s.future_intensity_all, ) else: - _one_time_processp( + _one_time_processp_cached( s.buf, s.G, s.past_intensity, @@ -574,6 +593,7 @@ def lazy_one_timep( buf_no, s.norm, s.lev_len, + intensity_buf, ) level += 1 @@ -606,6 +626,50 @@ def lazy_one_timep( return g2, s.lag_steps[:g_max] # , s +def _balance_roi_jobs(pixel_counts, worker_count): + """Assign ROI indices to workers using largest-pixel-count-first packing.""" + roi_count = len(pixel_counts) + if roi_count <= worker_count: + return [[index] for index in range(roi_count)] + + groups = [[] for _ in range(worker_count)] + group_pixels = np.zeros(worker_count, dtype=np.int64) + for index in sorted(range(roi_count), key=lambda item: (-pixel_counts[item], item)): + group = int(np.argmin(group_pixels)) + groups[group].append(index) + group_pixels[group] += pixel_counts[index] + return groups + + +def _run_one_time_group( + FD, + num_levels, + num_bufs, + ring_mask, + jobs, + bad_frame_list, + imgsum, + cal_error, +): + """Calculate one or more ROIs sequentially inside one worker.""" + group_results = [] + for index, label, norm in jobs: + label_mask = np.asarray(ring_mask == label, dtype=np.int64) + result = lazy_one_timep( + FD, + num_levels, + num_bufs, + label_mask, + None, + bad_frame_list, + imgsum, + norm, + cal_error, + ) + group_results.append((index, result)) + return group_results + + def cal_g2p( FD, ring_mask, @@ -622,120 +686,182 @@ def cal_g2p( if return_g2_details: return g2 with g2_denomitor, g2_past, g2_future """ FD.beg = max(FD.beg, good_start) - noframes = FD.end - FD.beg + 1 # number of frames, not "no frames" - for i in range(FD.beg, FD.end): - pass_FD(FD, i) + noframes = FD.end - FD.beg + 1 # preserve the historical level-selection convention if num_lev is None: num_lev = int(np.log(noframes / (num_buf - 1)) / np.log(2) + 1) + 1 print("In this g2 calculation, the buf and lev number are: %s--%s--" % (num_buf, num_lev)) - if bad_frame_list is not None: - if len(bad_frame_list) != 0: - print("%s Bad frames involved and will be discarded!" % len(bad_frame_list)) - noframes -= len(np.where(np.isin(bad_frame_list, range(good_start, FD.end)))[0]) + bad_frames = set() if bad_frame_list is None else set(np.atleast_1d(bad_frame_list).tolist()) + if bad_frames: + print("%s Bad frames involved and will be discarded!" % len(bad_frames)) + noframes -= len(np.where(np.isin(list(bad_frames), range(good_start, FD.end)))[0]) print("%s frames will be processed..." % (noframes - 1)) - roi_labels = np.unique(ring_mask) - roi_labels = roi_labels[roi_labels > 0] - ring_masks = [np.array(ring_mask == label, dtype=np.int64) for label in roi_labels] - qind, pixelist = roi.extract_label_indices(ring_mask) - nopr = np.array([np.count_nonzero(qind == label) for label in roi_labels]) - if norm is not None: - S = norm.shape - if len(S) > 1: - norms = [ - norm[:, np.isin(pixelist, extract_label_indices(np.array(ring_mask == i, dtype=np.int64))[1])] - for i in roi_labels - ] - else: - norms = [ - norm[np.isin(pixelist, extract_label_indices(np.array(ring_mask == i, dtype=np.int64))[1])] - for i in roi_labels - ] - inputs = range(len(ring_masks)) - pool = _make_pool(len(inputs)) - internal_state = None - print("Starting assign the tasks...") - results = {} - if norm is not None: - for i in tqdm(inputs): - results[i] = apply_async( - pool, - lazy_one_timep, - (FD, num_lev, num_buf, ring_masks[i], internal_state, bad_frame_list, imgsum, norms[i], cal_error), - ) - else: - # print ('for norm is None') - for i in tqdm(inputs): - results[i] = apply_async( - pool, - lazy_one_timep, - (FD, num_lev, num_buf, ring_masks[i], internal_state, bad_frame_list, imgsum, None, cal_error), + + qind, pixel_list = roi.extract_label_indices(ring_mask) + roi_labels = np.unique(qind) + if roi_labels.size == 0: + raise ValueError("ring_mask contains no positive ROI labels") + original_columns = [np.flatnonzero(qind == label) for label in roi_labels] + permutation = np.concatenate(original_columns) + grouped_pixels = np.asarray(pixel_list[permutation], dtype=np.int64) + pixel_counts = np.asarray([len(columns) for columns in original_columns], dtype=np.int64) + roi_starts = np.concatenate(([0], np.cumsum(pixel_counts))).astype(np.int64) + + lookup = np.full(FD.md["ncols"] * FD.md["nrows"], -1, dtype=np.int64) + lookup[grouped_pixels] = np.arange(grouped_pixels.size, dtype=np.int64) + norm_is_2d = norm is not None and np.ndim(norm) > 1 + norm_1d = np.asarray(norm if norm is not None and not norm_is_2d else np.ones(1), dtype=np.float64) + norm_2d = np.asarray(norm if norm_is_2d else np.ones((1, 1)), dtype=np.float64) + norm_columns = np.asarray(permutation if norm is not None else np.zeros(grouped_pixels.size), dtype=np.int64) + image_sums = np.asarray(imgsum if imgsum is not None else np.ones(1), dtype=np.float64) + normalization_flags = np.asarray( + [norm is not None and not norm_is_2d, norm_is_2d, imgsum is not None, False], dtype=np.bool_ + ) + dummy_means = np.ones((1, 1), dtype=np.float64) + dummy_qind = np.zeros(grouped_pixels.size, dtype=np.int64) + + _, lag_steps, dict_lags = multi_tau_lags(num_lev, num_buf) + level_lengths = np.asarray([len(dict_lags[key]) for key in dict_lags], dtype=np.int64) + level_offsets = np.zeros(num_lev, dtype=np.int64) + if num_lev > 1: + level_offsets[1:] = np.cumsum(level_lengths[:-1]) + lag_capacity = int((num_lev + 1) * num_buf / 2) + + states = [] + for pixel_count in pixel_counts: + error_shape = (lag_capacity, int(pixel_count)) if cal_error else (1, 1) + states.append( + { + "buf": np.zeros((num_lev, num_buf, int(pixel_count)), dtype=np.float64), + "G": np.zeros(lag_capacity, dtype=np.float64), + "past": np.zeros(lag_capacity, dtype=np.float64), + "future": np.zeros(lag_capacity, dtype=np.float64), + "images_per_level": np.zeros(num_lev, dtype=np.int64), + "track_level": np.zeros(num_lev, dtype=np.bool_), + "current": np.ones(num_lev, dtype=np.int64), + "bad_counts": np.zeros((num_lev, num_buf), dtype=np.int64), + "intensity_buf": np.zeros((num_lev, num_buf), dtype=np.float64), + "G_all": np.zeros(error_shape, dtype=np.float64), + "past_all": np.zeros(error_shape, dtype=np.float64), + "future_all": np.zeros(error_shape, dtype=np.float64), + } + ) + + worker_count = min(len(roi_labels), _available_cpu_count()) + groups = _balance_roi_jobs(pixel_counts, worker_count) + target_bytes = min(256 * 1024**2, max(8 * grouped_pixels.size, int(available_memory_bytes() * 0.10))) + block_frames = max(1, min(1024, target_bytes // max(8, 8 * grouped_pixels.size))) + + def process_group(group, block): + for roi_index in group: + start, stop = roi_starts[roi_index : roi_index + 2] + state = states[roi_index] + process_one_time_block( + block[:, start:stop], + state["buf"], + state["G"], + state["past"], + state["future"], + state["images_per_level"], + state["track_level"], + state["current"], + state["bad_counts"], + level_offsets, + state["intensity_buf"], + state["G_all"], + state["past_all"], + state["future_all"], + cal_error, ) - print("Starting running the tasks...") - res = _collect_pool_results(pool, results, show_progress=True) - len_lag = 10**10 - for i in inputs: # to get the smallest length of lag_step, - # ***************************** - # Here could result in problem for significantly cut useful data if some Q have very short tau list - # **************************** - if len_lag > len(res[i][1]): - lag_steps = res[i][1] - len_lag = len(lag_steps) - # lag_steps = res[0][1] - if not cal_error: - g2 = np.zeros([len(lag_steps), len(ring_masks)]) - else: - g2 = np.zeros([int((num_lev + 1) * num_buf / 2), len(ring_masks)]) - g2_err = np.zeros_like(g2) - # g2_G = np.zeros(( int( (num_lev + 1) * num_buf / 2), len(pixelist)) ) - # g2_P = np.zeros_like( g2_G ) - # g2_F = np.zeros_like( g2_G ) - Gmax = 0 - lag_steps_err = res[0][1] - for i in inputs: - # print( res[i][0][:,0].shape, g2.shape ) - if not cal_error: - g2[:, i] = res[i][0][:, 0][:len_lag] - else: - s_Gall_qi = res[i][2] # [:len_lag] - s_Pall_qi = res[i][3] # [:len_lag] - s_Fall_qi = res[i][4] # [:len_lag] - # print( s_Gall_qi.shape,s_Pall_qi.shape,s_Fall_qi.shape ) - avgGi = np.average(s_Gall_qi, axis=1) - devGi = np.std(s_Gall_qi, axis=1) - avgPi = np.average(s_Pall_qi, axis=1) - devPi = np.std(s_Pall_qi, axis=1) - avgFi = np.average(s_Fall_qi, axis=1) - devFi = np.std(s_Fall_qi, axis=1) - - if len(np.where(avgPi == 0)[0]) != 0: - g_max1 = np.where(avgPi == 0)[0][0] - else: - g_max1 = avgPi.shape[0] - if len(np.where(avgFi == 0)[0]) != 0: - g_max2 = np.where(avgFi == 0)[0][0] + if hasattr(FD, "_ensure_index"): + FD._ensure_index() + executor = ThreadPoolExecutor(max_workers=worker_count) if worker_count > 1 else None + try: + starts = range(FD.beg, FD.end, block_frames) + for block_start in tqdm(starts, desc="Correlating frame blocks"): + block_stop = min(FD.end, block_start + block_frames) + block = np.zeros((block_stop - block_start, grouped_pixels.size), dtype=np.float64) + for output_row, frame_index in enumerate(range(block_start, block_stop)): + if frame_index in bad_frames: + block[output_row].fill(np.nan) + continue + if hasattr(FD, "_raw_frame_view"): + positions, values = FD._raw_frame_view(frame_index) + else: + positions, values = FD.rdrawframe(frame_index) + sparse_scatter_normalized( + positions, + values, + lookup, + block, + output_row, + frame_index, + norm_1d, + norm_2d, + norm_columns, + image_sums, + dummy_means, + dummy_qind, + normalization_flags, + ) + if executor is None: + process_group(groups[0], block) else: - g_max2 = avgFi.shape[0] - g_max = min(g_max1, g_max2) - g2[:g_max, i] = avgGi[:g_max] / (avgPi[:g_max] * avgFi[:g_max]) - g2_err[:g_max, i] = np.sqrt( - (1 / (avgFi[:g_max] * avgPi[:g_max])) ** 2 * devGi[:g_max] ** 2 - + (avgGi[:g_max] / (avgFi[:g_max] ** 2 * avgPi[:g_max])) ** 2 * devFi[:g_max] ** 2 - + (avgGi[:g_max] / (avgFi[:g_max] * avgPi[:g_max] ** 2)) ** 2 * devPi[:g_max] ** 2 - ) - Gmax = max(g_max, Gmax) - lag_stepsi = res[i][1] - if len(lag_steps_err) < len(lag_stepsi): - lag_steps_err = lag_stepsi + futures = [executor.submit(process_group, group, block) for group in groups] + for future in futures: + future.result() + finally: + if executor is not None: + executor.shutdown() - del results - del res - if cal_error: - print("G2 with error bar calculation DONE!") - return g2[:Gmax, :], lag_steps_err[:Gmax], g2_err[:Gmax, :] / np.sqrt(nopr) - else: + if not cal_error: + roi_results = [] + valid_lengths = [] + for state in states: + zero_past = np.flatnonzero(state["past"] == 0) + zero_future = np.flatnonzero(state["future"] == 0) + g_max1 = int(zero_past[0]) if zero_past.size else state["past"].size + g_max2 = int(zero_future[0]) if zero_future.size else state["future"].size + valid_length = min(g_max1, g_max2) + valid_lengths.append(valid_length) + roi_results.append( + state["G"][:valid_length] / (state["past"][:valid_length] * state["future"][:valid_length]) + ) + common_length = min(valid_lengths) + g2 = np.column_stack([result[:common_length] for result in roi_results]) print("G2 calculation DONE!") - return g2, lag_steps + return g2, lag_steps[:common_length] + + g2 = np.zeros((lag_capacity, len(roi_labels)), dtype=np.float64) + g2_err = np.zeros_like(g2) + maximum_length = 0 + for roi_index, state in enumerate(states): + avg_g = np.average(state["G_all"], axis=1) + dev_g = np.std(state["G_all"], axis=1) + avg_past = np.average(state["past_all"], axis=1) + dev_past = np.std(state["past_all"], axis=1) + avg_future = np.average(state["future_all"], axis=1) + dev_future = np.std(state["future_all"], axis=1) + zero_past = np.flatnonzero(avg_past == 0) + zero_future = np.flatnonzero(avg_future == 0) + g_max1 = int(zero_past[0]) if zero_past.size else avg_past.size + g_max2 = int(zero_future[0]) if zero_future.size else avg_future.size + valid_length = min(g_max1, g_max2) + g2[:valid_length, roi_index] = avg_g[:valid_length] / (avg_past[:valid_length] * avg_future[:valid_length]) + g2_err[:valid_length, roi_index] = np.sqrt( + (1 / (avg_future[:valid_length] * avg_past[:valid_length])) ** 2 * dev_g[:valid_length] ** 2 + + (avg_g[:valid_length] / (avg_future[:valid_length] ** 2 * avg_past[:valid_length])) ** 2 + * dev_future[:valid_length] ** 2 + + (avg_g[:valid_length] / (avg_future[:valid_length] * avg_past[:valid_length] ** 2)) ** 2 + * dev_past[:valid_length] ** 2 + ) + maximum_length = max(maximum_length, valid_length) + print("G2 with error bar calculation DONE!") + return ( + g2[:maximum_length], + lag_steps[:maximum_length], + g2_err[:maximum_length] / np.sqrt(pixel_counts), + ) def cal_GPF( @@ -770,10 +896,7 @@ def cal_GPF( ring_masks = [np.array(ring_mask == label, dtype=np.int64) for label in roi_labels] qind, pixelist = roi.extract_label_indices(ring_mask) if norm is not None: - norms = [ - norm[np.isin(pixelist, extract_label_indices(np.array(ring_mask == i, dtype=np.int64))[1])] - for i in roi_labels - ] + norms = [norm[qind == label] for label in roi_labels] inputs = range(len(ring_masks)) pool = _make_pool(len(inputs)) diff --git a/pyCHX/tests/test_compression.py b/pyCHX/tests/test_compression.py index 5acb347..39ac1cd 100644 --- a/pyCHX/tests/test_compression.py +++ b/pyCHX/tests/test_compression.py @@ -1,3 +1,6 @@ +import pickle +import struct + import numpy as np import pytest @@ -93,6 +96,192 @@ def test_compressed_file_round_trip(tmp_path): assert compressed.FID.closed +@pytest.mark.portable +@pytest.mark.parametrize( + ("nobytes", "bins", "value_format"), + [(2, 1, "h"), (4, 1, "i"), (4, 2, "d")], +) +def test_compressed_payload_keeps_legacy_binary_layout(tmp_path, nobytes, bins, value_format): + from pyCHX.chx_compress import init_compress_eigerdata + + frames = np.array( + [ + [[0, 1, 2], [3, 0, 4]], + [[5, 0, 6], [0, 7, 8]], + [[9, 10, 0], [11, 12, 0]], + ], + dtype=np.int32, + ) + mask = np.ones(frames.shape[1:], dtype=bool) + filename = tmp_path / "layout.cmp" + + init_compress_eigerdata( + frames, + mask.copy(), + {"pixel_mask": mask.copy()}, + str(filename), + nobytes=nobytes, + bins=bins, + with_pickle=False, + ) + + expected = bytearray() + for start in range(0, len(frames), bins): + image = np.average(frames[start : start + bins], axis=0) + positions = np.flatnonzero(image.ravel() > 0) + values = image.ravel()[positions] + if bins == 1: + values = values.astype({2: np.int16, 4: np.int32}[nobytes]) + expected.extend(struct.pack("@I", len(positions))) + expected.extend(struct.pack("@{}i".format(len(positions)), *positions)) + expected.extend(struct.pack("@{}{}".format(len(positions), value_format), *values)) + + assert filename.read_bytes()[1024:] == bytes(expected) + + +@pytest.mark.portable +def test_compress_eigerdata_stages_then_publishes_serial_output(tmp_path): + from pyCHX.chx_compress import compress_eigerdata + + frames = np.arange(1, 25, dtype=np.int32).reshape(4, 2, 3) + mask = np.ones(frames.shape[1:], dtype=bool) + destination = tmp_path / "destination" / "serial.cmp" + staging = tmp_path / "staging" + destination.parent.mkdir() + staging.mkdir() + + compress_eigerdata( + frames, + mask.copy(), + {"pixel_mask": mask.copy()}, + str(destination), + force_compress=True, + dtypes="images", + direct_load_data=False, + with_pickle=False, + new_path=str(staging), + ) + + assert destination.is_file() + assert list(staging.iterdir()) == [] + + +@pytest.mark.portable +def test_parallel_compression_stages_and_publishes_identical_cmp(tmp_path): + from pyCHX.chx_compress import init_compress_eigerdata, para_compress_eigerdata + + frames = np.arange(1, 31, dtype=np.int32).reshape(5, 2, 3) + mask = np.ones(frames.shape[1:], dtype=bool) + metadata = { + "beam_center_x": 0, + "beam_center_y": 0, + "count_time": 0, + "detector_distance": 0, + "frame_time": 0, + "incident_wavelength": 0, + "pixel_mask": mask.copy(), + "x_pixel_size": 75, + "y_pixel_size": 75, + } + serial = tmp_path / "serial.cmp" + parallel = tmp_path / "destination" / "parallel.cmp" + staging = tmp_path / "staging" + parallel.parent.mkdir() + staging.mkdir() + + init_compress_eigerdata( + frames, + mask.copy(), + metadata.copy(), + str(serial), + with_pickle=False, + ) + para_compress_eigerdata( + frames, + mask.copy(), + metadata.copy(), + str(parallel), + num_sub=2, + dtypes="images", + cpu_core_number=2, + with_pickle=False, + copy_rawdata=False, + new_path=str(staging), + ) + + assert parallel.read_bytes() == serial.read_bytes() + assert list(staging.iterdir()) == [] + + +@pytest.mark.portable +def test_parallel_compression_keeps_partial_final_bin_byte_identical(tmp_path): + from pyCHX.chx_compress import init_compress_eigerdata, para_compress_eigerdata + + frames = np.arange(1, 31, dtype=np.int32).reshape(5, 2, 3) + mask = np.ones(frames.shape[1:], dtype=bool) + metadata = { + "beam_center_x": 0, + "beam_center_y": 0, + "count_time": 0, + "detector_distance": 0, + "frame_time": 0, + "incident_wavelength": 0, + "pixel_mask": mask.copy(), + "x_pixel_size": 0, + "y_pixel_size": 0, + } + serial = tmp_path / "serial-binned.cmp" + parallel = tmp_path / "parallel-binned.cmp" + + expected = init_compress_eigerdata( + frames, + mask.copy(), + metadata.copy(), + str(serial), + bins=2, + with_pickle=False, + ) + actual = para_compress_eigerdata( + frames, + mask.copy(), + metadata.copy(), + str(parallel), + num_sub=2, + bins=2, + dtypes="images", + cpu_core_number=2, + with_pickle=False, + copy_rawdata=False, + new_path=str(tmp_path), + ) + + assert parallel.read_bytes() == serial.read_bytes() + for actual_value, expected_value in zip(actual, expected): + np.testing.assert_allclose(actual_value, expected_value) + + +@pytest.mark.portable +def test_atomic_publish_preserves_existing_destination_on_copy_failure(monkeypatch, tmp_path): + from pyCHX import chx_compress + + source = tmp_path / "source.cmp" + destination = tmp_path / "destination.cmp" + source.write_bytes(b"new complete data") + destination.write_bytes(b"existing data") + + def fail_during_copy(_source, temporary): + with open(temporary, "wb") as stream: + stream.write(b"partial") + raise OSError("simulated copy failure") + + monkeypatch.setattr(chx_compress.shutil, "copyfile", fail_during_copy) + with pytest.raises(OSError, match="simulated copy failure"): + chx_compress._publish_file(source, destination) + + assert destination.read_bytes() == b"existing data" + assert list(tmp_path.glob(".destination.cmp.*.tmp")) == [] + + @pytest.mark.portable def test_compression_keeps_a_partial_final_frame_bin(tmp_path): from pyCHX.chx_compress import ( @@ -165,24 +354,51 @@ def test_compression_keeps_a_partial_final_frame_bin(tmp_path): @pytest.mark.portable -def test_parallel_compression_weights_segment_averages_by_valid_frames(monkeypatch, tmp_path): - from pyCHX import chx_compress +def test_unbinned_eiger_blocks_preserve_integer_frames_without_cast_warnings(tmp_path): + from pyCHX.chx_compress import _compress_segment + + invalid = np.iinfo(np.uint32).max + frames = np.array([[[1, invalid, 2]], [[3, invalid, 4]]], dtype=np.uint32) + + class DirectEigerFrames: + images_per_file = len(frames) + valid_keys = ["data_000001"] + _entry = {"data_000001": frames} + + detector_mask = np.array([[True, False, True]]) + with np.errstate(all="raise"): + final_mask, average, intensity, bad_frames = _compress_segment( + DirectEigerFrames(), + detector_mask.copy(), + str(tmp_path / "segment.cmp"), + bad_pixel_threshold=1e15, + hot_pixel_threshold=2**30, + bad_pixel_low_threshold=0, + nobytes=4, + bins=1, + start=0, + stop=len(frames), + ) + + np.testing.assert_array_equal(final_mask, detector_mask) + np.testing.assert_array_equal(average, [[2, 0, 3]]) + np.testing.assert_array_equal(intensity, [3, 7]) + np.testing.assert_array_equal(bad_frames, [False, False]) - class Result: - def __init__(self, value): - self.value = value - def get(self): - return self.value +@pytest.mark.portable +def test_parallel_compression_weights_segment_averages_by_valid_frames(monkeypatch, tmp_path): + from pyCHX import chx_compress mask = np.ones((1, 1), dtype=bool) - segment_results = { - 0: Result((mask.copy(), np.array([[2.0]]), np.array([1.0, 2.0, 3.0]), np.array([False, True, False]))), - 1: Result((mask.copy(), np.array([[8.0]]), np.array([4.0, 5.0]), np.array([False, False]))), - } - monkeypatch.setattr(chx_compress, "para_segment_compress_eigerdata", lambda **kwargs: segment_results) + segment_results = [ + (0, (mask.copy(), np.array([[2.0]]), np.array([1.0, 2.0, 3.0]), np.array([False, True, False]))), + (1, (mask.copy(), np.array([[8.0]]), np.array([4.0, 5.0]), np.array([False, False]))), + ] + monkeypatch.setattr(chx_compress, "_iter_parallel_segment_results", lambda **kwargs: segment_results) monkeypatch.setattr(chx_compress, "create_compress_header", lambda *args, **kwargs: None) monkeypatch.setattr(chx_compress, "combine_compressed", lambda *args, **kwargs: None) + monkeypatch.setattr(chx_compress, "_publish_file", lambda *args, **kwargs: None) _, average, intensity, bad_frames = chx_compress.para_compress_eigerdata( np.zeros((5, 1, 1)), @@ -229,6 +445,44 @@ def test_compression_handles_an_all_bad_segment_without_dividing_by_zero(tmp_pat np.testing.assert_array_equal(segment_bad, [True, True, True]) +@pytest.mark.portable +def test_parallel_hot_pixel_masking_preserves_segment_boundaries(tmp_path): + from pyCHX.chx_compress import Multifile, para_compress_eigerdata + + frames = np.array([[[200, 1]], [[2, 2]], [[5, 3]], [[6, 4]]], dtype=np.int32) + mask = np.ones((1, 2), dtype=bool) + metadata = { + "beam_center_x": 0, + "beam_center_y": 0, + "count_time": 0, + "detector_distance": 0, + "frame_time": 0, + "incident_wavelength": 0, + "pixel_mask": mask.copy(), + "x_pixel_size": 0, + "y_pixel_size": 0, + } + filename = tmp_path / "hot-pixel.cmp" + final_mask, _, _, _ = para_compress_eigerdata( + frames, + mask.copy(), + metadata, + str(filename), + num_sub=2, + hot_pixel_threshold=100, + dtypes="images", + cpu_core_number=2, + with_pickle=False, + copy_rawdata=False, + new_path=str(tmp_path), + ) + + assert not final_mask[0, 0] + with Multifile(str(filename), 0, len(frames)) as compressed: + assert compressed.rdframe(1)[0, 0] == 0 + assert compressed.rdframe(2)[0, 0] == 5 + + @pytest.mark.portable @pytest.mark.parametrize("bad_frame_list", [[], [3, 7, 18]]) def test_serial_and_parallel_g2_agree_for_compressed_data(tmp_path, bad_frame_list): @@ -282,17 +536,18 @@ def test_serial_and_parallel_g2_error_estimates_agree(tmp_path): filename, frames, ring_mask = _make_compressed_correlation_input(tmp_path) ring_mask = np.select([ring_mask == 1, ring_mask == 2], [2, 5], default=0) + norm = np.linspace(1.0, 2.0, np.count_nonzero(ring_mask)) serial_file = Multifile(str(filename), beg=0, end=len(frames)) parallel_file = Multifile(str(filename), beg=0, end=len(frames)) gpf_file = Multifile(str(filename), beg=0, end=len(frames)) try: serial_g2, serial_lags, serial_error, _ = cal_g2c( - serial_file, ring_mask, bad_frame_list=[], cal_error=True + serial_file, ring_mask, bad_frame_list=[], norm=norm, cal_error=True ) parallel_g2, parallel_lags, parallel_error = cal_g2p( - parallel_file, ring_mask, bad_frame_list=[], cal_error=True + parallel_file, ring_mask, bad_frame_list=[], norm=norm, cal_error=True ) - numerator, past, future = cal_GPF(gpf_file, ring_mask, bad_frame_list=[]) + numerator, past, future = cal_GPF(gpf_file, ring_mask, bad_frame_list=[], norm=norm) finally: serial_file.FID.close() parallel_file.FID.close() @@ -305,6 +560,46 @@ def test_serial_and_parallel_g2_error_estimates_agree(tmp_path): np.testing.assert_allclose(reconstructed_g2[: len(serial_g2)], serial_g2, rtol=1e-13, atol=0) +@pytest.mark.portable +@pytest.mark.parametrize( + ("norm_kind", "cal_error"), + [(None, False), ("1d", False), ("2d", True)], +) +def test_parallel_g2_worker_grouping_is_exact(tmp_path, monkeypatch, norm_kind, cal_error): + import skbeam.core.roi as roi + + from pyCHX import chx_correlationp + from pyCHX.chx_compress import Multifile + + filename, frames, ring_mask = _make_compressed_correlation_input(tmp_path) + ring_mask = np.select([ring_mask == 1, ring_mask == 2], [2, 5], default=0) + _, pixel_list = roi.extract_label_indices(ring_mask) + base_norm = np.linspace(1.0, 2.0, len(pixel_list)) + if norm_kind == "1d": + norm = base_norm + elif norm_kind == "2d": + norm = np.multiply.outer(np.linspace(1.0, 1.5, len(frames)), base_norm) + else: + norm = None + + def calculate(worker_count): + monkeypatch.setattr(chx_correlationp, "_available_cpu_count", lambda: worker_count) + with Multifile(str(filename), beg=0, end=len(frames)) as compressed: + return chx_correlationp.cal_g2p( + compressed, + ring_mask, + bad_frame_list=[3, 7], + imgsum=np.linspace(10.0, 20.0, len(frames)), + norm=norm, + cal_error=cal_error, + ) + + grouped = calculate(1) + one_roi_per_worker = calculate(8) + for grouped_value, ungrouped_value in zip(grouped, one_roi_per_worker): + np.testing.assert_array_equal(grouped_value, ungrouped_value) + + @pytest.mark.portable def test_serial_and_parallel_two_time_agree_for_compressed_data(tmp_path): from pyCHX.chx_compress import Multifile @@ -417,6 +712,50 @@ def test_frame_intensity_sampling_reports_source_frame_indices(tmp_path): np.testing.assert_array_equal(bad_frames, sample_indices[expected_intensity > threshold]) +@pytest.mark.portable +def test_read_compressed_reconstructs_average_and_bad_frames_in_one_pass(tmp_path, monkeypatch): + from pyCHX import chx_compress + + filename, frames, _ = _make_compressed_correlation_input(tmp_path) + threshold = float(frames[7].sum() - 1) + calls = [] + original = chx_compress.Multifile._raw_frame_view + + def counted(self, frame_index): + calls.append(frame_index) + return original(self, frame_index) + + monkeypatch.setattr(chx_compress.Multifile, "_raw_frame_view", counted) + _, average, intensity, bad_frames = chx_compress.read_compressed_eigerdata( + np.ones(frames.shape[1:], dtype=bool), + str(filename), + 0, + len(frames), + bad_pixel_threshold=threshold, + bad_pixel_low_threshold=-1, + bad_frame_list=[2], + with_pickle=False, + ) + expected_bad = np.unique(np.concatenate(([2], np.flatnonzero(frames.sum(axis=(1, 2)) > threshold)))) + + assert calls == list(range(len(frames))) + np.testing.assert_array_equal(intensity, frames.sum(axis=(1, 2))) + np.testing.assert_array_equal(bad_frames, expected_bad) + np.testing.assert_allclose(average, np.delete(frames, expected_bad, axis=0).mean(axis=0)) + + +@pytest.mark.portable +def test_waterfall_sparse_extraction_matches_dense_frames(tmp_path): + from pyCHX.chx_compress import Multifile + from pyCHX.chx_compress_analysis import cal_waterfallc + + filename, frames, ring_mask = _make_compressed_correlation_input(tmp_path) + with Multifile(str(filename), beg=0, end=len(frames)) as compressed: + actual = cal_waterfallc(compressed, ring_mask, qindex=2) + + np.testing.assert_array_equal(actual, frames[:, ring_mask == 2]) + + @pytest.mark.portable def test_collect_pool_results_always_reaps_workers(): from pyCHX.chx_compress import _collect_pool_results @@ -439,10 +778,170 @@ def test_pool_size_is_limited_by_available_cpus(monkeypatch): created_with = [] sentinel = object() - monkeypatch.setattr(chx_compress, "cpu_count", lambda: 4) + monkeypatch.setattr(chx_compress, "_available_cpu_count", lambda: 4) monkeypatch.setattr(chx_compress, "Pool", lambda processes: created_with.append(processes) or sentinel) assert chx_compress._make_pool(10) is sentinel assert created_with == [4] with pytest.raises(ValueError, match="at least one"): chx_compress._make_pool(0) + + +@pytest.mark.portable +def test_pool_size_respects_cpu_affinity(monkeypatch): + from pyCHX import chx_compress + + created_with = [] + sentinel = object() + monkeypatch.setattr(chx_compress, "cpu_count", lambda: 32) + monkeypatch.setattr(chx_compress.os, "sched_getaffinity", lambda _pid: set(range(3))) + monkeypatch.setattr(chx_compress, "physical_core_count", lambda cpu_ids: len(cpu_ids)) + monkeypatch.setattr(chx_compress, "Pool", lambda processes: created_with.append(processes) or sentinel) + + assert chx_compress._make_pool(10) is sentinel + assert created_with == [3] + + +@pytest.mark.portable +def test_pool_size_prefers_physical_cores_within_affinity(monkeypatch): + from pyCHX import chx_compress + + monkeypatch.setattr(chx_compress, "cpu_count", lambda: 256) + monkeypatch.setattr(chx_compress.os, "sched_getaffinity", lambda _pid: set(range(56))) + monkeypatch.setattr(chx_compress, "physical_core_count", lambda cpu_ids: 28) + assert chx_compress._available_cpu_count() == 28 + + +@pytest.mark.portable +@pytest.mark.parametrize(("nobytes", "dtype"), [(2, np.uint16), (4, np.uint32), (8, np.float64)]) +def test_multifile_indexed_views_support_legacy_value_widths_and_random_access(tmp_path, nobytes, dtype): + from pyCHX.chx_compress import Multifile, create_compress_header + + filename = tmp_path / f"legacy-{nobytes}.cmp" + metadata = {"img_shape": (2, 3)} + create_compress_header(metadata, str(filename), nobytes=nobytes) + values = [np.asarray([1, 3], dtype=dtype), np.asarray([], dtype=dtype), np.asarray([7], dtype=dtype)] + positions = [ + np.asarray([0, 5], dtype=np.int32), + np.asarray([], dtype=np.int32), + np.asarray([2], dtype=np.int32), + ] + with filename.open("ab") as stream: + for frame_positions, frame_values in zip(positions, values): + stream.write(np.asarray(len(frame_positions), dtype=np.uint32).tobytes()) + stream.write(frame_positions.tobytes()) + stream.write(frame_values.tobytes()) + + with Multifile(str(filename), beg=1, end=3) as compressed: + for frame_index in (2, 1, 2): + actual_positions, actual_values = compressed._raw_frame_view(frame_index) + np.testing.assert_array_equal(actual_positions, positions[frame_index]) + np.testing.assert_array_equal(actual_values, values[frame_index]) + assert not actual_positions.flags.writeable + assert not actual_values.flags.writeable + + restored = pickle.loads(pickle.dumps(compressed)) + try: + public_positions, public_values = restored.rdrawframe(2) + assert public_positions.flags.writeable + assert public_values.flags.writeable + public_values[:] = 0 + np.testing.assert_array_equal(restored.rdrawframe(2)[1], values[2]) + finally: + restored.close() + + compressed.reopen() + try: + np.testing.assert_array_equal(compressed.rdrawframe(1)[0], positions[1]) + finally: + compressed.close() + closed_copy = pickle.loads(pickle.dumps(compressed)) + assert closed_copy.FID.closed + closed_copy.reopen() + try: + np.testing.assert_array_equal(closed_copy.rdrawframe(2)[1], values[2]) + finally: + closed_copy.close() + + +@pytest.mark.portable +@pytest.mark.parametrize("corruption", ["header", "frame_header", "payload", "negative_length"]) +def test_multifile_rejects_malformed_or_truncated_input(tmp_path, corruption): + from pyCHX.chx_compress import Multifile, create_compress_header + + filename = tmp_path / "broken.cmp" + if corruption == "header": + filename.write_bytes(b"Version-COMP0001") + with pytest.raises(ValueError, match="header"): + Multifile(str(filename), 0, 1) + return + + create_compress_header({"img_shape": (2, 2)}, str(filename), nobytes=4) + if corruption == "frame_header": + with pytest.raises(ValueError, match="first frame"): + Multifile(str(filename), 0, 1) + return + with filename.open("ab") as stream: + stream.write(struct.pack("@i", -1 if corruption == "negative_length" else 2)) + if corruption == "payload": + stream.write(np.asarray([0], dtype=np.int32).tobytes()) + if corruption == "negative_length": + with pytest.raises(ValueError, match="negative"): + Multifile(str(filename), 0, 1) + else: + with Multifile(str(filename), 0, 1) as compressed: + with pytest.raises(ValueError, match="truncated"): + compressed._raw_frame_view(0) + + +@pytest.mark.portable +def test_cal_g2p_traverses_each_cmp_frame_once(tmp_path, monkeypatch): + from pyCHX.chx_compress import Multifile + from pyCHX.chx_correlationp import cal_g2p + + filename, frames, ring_mask = _make_compressed_correlation_input(tmp_path) + with Multifile(str(filename), 0, len(frames)) as compressed: + calls = [] + original = compressed._raw_frame_view + + def counted(frame_index): + calls.append(frame_index) + return original(frame_index) + + monkeypatch.setattr(compressed, "_raw_frame_view", counted) + cal_g2p(compressed, ring_mask, bad_frame_list=[]) + traversed = compressed._bytes_traversed + + assert calls == list(range(len(frames))) + assert traversed == filename.stat().st_size - 1024 + 4 * len(frames) + + +@pytest.mark.portable +def test_cal_g2p_supports_hundreds_of_sparse_rois(tmp_path, monkeypatch): + from pyCHX import chx_correlationp + from pyCHX.chx_compress import Multifile, init_compress_eigerdata + from pyCHX.chx_correlationc import cal_g2c + + roi_mask = np.arange(1, 201, dtype=np.int64).reshape(10, 20) + frame_number = np.arange(16, dtype=np.int32)[:, None, None] + frames = 1 + (frame_number + roi_mask[None, :, :]) % 23 + detector_mask = np.ones(roi_mask.shape, dtype=bool) + filename = tmp_path / "many-rois.cmp" + init_compress_eigerdata( + frames, + detector_mask.copy(), + {"pixel_mask": detector_mask.copy()}, + str(filename), + with_pickle=False, + ) + monkeypatch.setattr(chx_correlationp, "_available_cpu_count", lambda: 8) + monkeypatch.setattr(chx_correlationp, "available_memory_bytes", lambda: 1) + with ( + Multifile(str(filename), 0, len(frames)) as serial_file, + Multifile(str(filename), 0, len(frames)) as parallel_file, + ): + expected, expected_lags = cal_g2c(serial_file, roi_mask, bad_frame_list=[5]) + actual, actual_lags = chx_correlationp.cal_g2p(parallel_file, roi_mask, bad_frame_list=[5]) + + np.testing.assert_array_equal(actual_lags, expected_lags) + np.testing.assert_allclose(actual, expected, rtol=1e-13, atol=0) diff --git a/pyCHX/tests/test_import_compatibility.py b/pyCHX/tests/test_import_compatibility.py index 7608ef4..a40b473 100644 --- a/pyCHX/tests/test_import_compatibility.py +++ b/pyCHX/tests/test_import_compatibility.py @@ -38,6 +38,25 @@ def test_legacy_star_import_namespace_and_effective_bindings(): assert packages.save_lists is generic.save_lists +@pytest.mark.portable +def test_production_wildcard_import_order_keeps_optimized_bindings(): + from pyCHX.chx_compress import compress_eigerdata + from pyCHX.chx_correlationc import Get_Pixel_Arrayc, auto_two_Arrayc + from pyCHX.chx_correlationp import cal_g2p + from pyCHX.Two_Time_Correlation_Function import get_one_time_from_two_time + + namespace = {} + exec( + "from pyCHX.chx_packages import *\nfrom pyCHX.chx_xpcs_xsvs_jupyter_V1 import *", + namespace, + ) + assert namespace["compress_eigerdata"] is compress_eigerdata + assert namespace["cal_g2p"] is cal_g2p + assert namespace["Get_Pixel_Arrayc"] is Get_Pixel_Arrayc + assert namespace["auto_two_Arrayc"] is auto_two_Arrayc + assert namespace["get_one_time_from_two_time"] is get_one_time_from_two_time + + @pytest.mark.portable def test_final_marker_and_color_values_are_preserved(): import pyCHX.chx_packages as packages diff --git a/pyCHX/tests/test_numerical_regressions.py b/pyCHX/tests/test_numerical_regressions.py index 0dfffc1..60907b4 100644 --- a/pyCHX/tests/test_numerical_regressions.py +++ b/pyCHX/tests/test_numerical_regressions.py @@ -1,5 +1,6 @@ import ast import os +from contextlib import contextmanager from pathlib import Path import matplotlib.pyplot as plt @@ -151,6 +152,56 @@ def test_one_time_correlation_is_mean_of_two_time_diagonals(): np.testing.assert_allclose(get_one_time_from_two_time(two_time), expected) +@pytest.mark.portable +def test_one_time_from_two_time_preserves_nan_and_explicit_normalization_semantics(): + from pyCHX.Two_Time_Correlation_Function import get_one_time_from_two_time + + two_time = np.arange(72, dtype=float).reshape(6, 6, 2) + two_time[1, 3, 0] = np.nan + two_time[2, 2, 1] = np.nan + norms = np.linspace(1.0, 3.0, 12).reshape(6, 2) + pixel_counts = np.array([3, 7]) + expected = np.empty((6, 2)) + for delay in range(6): + for roi_index in range(2): + expected[delay, roi_index] = np.nanmean(np.diag(two_time[:, :, roi_index], delay)) / ( + np.average(norms[delay:, roi_index]) + * np.average(norms[: 6 - delay, roi_index]) + * pixel_counts[roi_index] + ) + + np.testing.assert_allclose( + get_one_time_from_two_time(two_time, norms=norms, nopr=pixel_counts), + expected, + rtol=1e-14, + atol=0, + ) + + +@pytest.mark.portable +def test_compiled_two_time_diagonal_reducer_matches_numpy(): + from pyCHX.Two_Time_Correlation_Function import get_one_time_from_two_time + + generator = np.random.default_rng(20260902) + two_time = generator.random((360, 360, 8)) + two_time[10, 17, 3] = np.nan + norms = 1.0 + generator.random((360, 8)) + pixel_counts = np.arange(2, 10) + expected = np.empty((360, 8)) + for delay in range(360): + diagonal = np.nanmean(two_time.diagonal(delay), axis=1) + expected[delay] = diagonal / ( + norms[delay:].mean(axis=0) * norms[: 360 - delay].mean(axis=0) * pixel_counts + ) + + np.testing.assert_allclose( + get_one_time_from_two_time(two_time, norms=norms, nopr=pixel_counts), + expected, + rtol=2e-13, + atol=0, + ) + + @pytest.mark.portable def test_legacy_delay_values_scale_once(): from skbeam.core.utils import multi_tau_lags @@ -368,14 +419,27 @@ def test_array_two_time_helpers_support_sparse_roi_labels(): ) expected = [] + expected_auto = [] expected_norm = [] for selected in (data[:, :2], data[:, 2:]): means = selected.mean(axis=1) expected.append(np.dot(selected, selected.T) / np.outer(means, means) / selected.shape[1]) + expected_auto.append( + np.dot(selected, selected.T) / means.reshape(1, -1) / means.reshape(-1, 1) / selected.shape[1] + ) expected_norm.append(means.mean()) expected = np.stack(expected, axis=2) + expected_auto = np.stack(expected_auto, axis=2) - np.testing.assert_allclose(auto_two_Arrayc(data, roi_mask), expected) + np.testing.assert_array_equal(auto_two_Arrayc(data, roi_mask), expected_auto) + integer_data = data.astype(np.int64) + integer_expected = [] + for selected in (integer_data[:, :2], integer_data[:, 2:]): + means = selected.mean(axis=1) + integer_expected.append( + np.dot(selected, selected.T) / means.reshape(1, -1) / means.reshape(-1, 1) / selected.shape[1] + ) + np.testing.assert_array_equal(auto_two_Arrayc(integer_data, roi_mask), np.stack(integer_expected, axis=2)) np.testing.assert_allclose(auto_two_Arrayc_ExplicitNorm(data, roi_mask, norm=data), expected) np.testing.assert_allclose(auto_two_Arrayp(data, roi_mask), expected) np.testing.assert_allclose(auto_two_Arrayp2(data, roi_mask), expected) @@ -385,12 +449,536 @@ def test_array_two_time_helpers_support_sparse_roi_labels(): assert set(mean_intensity) == {2, 5} np.testing.assert_allclose(mean_intensity[2], data[:, :2].mean(axis=1)) np.testing.assert_allclose(mean_intensity[5], data[:, 2:].mean(axis=1)) - np.testing.assert_allclose(auto_two_Arrayc(data, roi_mask, index=5), expected[:, :, 1:]) + np.testing.assert_array_equal(auto_two_Arrayc(data, roi_mask, index=5), expected_auto[:, :, 1:]) with pytest.raises(ValueError, match="ROI labels not present"): auto_two_Arrayc(data, roi_mask, index=3) +@pytest.mark.portable +def test_production_two_time_path_prenormalizes_before_symmetric_blas(monkeypatch): + from pyCHX import chx_correlationc + + generator = np.random.default_rng(9) + data = 1.0 + generator.random((512, 512)) + means = data.mean(axis=1) + observed = {} + original_dsyrk = chx_correlationc.dsyrk + + def recording_dsyrk(alpha, normalized, **kwargs): + observed["alpha"] = alpha + observed["means"] = normalized.mean(axis=1) + return original_dsyrk(alpha, normalized, **kwargs) + + monkeypatch.setattr(chx_correlationc, "dsyrk", recording_dsyrk) + actual = chx_correlationc._symmetric_two_time_product(data, means, data.shape[1]) + expected = np.dot(data, data.T) / np.outer(means, means) / data.shape[1] + + assert observed["alpha"] == pytest.approx(1 / data.shape[1]) + np.testing.assert_allclose(observed["means"], 1.0, rtol=1e-14, atol=1e-14) + np.testing.assert_allclose(actual, expected, rtol=2e-13, atol=1e-14) + + +@pytest.mark.portable +def test_array_two_time_preserves_zero_intensity_frame_results(): + from pyCHX.chx_correlationc import auto_two_Arrayc + + roi_mask = np.array([[1, 1]]) + data = np.array([[1.0, 2.0], [0.0, 0.0], [2.0, 4.0], [-1.0, 1.0]]) + means = np.average(data, axis=1) + with np.errstate(divide="ignore", invalid="ignore"): + expected = np.dot(data, data.T) / means.reshape(1, -1) / means.reshape(-1, 1) / data.shape[1] + + actual = auto_two_Arrayc(data, roi_mask)[:, :, 0] + + np.testing.assert_allclose(actual, expected, equal_nan=True) + + +@pytest.mark.portable +def test_two_time_uses_all_physical_cores_without_roi_threading(monkeypatch): + from pyCHX import chx_correlationc + + blas_limits = [] + numba_limits = [] + + def capture_limit(*, limits: int, user_api: str): + assert user_api == "blas" + blas_limits.append(limits) + + @contextmanager + def context(): + yield + + return context() + + @contextmanager + def capture_numba_limit(limit): + numba_limits.append(limit) + yield + + def reject_roi_thread_pool(*args, **kwargs): + raise AssertionError("two-time ROI calculations must not use Python threads around scipy BLAS") + + monkeypatch.setattr(chx_correlationc, "physical_core_count", lambda: 4) + monkeypatch.setattr(chx_correlationc, "threadpool_limits", capture_limit) + monkeypatch.setattr(chx_correlationc, "numba_thread_limit", capture_numba_limit) + monkeypatch.setattr(chx_correlationc, "ThreadPoolExecutor", reject_roi_thread_pool) + roi_mask = np.array([[1, 1], [2, 2]]) + data = np.arange(16, dtype=float).reshape(4, 4) + 1 + chx_correlationc.auto_two_Arrayc(data, roi_mask) + chx_correlationc.auto_two_Arrayc_ExplicitNorm(data, roi_mask, norm=data) + assert blas_limits == [4, 4] + assert numba_limits == [4, 4] + + +@pytest.mark.portable +def test_parallel_two_time_matrix_finishing_matches_serial_kernels(): + from pyCHX._performance import ( + mirror_and_normalize_two_time, + mirror_and_normalize_two_time_parallel, + mirror_two_time, + mirror_two_time_parallel, + store_symmetric_two_time_batch, + ) + + generator = np.random.default_rng(11) + upper_triangle = np.triu(generator.random((128, 128))) + expected_mirror = upper_triangle.copy() + actual_mirror = upper_triangle.copy() + mirror_two_time(expected_mirror) + mirror_two_time_parallel(actual_mirror) + np.testing.assert_array_equal(actual_mirror, expected_mirror) + + norms = generator.random(128) + expected_normalized = upper_triangle.copy() + actual_normalized = upper_triangle.copy() + mirror_and_normalize_two_time(expected_normalized, norms, 17) + mirror_and_normalize_two_time_parallel(actual_normalized, norms, 17) + np.testing.assert_array_equal(actual_normalized, expected_normalized) + + upper_batch = np.empty((128, 128, 2), dtype=np.float64, order="F") + upper_batch[:, :, 0] = upper_triangle + upper_batch[:, :, 1] = upper_triangle + row_norms = np.column_stack((np.ones(128), norms)) + output = np.empty((128, 128, 2), dtype=np.float64) + store_symmetric_two_time_batch( + upper_batch, + row_norms, + np.array([1, 17]), + np.array([True, False]), + output, + 0, + 2, + ) + np.testing.assert_array_equal(output[:, :, 0], expected_mirror) + np.testing.assert_array_equal(output[:, :, 1], expected_normalized) + + +@pytest.mark.portable +def test_two_time_batch_size_is_cache_and_memory_bounded(monkeypatch): + from pyCHX import chx_correlationc + + frame_count = 1_000 + matrix_bytes = frame_count * frame_count * np.dtype(np.float64).itemsize + monkeypatch.setattr(chx_correlationc, "available_memory_bytes", lambda: matrix_bytes * 40) + assert chx_correlationc._two_time_batch_size(frame_count, 100) == 2 + + monkeypatch.setattr(chx_correlationc, "available_memory_bytes", lambda: matrix_bytes * 1_000) + assert chx_correlationc._two_time_batch_size(frame_count, 100) == 16 + assert chx_correlationc._two_time_batch_size(frame_count, 7) == 7 + + +@pytest.mark.portable +def test_two_time_supports_hundreds_of_rois_and_returns_c_contiguous_output(monkeypatch): + from pyCHX import chx_correlationc + + roi_mask = np.arange(1, 201, dtype=np.int64).reshape(10, 20) + data = 1.0 + np.arange(8 * 200, dtype=np.float64).reshape(8, 200) % 31 + monkeypatch.setattr(chx_correlationc, "physical_core_count", lambda: 8) + actual = chx_correlationc.auto_two_Arrayc(data, roi_mask) + + assert actual.shape == (8, 8, 200) + assert actual.flags.c_contiguous + np.testing.assert_allclose(actual, 1.0) + + +@pytest.mark.portable +def test_two_time_multiple_output_batches_match_original_operations(monkeypatch): + from pyCHX import chx_correlationc + + generator = np.random.default_rng(81) + roi_count = 20 + pixels_per_roi = 3 + frame_count = 12 + roi_mask = np.repeat(np.arange(1, roi_count + 1), pixels_per_roi).reshape(6, 10) + data = 0.5 + generator.random((frame_count, roi_count * pixels_per_roi)) + explicit_norm = 0.5 + generator.random(data.shape) + monkeypatch.setattr(chx_correlationc, "_two_time_batch_size", lambda *_: 7) + + expected = [] + expected_explicit = [] + expected_without_norm = [] + for roi_index in range(roi_count): + start = roi_index * pixels_per_roi + selected = data[:, start : start + pixels_per_roi] + means = selected.mean(axis=1) + expected.append(np.dot(selected, selected.T) / np.outer(means, means) / pixels_per_roi) + explicit_means = explicit_norm[:, start : start + pixels_per_roi].mean(axis=1) + expected_explicit.append( + np.dot(selected, selected.T) / np.outer(explicit_means, explicit_means) / pixels_per_roi + ) + expected_without_norm.append(np.dot(selected, selected.T) / pixels_per_roi) + + actual = chx_correlationc.auto_two_Arrayc(data, roi_mask) + actual_explicit = chx_correlationc.auto_two_Arrayc_ExplicitNorm(data, roi_mask, norm=explicit_norm) + actual_without_norm = chx_correlationc.auto_two_Arrayc_ExplicitNorm(data, roi_mask) + np.testing.assert_allclose(actual, np.stack(expected, axis=2), rtol=1e-15, atol=1e-15) + np.testing.assert_allclose( + actual_explicit, + np.stack(expected_explicit, axis=2), + rtol=1e-15, + atol=1e-15, + ) + np.testing.assert_allclose( + actual_without_norm, + np.stack(expected_without_norm, axis=2), + rtol=1e-15, + atol=1e-15, + ) + + +@pytest.mark.portable +def test_diagonal_numba_limit_uses_physical_core_count(monkeypatch): + from pyCHX import Two_Time_Correlation_Function as two_time + + observed = [] + + @contextmanager + def recording_limit(limit): + observed.append(limit) + yield + + monkeypatch.setattr(two_time, "physical_core_count", lambda: 7) + monkeypatch.setattr(two_time, "numba_thread_limit", recording_limit) + data = np.ones((360, 360, 8)) + two_time.get_one_time_from_two_time(data) + assert observed == [7] + + +@pytest.mark.portable +def test_get_pixel_array_normalization_modes_are_exact(): + from pyCHX.chx_correlationc import Get_Pixel_Arrayc + + class SparseFrames: + beg = 1 + end = 4 + md = {"ncols": 2, "nrows": 4} + + def __init__(self, frames): + self.frames = frames + + def rdrawframe(self, index): + flattened = self.frames[index].ravel() + positions = np.flatnonzero(flattened).astype(np.int32) + return positions, flattened[positions] + + frames = np.arange(1, 33, dtype=np.float64).reshape(4, 2, 4) + pixel_list = np.array([0, 2, 5, 7]) + qind = np.array([1, 2, 1, 2]) + norm_1d = np.array([1.5, 2.0, 2.5, 4.0]) + norm_2d = np.multiply.outer(np.arange(1.0, 5.0), norm_1d) + imgsum = np.arange(10.0, 14.0) + mean_int_sets = np.array([[2.0, 3.0], [3.0, 4.0], [4.0, 5.0], [5.0, 6.0]]) + cases = [ + {}, + {"norm": norm_1d}, + {"norm": norm_2d}, + {"imgsum": imgsum}, + {"mean_int_sets": mean_int_sets, "qind": qind}, + {"norm": norm_2d, "imgsum": imgsum, "mean_int_sets": mean_int_sets, "qind": qind}, + ] + + selected_frames = frames[SparseFrames.beg : SparseFrames.end].reshape(3, -1)[:, pixel_list] + for kwargs in cases: + expected = np.zeros_like(selected_frames) + for output_index, frame_index in enumerate(range(SparseFrames.beg, SparseFrames.end)): + mean_norm = mean_int_sets[frame_index, qind - 1] if kwargs.get("mean_int_sets") is not None else 1.0 + sum_norm = imgsum[frame_index] if kwargs.get("imgsum") is not None else 1.0 + if kwargs.get("norm") is norm_2d: + pixel_norm = norm_2d[frame_index] + elif kwargs.get("norm") is norm_1d: + pixel_norm = norm_1d + else: + pixel_norm = 1.0 + expected[output_index] = selected_frames[output_index] / (mean_norm * sum_norm * pixel_norm) + + actual = Get_Pixel_Arrayc(SparseFrames(frames), pixel_list, **kwargs).get_data() + np.testing.assert_array_equal(actual, expected) + + +@pytest.mark.portable +@pytest.mark.parametrize( + ("cal_error", "use_intensity_cache"), + [(False, False), (False, True), (True, False)], +) +def test_optimized_one_time_kernel_matches_legacy_operations_exactly(cal_error, use_intensity_cache): + from pyCHX.chx_correlationc import _one_time_process, _one_time_process_cached, _one_time_process_error + + def legacy_process( + buf, + correlation, + past_norm, + future_norm, + labels, + num_bufs, + num_pixels, + images_per_level, + level, + buffer_number, + bad_counts, + level_lengths, + error_arrays=None, + ): + images_per_level[level] += 1 + minimum = num_bufs // 2 if level else 0 + for delay in range(minimum, min(images_per_level[level], num_bufs)): + time_index = int(level * num_bufs / 2 + delay) + past = buf[level, (buffer_number - delay) % num_bufs] + future = buf[level, buffer_number] + level_index = int(time_index - level_lengths[:level].sum()) + normalize = images_per_level[level] - delay - bad_counts[level + 1][level_index] + if np.isnan(past).any() or np.isnan(future).any(): + bad_counts[level + 1][level_index] += 1 + elif error_arrays is None: + for weights, output in zip( + [past * future, past, future], + [correlation, past_norm, future_norm], + ): + binned = np.bincount(labels, weights=weights)[1:] + output[time_index] += (binned / num_pixels - output[time_index]) / normalize + else: + for weights, output in zip([past * future, past, future], error_arrays): + output[time_index] += (weights - output[time_index]) / normalize + + buf = np.array( + [ + [ + [1.25, 2.5, 3.75, 5.0], + [np.nan, np.nan, np.nan, np.nan], + [2.0, 4.5, 7.0, 9.5], + [3.5, 5.25, 8.75, 11.0], + ] + ] + ) + labels = np.array([1, 1, 2, 2]) + num_pixels = np.array([2, 2]) + level_lengths = np.array([4]) + shape = (4, 2) + legacy_arrays = [np.zeros(shape), np.zeros(shape), np.zeros(shape)] + optimized_arrays = [array.copy() for array in legacy_arrays] + legacy_images_per_level = np.array([3]) + optimized_images_per_level = legacy_images_per_level.copy() + legacy_bad_counts = {1: np.zeros(4, dtype=np.int64)} + optimized_bad_counts = {1: np.zeros(4, dtype=np.int64)} + + if cal_error: + legacy_error_arrays = [np.zeros((4, 4)), np.zeros((4, 4)), np.zeros((4, 4))] + optimized_error_arrays = [array.copy() for array in legacy_error_arrays] + legacy_process( + buf, + *legacy_arrays, + labels, + 4, + num_pixels, + legacy_images_per_level, + 0, + 3, + legacy_bad_counts, + level_lengths, + legacy_error_arrays, + ) + _one_time_process_error( + buf, + *optimized_arrays, + labels, + 4, + num_pixels, + optimized_images_per_level, + 0, + 3, + optimized_bad_counts, + level_lengths, + *optimized_error_arrays, + ) + for actual, expected in zip(optimized_error_arrays, legacy_error_arrays): + np.testing.assert_array_equal(actual, expected) + else: + legacy_process( + buf, + *legacy_arrays, + labels, + 4, + num_pixels, + legacy_images_per_level, + 0, + 3, + legacy_bad_counts, + level_lengths, + ) + optimized_arguments = [ + buf, + *optimized_arrays, + labels, + 4, + num_pixels, + optimized_images_per_level, + 0, + 3, + optimized_bad_counts, + level_lengths, + ] + if use_intensity_cache: + intensity_cache = np.zeros((1, 4, 2), dtype=np.float64) + for buffer_index, image in enumerate(buf[0]): + if not np.isnan(image).any(): + intensity_cache[0, buffer_index] = np.bincount(labels, weights=image)[1:] + optimized_arguments.append(intensity_cache) + _one_time_process_cached(*optimized_arguments) + else: + _one_time_process(*optimized_arguments) + for actual, expected in zip(optimized_arrays, legacy_arrays): + np.testing.assert_array_equal(actual, expected) + + np.testing.assert_array_equal(optimized_images_per_level, legacy_images_per_level) + np.testing.assert_array_equal(optimized_bad_counts[1], legacy_bad_counts[1]) + + +@pytest.mark.portable +@pytest.mark.parametrize("use_intensity_cache", [False, True]) +@pytest.mark.parametrize(("level", "current_time"), [(0, 8), (1, 8.5)]) +def test_optimized_two_time_kernel_matches_legacy_operations_exactly(level, current_time, use_intensity_cache): + from pyCHX.chx_correlationc import _create_intensity_buffer, _two_time_process, _two_time_process_cached + + def legacy_process( + buf, + correlation, + labels, + num_bufs, + num_pixels, + images_per_level, + lag_steps, + current_time, + level, + buffer_number, + ): + images_per_level[level] += 1 + minimum = 0 if level == 0 else num_bufs // 2 + for delay in range(minimum, min(images_per_level[level], num_bufs)): + time_index = level * num_bufs / 2 + delay + past = buf[level, (buffer_number - delay) % num_bufs] + future = buf[level, buffer_number] + product_sum = np.bincount(labels, weights=past * future)[1:] + past_sum = np.bincount(labels, weights=past)[1:] + future_sum = np.bincount(labels, weights=future)[1:] + first_time = current_time - 1 + second_time = current_time - lag_steps[int(time_index)] - 1 + values = product_sum / (past_sum * future_sum) * num_pixels + if not isinstance(current_time, int): + shift = 2 ** (level - 1) + for offset in range(-shift + 1, shift + 1): + correlation[:, int(first_time + offset), int(second_time + offset)] = values + else: + correlation[:, int(first_time), int(second_time)] = values + + buf = np.array( + [ + [ + [1.25, 2.5, 3.75, 5.0], + [2.0, 4.5, 7.0, 9.5], + [3.5, 5.25, 8.75, 11.0], + [4.25, 6.5, 9.25, 12.5], + ], + [ + [1.625, 3.5, 5.375, 7.25], + [2.75, 4.875, 7.875, 10.25], + [3.875, 5.875, 9.0, 11.75], + [2.9375, 4.5, 6.5625, 8.875], + ], + ] + ) + labels = np.array([1, 1, 2, 2]) + num_pixels = np.array([2, 2]) + lag_steps = np.array([0, 1, 2, 3, 4, 6]) + legacy_correlation = np.zeros((2, 12, 12)) + optimized_correlation = legacy_correlation.copy() + legacy_images_per_level = np.array([3, 3]) + optimized_images_per_level = legacy_images_per_level.copy() + arguments = ( + buf, + labels, + 4, + num_pixels, + lag_steps, + current_time, + level, + 3, + ) + + legacy_process( + arguments[0], + legacy_correlation, + arguments[1], + arguments[2], + arguments[3], + legacy_images_per_level, + *arguments[4:], + ) + optimized_arguments = [ + arguments[0], + optimized_correlation, + arguments[1], + arguments[2], + arguments[3], + optimized_images_per_level, + *arguments[4:], + ] + if use_intensity_cache: + optimized_arguments.append(_create_intensity_buffer(buf, labels, len(num_pixels))) + _two_time_process_cached(*optimized_arguments) + else: + _two_time_process(*optimized_arguments) + + np.testing.assert_array_equal(optimized_correlation, legacy_correlation) + np.testing.assert_array_equal(optimized_images_per_level, legacy_images_per_level) + + +@pytest.mark.portable +def test_two_time_intensity_cache_tracks_frames_before_the_first_correlated_lag(): + from pyCHX.chx_correlationc import _create_intensity_buffer, _two_time_process_cached + + buf = np.zeros((2, 4, 4), dtype=np.float64) + buf[1, 0] = [1.0, 2.0, 3.0, 4.0] + labels = np.array([1, 1, 2, 2]) + intensity_buf = _create_intensity_buffer(buf, labels, 2) + intensity_buf[1, 0] = 0 + + _two_time_process_cached( + buf, + np.zeros((2, 8, 8)), + labels, + 4, + np.array([2, 2]), + np.zeros(2, dtype=np.int64), + np.array([0, 1, 2, 3, 4, 6]), + 1.5, + 1, + 0, + intensity_buf, + ) + + np.testing.assert_array_equal(intensity_buf[1, 0], [3.0, 7.0]) + + @pytest.mark.portable def test_bad_frames_mask_rows_and_columns(): from pyCHX.Two_Time_Correlation_Function import make_g12_mask diff --git a/pyproject.toml b/pyproject.toml index 76ddd02..bbb7e37 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -32,6 +32,7 @@ dependencies = [ "ipython>=8", "lmfit>=1.1,<2", "matplotlib>=3.6,<4", + "numba>=0.61", "numpy>=1.23,<3", "pandas>=1.5,<4", "pillow>=9", @@ -42,6 +43,7 @@ dependencies = [ "scikit-image>=0.20,<1", "scipy>=1.9,<2", "tifffile>=2022.8", + "threadpoolctl>=3.1", "tqdm>=4.64", "xray-vision>=0.1.1,<0.2", ]