import h5py
import numpy as np
import os
import shutil
import traceback
from coseda.io import get_free_space_windows, get_free_space_unix, handle_input, read_config, config_to_paths
from coseda.logging_utils import log_result, log_start
from coseda.pipeline.frame_intensities import calculate_mean_intensities_dask_file_inprocess
from coseda.nexus.logs import ensure_dense_logs
from coseda.nexus.goniometer import ensure_goniometer_transforms, get_goniometer_transform_order
from coseda.nexus.images import ensure_image_nxdata
from coseda.nexus.paths import DETECTOR_GROUP_PATH, LEGACY_MASK_PATH, NEXUS_MASK_PATH, get_mask_dataset
from coseda.nexus.peaks import ensure_peak_nxdata
from coseda.nexus.process import write_nxprocess_stripping
from coseda.logging_utils import get_logger
_LOGGER = get_logger(__name__)
_DENSE_INTENSITY_DATASETS = {
# Current names.
"frame_mean_intensities",
"frame_total_intensities",
"frame_radial_intensities",
# Legacy importer names.
"mean_intensities",
"total_intensities",
"radial_intensities",
}
_DENSE_METADATA_DATASETS = {
"alphatilt",
"betatilt",
"frame_id",
"streak_id",
"streak_frame",
"streak_endpoints",
}
_DENSE_STAGE_PREFIXES = (
"stagepos",
"stage_position",
"stageposition",
)
def _preserve_full_length_dataset(dataset_name):
"""Return True for dense per-original-frame arrays that must not be stripped."""
name = dataset_name.lower()
return (
name in _DENSE_INTENSITY_DATASETS
or name in _DENSE_METADATA_DATASETS
or name.startswith(_DENSE_STAGE_PREFIXES)
)
def _copy_detector_mask(src_file, dst_file, logfile_path=None):
"""Copy the 2-D detector mask and recreate the legacy /mask hardlink."""
mask_ds = get_mask_dataset(src_file)
if mask_ds is None:
return
for path in (LEGACY_MASK_PATH, NEXUS_MASK_PATH):
if path in dst_file:
del dst_file[path]
detector_group = dst_file.require_group(DETECTOR_GROUP_PATH.lstrip("/"))
if "mask" in detector_group:
del detector_group["mask"]
detector_group.copy(mask_ds, "mask")
dst_file[LEGACY_MASK_PATH] = dst_file[NEXUS_MASK_PATH]
log_start(logfile_path, f"Copied detector mask from {mask_ds.name}.")
[docs]
def check_if_index_dataset_exists(hdf5_file_path):
"""Return True if entry/data/index exists in the HDF5 file."""
index_dataset_name = 'index'
try:
# Open the HDF5 file in append mode (since we might need to create the index dataset)
with h5py.File(hdf5_file_path, 'a') as hdf5_file:
data_group = hdf5_file['entry/data']
# Check if the index dataset already exists
if index_dataset_name in data_group:
_LOGGER.info(f"Index dataset '{index_dataset_name}' exists.")
return True
else:
_LOGGER.info(f"Index dataset '{index_dataset_name}' does not exist.")
return False
except Exception as e:
_LOGGER.error(f"Error accessing HDF5 file: {e}")
return False
[docs]
def check_if_stripped(hdf5_file_path):
"""Heuristic: compare images vs index length; if they differ, the file was stripped."""
images_dataset_name = 'images'
index_dataset_name = 'index'
try:
with h5py.File(hdf5_file_path, 'r') as hdf5_file:
data_group = hdf5_file['entry/data']
# Check if both datasets exist
if images_dataset_name not in data_group or index_dataset_name not in data_group:
_LOGGER.warning(
f"One or both datasets do not exist: '{images_dataset_name}', '{index_dataset_name}'"
)
return False
# Get the lengths of both datasets
images_length = data_group[images_dataset_name].shape[0]
index_length = data_group[index_dataset_name].shape[0]
# Compare the lengths
if images_length == index_length:
return False # Not stripped
else:
return True # Stripped
except Exception as e:
_LOGGER.error(f"Error accessing HDF5 file: {e}")
return False
[docs]
def check_if_peaks_and_intensities_datasets_exist(hdf5_file_path):
"""Return booleans for presence of nPeaks and frame intensity datasets."""
peaks_dataset_name = 'nPeaks'
peaks_dataset_exists = False
intensities_exist = False
try:
with h5py.File(hdf5_file_path, 'r') as hdf5_file:
data_group = hdf5_file['entry/data']
# Check if 'nPeaks' exists
if peaks_dataset_name in data_group:
_LOGGER.info(f"Dataset '{peaks_dataset_name}' exists.")
peaks_dataset_exists = True
else:
_LOGGER.info(f"Dataset '{peaks_dataset_name}' does not exist.")
intensity_names = [
name for name in data_group
if isinstance(data_group[name], h5py.Dataset)
and name.lower() in _DENSE_INTENSITY_DATASETS
]
if intensity_names:
_LOGGER.info(f"Intensity dataset(s) exist: {', '.join(sorted(intensity_names))}.")
intensities_exist = True
else:
_LOGGER.info("No frame intensity datasets exist.")
except Exception as e:
_LOGGER.error(f"Error accessing HDF5 file: {e}")
return None
return peaks_dataset_exists, intensities_exist
[docs]
def create_index_dataset(hdf5_file_path):
"""Create entry/data/index = arange(len(images)) if it does not exist."""
images_dataset_name = 'images'
index_dataset_name = 'index'
try:
with h5py.File(hdf5_file_path, 'a') as hdf5_file:
data_group = hdf5_file['entry/data']
# Check if the images dataset exists
if images_dataset_name not in data_group:
_LOGGER.warning(f"Images dataset '{images_dataset_name}' not found.")
return False
images_length = data_group[images_dataset_name].shape[0]
# Check if the index dataset already exists
if index_dataset_name in data_group:
_LOGGER.info(f"Index dataset '{index_dataset_name}' already exists.")
return False
# Create the index dataset with integer values from 0 to images_length - 1
index_dataset = data_group.create_dataset(index_dataset_name, shape=(images_length,), dtype='i4')
index_dataset[:] = np.arange(images_length)
_LOGGER.info(f"Index dataset '{index_dataset_name}' created with {images_length} entries.")
return True
except Exception as e:
_LOGGER.error(f"Error creating index dataset: {e}")
return False
[docs]
def check_disk_space(hdf5_file_path, threshold):
"""Check free space and estimate fraction of frames kept; returns (ok, message, fraction_kept)."""
file_directory = os.path.dirname(hdf5_file_path)
# Get free space based on the operating system
if os.name == 'nt':
free_space = get_free_space_windows(file_directory) / (1024 ** 3) # Convert to GB
else:
free_space = get_free_space_unix(file_directory) / (1024 ** 3) # Convert to GB
try:
# Get total file size in bytes
total_file_size_bytes = os.path.getsize(hdf5_file_path)
total_file_size_gb = total_file_size_bytes / (1024 ** 3)
with h5py.File(hdf5_file_path, 'r') as hdf5_file:
data_group = hdf5_file['entry/data']
# Access the datasets
n_peaks_dataset = data_group['nPeaks']
frames_dataset = data_group['images']
index_dataset = data_group['index']
total_frames = index_dataset.shape[0] # Total number of frames (from index)
# Find the valid frames in the current dataset (frames that were not stripped)
valid_indices = np.where(index_dataset[:] != -1)[0]
# Calculate the number of valid frames that meet the threshold
n_peaks_values = n_peaks_dataset[:]
frames_to_keep = np.sum(n_peaks_values >= threshold)
# Fraction of valid frames that will be kept
fraction_kept = frames_to_keep / len(valid_indices) if len(valid_indices) > 0 else 0
totalfraction_kept = frames_to_keep / len(index_dataset) if len(index_dataset) > 0 else 0
# Read the actual size of the frames dataset from the HDF5 file
frames_dataset_size_bytes = frames_dataset.id.get_storage_size()
# Convert the dataset size to GB
frames_dataset_size_gb = frames_dataset_size_bytes / (1024 ** 3)
# Calculate overhead size (size of other datasets and metadata)
overhead_size_gb = total_file_size_gb - frames_dataset_size_gb
# Estimate the size of the stripped frames dataset
stripped_frames_size_gb = frames_dataset_size_gb * fraction_kept
# Estimate the new total file size after stripping
estimated_new_file_size_gb = stripped_frames_size_gb + overhead_size_gb
# Estimate the space required for repacking (allowing overhead for temp space and metadata)
estimated_required_space_gb = estimated_new_file_size_gb * 2 # Double the space for temporary storage
except Exception as e:
result = f"Error accessing HDF5 file: {e}"
return False, result, None
# Check if there is sufficient free space for the repacking operation
if free_space >= estimated_required_space_gb:
result = f"Sufficient disk space available. Estimated file size after repacking: {round(estimated_new_file_size_gb, 2)} GB."
return True, result, totalfraction_kept
else:
result = f"Not enough disk space for repacking procedure. Required: {round(estimated_required_space_gb, 2)} GB, available: {round(free_space, 2)} GB."
return False, result, totalfraction_kept
[docs]
def repack_hdf5_with_backup(hdf5_file_path, logfile_path=None):
backup_file_path = hdf5_file_path + ".backup"
temp_repacked_file = hdf5_file_path + ".repacked"
try:
# Step 1: Rename the original file to .backup
os.rename(hdf5_file_path, backup_file_path)
log_start(logfile_path, f"Backup created: {backup_file_path}")
# Step 2: Open the backup file and create a new repacked file
with h5py.File(backup_file_path, 'r') as src_file, h5py.File(temp_repacked_file, 'w') as dst_file:
# Copy all groups, datasets, and attributes from the backup to the new file
def recursive_copy(src, dst):
for key in src:
item = src[key]
if isinstance(item, h5py.Group):
# Create the group in the destination and copy recursively
dst_group = dst.create_group(key)
# Copy attributes
for attr_name, attr_value in item.attrs.items():
dst_group.attrs[attr_name] = attr_value
recursive_copy(item, dst_group)
elif isinstance(item, h5py.Dataset):
# Copy dataset
dst.copy(item, key)
# Copy attributes
dst[key].attrs.update(item.attrs)
# Start copying recursively from the root group
recursive_copy(src_file, dst_file)
# Step 3: Check if repacking was successful
if os.path.exists(temp_repacked_file):
# Delete the backup file after successful repacking
os.remove(backup_file_path)
# Rename the repacked file to the original file name
os.rename(temp_repacked_file, hdf5_file_path)
log_start(logfile_path, f"Repacking successful, backup deleted. File restored to original name: {hdf5_file_path}")
return True
else:
raise Exception("Repacked file was not created successfully.")
except Exception as e:
log_start(logfile_path, f"Error during repacking process: {e}")
# Rollback: Delete the new file if it exists
if os.path.exists(temp_repacked_file):
os.remove(temp_repacked_file)
log_start(logfile_path, f"Temporary repacked file deleted: {temp_repacked_file}")
# Restore the backup by renaming it to the original file name
if os.path.exists(backup_file_path):
os.rename(backup_file_path, hdf5_file_path)
log_start(logfile_path, f"Backup restored to original file: {hdf5_file_path}")
return False
[docs]
def strip_h5(
hdf5_file_path,
threshold=1,
force=False,
logfile_path=None,
progress_callback=None,
goniometer_transform_order=None,
):
"""Strip frames from a single HDF5 based on peak counts, preserving metadata and backups."""
# Rename original file to backup
backup_path = hdf5_file_path + ".backup"
os.rename(hdf5_file_path, backup_path)
try:
# Check if the nPeaks and mean_intensities datasets exist
peaks_exist, intensities_exist = check_if_peaks_and_intensities_datasets_exist(backup_path)
if not intensities_exist:
log_start(logfile_path, 'File does not contain frame intensities - calculating them now.')
calculate_mean_intensities_dask_file_inprocess(backup_path, logfile_path)
if not peaks_exist:
log_start(logfile_path, "Peak dataset is missing. Run peakfinding and try again.")
raise Exception("Peak dataset missing")
# Ensure the index dataset exists, or create one
if not check_if_index_dataset_exists(backup_path):
if not create_index_dataset(backup_path):
raise Exception("Failed to create index dataset")
# Check disk space
disk_space_check, result, fraction_kept = check_disk_space(backup_path, threshold)
log_start(logfile_path, result)
if not disk_space_check:
raise Exception("Not enough disk space to proceed")
if fraction_kept < 0.1 and not force:
log_start(logfile_path, f'You are attempting to delete {round(((1 - fraction_kept) * 100), 2)}% of your dataset! Please review your peakfinding results and rerun this function in forced mode if correct.')
raise Exception("Too many frames would be deleted, aborting (not forced).")
# Open backup for reading and new file for writing
with h5py.File(backup_path, 'r') as src_file, h5py.File(hdf5_file_path, 'w') as dst_file:
src_group = src_file['entry/data']
dst_group = dst_file.create_group('entry/data')
# Now build and write the stripped images dataset as before, but using src_group and dst_group
images_dataset = src_group['images']
n_peaks_dataset = src_group['nPeaks']
index_dataset = src_group['index']
images_length = images_dataset.shape[0]
n_peaks_length = n_peaks_dataset.shape[0]
index_length = index_dataset.shape[0]
# Get the mapping from original indices to current images indices
current_mapping = index_dataset[:]
# Get the valid indices where frames haven't been stripped (i.e., index != -1)
valid_indices = np.where(current_mapping != -1)[0]
if len(valid_indices) == 0:
log_start(logfile_path, "All frames have been stripped. No further action required.")
# On success, delete backup
os.remove(backup_path)
return True
# Retrieve the nPeaks values for valid frames
n_peaks_values = n_peaks_dataset[:]
# Initialize new mapping with -1 (assuming all frames are stripped)
new_mapping = np.full(index_length, -1, dtype='i4')
# Determine which valid frames meet the threshold
# Map original indices to their current image indices
valid_img_indices = current_mapping[valid_indices]
# Boolean mask for frames to keep based on peak counts
frames_to_keep_mask = n_peaks_values[valid_img_indices] >= threshold
# Filter original and image indices
frames_to_keep_orig = valid_indices[frames_to_keep_mask]
frames_to_keep_img = valid_img_indices[frames_to_keep_mask]
# Update the index mapping for frames to keep
index_counter = 0 # New images index
frames_to_keep = [] # List of tuples (old_image_index, new_image_index, original_index)
log_start(logfile_path, 'Identifying frames to keep.')
for original_idx, current_image_index in zip(frames_to_keep_orig, frames_to_keep_img):
# Ensure current_image_index is valid
if current_image_index < 0 or current_image_index >= images_length:
log_start(logfile_path, f"Warning: Skipping frame {original_idx} with invalid image index {current_image_index}.")
continue # Skip invalid index
# Update the new_mapping and keep track of frames to keep
new_mapping[original_idx] = index_counter
frames_to_keep.append((current_image_index, index_counter, original_idx))
index_counter += 1
num_frames_to_keep = len(frames_to_keep)
num_frames_to_strip = np.sum(new_mapping == -1)
log_start(logfile_path, f"Frames to keep: {num_frames_to_keep}")
log_start(logfile_path, f"Frames to strip: {num_frames_to_strip}")
if num_frames_to_keep == 0:
log_start(logfile_path, "No frames meet the criteria for stripping. All valid frames will be discarded.")
# Remove the images dataset since no frames are kept
# (Do not create images dataset at all)
log_start(logfile_path, "All frames have been stripped.")
# On success, delete backup
os.remove(backup_path)
return True
# Write the new index dataset
dst_group.create_dataset('index', data=new_mapping, dtype='i4')
# Write the compact->original mapping
image_key = np.array([orig_idx for _, _, orig_idx in frames_to_keep], dtype='i4')
dst_group.create_dataset('image_key', data=image_key, dtype='i4')
# Datasets to be reduced are image-indexed arrays. Dense metadata
# streams such as stage positions and frame intensity summaries stay
# full-length so maps/atlases can be rebuilt after images are stripped.
datasets_to_reduce = [
dataset_name for dataset_name in src_group
if dataset_name not in {'index', 'image_key', 'images'}
and not _preserve_full_length_dataset(dataset_name)
and isinstance(src_group[dataset_name], h5py.Dataset)
]
# Determine names to skip during initial bulk copy (will be written reduced later)
# Also skip HDF5 groups (e.g. NXdata groups like 'peaks', 'peak_counts')
# as they contain soft links that would dereference and duplicate data.
# These groups are rebuilt by ensure_peak_nxdata after stripping.
skip_names = set(datasets_to_reduce) | {'images', 'index', 'image_key'}
# Copy only metadata datasets that won't be reduced
for name in src_group:
if name in skip_names or isinstance(src_group[name], h5py.Group):
continue
dst_group.copy(src_group[name], name)
# Preserve the atlas group if present. It is built from the (kept) dense
# stage trajectory, so it stays valid after stripping.
if 'atlas' in src_file['entry']:
dst_file['entry'].copy(src_file['entry/atlas'], 'atlas')
log_start(logfile_path, "Copied existing atlas group.")
log_start(logfile_path, f"Datasets to reduce: {datasets_to_reduce}")
# All reducible datasets are image-indexed (their length matches images.shape[0])
frames_to_keep_by_image = np.array([img_idx for img_idx, _, _ in frames_to_keep], dtype='i4')
# Create new reduced datasets
for dataset_name in datasets_to_reduce:
dataset = src_group[dataset_name]
dataset_length = dataset.shape[0]
if dataset_length != images_length:
log_start(logfile_path, f"Warning: Dataset '{dataset_name}' length ({dataset_length}) does not match images length ({images_length}). Skipping.")
continue
new_shape = (num_frames_to_keep,) + dataset.shape[1:]
dataset_chunks = dataset.chunks
if dataset_chunks is not None:
new_chunks = tuple(
min(chunk_dim, shape_dim)
for chunk_dim, shape_dim in zip(dataset_chunks, new_shape)
)
else:
new_chunks = None
new_dataset = dst_group.create_dataset(
dataset_name,
shape=new_shape,
dtype=dataset.dtype,
chunks=new_chunks,
compression=dataset.compression
)
new_dataset[:] = dataset[frames_to_keep_by_image]
log_start(logfile_path, f"Dataset '{dataset_name}' reduced and updated.")
log_start(logfile_path, 'Processing images dataset...')
chunk_size = images_dataset.chunks
if chunk_size is not None:
adjusted_chunk_size = min(chunk_size[0], num_frames_to_keep)
new_chunk_size = (adjusted_chunk_size,) + chunk_size[1:]
else:
new_chunk_size = (min(1000, num_frames_to_keep),) + images_dataset.shape[1:]
new_images_dataset = dst_group.create_dataset(
'images',
shape=(num_frames_to_keep, *images_dataset.shape[1:]),
dtype=images_dataset.dtype,
chunks=new_chunk_size,
compression=images_dataset.compression
)
log_start(logfile_path, 'Writing new images dataset...')
for i, (old_image_index, new_image_index, _) in enumerate(frames_to_keep):
new_images_dataset[new_image_index] = images_dataset[old_image_index]
if progress_callback:
progress_callback(int((i+1) / num_frames_to_keep * 100))
if new_image_index % 1000 == 0:
log_start(logfile_path, f"Copied frame {new_image_index}.")
new_images_dataset.flush()
dst_file.flush()
log_start(logfile_path, f"Copied all {num_frames_to_keep} frames to the new images dataset.")
log_start(logfile_path, "Images dataset reduced and updated.")
_copy_detector_mask(src_file, dst_file, logfile_path)
ensure_dense_logs(dst_file)
ensure_image_nxdata(dst_file)
ensure_peak_nxdata(dst_file)
ensure_goniometer_transforms(dst_file, goniometer_transform_order)
write_nxprocess_stripping(
dst_file,
threshold=threshold,
force=force,
input_path=backup_path,
output_path=hdf5_file_path,
kept_frames=num_frames_to_keep,
stripped_frames=num_frames_to_strip,
)
dst_file.flush()
# On success, delete backup
os.remove(backup_path)
return True
except Exception as e:
log_start(logfile_path, f"Strip aborted: {e}\n{traceback.format_exc()}. Restoring original file.")
# Remove partially modified file
if os.path.exists(hdf5_file_path):
os.remove(hdf5_file_path)
# Restore snapshot
shutil.move(backup_path, hdf5_file_path)
return False
[docs]
def translate_to_stripped_index(hdf5_file_path, unstripped_frame_index):
index_dataset_name = 'entry/data/index'
try:
# Open the HDF5 file and access the index dataset
with h5py.File(hdf5_file_path, 'a', swmr=True) as hdf5_file:
data_group = hdf5_file['entry/data']
if index_dataset_name not in data_group:
_LOGGER.warning(f"Index dataset '{index_dataset_name}' not found.")
return None
index_dataset = data_group[index_dataset_name]
index_length = index_dataset.shape[0]
# Check if the unstripped frame index is within valid range
if unstripped_frame_index >= index_length:
_LOGGER.warning(
f"Unstripped frame index {unstripped_frame_index} is out of bounds (max index: {index_length - 1})."
)
return None
# Retrieve the corresponding stripped index
stripped_index = index_dataset[unstripped_frame_index]
if stripped_index == -1:
_LOGGER.info(f"Frame {unstripped_frame_index} was stripped.")
return None
else:
return stripped_index
except Exception as e:
_LOGGER.error(f"Error accessing HDF5 file: {e}")
return None
[docs]
def strip_h5_batch(input_path, progress_callback=None):
"""
Strip one or more INI-defined datasets. Accepts a list of INI paths (can span
multiple directories) or a single path/dir handled by handle_input.
"""
# If caller already provided an explicit list, use it as-is to allow mixed directories
if isinstance(input_path, list):
configfiles = list(input_path)
else:
# Handle input path and retrieve .ini files (enforces same directory)
configfiles, input_path = handle_input(input_path)
if not configfiles:
_LOGGER.info("No .ini files found to process.")
return
_LOGGER.info("The following .ini files were found and will be processed:")
for configfile in configfiles:
_LOGGER.info(f"- {configfile}")
filecount = 1
for configfile in configfiles:
# Get paths
outputfolder, outputfolder_path, logfile, logfile_path, h5file, h5file_path = config_to_paths(configfile)
# Get config
config = read_config(configfile)
# Check if necessary parameters are defined in the config file
required_params = ['strip_threshold', 'strip_force']
missing_params = [param for param in required_params if not config.has_option('Parameters', param)]
if missing_params:
log_start(logfile_path, f"Error: Missing parameters {', '.join(missing_params)} in config file {configfile}")
continue # Skip to the next config file
# Load parameters
try:
strip_threshold = float(config.get('Parameters', 'strip_threshold'))
strip_force = config.getboolean('Parameters', 'strip_force')
except Exception as e:
log_start(logfile_path, f"Error parsing parameters in config file {configfile}: {e}")
continue # Skip to the next config file
goniometer_transform_order = get_goniometer_transform_order(config)
# Log the starting of the stripping process
log_start(logfile_path, f"Starting strip_h5 with threshold={strip_threshold}, force={strip_force}")
# Call the strip_h5 function
try:
result = strip_h5(
h5file_path,
threshold=strip_threshold,
force=strip_force,
logfile_path=logfile_path,
progress_callback=progress_callback,
goniometer_transform_order=goniometer_transform_order,
)
if result is True:
log_start(logfile_path, "Stripping completed successfully.")
else:
log_start(logfile_path, "Stripping failed.")
except Exception as e:
log_start(logfile_path, f"Error during stripping: {e}")
if filecount < len(configfiles):
_LOGGER.info("")
_LOGGER.info(f"Proceeding to next task (file {filecount}/{len(configfiles)})")
_LOGGER.info("")
filecount += 1
_LOGGER.info(f"Batch processing finished ({filecount-1}/{len(configfiles)})")