Merge branch 'feat/LPD-max-nodes-parameter' into 'master'

[LPD] [CORRECT] Expose max-nodes parameter See merge request detectors/pycalibration!678

Merge branch 'feat/LPD-max-nodes-parameter' into 'master'
[LPD] [CORRECT] Expose max-nodes parameter See merge request detectors/pycalibration!678
331d7eeb · Philipp Schmidt · 485f6c8d · 1055cbb9 · 331d7eeb
Commit 331d7eeb authored 2 years ago by Philipp Schmidt
--- a/notebooks/LPD/LPD_Correct_Fast.ipynb
+++ b/notebooks/LPD/LPD_Correct_Fast.ipynb
@@ -61,12 +61,13 @@
    "\n",
    "# Parallelization options\n",
    "sequences_per_node = 1  # Sequence files to process per node\n",
+    "max_nodes = 8  # Maximum number of SLURM jobs to split correction work into\n",
    "num_workers = 8  # Worker processes per node, 8 is safe on 768G nodes but won't work on 512G.\n",
    "num_threads_per_worker = 32  # Number of threads per worker.\n",
    "\n",
-    "def balance_sequences(in_folder, run, sequences, sequences_per_node, karabo_da):\n",
+    "def balance_sequences(in_folder, run, sequences, sequences_per_node, karabo_da, max_nodes):\n",
    "    from xfel_calibrate.calibrate import balance_sequences as bs\n",
-    "    return bs(in_folder, run, sequences, sequences_per_node, karabo_da)"
+    "    return bs(in_folder, run, sequences, sequences_per_node, karabo_da, max_nodes=max_nodes)"
   ]
  },
  {

 %% Cell type:markdown id: tags:

 # LPD Offline Correction #

 Author: European XFEL Data Analysis Group

 %% Cell type:code id: tags:

 ``` python
 # Input parameters
 in_folder = "/gpfs/exfel/exp/FXE/202201/p003073/raw/"  # the folder to read data from, required
 out_folder = "/gpfs/exfel/data/scratch/schmidtp/random/LPD_test"  # the folder to output to, required
 metadata_folder = ''  # Directory containing calibration_metadata.yml when run by xfel-calibrate.
 sequences = [-1]  # Sequences to correct, use [-1] for all
 modules = [-1]  # Modules indices to correct, use [-1] for all, only used when karabo_da is empty
 karabo_da = ['']  # Data aggregators names to correct, use [''] for all
 run = 10 # runs to process, required

 # Source parameters
 karabo_id = 'FXE_DET_LPD1M-1'  # Karabo domain for detector.
 input_source = '{karabo_id}/DET/{module_index}CH0:xtdf'  # Input fast data source.
 output_source = ''  # Output fast data source, empty to use same as input.

 # CalCat parameters
 use_dir_creation_date = True  # Use the creation date of the directory for database time derivation.
 cal_db_interface = ''  # Not needed, compatibility with current webservice.
 cal_db_timeout = 0  # Not needed, compatbility with current webservice.
 cal_db_root = '/gpfs/exfel/d/cal/caldb_store'

 # Operating conditions
 mem_cells = 512  # Memory cells, LPD constants are always taken with 512 cells.
 bias_voltage = 250.0  # Detector bias voltage.
 capacitor = '5pF'  # Capacitor setting: 5pF or 50pF
 photon_energy = 9.2  # Photon energy in keV.
 category = 0  # Whom to blame.

 # Correction parameters
 offset_corr = True  # Offset correction.
 rel_gain = True  # Gain correction based on RelativeGain constant.
 ff_map = True  # Gain correction based on FFMap constant.
 gain_amp_map = True  # Gain correction based on GainAmpMap constant.

 # Output options
 overwrite = True  # set to True if existing data should be overwritten
 chunks_data = 1  # HDF chunk size for pixel data in number of frames.
 chunks_ids = 32  # HDF chunk size for cellId and pulseId datasets.
 create_virtual_cxi_in = ''  # Folder to create virtual CXI files in (for each job's chunk)

 # Parallelization options
 sequences_per_node = 1  # Sequence files to process per node
+max_nodes = 8  # Maximum number of SLURM jobs to split correction work into
 num_workers = 8  # Worker processes per node, 8 is safe on 768G nodes but won't work on 512G.
 num_threads_per_worker = 32  # Number of threads per worker.

-def balance_sequences(in_folder, run, sequences, sequences_per_node, karabo_da):
+def balance_sequences(in_folder, run, sequences, sequences_per_node, karabo_da, max_nodes):
    from xfel_calibrate.calibrate import balance_sequences as bs
-    return bs(in_folder, run, sequences, sequences_per_node, karabo_da)
+    return bs(in_folder, run, sequences, sequences_per_node, karabo_da, max_nodes=max_nodes)
 ```

 %% Cell type:code id: tags:

 ``` python
 from collections import OrderedDict
 from pathlib import Path
 from time import perf_counter
 import gc
 import re
 import warnings

 import numpy as np
 import h5py

 import matplotlib
 matplotlib.use('agg')
 import matplotlib.pyplot as plt
 %matplotlib inline

 from calibration_client import CalibrationClient
 from calibration_client.modules import CalibrationConstantVersion
 import extra_data as xd
 import extra_geom as xg
 import pasha as psh

 from extra_data.components import LPD1M

 from cal_tools.lpdalgs import correct_lpd_frames
 from cal_tools.tools import CalibrationMetadata, get_dir_creation_date, write_compressed_frames
 from cal_tools.files import DataFile
 from cal_tools.restful_config import restful_config
 ```

 %% Cell type:markdown id: tags:

 # Prepare environment

 %% Cell type:code id: tags:

 ``` python
 file_re = re.compile(r'^RAW-R(\d{4})-(\w+\d+)-S(\d{5})$')  # This should probably move to cal_tools

 run_folder = Path(in_folder) / f'r{run:04d}'
 out_folder = Path(out_folder)
 out_folder.mkdir(exist_ok=True)

 output_source = output_source or input_source

 cal_db_root = Path(cal_db_root)

 metadata = CalibrationMetadata(metadata_folder or out_folder)

 if use_dir_creation_date:
    creation_time = get_dir_creation_date(in_folder, run)
 else:
    from datetime import datetime
    creation_time = datetime.now()

 print(f'Using {creation_time.isoformat()} as creation time')

 # Pick all modules/aggregators or those selected.
 if not karabo_da or karabo_da == ['']:
    if not modules or modules == [-1]:
        modules = list(range(16))

    karabo_da = [f'LPD{i:02d}' for i in modules]

 # Pick all sequences or those selected.
 if not sequences or sequences == [-1]:
    do_sequence = lambda seq: True
 else:
    do_sequence = [int(x) for x in sequences].__contains__

 # List of detector sources.
 det_inp_sources = [input_source.format(karabo_id=karabo_id, module_index=int(da[-2:])) for da in karabo_da]
 ```

 %% Cell type:markdown id: tags:

 # Select data to process

 %% Cell type:code id: tags:

 ``` python
 data_to_process = []

 for inp_path in run_folder.glob('RAW-*.h5'):
    match = file_re.match(inp_path.stem)

    if match[2] not in karabo_da or not do_sequence(int(match[3])):
        continue

    outp_path = out_folder / 'CORR-R{run:04d}-{aggregator}-S{seq:05d}.h5'.format(
        run=int(match[1]), aggregator=match[2], seq=int(match[3]))

    data_to_process.append((match[2], inp_path, outp_path))

 print('Files to process:')
 for data_descr in sorted(data_to_process, key=lambda x: f'{x[0]}{x[1]}'):
    print(f'{data_descr[0]}\t{data_descr[1]}')
 ```

 %% Cell type:markdown id: tags:

 # Obtain and prepare calibration constants

 %% Cell type:code id: tags:

 ``` python
 # Connect to CalCat.
 calcat_config = restful_config['calcat']
 client = CalibrationClient(
    client_id=calcat_config['user-id'],
    client_secret=calcat_config['user-secret'],
    user_email=calcat_config['user-email'],
    base_api_url=calcat_config['base-api-url'],
    token_url=calcat_config['token-url'],
    refresh_url=calcat_config['refresh-url'],
    auth_url=calcat_config['auth-url'],
    scope='')
 ```

 %% Cell type:code id: tags:

 ``` python
 dark_calibrations = {
    1: 'Offset',  # np.float32
    14: 'BadPixelsDark'  # should be np.uint32, but is np.float64
 }

 dark_condition = [
    dict(parameter_id=1, value=bias_voltage),  # Sensor bias voltage
    dict(parameter_id=7, value=mem_cells),  # Memory cells
    dict(parameter_id=15, value=capacitor),  # Feedback capacitor
    dict(parameter_id=13, value=256),  # Pixels X
    dict(parameter_id=14, value=256),  # Pixels Y
 ]

 illuminated_calibrations = {
    20: 'BadPixelsFF',  # np.uint32
    42: 'GainAmpMap',  # np.float32
    43: 'FFMap',  # np.float32
    44: 'RelativeGain'  # np.float32
 }

 illuminated_condition = dark_condition.copy()
 illuminated_condition += [
    dict(parameter_id=3, value=photon_energy),  # Source energy
    dict(parameter_id=25, value=category)  # category
 ]

 const_data = {}
 const_load_mp = psh.ProcessContext(num_workers=24)

 print('Querying calibration database', end='', flush=True)
 start = perf_counter()
 for calibrations, condition in [
    (dark_calibrations, dark_condition),
    (illuminated_calibrations, illuminated_condition)
 ]:
    resp = CalibrationConstantVersion.get_closest_by_time_by_detector_conditions(
        client, karabo_id, list(calibrations.keys()),
        {'parameters_conditions_attributes': condition},
        karabo_da='', event_at=creation_time.isoformat(), snapshot_at=None)

    if not resp['success']:
        raise RuntimeError(resp)

    for ccv in resp['data']:
        cc = ccv['calibration_constant']
        da = ccv['physical_detector_unit']['karabo_da']
        calibration_name = calibrations[cc['calibration_id']]

        dtype = np.uint32 if calibration_name.startswith('BadPixels') else np.float32

        const_data[(da, calibration_name)] = dict(
            path=Path(ccv['path_to_file']) / ccv['file_name'],
            dataset=ccv['data_set_name'],
            data=const_load_mp.alloc(shape=(256, 256, mem_cells, 3), dtype=dtype)
        )
    print('.', end='', flush=True)

 total_time = perf_counter() - start
 print(f'{total_time:.1f}s')
 ```

 %% Cell type:code id: tags:

 ``` python
 def load_constant_dataset(wid, index, const_descr):
    ccv_entry = const_data[const_descr]

    with h5py.File(cal_db_root / ccv_entry['path'], 'r') as fp:
        fp[ccv_entry['dataset'] + '/data'].read_direct(ccv_entry['data'])

    print('.', end='', flush=True)

 print('Loading calibration data', end='', flush=True)
 start = perf_counter()
 const_load_mp.map(load_constant_dataset, list(const_data.keys()))
 total_time = perf_counter() - start

 print(f'{total_time:.1f}s')
 ```

 %% Cell type:code id: tags:

 ``` python
 # These are intended in order cell, X, Y, gain
 ccv_offsets = {}
 ccv_gains = {}
 ccv_masks = {}

 ccv_shape = (mem_cells, 256, 256, 3)

 constant_order = {
    'Offset':        (2, 1, 0, 3),
    'BadPixelsDark': (2, 1, 0, 3),
    'RelativeGain':  (2, 1, 0, 3),
    'FFMap':         (2, 0, 1, 3),
    'BadPixelsFF':   (2, 0, 1, 3),
    'GainAmpMap':    (2, 0, 1, 3),
 }

 def prepare_constants(wid, index, aggregator):
    consts = {calibration_name: entry['data']
              for (aggregator_, calibration_name), entry
              in const_data.items()
              if aggregator == aggregator_}

    def _prepare_data(calibration_name, dtype):
        return consts[calibration_name] \
            .transpose(constant_order[calibration_name]) \
            .astype(dtype, copy=True)  # Make sure array is contiguous.

    if offset_corr and 'Offset' in consts:
        ccv_offsets[aggregator] = _prepare_data('Offset', np.float32)
    else:
        ccv_offsets[aggregator] = np.zeros(ccv_shape, dtype=np.float32)

    ccv_gains[aggregator] = np.ones(ccv_shape, dtype=np.float32)

    if 'BadPixelsDark' in consts:
        ccv_masks[aggregator] = _prepare_data('BadPixelsDark', np.uint32)
    else:
        ccv_masks[aggregator] = np.zeros(ccv_shape, dtype=np.uint32)

    if rel_gain and 'RelativeGain' in consts:
        ccv_gains[aggregator] *= _prepare_data('RelativeGain', np.float32)

    if ff_map and 'FFMap' in consts:
        ccv_gains[aggregator] *= _prepare_data('FFMap', np.float32)

        if 'BadPixelsFF' in consts:
            np.bitwise_or(ccv_masks[aggregator], _prepare_data('BadPixelsFF', np.uint32),
                          out=ccv_masks[aggregator])

    if gain_amp_map and 'GainAmpMap' in consts:
        ccv_gains[aggregator] *= _prepare_data('GainAmpMap', np.float32)

    print('.', end='', flush=True)


 print('Preparing constants', end='', flush=True)
 start = perf_counter()
 psh.ThreadContext(num_workers=len(karabo_da)).map(prepare_constants, karabo_da)
 total_time = perf_counter() - start
 print(f'{total_time:.1f}s')

 const_data.clear()  # Clear raw constants data now to save memory.
 gc.collect();
 ```

 %% Cell type:code id: tags:

 ``` python
 def correct_file(wid, index, work):
    aggregator, inp_path, outp_path = work
    module_index = int(aggregator[-2:])

    start = perf_counter()
    dc = xd.H5File(inp_path).select('*', 'image.*', require_all=True)
    inp_source = dc[input_source.format(karabo_id=karabo_id, module_index=module_index)]
    open_time = perf_counter() - start

    # Load raw data for this file.
    start = perf_counter()
    in_raw = inp_source['image.data'].ndarray().squeeze()
    in_cell = inp_source['image.cellId'].ndarray().squeeze()
    in_pulse = inp_source['image.pulseId'].ndarray().squeeze()
    read_time = perf_counter() - start

    # Allocate output arrays.
    out_data = np.zeros((in_raw.shape[0], 256, 256), dtype=np.float32)
    out_gain = np.zeros((in_raw.shape[0], 256, 256), dtype=np.uint8)
    out_mask = np.zeros((in_raw.shape[0], 256, 256), dtype=np.uint32)

    start = perf_counter()
    correct_lpd_frames(in_raw, in_cell,
                       out_data, out_gain, out_mask,
                       ccv_offsets[aggregator], ccv_gains[aggregator], ccv_masks[aggregator],
                       num_threads=num_threads_per_worker)
    correct_time = perf_counter() - start

    image_counts = inp_source['image.data'].data_counts(labelled=False)

    start = perf_counter()
    if (not outp_path.exists() or overwrite) and image_counts.sum() > 0:
        fa = dc.files[0]
        sel_trains = np.isin(fa.train_ids, dc.train_ids)

        outp_source_name = output_source.format(karabo_id=karabo_id, module_index=module_index)

        with DataFile(outp_path, 'w') as outp_file:
            outp_file.create_index(
                train_ids=dc.train_ids,
                timestamps=fa.file['INDEX/timestamp'][sel_trains],
                flags=fa.validity_flag[sel_trains])
            outp_file.create_metadata(like=dc, instrument_channels=(f'{outp_source_name}/image',))

            outp_source = outp_file.create_instrument_source(outp_source_name)

            outp_source.create_index(image=image_counts)
            outp_source.create_key('image.cellId', data=in_cell,
                                   chunks=(chunks_ids,))
            outp_source.create_key('image.pulseId', data=in_pulse,
                                   chunks=(chunks_ids,))
            outp_source.create_key('image.data', data=out_data,
                                   chunks=(chunks_data, 256, 256))
            write_compressed_frames(
                out_gain, outp_file, f'INSTRUMENT/{outp_source_name}/image/gain', comp_threads=8)
            write_compressed_frames(
                out_mask, outp_file, f'INSTRUMENT/{outp_source_name}/image/mask', comp_threads=8)
    write_time = perf_counter() - start

    total_time = open_time + read_time + correct_time + write_time
    frame_rate = in_raw.shape[0] / total_time

    print('{}\t{}\t{:.3f}\t{:.3f}\t{:.3f}\t{:.3f}\t{:.3f}\t{}\t{:.1f}'.format(
        wid, aggregator, open_time, read_time, correct_time, write_time, total_time,
        in_raw.shape[0], frame_rate))

    in_raw = None
    in_cell = None
    in_pulse = None
    out_data = None
    out_gain = None
    out_mask = None
    gc.collect()

 print('worker\tDA\topen\tread\tcorrect\twrite\ttotal\tframes\trate')
 start = perf_counter()
 psh.ProcessContext(num_workers=num_workers).map(correct_file, data_to_process)
 total_time = perf_counter() - start
 print(f'Total time: {total_time:.1f}s')
 ```

 %% Cell type:markdown id: tags:

 # Data preview for first train

 %% Cell type:code id: tags:

 ``` python
 geom = xg.LPD_1MGeometry.from_quad_positions(
    [(11.4, 299), (-11.5, 8), (254.5, -16), (278.5, 275)])

 output_paths = [outp_path for _, _, outp_path in data_to_process if outp_path.exists()]
 dc = xd.DataCollection.from_paths(output_paths).select_trains(np.s_[0])

 det = LPD1M(dc, detector_name=karabo_id)
 data = det.get_array('image.data')
 ```

 %% Cell type:markdown id: tags:

 ### Intensity histogram across all cells

 %% Cell type:code id: tags:

 ``` python
 left_edge_ratio = 0.01
 right_edge_ratio = 0.99

 fig, ax = plt.subplots(num=1, clear=True, figsize=(15, 6))
 values, bins, _ = ax.hist(np.ravel(data.data), bins=2000, range=(-1500, 2000))

 def find_nearest_index(array, value):
    return (np.abs(array - value)).argmin()

 cum_values = np.cumsum(values)
 vmin = bins[find_nearest_index(cum_values, cum_values[-1]*left_edge_ratio)]
 vmax = bins[find_nearest_index(cum_values, cum_values[-1]*right_edge_ratio)]

 max_value = values.max()
 ax.vlines([vmin, vmax], 0, max_value, color='red', linewidth=5, alpha=0.2)
 ax.text(vmin, max_value, f'{left_edge_ratio*100:.0f}%',
        color='red', ha='center', va='bottom', size='large')
 ax.text(vmax, max_value, f'{right_edge_ratio*100:.0f}%',
        color='red', ha='center', va='bottom', size='large')
 ax.text(vmax+(vmax-vmin)*0.01, max_value/2, 'Colormap interval',
        color='red', rotation=90, ha='left', va='center', size='x-large')

 ax.set_xlim(vmin-(vmax-vmin)*0.1, vmax+(vmax-vmin)*0.1)
 ax.set_ylim(0, max_value*1.1)
 pass
 ```

 %% Cell type:markdown id: tags:

 ### First memory cell

 %% Cell type:code id: tags:

 ``` python
 fig, ax = plt.subplots(num=2, figsize=(15, 15), clear=True, nrows=1, ncols=1)
 geom.plot_data_fast(data[:, 0, 0], ax=ax, vmin=vmin, vmax=vmax)
 pass
 ```

 %% Cell type:markdown id: tags:

 ### Train average

 %% Cell type:code id: tags:

 ``` python
 fig, ax = plt.subplots(num=3, figsize=(15, 15), clear=True, nrows=1, ncols=1)
 geom.plot_data_fast(data[:, 0].mean(axis=1), ax=ax, vmin=vmin, vmax=vmax)
 pass
 ```

 %% Cell type:markdown id: tags:

 ### Lowest gain stage per pixel

 %% Cell type:code id: tags:

 ``` python
 highest_gain_stage = det.get_array('image.gain', pulses=np.s_[:]).max(axis=(1, 2))

 fig, ax = plt.subplots(num=4, figsize=(15, 15), clear=True, nrows=1, ncols=1)
 p = geom.plot_data_fast(highest_gain_stage, ax=ax, vmin=0, vmax=2);

 cb = ax.images[0].colorbar
 cb.set_ticks([0, 1, 2])
 cb.set_ticklabels(['High gain', 'Medium gain', 'Low gain'])
 ```

 %% Cell type:markdown id: tags:

 ### Create virtual CXI file

 %% Cell type:code id: tags:

 ``` python
 if create_virtual_cxi_in:
    vcxi_folder = Path(create_virtual_cxi_in)
    vcxi_folder.mkdir(parents=True, exist_ok=True)

    if not sequences or sequences == [-1]:
        seqs = 'all'
    else:
        seqs = '_'.join(str(s) for s in sequences)

    det.write_virtual_cxi(vcxi_folder / f'r{run}_seqs_{seqs}.cxi')
 ```