From 2b1b47f3d169e0532ae275cfa7efea756b20e0a3 Mon Sep 17 00:00:00 2001
From: ahmedk <karim.ahmed@xfel.eu>
Date: Wed, 14 Dec 2022 11:20:03 +0100
Subject: [PATCH] LPD calcat interface

---
 notebooks/LPD/LPD_Correct_Fast.ipynb          | 194 +++++++-----------
 ...LPD_retrieve_constants_precorrection.ipynb | 146 ++++++-------
 2 files changed, 140 insertions(+), 200 deletions(-)

diff --git a/notebooks/LPD/LPD_Correct_Fast.ipynb b/notebooks/LPD/LPD_Correct_Fast.ipynb
index 184d6ce86..175d1b3ef 100644
--- a/notebooks/LPD/LPD_Correct_Fast.ipynb
+++ b/notebooks/LPD/LPD_Correct_Fast.ipynb
@@ -108,6 +108,7 @@
     "\n",
     "from extra_data.components import LPD1M\n",
     "\n",
+    "from cal_tools.calcat_interface import LPD_CalibrationData\n",
     "from cal_tools.lpdalgs import correct_lpd_frames\n",
     "from cal_tools.lpdlib import get_mem_cell_pattern, make_cell_order_condition\n",
     "from cal_tools.tools import CalibrationMetadata, calcat_creation_time\n",
@@ -139,16 +140,19 @@
     "cal_db_root = Path(cal_db_root)\n",
     "\n",
     "metadata = CalibrationMetadata(metadata_folder or out_folder)\n",
+    "# Constant paths & timestamps are saved under retrieved-constants in calibration_metadata.yml\n",
+    "retrieved_constants = metadata.setdefault(\"retrieved-constants\", {})\n",
     "\n",
     "creation_time = calcat_creation_time(in_folder, run, creation_time)\n",
     "print(f'Using {creation_time.isoformat()} as creation time')\n",
     "\n",
     "# Pick all modules/aggregators or those selected.\n",
-    "if not karabo_da or karabo_da == ['']:\n",
-    "    if not modules or modules == [-1]:\n",
+    "if karabo_da == ['']:\n",
+    "    if modules == [-1]:\n",
     "        modules = list(range(16))\n",
-    "\n",
     "    karabo_da = [f'LPD{i:02d}' for i in modules]\n",
+    "else:\n",
+    "    modules = [int(x[-2:]) for x in karabo_da]\n",
     "    \n",
     "# Pick all sequences or those selected.\n",
     "if not sequences or sequences == [-1]:\n",
@@ -236,6 +240,17 @@
     "# Obtain and prepare calibration constants"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "metadata = CalibrationMetadata(metadata_folder or out_folder)\n",
+    "# Constant paths & timestamps are saved under retrieved-constants in calibration_metadata.yml\n",
+    "const_yaml = metadata.setdefault(\"retrieved-constants\", {})"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -259,53 +274,36 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "metadata = CalibrationMetadata(metadata_folder or out_folder)\n",
-    "# Constant paths & timestamps are saved under retrieved-constants in calibration_metadata.yml\n",
-    "const_yaml = metadata.setdefault(\"retrieved-constants\", {})"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "tags": []
+   },
    "outputs": [],
    "source": [
-    "const_data = {}\n",
-    "const_load_mp = psh.ProcessContext(num_workers=24)\n",
+    "const_data = dict()  # {\"ModuleName\": {\"ConstantName\": ndarray}}\n",
+    "start = perf_counter()\n",
+    "if const_yaml:\n",
     "\n",
-    "if const_yaml:  # Read constants from YAML file.\n",
-    "    start = perf_counter()\n",
-    "    for da, ccvs in const_yaml.items():\n",
+    "    const_load_mp = psh.ProcessContext(num_workers=24)\n",
     "\n",
-    "        for calibration_name, ccv in ccvs['constants'].items():\n",
-    "            if ccv['file-path'] is None:\n",
-    "                warning(f\"Missing {calibration_name} for {da}\")\n",
+    "    for mod, constants in const_yaml.items():\n",
+    "        for cname, cmdata in constants[\"constants\"].items():\n",
+    "            if cmdata['file-path'] is None:  # TODO: CHECK AFTER REBASING\n",
+    "                warning(f\"Missing {cname} for {mod}\")\n",
     "                continue\n",
+    "            const_data.setdefault(mod, {})[cname] = const_load_mp.alloc(\n",
+    "                    shape=(256, 256, mem_cells, 3),  # All LPD constants have the same shape.\n",
+    "                    dtype=np.uint32 if cname.startswith('BadPixels') else np.float32\n",
+    
+    "            )\n",
+    "    def load_constant_dataset(wid, index, mod):\n",
+    "        for cname, mdata in const_yaml[mod][\"constants\"].items():\n",
+    "            with h5py.File(mdata[\"path\"], \"r\") as cf:\n",
+    "                cf[f\"{mdata['dataset']}/data\"].read_direct(const_data[mod][cname])\n",
     "\n",
-    "            dtype = np.uint32 if calibration_name.startswith('BadPixels') else np.float32\n",
+    "    const_load_mp.map(load_constant_dataset, karabo_da)\n",
     "\n",
-    "            const_data[(da, calibration_name)] = dict(\n",
-    "                path=Path(ccv['file-path']),\n",
-    "                dataset=ccv['dataset-name'],\n",
-    "                data=const_load_mp.alloc(shape=(256, 256, mem_cells, 3), dtype=dtype)\n",
-    "            )\n",
-    "else:  # Retrieve constants from CALCAT.\n",
-    "    dark_calibrations = {\n",
-    "        1: 'Offset',  # np.float32\n",
-    "        14: 'BadPixelsDark'  # should be np.uint32, but is np.float64\n",
-    "    }\n",
-    "\n",
-    "    base_condition = [\n",
-    "        dict(parameter_name='Sensor Bias Voltage', value=bias_voltage),\n",
-    "        dict(parameter_name='Memory cells', value=mem_cells),\n",
-    "        dict(parameter_name='Feedback capacitor', value=capacitor),\n",
-    "        dict(parameter_name='Pixels X', value=256),\n",
-    "        dict(parameter_name='Pixels Y', value=256),\n",
-    "    ]\n",
-    "    cell_ids_pattern_s = None\n",
+    "else:\n",
+    "    cell_ids_pattern_s = None  # TODO: FIX AFTER REBASING\n",
     "    if use_cell_order != 'never':\n",
     "        # Read the order of memory cells used\n",
     "        raw_data = xd.DataCollection.from_paths([e[1] for e in data_to_process])\n",
@@ -320,71 +318,35 @@
     "        ]\n",
     "    else:\n",
     "        dark_condition = base_condition.copy()\n",
-    "\n",
-    "    illuminated_calibrations = {\n",
-    "        20: 'BadPixelsFF',  # np.uint32\n",
-    "        42: 'GainAmpMap',  # np.float32\n",
-    "        43: 'FFMap',  # np.float32\n",
-    "        44: 'RelativeGain'  # np.float32\n",
-    "    }\n",
-    "\n",
-    "    illuminated_condition = base_condition + [\n",
-    "        dict(parameter_name='Source Energy', value=photon_energy),\n",
-    "        dict(parameter_name='category', value=category)\n",
-    "    ]\n",
-    "\n",
-    "    print('Querying calibration database', end='', flush=True)\n",
-    "    start = perf_counter()\n",
-    "    for calibrations, condition in [\n",
-    "        (dark_calibrations, dark_condition),\n",
-    "        (illuminated_calibrations, illuminated_condition)\n",
-    "    ]:\n",
-    "        resp = CalibrationConstantVersion.get_closest_by_time_by_detector_conditions(\n",
-    "            client, karabo_id, list(calibrations.keys()),\n",
-    "            {'parameters_conditions_attributes': condition},\n",
-    "            karabo_da='', event_at=creation_time.isoformat()\n",
-    "        )\n",
-    "\n",
-    "        if not resp['success']:\n",
-    "            raise RuntimeError(resp)\n",
-    "\n",
-    "        for ccv in resp['data']:\n",
-    "            cc = ccv['calibration_constant']\n",
-    "            da = ccv['physical_detector_unit']['karabo_da']\n",
-    "            calibration_name = calibrations[cc['calibration_id']]\n",
-    "            \n",
-    "            dtype = np.uint32 if calibration_name.startswith('BadPixels') else np.float32\n",
-    "            \n",
-    "            const_data[(da, calibration_name)] = dict(\n",
-    "                path=Path(ccv['path_to_file']) / ccv['file_name'],\n",
-    "                dataset=ccv['data_set_name'],\n",
-    "                data=const_load_mp.alloc(shape=(256, 256, mem_cells, 3), dtype=dtype)\n",
-    "            )\n",
-    "        print('.', end='', flush=True)\n",
-    "            \n",
-    "total_time = perf_counter() - start\n",
-    "print(f'{total_time:.1f}s')"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "def load_constant_dataset(wid, index, const_descr):\n",
-    "    ccv_entry = const_data[const_descr]\n",
-    "    \n",
-    "    with h5py.File(cal_db_root / ccv_entry['path'], 'r') as fp:\n",
-    "        fp[ccv_entry['dataset'] + '/data'].read_direct(ccv_entry['data'])\n",
-    "        \n",
-    "    print('.', end='', flush=True)\n",
-    "\n",
-    "print('Loading calibration data', end='', flush=True)\n",
-    "start = perf_counter()\n",
-    "const_load_mp.map(load_constant_dataset, list(const_data.keys()))\n",
+    "    lpd_cal = LPD_CalibrationData(\n",
+    "        detector_name=karabo_id,\n",
+    "        modules=karabo_da,\n",
+    "        sensor_bias_voltage=bias_voltage,\n",
+    "        memory_cells=mem_cells,  # TODO: ADD MISSING CONDITION AFTER REBASING\n",
+    "        feedback_capacitor=capacitor,\n",
+    "        source_energy=photon_energy,\n",
+    "        category=category,\n",
+    "        event_at=creation_time,\n",
+    "        snapshot_at=creation_time,\n",
+    "        client=client,\n",
+    "    )\n",
+    "    const_data = lpd_cal.ndarray_map()\n",
+    "\n",
+    "# Validate the constants availability and raise/warn correspondingly. \n",
+    "for mod, calibrations in const_data.items():\n",
+    "    missing_dark_constants = set(\n",
+    "        c for c in [\"Offset\", \"BadPixelsDark\"] if c not in calibrations.keys())\n",
+    "    missing_gain_constants = set(\n",
+    "        c for c in [\"BadPixelsFF\", \"GainAmpMap\", \"FFMap\", \"RelativeGain\"] if c not in calibrations.keys())  # noqa\n",
+    "    if missing_dark_constants:\n",
+    "        raise KeyError(\n",
+    "            f\"Dark constants {missing_dark_constants} are not available for correction. Module: {mod}\")  # noqa\n",
+    "    if missing_gain_constants:\n",
+    "        warning(\n",
+    "            f\"Gain constants {missing_gain_constants} were not retrieved. Module: {mod}\")\n",
+    "    if calibrations[\"BadPixelsDark\"].dtype != np.uint32:  # Old LPD constants are stored as float32.\n",
+    "        calibrations[\"BadPixelsDark\"] = calibrations[\"BadPixelsDark\"].astype(np.uint32)\n",
     "total_time = perf_counter() - start\n",
-    "\n",
     "print(f'{total_time:.1f}s')"
    ]
   },
@@ -411,39 +373,35 @@
     "}\n",
     "\n",
     "def prepare_constants(wid, index, aggregator):\n",
-    "    consts = {calibration_name: entry['data']\n",
-    "              for (aggregator_, calibration_name), entry\n",
-    "              in const_data.items()\n",
-    "              if aggregator == aggregator_}\n",
     "    \n",
     "    def _prepare_data(calibration_name, dtype):\n",
-    "        return consts[calibration_name] \\\n",
+    "        return const_data[calibration_name] \\\n",
     "            .transpose(constant_order[calibration_name]) \\\n",
     "            .astype(dtype, copy=True)  # Make sure array is contiguous.\n",
     "    \n",
-    "    if offset_corr and 'Offset' in consts:\n",
+    "    if offset_corr and 'Offset' in const_data:\n",
     "        ccv_offsets[aggregator] = _prepare_data('Offset', np.float32)\n",
     "    else:\n",
     "        ccv_offsets[aggregator] = np.zeros(ccv_shape, dtype=np.float32)\n",
     "        \n",
     "    ccv_gains[aggregator] = np.ones(ccv_shape, dtype=np.float32)\n",
     "    \n",
-    "    if 'BadPixelsDark' in consts:\n",
+    "    if 'BadPixelsDark' in const_data:\n",
     "        ccv_masks[aggregator] = _prepare_data('BadPixelsDark', np.uint32)\n",
     "    else:\n",
     "        ccv_masks[aggregator] = np.zeros(ccv_shape, dtype=np.uint32)\n",
     "    \n",
-    "    if rel_gain and 'RelativeGain' in consts:\n",
+    "    if rel_gain and 'RelativeGain' in const_data:\n",
     "        ccv_gains[aggregator] *= _prepare_data('RelativeGain', np.float32)\n",
     "        \n",
-    "    if ff_map and 'FFMap' in consts:\n",
+    "    if ff_map and 'FFMap' in const_data:\n",
     "        ccv_gains[aggregator] *= _prepare_data('FFMap', np.float32)\n",
     "        \n",
-    "        if 'BadPixelsFF' in consts:\n",
+    "        if 'BadPixelsFF' in const_data:\n",
     "            np.bitwise_or(ccv_masks[aggregator], _prepare_data('BadPixelsFF', np.uint32),\n",
     "                          out=ccv_masks[aggregator])\n",
     "        \n",
-    "    if gain_amp_map and 'GainAmpMap' in consts:\n",
+    "    if gain_amp_map and 'GainAmpMap' in const_data:\n",
     "        ccv_gains[aggregator] *= _prepare_data('GainAmpMap', np.float32)\n",
     "        \n",
     "    print('.', end='', flush=True)\n",
diff --git a/notebooks/LPD/LPD_retrieve_constants_precorrection.ipynb b/notebooks/LPD/LPD_retrieve_constants_precorrection.ipynb
index 419a9a43f..1ff96e641 100644
--- a/notebooks/LPD/LPD_retrieve_constants_precorrection.ipynb
+++ b/notebooks/LPD/LPD_retrieve_constants_precorrection.ipynb
@@ -47,8 +47,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "from logging import warning\n",
     "from pathlib import Path\n",
-    "from time import perf_counter\n",
     "\n",
     "import numpy as np\n",
     "\n",
@@ -57,12 +57,14 @@
     "import extra_data as xd\n",
     "\n",
     "from cal_tools.lpdlib import get_mem_cell_pattern, make_cell_order_condition\n",
+    "from cal_tools.calcat_interface import LPD_CalibrationData\n",
     "from cal_tools.tools import (\n",
     "    CalibrationMetadata,\n",
     "    calcat_creation_time,\n",
-    "    save_constant_metadata,\n",
     ")\n",
-    "from cal_tools.restful_config import restful_config"
+    "from calibration_client import CalibrationClient\n",
+    "from cal_tools.restful_config import restful_config\n",
+    "from cal_tools.step_timing import StepTimer"
    ]
   },
   {
@@ -82,14 +84,21 @@
     "print(f'Using {creation_time.isoformat()} as creation time')\n",
     "\n",
     "# Pick all modules/aggregators or those selected.\n",
-    "if not karabo_da or karabo_da == ['']:\n",
-    "    if not modules or modules == [-1]:\n",
+    "if karabo_da == ['']:\n",
+    "    if modules == [-1]:\n",
     "        modules = list(range(16))\n",
-    "\n",
     "    karabo_da = [f'LPD{i:02d}' for i in modules]\n",
-    "\n",
-    "# List of detector sources.\n",
-    "det_inp_sources = [input_source.format(karabo_id=karabo_id, module_index=int(da[-2:])) for da in karabo_da]"
+    "else:\n",
+    "    modules = [int(x[-2:]) for x in karabo_da]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "step_timer = StepTimer()"
    ]
   },
   {
@@ -118,17 +127,23 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "dark_calibrations = {\n",
-    "    1: 'Offset',\n",
-    "    14: 'BadPixelsDark',\n",
-    "}\n",
-    "\n",
-    "base_condition = [\n",
-    "    dict(parameter_name='Sensor Bias Voltage', value=bias_voltage),\n",
-    "    dict(parameter_name='Memory cells', value=mem_cells),\n",
-    "    dict(parameter_name='Feedback capacitor', value=capacitor),\n",
-    "    dict(parameter_name='Pixels X', value=256),\n",
-    "    dict(parameter_name='Pixels Y', value=256),\n",
+    "step_timer.start()\n",
+    "lpd_cal = LPD_CalibrationData(\n",
+    "    detector_name=karabo_id,\n",
+    "    modules=karabo_da,\n",
+    "    sensor_bias_voltage=bias_voltage,\n",
+    "    memory_cells=mem_cells,\n",
+    "    feedback_capacitor=capacitor,\n",
+    "    source_energy=photon_energy,\n",
+    "    category=category,\n",
+    "    event_at=creation_time,\n",
+    "    snapshot_at=creation_time,\n",
+    "    client=client,\n",
+    ")\n",
+    "constant_names = [\n",
+    "    \"Offset\", \"BadPixelsDark\",\n",
+    "    \"BadPixelsFF\", \"GainAmpMap\",\n",
+    "    \"FFMap\", \"RelativeGain\",\n",
     "]\n",
     "cell_ids_pattern_s = None\n",
     "if use_cell_order != 'never':\n",
@@ -146,70 +161,37 @@
     "else:\n",
     "    dark_condition = base_condition.copy()\n",
     "\n",
-    "illuminated_calibrations = {\n",
-    "    20: 'BadPixelsFF',\n",
-    "    42: 'GainAmpMap',\n",
-    "    43: 'FFMap',\n",
-    "    44: 'RelativeGain',\n",
-    "}\n",
-    "\n",
-    "illuminated_condition = base_condition + [\n",
-    "    dict(parameter_name='Source Energy', value=photon_energy),\n",
-    "    dict(parameter_name='category', value=category)\n",
-    "]"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "const_data = {}\n",
-    "\n",
-    "print('Querying calibration database', end='', flush=True)\n",
-    "start = perf_counter()\n",
-    "for k_da in karabo_da:\n",
-    "    pdu = None\n",
-    "    retrieved_constants[k_da] = dict()\n",
-    "    const_mdata = retrieved_constants[k_da][\"constants\"] = dict()\n",
-    "    for calibrations, condition in [\n",
-    "        (dark_calibrations, dark_condition),\n",
-    "        (illuminated_calibrations, illuminated_condition)\n",
-    "    ]:\n",
-    "        resp = CalibrationConstantVersion.get_closest_by_time_by_detector_conditions(\n",
-    "            client, karabo_id, list(calibrations.keys()),\n",
-    "            {'parameters_conditions_attributes': condition},\n",
-    "            karabo_da=k_da, event_at=creation_time.isoformat())\n",
-    "\n",
-    "        if not resp[\"success\"]:\n",
-    "            print(f\"ERROR: Constants {list(calibrations.values())} \"\n",
-    "            f\"were not retrieved, {resp['app_info']}\")\n",
-    "            for cname in calibrations.values():\n",
-    "                if cname == 'Offset':\n",
-    "                    raise Exception(\"Could not find offset constant, will not correct data\")\n",
-    "                const_mdata[cname] = dict()\n",
-    "                const_mdata[cname][\"file-path\"] = None\n",
-    "                const_mdata[cname][\"dataset-name\"] = None\n",
-    "                const_mdata[cname][\"creation-time\"] = None     \n",
-    "            continue\n",
-    "\n",
-    "        for ccv in resp[\"data\"]:\n",
-    "            cc = ccv['calibration_constant']\n",
-    "            cname = calibrations[cc['calibration_id']]\n",
-    "            const_mdata[cname] = dict()\n",
-    "            const_mdata[cname][\"file-path\"] = str(Path(ccv['path_to_file']) / ccv['file_name'])\n",
-    "            const_mdata[cname][\"dataset-name\"] = ccv['data_set_name']\n",
-    "            const_mdata[cname][\"creation-time\"] = ccv['begin_at']\n",
-    "            pdu = ccv['physical_detector_unit']['physical_name']\n",
-    "\n",
-    "        print('.', end='', flush=True)\n",
-    "    retrieved_constants[k_da][\"physical-detector-unit\"] = pdu\n",
+    "# for cname in constant_names:\n",
+    "# Error out if dark constants are not retrieved.\n",
+    "\n",
+    "lpd_cal_metadata = lpd_cal.metadata()\n",
+    "\n",
+    "# Validate the constants availability and raise/warn correspondingly. \n",
+    "for mod, ccv_dict in lpd_cal_metadata.items():\n",
+    "    missing_dark_constants = set(\n",
+    "        c for c in [\"Offset\", \"BadPixelsDark\"] if c not in ccv_dict.keys())\n",
+    "    missing_gain_constants = set(\n",
+    "        c for c in [\"BadPixelsFF\", \"GainAmpMap\", \"FFMap\", \"RelativeGain\"] if c not in ccv_dict.keys())  # noqa\n",
+    "    if missing_dark_constants:\n",
+    "        raise KeyError(\n",
+    "            f\"Dark constants {missing_dark_constants} are not available for correction. Module: {mod}\")  # noqa\n",
+    "    if missing_gain_constants:\n",
+    "        warning(\n",
+    "            f\"Gain constants {missing_gain_constants} were not retrieved. Module: {mod}\")\n",
+    "\n",
+    "for mod, ccv_dict in lpd_cal_metadata.items():\n",
+    "    mdata_dict = {\"constants\": dict()}\n",
+    "    for cname, ccv_metadata in ccv_dict.items():\n",
+    "        mdata_dict[\"constants\"][cname] = {\n",
+    "                \"path\": str(lpd_cal.caldb_root / ccv_metadata[\"path\"]),\n",
+    "                \"dataset\": ccv_metadata[\"dataset\"],\n",
+    "                \"creation-time\": ccv_metadata[\"begin_validity_at\"],\n",
+    "        }\n",
+    "    mdata_dict[\"physical-name\"] = ccv_metadata[\"physical_name\"]\n",
+    "    retrieved_constants[mod] = mdata_dict\n",
     "metadata.save()\n",
     "\n",
-    "total_time = perf_counter() - start\n",
-    "print(f'{total_time:.1f}s')\n",
-    "print(f\"Stored retrieved constants in {metadata.filename}\")"
+    "step_timer.done_step(f\"Stored retrieved constants in {metadata.filename}.\")"
    ]
   }
  ],
-- 
GitLab