make sure gain constant is converted to float32

6a963304 · Karim Ahmed · 8a5086bd · 6a963304
Commit 6a963304 authored 2 years ago by Karim Ahmed
--- a/notebooks/Gotthard2/Correction_Gotthard2_NBC.ipynb
+++ b/notebooks/Gotthard2/Correction_Gotthard2_NBC.ipynb
@@ -326,7 +326,7 @@
    "        mask[index, ...],\n",
    "        g,\n",
    "        const_data[mod][\"Offset\"],\n",
-    "        const_data[mod][\"RelativeGain\"],\n",
+    "        const_data[mod][\"RelativeGain\"].astype(np.float32),\n",
    "        const_data[mod][\"Mask\"],\n",
    "        apply_offset=offset_correction,\n",
    "        apply_gain=gain_correction,\n",

 %% Cell type:markdown id:bed7bd15-21d9-4735-82c1-c27c1a5e3346 tags:

 # Gotthard2 Offline Correction #

 Author: European XFEL Detector Group, Version: 1.0

 Offline Calibration for the Gothard2 Detector

 %% Cell type:code id:570322ed-f611-4fd1-b2ec-c12c13d55843 tags:

 ``` python
 in_folder = "/gpfs/exfel/exp/FXE/202221/p003225/raw"  # the folder to read data from, required
 out_folder = "/gpfs/exfel/data/scratch/ahmedk/test/gotthard2"  # the folder to output to, required
 metadata_folder = ""  # Directory containing calibration_metadata.yml when run by xfel-calibrate
 run = 50  # run to process, required
 sequences = [-1]  # sequences to correct, set to [-1] for all, range allowed
 sequences_per_node = 1  # number of sequence files per node if notebook executed through xfel-calibrate, set to 0 to not run SLURM parallel

 # Parameters used to access raw data.
 karabo_id = "FXE_XAD_G2XES"  # karabo prefix of Gotthard-II devices
 karabo_da = ["GH201"]  # data aggregators
 receiver_template = "RECEIVER"  # receiver template used to read INSTRUMENT keys.
 control_template = "CONTROL"  # control template used to read CONTROL keys.
 instrument_source_template = "{}/DET/{}:daqOutput"  # template for source name (filled with karabo_id & receiver_id). e.g. 'SPB_IRDA_JF4M/DET/JNGFR01:daqOutput'
 ctrl_source_template = "{}/DET/{}"  # template for control source name (filled with karabo_id_control)
 karabo_id_control = ""  # Control karabo ID. Set to empty string to use the karabo-id

 # Parameters for calibration database.
 use_dir_creation_date = True  # use the creation data of the input dir for database queries.
 cal_db_interface = "tcp://max-exfl016:8016#8025"  # the database interface to use.
 cal_db_timeout = 180000  # timeout on caldb requests.
 overwrite_creation_time = ""  # To overwrite the measured creation_time. Required Format: YYYY-MM-DD HR:MN:SC.00 e.g. "2022-06-28 13:00:00.00"

 # Parameters affecting corrected data.
 constants_file = ""  # Use constants in given constant file path. /gpfs/exfel/data/scratch/ahmedk/dont_remove/gotthard2/constants/calibration_constants_GH2.h5
 offset_correction = True  # apply offset correction. This can be disabled to only apply LUT or apply LUT and gain correction for non-linear differential results.
 gain_correction = True  # apply gain correction.

 # Parameter conditions.
 bias_voltage = -1  # Detector bias voltage, set to -1 to use value in raw file.
 exposure_time = -1.  # Detector exposure time, set to -1 to use value in raw file.
 exposure_period = -1.  # Detector exposure period, set to -1 to use value in raw file.
 acquisition_rate = -1.  # Detector acquisition rate (1.1/4.5), set to -1 to use value in raw file.
 single_photon = -1  # Detector single photon mode (High/Low CDS), set to -1 to use value in raw file.

 # Parameters for plotting
 skip_plots = False  # exit after writing corrected files
 pulse_idx_preview = 3  # pulse index to preview. The following even/odd pulse index is used for preview. # TODO: update to pulseId preview.


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

 %% Cell type:code id:6e9730d8-3908-41d7-abe2-d78e046d5de2 tags:

 ``` python
 import datetime
 import warnings
 from functools import partial

 import h5py
 import pasha as psh
 import numpy as np
 import matplotlib.pyplot as plt
 from IPython.display import Markdown, display
 from extra_data import RunDirectory, H5File
 from pathlib import Path

 from cal_tools import h5_copy_except
 from cal_tools.gotthard2 import gotthard2algs, gotthard2lib
 from cal_tools.step_timing import StepTimer
 from cal_tools.tools import (
    get_constant_from_db_and_time,
    get_dir_creation_date,
    get_pdu_from_db,
    CalibrationMetadata,
 )
 from iCalibrationDB import Conditions, Constants
 from XFELDetAna.plotting.heatmap import heatmapPlot

 warnings.filterwarnings('ignore')

 %matplotlib inline
 ```

 %% Cell type:code id:d7c02c48-4429-42ea-a42e-de45366d7fa3 tags:

 ``` python
 in_folder = Path(in_folder)
 run_folder = in_folder / f"r{run:04d}"
 out_folder = Path(out_folder)
 out_folder.mkdir(parents=True, exist_ok=True)

 metadata = CalibrationMetadata(metadata_folder or out_folder)
 # NOTE: this notebook will not overwrite calibration metadata file
 const_yaml = metadata.get("retrieved-constants", {})

 if not karabo_id_control:
    karabo_id_control = karabo_id

 instrument_src = instrument_source_template.format(karabo_id, receiver_template)
 ctrl_src = ctrl_source_template.format(karabo_id_control, control_template)

 print(f"Process modules: {karabo_da} for run {run}")

 creation_time = None
 if overwrite_creation_time:
    creation_time = datetime.datetime.strptime(
        overwrite_creation_time, "%Y-%m-%d %H:%M:%S.%f"
    )
 elif use_dir_creation_date:
    creation_time = get_dir_creation_date(in_folder, run)
    print(f"Using {creation_time} as creation time")
 ```

 %% Cell type:code id:b5eb816e-b5f2-44ce-9907-0273d82341b6 tags:

 ``` python
 # Select only sequence files to process for the selected detector.
 if sequences == [-1]:
    possible_patterns = list(f"*{mod}*.h5" for mod in karabo_da)
 else:
    possible_patterns = list(
        f"*{mod}-S{s:05d}.h5" for mod in karabo_da for s in sequences
    )

 run_folder = Path(in_folder / f"r{run:04d}")
 seq_files = [
    f for f in run_folder.glob("*.h5") if any(f.match(p) for p in possible_patterns)
 ]

 seq_files = sorted(seq_files)

 if not seq_files:
    raise IndexError("No sequence files available for the selected sequences.")

 print(f"Processing a total of {len(seq_files)} sequence files")
 ```

 %% Cell type:code id:f9a8d1eb-ce6a-4ed0-abf4-4a6029734672 tags:

 ``` python
 step_timer = StepTimer()
 ```

 %% Cell type:code id:892172d8 tags:

 ``` python
 # Read slow data
 run_dc = RunDirectory(run_folder)
 g2ctrl = gotthard2lib.Gotthard2Ctrl(run_dc=run_dc, ctrl_src=ctrl_src)

 if bias_voltage == -1:
    bias_voltage = g2ctrl.get_bias_voltage()
 if exposure_time == -1:
    exposure_time = g2ctrl.get_exposure_time()
 if exposure_period == -1:
    exposure_period = g2ctrl.get_exposure_period()
 if acquisition_rate == -1:
    acquisition_rate = g2ctrl.get_acquisition_rate()
 if single_photon == -1:
    single_photon = g2ctrl.get_single_photon()

 print("Bias Voltage:", bias_voltage)
 print("Exposure Time:", exposure_time)
 print("Exposure Period:", exposure_period)
 print("Acquisition Rate:", acquisition_rate)
 print("Single Photon:", single_photon)
 ```

 %% Cell type:markdown id:8c852392-bb19-4c40-b2ce-3b787538a92d tags:

 ### Retrieving calibration constants

 %% Cell type:code id:5717d722 tags:

 ``` python
 # Used for old FXE (p003225) runs before adding Gotthard2 to CALCAT
 const_data = dict()
 if constants_file:
    for mod in karabo_da:
        const_data[mod] = dict()
        # load constants temporarily using defined local paths.
        with h5py.File(constants_file, "r") as cfile:
            const_data[mod]["LUT"] = cfile["LUT"][()]
            const_data[mod]["Offset"] = cfile["offset_map"][()].astype(np.float32)
            const_data[mod]["RelativeGain"] = cfile["gain_map"][()].astype(np.float32)
            const_data[mod]["Mask"] = cfile["bpix_ff"][()].astype(np.uint32)
 ```

 %% Cell type:code id:1cdbe818 tags:

 ``` python
 # Conditions iCalibrationDB object.
 condition = Conditions.Dark.Gotthard2(
    bias_voltage=bias_voltage,
    exposure_time=exposure_time,
    exposure_period=exposure_period,
    single_photon=single_photon,
    acquisition_rate=acquisition_rate,
 )

 # TODO: Maybe this condition and previous cell can be removed later after the initial phase.
 if not constants_file:
    # Prepare a dictionary of empty constants to loop on
    # it's keys and initiate non-retrieved constants.
    empty_lut = (np.arange(2 ** 12).astype(np.float64) * 2 ** 10 / 2 ** 12).astype(
        np.uint16
    )
    empty_lut = np.stack(1280 * [np.stack([empty_lut] * 2)], axis=0)
    empty_constants = {
        "LUT": empty_lut,
        "Offset": np.zeros((1280, 2, 3), dtype=np.float32),
        "BadPixelsDark": np.zeros((1280, 2, 3), dtype=np.uint32),
        "RelativeGain": np.ones((1280, 2, 3), dtype=np.float32),
        "BadPixelsFF": np.zeros((1280, 2, 3), dtype=np.uint32),
    }

    for mod in karabo_da:
        const_data[mod] = dict()
        # Only used for printing timestamps within the loop.
        when = dict()
        # Check YAML file for constant metadata of file path and creation-time
        if const_yaml:
            for cname, mdata in const_yaml[mod]["constants"].items():
                const_data[mod][cname] = dict()
                when[cname] = mdata["creation-time"]
                if when[cname]:
                    with h5py.File(mdata["file-path"], "r") as cf:
                        const_data[mod][cname] = np.copy(
                            cf[f"{mdata['dataset-name']}/data"]
                        )
                else:
                    const_data[mod][cname] = empty_constants[cname]
        else:  # Retrieve constants from CALCAT. Missing YAML file or running notebook interactively.
            for cname, cempty in empty_constants.items():
                const_data[mod][cname] = dict()
                const_data[mod][cname], when[cname] = get_constant_from_db_and_time(
                    karabo_id=karabo_id,
                    karabo_da=mod,
                    cal_db_interface=cal_db_interface,
                    creation_time=creation_time,
                    timeout=cal_db_timeout,
                    print_once=False,
                    condition=condition,
                    constant=getattr(Constants.Gotthard2, cname)(),
                    empty_constant=cempty,
                )
        bpix = const_data[mod]["BadPixelsDark"]
        bpix |= const_data[mod]["BadPixelsFF"]
        const_data[mod]["Mask"] = bpix

        # Print timestamps for the retrieved constants.
        print(f"Constants for module {mod}:")
        for cname, ctime in when.items():
            print(f"  {cname} injected at {ctime}")
        del when
 ```

 %% Cell type:code id:23fcf7f4-351a-4df7-8829-d8497d94fecc tags:

 ``` python
 context = psh.ProcessContext(num_workers=23)
 ```

 %% Cell type:code id:daecd662-26d2-4cb8-aa70-383a579cf9f9 tags:

 ``` python
 def correct_train(wid, index, d):
    g = gain[index]
    gotthard2algs.convert_to_10bit(d, const_data[mod]["LUT"], data_corr[index, ...])
    gotthard2algs.correct_train(
        data_corr[index, ...],
        mask[index, ...],
        g,
        const_data[mod]["Offset"],
-        const_data[mod]["RelativeGain"],
+        const_data[mod]["RelativeGain"].astype(np.float32),
        const_data[mod]["Mask"],
        apply_offset=offset_correction,
        apply_gain=gain_correction,
    )
 ```

 %% Cell type:code id:f88c1aa6-a735-4b72-adce-b30162f5daea tags:

 ``` python
 for mod in karabo_da:
    # This is used in case receiver template consists of
    # karabo data aggregator index. e.g. detector at DETLAB
    instr_mod_src = instrument_src.format(mod[-2:])
    data_path = "INSTRUMENT/" + instr_mod_src + "/data"
    for raw_file in seq_files:
        step_timer.start()

        dc = H5File(raw_file)
        out_file = out_folder / raw_file.name.replace("RAW", "CORR")

        # Select module INSTRUMENT source and deselect empty trains.
        dc = dc.select(instr_mod_src, require_all=True)
        data = dc[instr_mod_src, "data.adc"].ndarray()
        gain = dc[instr_mod_src, "data.gain"].ndarray()
        step_timer.done_step("preparing raw data")
        dshape = data.shape

        step_timer.start()

        # Allocate shared arrays.
        data_corr = context.alloc(shape=dshape, dtype=np.float32)
        mask = context.alloc(shape=dshape, dtype=np.uint32)
        context.map(correct_train, data)
        step_timer.done_step("Correcting one sequence file")

        step_timer.start()

        # Provided PSI gain map has 0 values. Set inf values to nan.
        # TODO: This can maybe be removed after creating XFEL gain maps.?
        data_corr[np.isinf(data_corr)] = np.nan
        # Create CORR files and add corrected data sources.
        # Exclude raw data images (data/adc)
        with h5py.File(out_file, "w") as ofile:
            # Copy RAW non-calibrated sources.
            with h5py.File(raw_file, "r") as sfile:
                h5_copy_except.h5_copy_except_paths(sfile, ofile, [f"{data_path}/adc"])
            # Create datasets with the available corrected data
            ddset = ofile.create_dataset(
                f"{data_path}/adc",
                data=data_corr,
                chunks=((1,) + dshape[1:]),  # 1 chunk == 1 image
                dtype=np.float32,
            )
            # Create datasets with the available corrected data
            ddset = ofile.create_dataset(
                f"{data_path}/mask",
                data=mask,
                chunks=((1,) + dshape[1:]),  # 1 chunk == 1 image
                dtype=np.uint32,
                compression="gzip",
                compression_opts=1,
                shuffle=True,
            )
        step_timer.done_step("Storing data")
 ```

 %% Cell type:code id:94b8e4d2-9f8c-4c23-a509-39238dd8435c tags:

 ``` python
 print(f"Total processing time {step_timer.timespan():.01f} s")
 step_timer.print_summary()
 ```

 %% Cell type:code id:0ccc7f7e-2a3f-4ac0-b854-7d505410d2fd tags:

 ``` python
 if skip_plots:
    print("Skipping plots")
    import sys

    sys.exit(0)
 ```

 %% Cell type:code id:ff203f77-3811-46f3-bf7d-226d2dcab13f tags:

 ``` python
 mod_dcs = {}
 first_seq_raw = seq_files[0]
 first_seq_corr = out_folder / first_seq_raw.name.replace("RAW", "CORR")
 for mod in karabo_da:
    mod_dcs[mod] = {}
    with H5File(first_seq_corr) as out_dc:
        tid, mod_dcs[mod]["train_corr_data"] = next(
            out_dc[instr_mod_src, "data.adc"].trains()
        )
    with H5File(first_seq_raw) as in_dc:
        train_dict = in_dc.train_from_id(tid)[1][instr_mod_src]
        mod_dcs[mod]["train_raw_data"] = train_dict["data.adc"]
        mod_dcs[mod]["train_raw_gain"] = train_dict["data.gain"]
 ```

 %% Cell type:code id:494b453a tags:

 ``` python
 # Keep as long as it is essential to correct
 # RAW data (FXE p003225) before the data mapping was added to CALCAT.
 try:
    db_modules = get_pdu_from_db(
        karabo_id=karabo_id,
        karabo_da=karabo_da,
        constant=Constants.jungfrau.Offset(),
        condition=condition,
        cal_db_interface=cal_db_interface,
        snapshot_at=creation_time,
    )
 except RuntimeError:
    print(
        "No Physical detector units found for this"
        " detector mapping at the RAW data creation time."
    )
    db_modules = [None] * len(karabo_da)
 ```

 %% Cell type:code id:1b379438-eb1d-42b2-ac83-eb8cf88c46db tags:

 ``` python
 display(Markdown("### Mean RAW and CORRECTED across pulses for one train:"))
 display(Markdown(f"Train: {tid}"))

 step_timer.start()
 for mod, pdu in zip(karabo_da, db_modules):

    fig, ax = plt.subplots(figsize=(20, 10))
    raw_data = mod_dcs[mod]["train_raw_data"]
    im = ax.plot(np.mean(raw_data, axis=0))
    ax.set_title(f"RAW module {mod} ({pdu})")
    ax.set_xlabel("Strip #", size=20)
    ax.set_ylabel("12-bit ADC output", size=20)
    plt.xticks(fontsize=20)
    plt.yticks(fontsize=20)
    pass

    fig, ax = plt.subplots(figsize=(20, 10))
    corr_data = mod_dcs[mod]["train_corr_data"]
    im = ax.plot(np.mean(corr_data, axis=0))
    ax.set_title(f"CORRECTED module {mod} ({pdu})")
    ax.set_xlabel("Strip #", size=20)
    ax.set_ylabel("10-bit KeV. output", size=20)
    plt.xticks(fontsize=20)
    plt.yticks(fontsize=20)
    pass
 step_timer.done_step("Plotting mean data")
 ```

 %% Cell type:code id:58a6a276 tags:

 ``` python
 display(Markdown(f"### RAW and CORRECTED strips across pulses for train {tid}"))

 step_timer.start()
 for mod, pdu in zip(karabo_da, db_modules):
    for plt_data, dname in zip(
        ["train_raw_data", "train_corr_data"], ["RAW", "CORRECTED"]
    ):
        fig, ax = plt.subplots(figsize=(15, 20))
        plt.rcParams.update({"font.size": 20})

        heatmapPlot(
            mod_dcs[mod][plt_data],
            y_label="Pulses",
            x_label="Strips",
            title=f"{dname} module {mod} ({pdu})",
            use_axis=ax,
        )
        pass
 step_timer.done_step("Plotting RAW and CORRECTED data for one train")
 ```

 %% Cell type:code id:cd8f5e08-fcee-4bff-ba63-6452b3d892a2 tags:

 ``` python
 # Validate given "pulse_idx_preview"

 if pulse_idx_preview + 1 > data.shape[1]:
    print(
        f"WARNING: selected pulse_idx_preview {pulse_idx_preview} is not available in data."
        " Previewing 1st pulse."
    )
    pulse_idx_preview = 1

 if data.shape[1] == 1:
    odd_pulse = 1
    even_pulse = None
 else:
    odd_pulse = pulse_idx_preview if pulse_idx_preview % 2 else pulse_idx_preview + 1
    even_pulse = (
        pulse_idx_preview if not (pulse_idx_preview % 2) else pulse_idx_preview + 1
    )

 if pulse_idx_preview + 1 > data.shape[1]:
    pulse_idx_preview = 1
    if data.shape[1] > 1:
        pulse_idx_preview = 2
 ```

 %% Cell type:code id:e5f0d4d8-e32c-4f2c-8469-4ebbfd3f644c tags:

 ``` python
 display(Markdown("### RAW and CORRECTED even/odd pulses for one train:"))
 display(Markdown(f"Train: {tid}"))
 for mod, pdu in zip(karabo_da, db_modules):
    fig, ax = plt.subplots(figsize=(20, 20))
    raw_data = mod_dcs[mod]["train_raw_data"]
    corr_data = mod_dcs[mod]["train_corr_data"]

    ax.plot(raw_data[odd_pulse], label=f"Odd Pulse {odd_pulse}")
    if even_pulse:
        ax.plot(raw_data[even_pulse], label=f"Even Pulse {even_pulse}")

    ax.set_title(f"RAW module {mod} ({pdu})")
    ax.set_xlabel("Strip #", size=20)
    ax.set_ylabel("12-bit ADC RAW", size=20)
    plt.xticks(fontsize=20)
    plt.yticks(fontsize=20)
    ax.legend()
    pass

    fig, ax = plt.subplots(figsize=(20, 20))
    ax.plot(corr_data[odd_pulse], label=f"Odd Pulse {odd_pulse}")
    if even_pulse:
        ax.plot(corr_data[even_pulse], label=f"Even Pulse {even_pulse}")
    ax.set_title(f"CORRECTED module {mod} ({pdu})")
    ax.set_xlabel("Strip #", size=20)
    ax.set_ylabel("10-bit KeV CORRECTED", size=20)
    plt.xticks(fontsize=20)
    plt.yticks(fontsize=20)
    ax.legend()
    pass
 step_timer.done_step("Plotting RAW and CORRECTED odd/even pulses.")
 ```