from coseda.logging_utils import log_print
from dask import delayed
from dask.diagnostics import ProgressBar
import h5py
import numpy as np
from coseda.initialize import find_configfiles
from coseda.io import handle_input, parse_config
from coseda.logging_utils import log_start, log_result, shoutout
from coseda.dask_client_manager import DaskClientManager
from coseda.nexus.paths import get_mask_dataset
@delayed
def calculate_mean_intensities_batch(h5file_path, indices):
"""Calculate mean intensities for a batch of frames."""
mean_intensities = []
with h5py.File(h5file_path, 'r') as f:
images_dataset = f['entry/data/images']
for frame_idx in indices:
image_data = images_dataset[frame_idx]
mean_intensity = image_data.mean()
mean_intensities.append(mean_intensity)
return mean_intensities
[docs]
def calculate_mean_intensities_dask_file(client, h5file_path, batch_size=100):
with h5py.File(h5file_path, 'r+') as f:
num_images = f['entry/data/images'].shape[0]
num_batches = (num_images + batch_size - 1) // batch_size # Calculate the number of batches
# Create a list to collect all delayed tasks for calculating mean intensities
delayed_tasks = []
for batch_idx in range(num_batches):
start_idx = batch_idx * batch_size
end_idx = min(start_idx + batch_size, num_images)
indices = range(start_idx, end_idx)
# Create a delayed task for each batch
delayed_task = calculate_mean_intensities_batch(h5file_path, indices)
delayed_tasks.append(delayed_task)
# Submit all tasks to Dask and gather results
futures = client.compute(delayed_tasks)
# Display progress bar using ProgressBar
with ProgressBar():
results = client.gather(futures)
# Flatten the results list and write to the HDF5 file
all_mean_intensities = [item for batch_results in results for item in batch_results]
# Write results back to the HDF5 file
with h5py.File(h5file_path, 'r+') as f:
intensities_dataset = 'entry/data/frame_mean_intensities'
# Save the results
if intensities_dataset in f:
f[intensities_dataset][...] = all_mean_intensities
else:
f.create_dataset(intensities_dataset, data=all_mean_intensities)
[docs]
def calculate_mean_intensities_dask_file_inprocess(h5file_path, logfile_path):
client = DaskClientManager.get_client()
dashboard_address, num_workers, threads_per_worker, memory_per_worker = DaskClientManager.get_client_info()
DaskClientManager.log_client_info(logfile_path)
log_print('')
log_print(f'Click here to monitor progress: {dashboard_address}')
log_print('')
log_start(logfile_path, f'Start calculating mean intensities using dask, batch size = default')
try:
calculate_mean_intensities_dask_file(client, h5file_path)
log_start(logfile_path, f'Mean intensity calculation finished')
except Exception as e:
log_result(logfile_path, "Mean intensity calculation finished with errors", str(e))
[docs]
def calculate_mean_intensities_dask_batch(input_path, batch_size=100):
configfiles, input_path = handle_input(input_path)
filecount = 1
client = DaskClientManager.get_client()
dashboard_address, num_workers, threads_per_worker, memory_per_worker = DaskClientManager.get_client_info()
log_print('')
log_print(f'Click here to monitor progress: {dashboard_address}')
log_print('')
for configfile in configfiles:
log_print(f'Processing: {configfile}')
config, outputfolder, originalfile, logfile, path, outputfolder_path, originalfile_path, logfile_path, framepath, h5file, h5file_path = parse_config(configfile)
DaskClientManager.log_client_info(logfile_path)
log_start(logfile_path, f'Start calculating mean intensities using dask, batch size = {batch_size}')
try:
calculate_mean_intensities_dask_file(client, h5file_path, batch_size)
except Exception as e:
log_result(logfile_path, "Mean intensity calculation finished with errors", str(e))
if filecount < len(configfiles):
log_print("")
log_print(f"Proceeding to next task (file {filecount}/{len(configfiles)})")
log_print("")
filecount += 1
log_print(f"Finished intensity calculation ({filecount-1}/{len(configfiles)} file(s))")
# ---------------------------------------------------------------------------
# Radial intensity profiles (one per frame, computed for every frame so they
# survive the later strip step -- frame_radial_intensities is in strip_h5's
# keep-list). Bin r holds the azimuthal SUM of intensity at round(radius) == r
# around the frame's beam center, so summing annuli reproduces virtual
# BF/DF/HAADF signals (same layout as the gridsweeper radial_sum_intensity).
# ---------------------------------------------------------------------------
def _radial_profile_row(image, cx, cy, nbins, X, Y, mask=None):
"""Azimuthal radial sum of one image around (cx, cy), 1-px integer radius bins.
Returns a length-``nbins`` float32 array; masked-out pixels (mask == 0)
contribute 0. ``X``/``Y`` are the reusable pixel-coordinate meshgrids.
"""
R = np.sqrt((X - cx) ** 2 + (Y - cy) ** 2).astype(np.int32)
weights = image.astype(np.float64)
if mask is not None:
weights = weights * mask
r_flat = R.ravel()
w_flat = weights.ravel()
keep = r_flat < nbins
prof = np.bincount(r_flat[keep], weights=w_flat[keep], minlength=nbins)
return prof[:nbins].astype(np.float32)
@delayed
def calculate_radial_profiles_batch(h5file_path, indices, cx, cy, nbins):
"""Radial profiles for a batch of frames -> (start_index, block[n, nbins])."""
indices = list(indices)
with h5py.File(h5file_path, 'r') as f:
images = f['entry/data/images']
H, W = images.shape[1], images.shape[2]
Y, X = np.mgrid[0:H, 0:W]
mask = None
pm = get_mask_dataset(f)
if pm is not None and pm.ndim == 2 and pm.shape == (H, W):
mask = pm[()].astype(np.float32)
block = np.zeros((len(indices), nbins), dtype=np.float32)
for k, idx in enumerate(indices):
block[k] = _radial_profile_row(images[idx], cx[idx], cy[idx], nbins, X, Y, mask)
return indices[0], block
[docs]
def calculate_radial_profiles_dask_file(client, h5file_path, batch_size=64, max_radius=None):
"""Compute a radial intensity profile per frame and save frame_radial_intensities.
Uses the per-frame beam center (entry/data/center_x / center_y) when present,
otherwise the image center. ``max_radius`` (px) caps the number of bins; by
default it reaches the farthest corner from any frame's center.
"""
with h5py.File(h5file_path, 'r') as f:
g = f['entry/data']
n = g['images'].shape[0]
H, W = g['images'].shape[1], g['images'].shape[2]
if 'center_x' in g and 'center_y' in g and g['center_x'].shape[0] == n:
cx = g['center_x'][:].astype(np.float64)
cy = g['center_y'][:].astype(np.float64)
center_source = 'center_x/center_y'
else:
cx = np.full(n, (W - 1) / 2.0)
cy = np.full(n, (H - 1) / 2.0)
center_source = 'image center (center_x/center_y not found)'
# Bins reach the farthest image corner from the *typical* (median) beam center.
# Using the median rather than the per-frame maximum keeps a stray / failed
# center from inflating the profile length (and memory) for the whole file;
# a rare off-center frame simply drops the pixels beyond nbins.
cx_med, cy_med = float(np.median(cx)), float(np.median(cy))
dx = max(cx_med, (W - 1) - cx_med)
dy = max(cy_med, (H - 1) - cy_med)
nbins = int(np.ceil(np.hypot(dx, dy))) + 1
if max_radius is not None:
nbins = min(nbins, int(max_radius) + 1)
log_print(f"Radial profiles: {n} frames x {nbins} bins "
f"(median center {cx_med:.0f},{cy_med:.0f}; source {center_source})")
tasks = [calculate_radial_profiles_batch(h5file_path, range(s, min(s + batch_size, n)), cx, cy, nbins)
for s in range(0, n, batch_size)]
futures = client.compute(tasks)
with ProgressBar():
results = client.gather(futures)
profiles = np.zeros((n, nbins), dtype=np.float32)
for start_idx, block in results:
profiles[start_idx:start_idx + block.shape[0]] = block
with h5py.File(h5file_path, 'r+') as f:
path = 'entry/data/frame_radial_intensities'
if path in f:
del f[path]
ds = f.create_dataset(path, data=profiles)
ds.attrs['center_source'] = center_source
ds.attrs['radial_bin_width_px'] = 1.0
ds.attrs['reduction'] = 'sum'
log_print(f"Saved entry/data/frame_radial_intensities {profiles.shape} to {h5file_path}")
[docs]
def calculate_radial_profiles_dask_file_inprocess(h5file_path, logfile_path, batch_size=64, max_radius=None):
"""In-process entry point mirroring the mean-intensity variant."""
client = DaskClientManager.get_client()
dashboard_address, num_workers, threads_per_worker, memory_per_worker = DaskClientManager.get_client_info()
DaskClientManager.log_client_info(logfile_path)
log_print('')
log_print(f'Click here to monitor progress: {dashboard_address}')
log_print('')
log_start(logfile_path, 'Start calculating radial intensity profiles using dask')
try:
calculate_radial_profiles_dask_file(client, h5file_path, batch_size=batch_size, max_radius=max_radius)
log_start(logfile_path, 'Radial profile calculation finished')
except Exception as e:
log_result(logfile_path, "Radial profile calculation finished with errors", str(e))