diff --git a/notebooks/Jungfrau/Jungfrau_Gain_Correct_and_Verify_NBC.ipynb b/notebooks/Jungfrau/Jungfrau_Gain_Correct_and_Verify_NBC.ipynb
index f7ccbfff05bcdba173017293ee2e0338f8c0986c..5dc2a73250f74a1c1b2e3f3cc73806ba0bac71c1 100644
--- a/notebooks/Jungfrau/Jungfrau_Gain_Correct_and_Verify_NBC.ipynb
+++ b/notebooks/Jungfrau/Jungfrau_Gain_Correct_and_Verify_NBC.ipynb
@@ -38,7 +38,8 @@
"cal_db_timeout = 180000 # timeout on caldb requests\n",
"\n",
"# Parameters affecting corrected data.\n",
- "relative_gain = True # do relative gain correction\n",
+ "relative_gain = True # do relative gain correction.\n",
+ "strixel_sensor = False # reordering for strixel detector layout.\n",
"limit_trains = 0 # ONLY FOR TESTING. process only first N trains, Use 0 to process all.\n",
"chunks_ids = 32 # HDF chunk size for memoryCell and frameNumber.\n",
"chunks_data = 1 # HDF chunk size for pixel data in number of frames.\n",
@@ -208,6 +209,21 @@
"print(f\"Number of memory cells are {memory_cells}\")"
]
},
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "strixel_transform = None\n",
+ "output_frame_shape = None\n",
+ "\n",
+ "if strixel_sensor:\n",
+ " from cal_tools.jfalgs import strixel_transform, strixel_shape\n",
+ " output_frame_shape = strixel_shape()\n",
+ " print('Strixel sensor transformation enabled')"
+ ]
+ },
{
"cell_type": "markdown",
"metadata": {},
@@ -340,8 +356,15 @@
"def correct_train(wid, index, d):\n",
" d = d.astype(np.float32) # [cells, x, y]\n",
" g = gain[index]\n",
+ " \n",
+ " # Copy gain over first to keep it at the original 3 for low gain.\n",
+ " if strixel_transform is not None:\n",
+ " strixel_transform(g, out=gain_corr[index, ...])\n",
+ " else:\n",
+ " gain_corr[index, ...] = g\n",
"\n",
" # Jungfrau gains 0[00], 1[01], 3[11]\n",
+ " # Change low gain to 2 for indexing purposes.\n",
" g[g==3] = 2\n",
"\n",
" # Select memory cells\n",
@@ -380,8 +403,14 @@
"\n",
" msk = np.choose(g, np.moveaxis(mask_cell, -1, 0))\n",
"\n",
- " data_corr[index, ...] = d\n",
- " mask_corr[index, ...] = msk"
+ " if strixel_transform is not None:\n",
+ " strixel_transform(d, out=data_corr[index, ...])\n",
+ " data_corr[index, :, Ydouble, Xdouble] /= strixel_double_norm\n",
+ "\n",
+ " strixel_transform(msk, out=mask_corr[index, ...])\n",
+ " else:\n",
+ " data_corr[index, ...] = d\n",
+ " mask_corr[index, ...] = msk"
]
},
{
@@ -464,8 +493,8 @@
" # and number of available trains to correct.\n",
" seq_dc = H5File(sequence_file)\n",
" seq_dc_adc = seq_dc[instrument_src_kda, \"data.adc\"]\n",
- " dshape = seq_dc_adc.shape\n",
- " corr_ntrains = seq_dc_adc.shape[0] # number of available trains to correct.\n",
+ " ishape = seq_dc_adc.shape # input shape.\n",
+ " corr_ntrains = ishape[0] # number of available trains to correct.\n",
" all_train_ids = seq_dc_adc.train_ids\n",
"\n",
" # Raise a WARNING if this sequence has no trains to correct.\n",
@@ -488,10 +517,17 @@
" # Load constants from the constants dictionary.\n",
" # These arrays are used by `correct_train()` function\n",
" offset_map, mask, gain_map = constants[local_karabo_da]\n",
+ " \n",
+ " # Determine total output shape.\n",
+ " if output_frame_shape is not None:\n",
+ " oshape = (*ishape[:-2], *output_frame_shape)\n",
+ " else:\n",
+ " oshape = ishape\n",
"\n",
" # Allocate shared arrays for corrected data. Used in `correct_train()`\n",
- " data_corr = context.alloc(shape=(corr_ntrains, *dshape[1:]), dtype=np.float32)\n",
- " mask_corr = context.alloc(shape=(corr_ntrains, *dshape[1:]), dtype=np.uint32)\n",
+ " data_corr = context.alloc(shape=oshape, dtype=np.float32)\n",
+ " gain_corr = context.alloc(shape=oshape, dtype=np.uint8)\n",
+ " mask_corr = context.alloc(shape=oshape, dtype=np.uint32)\n",
"\n",
" step_timer.start()\n",
" # Overwrite seq_dc after eliminating empty trains or/and applying limited images.\n",
@@ -551,10 +587,10 @@
" # Add main corrected `data.adc`` dataset and store corrected data.\n",
" outp_source.create_key(\n",
" \"data.adc\", data=data_corr,\n",
- " chunks=(min(chunks_data, data_corr.shape[0]), *dshape[1:]))\n",
+ " chunks=(min(chunks_data, data_corr.shape[0]), *oshape[1:]))\n",
"\n",
" write_compressed_frames(\n",
- " gain, outp_file, f\"{outp_source.name}/data/gain\", comp_threads=8)\n",
+ " gain_corr, outp_file, f\"{outp_source.name}/data/gain\", comp_threads=8)\n",
" write_compressed_frames(\n",
" mask_corr, outp_file, f\"{outp_source.name}/data/mask\", comp_threads=8)\n",
"\n",
@@ -739,14 +775,21 @@
"corrected_mean = np.mean(corrected, axis=1)\n",
"_corrected_vmin = min(0.75*np.median(corrected_mean[corrected_mean > 0]), -0.5)\n",
"_corrected_vmax = max(2.*np.median(corrected_mean[corrected_mean > 0]), 100)\n",
- "geom.plot_data_fast(\n",
- " corrected_mean,\n",
- " ax=ax,\n",
- " vmin=_corrected_vmin,\n",
- " vmax=_corrected_vmax,\n",
- " cmap=\"jet\",\n",
- " colorbar={'shrink': 1, 'pad': 0.01},\n",
+ "\n",
+ "mean_plot_kwargs = dict(\n",
+ " vmin=_corrected_vmin, vmax=_corrected_vmax, cmap=\"jet\"\n",
")\n",
+ "\n",
+ "if not strixel_sensor:\n",
+ " geom.plot_data_fast(\n",
+ " corrected_mean,\n",
+ " ax=ax,\n",
+ " colorbar={'shrink': 1, 'pad': 0.01},\n",
+ " **mean_plot_kwargs\n",
+ " )\n",
+ "else:\n",
+ " ax.imshow(corrected_mean.squeeze(), aspect=10, **mean_plot_kwargs)\n",
+ " \n",
"ax.set_title(f'{karabo_id} - Mean CORRECTED', size=18)\n",
"\n",
"plt.show()"
@@ -763,14 +806,17 @@
"corrected_masked[mask != 0] = np.nan\n",
"corrected_masked_mean = np.nanmean(corrected_masked, axis=1)\n",
"del corrected_masked\n",
- "geom.plot_data_fast(\n",
- " corrected_masked_mean,\n",
- " ax=ax,\n",
- " vmin=_corrected_vmin,\n",
- " vmax=_corrected_vmax,\n",
- " cmap=\"jet\",\n",
- " colorbar={'shrink': 1, 'pad': 0.01},\n",
- ")\n",
+ "\n",
+ "if not strixel_sensor:\n",
+ " geom.plot_data_fast(\n",
+ " corrected_masked_mean,\n",
+ " ax=ax,\n",
+ " colorbar={'shrink': 1, 'pad': 0.01},\n",
+ " **mean_plot_kwargs\n",
+ " )\n",
+ "else:\n",
+ " ax.imshow(corrected_mean.squeeze(), aspect=10, **mean_plot_kwargs)\n",
+ "\n",
"ax.set_title(f'{karabo_id} - Mean CORRECTED with mask', size=18)\n",
"\n",
"plt.show()"
@@ -785,14 +831,23 @@
"display(Markdown((f\"#### A single image from train {tid}\")))\n",
"\n",
"fig, ax = plt.subplots(figsize=(18, 10))\n",
- "geom.plot_data_fast(\n",
- " corrected_train,\n",
- " ax=ax,\n",
+ "\n",
+ "single_plot_kwargs = dict(\n",
" vmin=min(0.75 * np.median(corrected_train[corrected_train > 0]), -0.5),\n",
" vmax=max(2.0 * np.median(corrected_train[corrected_train > 0]), 100),\n",
- " cmap=\"jet\",\n",
- " colorbar={\"shrink\": 1, \"pad\": 0.01},\n",
+ " cmap=\"jet\"\n",
")\n",
+ "\n",
+ "if not strixel_sensor:\n",
+ " geom.plot_data_fast(\n",
+ " corrected_train,\n",
+ " ax=ax,\n",
+ " colorbar={\"shrink\": 1, \"pad\": 0.01},\n",
+ " **single_plot_kwargs\n",
+ " )\n",
+ "else:\n",
+ " ax.imshow(corrected_train.squeeze(), aspect=10, **single_plot_kwargs)\n",
+ "\n",
"ax.set_title(f\"{karabo_id} - CORRECTED train: {tid}\", size=18)\n",
"\n",
"plt.show()"
@@ -951,12 +1006,16 @@
"display(Markdown(f\"#### Bad pixels image for train {tid}\"))\n",
"\n",
"fig, ax = plt.subplots(figsize=(18, 10))\n",
- "geom.plot_data_fast(\n",
- " np.log2(mask_train),\n",
- " ax=ax,\n",
- " vmin=0, vmax=32, cmap=\"jet\",\n",
- " colorbar={'shrink': 1, 'pad': 0.01},\n",
- ")\n",
+ "if not strixel_sensor:\n",
+ " geom.plot_data_fast(\n",
+ " np.log2(mask_train),\n",
+ " ax=ax,\n",
+ " vmin=0, vmax=32, cmap=\"jet\",\n",
+ " colorbar={'shrink': 1, 'pad': 0.01},\n",
+ " )\n",
+ "else:\n",
+ " ax.imshow(np.log2(mask_train).squeeze(), vmin=0, vmax=32, cmap='jet', aspect=10)\n",
+ "\n",
"plt.show()"
]
}