cookbooks/visualizing_scalar_results.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "97de1512",
   "metadata": {},
   "source": [
    "# Postprocessing/Visualizing scalar results using reduce functions\n",
    "\n",
    "## Introduction\n",
    "\n",
    "This cookbook will guide you through the creation of a simple visualization from an existing trial in jinko.  \n",
    "In particular, you will be able to retrieve scalar results and plot them using plotly.  \n",
    "\n",
    "Linked resources: [Jinko](https://jinko.ai/project/e0fbb5bb-8929-439a-bad6-9e12d19d9ae4?labels=d59e57e6-889b-427a-b4ee-bed52a6c6ca2)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4a8a8c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Jinko specifics imports & initialization\n",
    "# Please fold this section and do not change\n",
    "import jinko_helpers as jinko\n",
    "\n",
    "# Connect to Jinko (see README.md for more options)\n",
    "jinko.initialize()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8a8364c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cookbook-specific imports\n",
    "\n",
    "import datetime\n",
    "import io\n",
    "import json\n",
    "import pandas as pd\n",
    "import plotly.express as px\n",
    "import re\n",
    "import zipfile\n",
    "import plotly.graph_objects as go\n",
    "\n",
    "# Cookbook-specific constants\n",
    "\n",
    "# @param {\"name\":\"trialId\", \"type\": \"string\"}\n",
    "# Fill the short Id of your Trial (ex: tr-EKRx-3HRt)\n",
    "trialId = \"tr-Fzt9-uO98\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1f6ec429",
   "metadata": {},
   "source": [
    "## Let's use the API and plot the data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2fc5ed2a",
   "metadata": {},
   "source": [
    "### Load the trial"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93bf3446",
   "metadata": {},
   "outputs": [],
   "source": [
    "if trialId is None:\n",
    "    raise Exception(\"Please specify a Trial Id\")\n",
    "else:\n",
    "    print(f\"Using Trial ID: {trialId}\")\n",
    "\n",
    "# Convert short Id to coreItemId\n",
    "try:\n",
    "    coreItemId = jinko.get_core_item_id(trialId, 1)\n",
    "except Exception as e:\n",
    "    print(f\"Failed to find corresponding trial, check the trialId\")\n",
    "    raise\n",
    "\n",
    "# List all Trial versions (https://doc.jinko.ai/api/#/paths/core-v2-trial_manager-trial-status/post)\n",
    "try:\n",
    "    trialVersions = jinko.make_request(\n",
    "        f'/core/v2/trial_manager/trial/{coreItemId[\"id\"]}/status'\n",
    "    ).json()\n",
    "    print(f\"Fetched {len(trialVersions)} versions for the trial.\")\n",
    "except Exception as e:\n",
    "    print(f\"Error fetching trial versions: {e}\")\n",
    "    raise\n",
    "\n",
    "# Get the latest completed version\n",
    "try:\n",
    "    latestCompletedVersion = next(\n",
    "        (item for item in trialVersions if item[\"status\"] == \"completed\"), None\n",
    "    )\n",
    "    if latestCompletedVersion is None:\n",
    "        raise Exception(\"No completed Trial version found\")\n",
    "    else:\n",
    "        print(\n",
    "            \"Successfully fetched this simulation:\\n\",\n",
    "            json.dumps(latestCompletedVersion, indent=1),\n",
    "        )\n",
    "        # Store the trial Id and the snapshot Id to use in the API requests\n",
    "        simulationId = latestCompletedVersion[\"simulationId\"]\n",
    "        trialId = simulationId[\"coreItemId\"]\n",
    "        trialSnapshotId = simulationId[\"snapshotId\"]\n",
    "except Exception as e:\n",
    "    print(f\"Error processing trial versions: {e}\")\n",
    "    raise"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f115b8af",
   "metadata": {},
   "source": [
    "### Display a result summary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc0b3f84",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Retrieve result summary (https://doc.jinko.ai/api/#/paths/core-v2-trial_manager-trial-trialId--snapshots--trialIdSnapshot--results_summary/get)\n",
    "response = jinko.make_request(\n",
    "    \"/core/v2/trial_manager/trial/%s/snapshots/%s/results_summary\"\n",
    "    % (trialId, trialSnapshotId),\n",
    "    method=\"GET\",\n",
    ")\n",
    "responseSummary = json.loads(response.content)\n",
    "\n",
    "defaultScalars = [\n",
    "    \"__jinkoAllocationMiB.tmax\",\n",
    "    \"__jinkoSolvingTime.tmax\",\n",
    "    \"SimulationTMax\",\n",
    "    \"SimulationTMin\",\n",
    "]\n",
    "\n",
    "# Print a summary of the results content\n",
    "print(\"Keys in the results summary:\\n\", list(responseSummary.keys()), \"\\n\")\n",
    "print(\"Available patients:\\n\", responseSummary[\"patients\"], \"\\n\")\n",
    "print(\"Available arms:\\n\", responseSummary[\"arms\"], \"\\n\")\n",
    "print(\n",
    "    \"Available scalars:\\n\",\n",
    "    [scalar[\"id\"] for scalar in responseSummary[\"scalars\"]],\n",
    "    \"\\n\",\n",
    ")\n",
    "print(\n",
    "    \"Available cross-arm scalars:\\n\",\n",
    "    [scalar[\"id\"] for scalar in responseSummary[\"scalarsCrossArm\"]],\n",
    "    \"\\n\",\n",
    ")\n",
    "print(\n",
    "    \"Available categorical parameters:\\n\",\n",
    "    [scalar[\"id\"] for scalar in responseSummary[\"categoricals\"]],\n",
    "    \"\\n\",\n",
    ")\n",
    "print(\n",
    "    \"Available cross-arm categorical parameters:\\n\",\n",
    "    [scalar[\"id\"] for scalar in responseSummary[\"categoricalsCrossArm\"]],\n",
    "    \"\\n\",\n",
    ")\n",
    "\n",
    "# Store the list of scenario overrides to use them later\n",
    "scenarioOverrides = [\n",
    "    scalar[\"id\"]\n",
    "    for scalar in (responseSummary[\"scalars\"] + responseSummary[\"categoricals\"])\n",
    "    if \"ScenarioOverride\" in scalar[\"type\"][\"labels\"]\n",
    "]\n",
    "print(\"List of scenario overrides:\\n\", scenarioOverrides, \"\\n\")\n",
    "\n",
    "# Store the list of scalars that are not scenario descriptors to use them later\n",
    "resultScalars = [\n",
    "    scalar[\"id\"]\n",
    "    for scalar in responseSummary[\"scalars\"]\n",
    "    if \"Custom\" in scalar[\"type\"][\"labels\"]\n",
    "]\n",
    "print(\"List of result scalars:\\n\", resultScalars, \"\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cb7a6106",
   "metadata": {},
   "source": [
    "### Download scalar results data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7436ffb3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Retrieve scalar results (https://doc.jinko.ai/api/#/paths/core-v2-result_manager-scalars_summary/post)\n",
    "\n",
    "# replace here by the scalar ids list you want\n",
    "idsForScalars = {\"scalars\": resultScalars, \"scenarioOverrides\": scenarioOverrides}\n",
    "csvScalars = {}\n",
    "\n",
    "\n",
    "def retrieve_scalars(scalar_type):\n",
    "    try:\n",
    "        print(\"Retrieving scalar results...\")\n",
    "        response = jinko.make_request(\n",
    "            \"/core/v2/result_manager/scalars_summary\",\n",
    "            method=\"POST\",\n",
    "            json={\n",
    "                \"select\": idsForScalars[scalar_type],\n",
    "                \"trialId\": latestCompletedVersion[\"simulationId\"],\n",
    "            },\n",
    "        )\n",
    "        if response.status_code == 200:\n",
    "            print(\"Scalar results retrieved successfully.\")\n",
    "            archive = zipfile.ZipFile(io.BytesIO(response.content))\n",
    "            filename = archive.namelist()[0]\n",
    "            print(f\"Extracted scalar results file: {filename}\")\n",
    "            csvScalars[scalar_type] = archive.read(filename).decode(\"utf-8\")\n",
    "        else:\n",
    "            print(\n",
    "                f\"Failed to retrieve scalar results: {response.status_code} - {response.reason}\"\n",
    "            )\n",
    "            response.raise_for_status()\n",
    "    except Exception as e:\n",
    "        print(f\"Error during scalar results retrieval or processing: {e}\")\n",
    "        raise\n",
    "\n",
    "\n",
    "retrieve_scalars(\"scalars\")\n",
    "retrieve_scalars(\"scenarioOverrides\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "42410425",
   "metadata": {},
   "source": [
    "### Postprocess the data in a pandas dataframe\n",
    "\n",
    "The data is post-processed using the pandas library, and transformed into a table separating the information necessary to plot. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c28a055c-280b-4497-bd43-e4f8429232c2",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# All currently available reduce functions\n",
    "reduceFuncs = [\n",
    "    \"time-of-max\",\n",
    "    \"time-of-min\",\n",
    "    \"max-slope\",\n",
    "    \"min-slope\",\n",
    "    \"amplitude\",\n",
    "    \"end-minus-start\",\n",
    "    \"at\",\n",
    "    \"auc\",\n",
    "    \"avg\",\n",
    "    \"max\",\n",
    "    \"min\",\n",
    "]\n",
    "reduceFuncRegex = \"(\" + \"|\".join(reduceFuncs) + \")\"\n",
    "\n",
    "# Load scalars into a dataframe\n",
    "dfScalars = pd.read_csv(io.StringIO(csvScalars[\"scalars\"]))\n",
    "print(\"\\nRaw scalar data (first rows):\\n\")\n",
    "display(dfScalars.head())\n",
    "\n",
    "# Postprocessing - Split the scalarId column into pieces depending on the applicable pattern\n",
    "dfScalars[[\"generalScalarId\", \"reduceFunction\", \"startTime\"]] = dfScalars[\n",
    "    \"scalarId\"\n",
    "].str.extract(r\"(.*?)-\" + reduceFuncRegex + \"(?:-from-|-)(.*)\")\n",
    "\n",
    "\n",
    "# Postprocessing - Extract start and end times for each measure created on jinkō (for point measures like `at`, start and end time are equal)\n",
    "def extract_endtime(value):\n",
    "    match = re.match(r\"(?:.*?-to-(.*)|^(.*)$)\", value)\n",
    "    if match:\n",
    "        return match.group(1) or match.group(2)\n",
    "    return None\n",
    "\n",
    "\n",
    "def extract_starttime(value):\n",
    "    match = re.match(r\"(?:(.*?)-to-.*|^(.*)$)\", value)\n",
    "    if match:\n",
    "        return match.group(1) or match.group(2)\n",
    "    return None\n",
    "\n",
    "\n",
    "dfScalars[\"endTime\"] = dfScalars[\"startTime\"].apply(extract_endtime)\n",
    "\n",
    "dfScalars[\"startTime\"] = dfScalars[\"startTime\"].apply(extract_starttime)\n",
    "\n",
    "\n",
    "# Load scenarioOverrides into a dataframe\n",
    "dfOverrides = pd.read_csv(io.StringIO(csvScalars[\"scenarioOverrides\"]))\n",
    "\n",
    "# Remove unneccessary trailing tmins\n",
    "dfOverrides[\"scalarId\"] = [x.strip(\".tmin\") for x in dfOverrides[\"scalarId\"]]\n",
    "\n",
    "# Pivot to a wide format to obtain protocol overrides in columns\n",
    "dfOverrides = dfOverrides.pivot(\n",
    "    index=[\"armId\", \"patientId\"], columns=\"scalarId\", values=\"value\"\n",
    ")\n",
    "\n",
    "print(\"\\nPivotted scenario override table (first rows): \\n\")\n",
    "display(dfOverrides.head())\n",
    "\n",
    "# Merge dataframes to have scenario overrides as columns in the results dataframe\n",
    "dfScalars = dfScalars.merge(dfOverrides, how=\"inner\", on=[\"patientId\", \"armId\"])\n",
    "print(\"\\nScalar result table with overrides (first rows): \\n\")\n",
    "display(dfScalars.head())\n",
    "\n",
    "\n",
    "def parse_isoduration(str):\n",
    "    ## Taken from https://stackoverflow.com/questions/36976138/is-there-an-easy-way-to-convert-iso-8601-duration-to-timedelta\n",
    "    ## Parse the ISO8601 duration as years,months,weeks,days, hours,minutes,seconds\n",
    "    ## Returns: time in days with 86400 seconds, converting units higher than days into seconds used on jinkō\n",
    "    ## Examples: \"PT1H30M15.460S\", \"P5DT4M\", \"P2WT3H\"\n",
    "    def get_isosplit(str, split):\n",
    "        if split in str:\n",
    "            n, str = str.split(split, 1)\n",
    "        else:\n",
    "            n = \"0\"\n",
    "        return n.replace(\",\", \".\"), str  # to handle like \"P0,5Y\"\n",
    "\n",
    "    str = str.split(\"P\", 1)[-1]  # Remove prefix\n",
    "    s_yr, str = get_isosplit(str, \"Y\")  # Step through letter dividers\n",
    "    s_mo, str = get_isosplit(str, \"M\")\n",
    "    s_wk, str = get_isosplit(str, \"W\")\n",
    "    s_dy, str = get_isosplit(str, \"D\")\n",
    "    _, str = get_isosplit(str, \"T\")\n",
    "    s_hr, str = get_isosplit(str, \"H\")\n",
    "    s_mi, str = get_isosplit(str, \"M\")\n",
    "    s_sc, str = get_isosplit(str, \"S\")\n",
    "    n_yr = float(s_yr) * 31557600  # approx seconds for year, month, week\n",
    "    n_mo = float(s_mo) * 2629800\n",
    "    n_wk = float(s_wk) * 604800\n",
    "    dt = datetime.timedelta(\n",
    "        days=float(s_dy),\n",
    "        hours=float(s_hr),\n",
    "        minutes=float(s_mi),\n",
    "        seconds=float(s_sc) + n_yr + n_mo + n_wk,\n",
    "    )\n",
    "    return dt.total_seconds() / 86400\n",
    "\n",
    "\n",
    "# Parse iso durations into time in weeks\n",
    "dfScalars[\"endTime\"] = dfScalars[\"endTime\"].apply(parse_isoduration)\n",
    "dfScalars[\"startTime\"] = dfScalars[\"startTime\"].apply(parse_isoduration)\n",
    "\n",
    "# parse tiem in seconds (default) into time in weeks for observables containing 'time'\n",
    "dfScalars.loc[dfScalars[\"reduceFunction\"].str.contains(\"time\"), \"value\"] /= 86400\n",
    "dfScalars.loc[dfScalars[\"reduceFunction\"].str.contains(\"time\"), \"unit\"] = \"days\"\n",
    "\n",
    "# Sort values by starting time\n",
    "dfScalars = dfScalars.sort_values(by=\"startTime\", ascending=True, na_position=\"first\")\n",
    "\n",
    "print(\"\\n Final scalar result table with overrides (first rows): \\n\")\n",
    "display(dfScalars.head())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "720e4a1f",
   "metadata": {},
   "source": [
    "### Plot the data\n",
    "\n",
    "Finally we plot the time series data by facetting over scenario overrides. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e768c432-016a-49a1-9e83-fdccdbae13bc",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Define constants for color mapping\n",
    "COLOR_MAP = {\n",
    "    \"intravenous\": \"#1f77b4\",  # blue\n",
    "    \"subcutaneous\": \"#ff7f0e\",  # orange\n",
    "}\n",
    "\n",
    "# Define offsets for annotation positioning\n",
    "OFFSET_MAP = {\n",
    "    \"intravenous\": 0,\n",
    "    \"subcutaneous\": 1,\n",
    "}\n",
    "\n",
    "# Define descriptions for summary measures\n",
    "SUMMARY_MEASURE_MAP = {\n",
    "    \"auc\": \"AUC from day 7 to day 10\",\n",
    "    \"avg\": \"Average from day 7 to day 10\",\n",
    "    \"max\": \"Maximum from day 7 to day 10\",\n",
    "}\n",
    "\n",
    "# Define offsets for summary measures\n",
    "MEASURE_OFFSET_MAP = {\n",
    "    \"auc\": 1,\n",
    "    \"avg\": 0,\n",
    "    \"max\": 2,\n",
    "}\n",
    "\n",
    "# Lists of measures to summarize and plot\n",
    "SUMMARY_MEASURES = [\"avg\", \"auc\"]\n",
    "PLOT_MEASURES = [\"max\"]\n",
    "\n",
    "# Columns used for faceting and coloring the plot\n",
    "FACET_COLUMN = \"fullDose\"\n",
    "FACET_ROW = \"primingDose\"\n",
    "COLOR_COLUMN = \"administrationMode\"\n",
    "\n",
    "# Base query to filter data\n",
    "QUERY_BASE = 'generalScalarId == \"Blood.Drug\" and reduceFunction == \"at\"'\n",
    "\n",
    "\n",
    "def create_line_plot(data):\n",
    "    \"\"\"Create a line plot using the provided data.\"\"\"\n",
    "    return px.line(\n",
    "        data,\n",
    "        x=\"endTime\",\n",
    "        y=\"value\",\n",
    "        facet_col=FACET_COLUMN,\n",
    "        facet_row=FACET_ROW,\n",
    "        color=COLOR_COLUMN,\n",
    "        labels={\n",
    "            \"endTime\": \"Time [days]\",\n",
    "            \"value\": \"Concentration of Drug in Blood [mg/ml]\",\n",
    "            \"fullDose\": \"Full dose [mg]\",\n",
    "            \"primingDose\": \"Priming dose [mg]\",\n",
    "            \"administrationMode\": \"Administration mode\",\n",
    "        },\n",
    "        height=600,\n",
    "        color_discrete_map=COLOR_MAP,\n",
    "    )\n",
    "\n",
    "\n",
    "def get_facet_indices(data, row):\n",
    "    \"\"\"Get facet indices for the given row.\"\"\"\n",
    "    col_idx = data[FACET_COLUMN].unique().tolist().index(row[FACET_COLUMN])\n",
    "    row_idx = data[FACET_ROW].unique().tolist().index(row[FACET_ROW])\n",
    "    return col_idx, row_idx\n",
    "\n",
    "\n",
    "def add_summary_annotations(annotations, data, key, max_value):\n",
    "    \"\"\"Add summary measure annotations.\"\"\"\n",
    "    query_key = f'generalScalarId == \"Blood.Drug\" and reduceFunction == \"{key}\"'\n",
    "    for i, row in data.query(query_key).iterrows():\n",
    "        # Get indices for the facet grid\n",
    "        col_idx, row_idx = get_facet_indices(data.query(QUERY_BASE), row)\n",
    "        xref = f\"x{col_idx + 1}\"\n",
    "        yref = f\"y{row_idx + 1}\"\n",
    "\n",
    "        # Add the annotation\n",
    "        annotations.append(\n",
    "            dict(\n",
    "                x=10,\n",
    "                y=(\n",
    "                    max_value\n",
    "                    - max_value / 5 * OFFSET_MAP[row[COLOR_COLUMN]]\n",
    "                    - max_value / 10 * MEASURE_OFFSET_MAP[key]\n",
    "                ),\n",
    "                xref=xref if col_idx != 0 else \"x\",\n",
    "                yref=yref if row_idx != 0 else \"y\",\n",
    "                text=f'{row[COLOR_COLUMN]}: {SUMMARY_MEASURE_MAP[key]} [{row[\"unit\"]}]: {\"%.2E\" % row[\"value\"]}',\n",
    "                showarrow=False,\n",
    "                font=dict(color=COLOR_MAP[row[COLOR_COLUMN]]),\n",
    "                bgcolor=\"white\",\n",
    "                bordercolor=\"black\",\n",
    "                textangle=0,\n",
    "                xanchor=\"right\",\n",
    "                yanchor=\"top\",\n",
    "            )\n",
    "        )\n",
    "\n",
    "\n",
    "def add_plot_annotations(fig, annotations, data, key):\n",
    "    \"\"\"Add plot measure points and annotations.\"\"\"\n",
    "    query_key = f'generalScalarId == \"Blood.Drug\" and reduceFunction == \"{key}\"'\n",
    "    for i, row in data.query(query_key).iterrows():\n",
    "        # Get indices for the facet grid\n",
    "        col_idx, row_idx = get_facet_indices(data.query(QUERY_BASE), row)\n",
    "        xref = f\"x{col_idx + 1}\"\n",
    "        yref = f\"y{row_idx + 1}\"\n",
    "\n",
    "        # Get the x and y values for the annotation\n",
    "        x_value = float(\n",
    "            data.query(\n",
    "                f'generalScalarId == \"Blood.Drug\" and reduceFunction == \"time-of-{key}\" '\n",
    "                f'and {COLOR_COLUMN} == \"{row[COLOR_COLUMN]}\" '\n",
    "                f'and {FACET_COLUMN} == \"{row[FACET_COLUMN]}\" '\n",
    "                f'and {FACET_ROW} == \"{row[FACET_ROW]}\"'\n",
    "            )[\"value\"].iloc[0]\n",
    "        )\n",
    "        y_value = row[\"value\"]\n",
    "\n",
    "        # Add a scatter plot point for the measure\n",
    "        fig.add_trace(\n",
    "            go.Scatter(\n",
    "                x=[x_value],\n",
    "                y=[y_value],\n",
    "                mode=\"markers\",\n",
    "                marker_symbol=\"x\",\n",
    "                marker=dict(color=COLOR_MAP[row[COLOR_COLUMN]]),\n",
    "                showlegend=False,\n",
    "            ),\n",
    "            row=row_idx + 1,\n",
    "            col=col_idx + 1,\n",
    "        )\n",
    "\n",
    "        # Add an annotation for the measure\n",
    "        annotations.append(\n",
    "            dict(\n",
    "                x=x_value,\n",
    "                y=y_value,\n",
    "                xref=xref if col_idx != 0 else \"x\",\n",
    "                yref=yref if row_idx != 0 else \"y\",\n",
    "                text=f'{row[COLOR_COLUMN]}: {SUMMARY_MEASURE_MAP[key]} [{row[\"unit\"]}]:<br> {\"%.2E\" % y_value} at {x_value} days',\n",
    "                showarrow=True,\n",
    "                font=dict(color=COLOR_MAP[row[COLOR_COLUMN]]),\n",
    "                bgcolor=\"white\",\n",
    "                bordercolor=\"black\",\n",
    "                textangle=0,\n",
    "                xanchor=\"left\",\n",
    "                yanchor=\"bottom\",\n",
    "            )\n",
    "        )\n",
    "\n",
    "\n",
    "def plot_here(data):\n",
    "    \"\"\"Main function to create the plot and add annotations.\"\"\"\n",
    "    filtered_data = data.query(QUERY_BASE)\n",
    "    fig = create_line_plot(filtered_data)\n",
    "\n",
    "    # Get the maximum value in the filtered data for annotation positioning\n",
    "    max_value = max(filtered_data[\"value\"])\n",
    "    annotations = []\n",
    "\n",
    "    # Add annotations for summary and plot measures\n",
    "    for key in SUMMARY_MEASURE_MAP.keys():\n",
    "        if key in SUMMARY_MEASURES:\n",
    "            add_summary_annotations(annotations, data, key, max_value)\n",
    "        if key in PLOT_MEASURES:\n",
    "            add_plot_annotations(fig, annotations, data, key)\n",
    "\n",
    "    # Assign annotations to the plot\n",
    "    fig[\"layout\"][\"annotations\"] = annotations\n",
    "    fig.show()\n",
    "\n",
    "\n",
    "# Call the main function with the dataframe `dfScalars`\n",
    "plot_here(dfScalars)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2caedec7-5d40-4150-809d-78f310e9d6a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define a color map for different administration modes\n",
    "color_map = {\n",
    "    \"intravenous\": \"#1f77b4\",  # blue\n",
    "    \"subcutaneous\": \"#ff7f0e\",  # orange\n",
    "}\n",
    "\n",
    "# Define an offset map for annotation handling\n",
    "offset_map = {\n",
    "    \"intravenous\": 0,\n",
    "    \"subcutaneous\": 1,\n",
    "}\n",
    "\n",
    "# Define descriptions for various summary measures\n",
    "summary_measure_map = {\n",
    "    \"end-minus-start\": \"Value after three weeks minus baseline\",\n",
    "    \"amplitude\": \"Amplitude from baseline to week 3\",\n",
    "    \"min\": \"Minimum from baseline to week 3\",\n",
    "    \"min-slope\": \"Minimum slope from baseline to week 3\",\n",
    "    \"max-slope\": \"Maximum slope from baseline to week 3\",\n",
    "}\n",
    "\n",
    "# Define offsets for summary measures\n",
    "measure_offset_map = {\n",
    "    \"end-minus-start\": 1,\n",
    "    \"amplitude\": 0,\n",
    "    \"min-slope\": 2,\n",
    "    \"max-slope\": 3,\n",
    "    \"min\": 4,\n",
    "}\n",
    "\n",
    "# Lists of measures to summarize and plot\n",
    "summary_measures = [\"min-slope\", \"max-slope\", \"amplitude\", \"end-minus-start\"]\n",
    "plot_measures = [\"min\"]\n",
    "\n",
    "# Columns used for faceting and coloring the plot\n",
    "facet_column = \"fullDose\"\n",
    "facet_row = \"primingDose\"\n",
    "color_column = \"administrationMode\"\n",
    "\n",
    "# Base query to filter data for the main plot\n",
    "query_here = 'generalScalarId == \"Tumor.CancerCell\" and reduceFunction == \"at\"'\n",
    "\n",
    "# Create the main line plot using Plotly Express\n",
    "fig = px.line(\n",
    "    dfScalars.query(query_here),\n",
    "    x=\"endTime\",\n",
    "    y=\"value\",\n",
    "    facet_col=facet_column,\n",
    "    facet_row=facet_row,\n",
    "    color=color_column,\n",
    "    labels={\n",
    "        \"endTime\": \"Time [days]\",\n",
    "        \"value\": \"Amount of Tumor.CancerCell [cell count]\",\n",
    "        \"fullDose\": \"Full dose [mg]\",\n",
    "        \"primingDose\": \"Priming dose [mg]\",\n",
    "        \"administrationMode\": \"Administration mode\",\n",
    "    },\n",
    "    height=600,\n",
    "    color_discrete_map=color_map,\n",
    ")\n",
    "\n",
    "# Determine the maximum and minimum values in the filtered data for annotation positioning\n",
    "max_value = max(dfScalars.query(query_here)[\"value\"])\n",
    "min_value = min(dfScalars.query(query_here)[\"value\"])\n",
    "\n",
    "annotations = []\n",
    "\n",
    "# Iterate over each summary measure to add annotations\n",
    "for key in summary_measure_map.keys():\n",
    "    query_here_key = (\n",
    "        f'generalScalarId == \"Tumor.CancerCell\" and reduceFunction == \"{key}\"'\n",
    "    )\n",
    "\n",
    "    # Iterate over each row in the filtered data for the current summary measure\n",
    "    for i, row in dfScalars.query(query_here_key).iterrows():\n",
    "        # Determine the xref and yref for the correct subplot\n",
    "        col_idx = (\n",
    "            dfScalars.query(query_here)[facet_column]\n",
    "            .unique()\n",
    "            .tolist()\n",
    "            .index(row[facet_column])\n",
    "        )\n",
    "        row_idx = (\n",
    "            dfScalars.query(query_here)[facet_row]\n",
    "            .unique()\n",
    "            .tolist()\n",
    "            .index(row[facet_row])\n",
    "        )\n",
    "        xref = f\"x{col_idx+1}\"\n",
    "        yref = f\"y{row_idx+1}\"\n",
    "\n",
    "        # Add annotations for summary measures\n",
    "        if key in summary_measures:\n",
    "            annotations.append(\n",
    "                dict(\n",
    "                    x=35,\n",
    "                    y=(\n",
    "                        max_value\n",
    "                        - (max_value - min_value) / 5 * offset_map[row[color_column]]\n",
    "                        - (max_value - min_value) / 20 * measure_offset_map[key]\n",
    "                    ),\n",
    "                    xref=(\n",
    "                        xref if col_idx != 0 else \"x\"\n",
    "                    ),  # Use 'x' for the first column subplot\n",
    "                    yref=(\n",
    "                        yref if row_idx != 0 else \"y\"\n",
    "                    ),  # Use 'y' for the first row subplot\n",
    "                    text=f'{row[color_column]}: {summary_measure_map[key]} [{row[\"unit\"]}]: {\"%.2E\" % (row[\"value\"])}',\n",
    "                    showarrow=False,\n",
    "                    font=dict(\n",
    "                        color=color_map[row[color_column]]\n",
    "                    ),  # Use the color from the color map\n",
    "                    bgcolor=\"white\",\n",
    "                    bordercolor=\"black\",\n",
    "                    textangle=0,\n",
    "                    xanchor=\"left\",\n",
    "                    yanchor=\"top\",\n",
    "                )\n",
    "            )\n",
    "\n",
    "        # Add annotations and plot points for specific measures\n",
    "        if key in plot_measures:\n",
    "            # Determine the x and y values for the plot points\n",
    "            x_value = float(\n",
    "                dfScalars.query(\n",
    "                    f'generalScalarId == \"Tumor.CancerCell\" and reduceFunction == \"time-of-{key}\" and '\n",
    "                    f'{color_column} == \"{row[color_column]}\" and {facet_column} == \"{row[facet_column]}\" and '\n",
    "                    f'{facet_row} == \"{row[facet_row]}\"'\n",
    "                )[\"value\"].iloc[0]\n",
    "            )\n",
    "            y_value = [row[\"value\"]]\n",
    "\n",
    "            # Add scatter plot points to the main plot\n",
    "            fig.add_trace(\n",
    "                go.Scatter(\n",
    "                    x=[x_value],\n",
    "                    y=y_value,\n",
    "                    mode=\"markers\",\n",
    "                    marker_symbol=\"x\",\n",
    "                    marker=dict(color=color_map[row[color_column]]),\n",
    "                    showlegend=False,\n",
    "                ),\n",
    "                row=row_idx + 1,\n",
    "                col=col_idx + 1,\n",
    "            )\n",
    "\n",
    "            # Add annotations for the plot points\n",
    "            annotations.append(\n",
    "                dict(\n",
    "                    x=x_value,\n",
    "                    y=y_value[0],\n",
    "                    xref=(\n",
    "                        xref if col_idx != 0 else \"x\"\n",
    "                    ),  # Use 'x' for the first column subplot\n",
    "                    yref=(\n",
    "                        yref if row_idx != 0 else \"y\"\n",
    "                    ),  # Use 'y' for the first row subplot\n",
    "                    text=f'{row[color_column]}: {summary_measure_map[key]} [{row[\"unit\"]}]:<br> {\"%.2E\" % y_value[0]} at {x_value} days',\n",
    "                    showarrow=True,\n",
    "                    font=dict(\n",
    "                        color=color_map[row[color_column]]\n",
    "                    ),  # Use the color from the color map\n",
    "                    bgcolor=\"white\",\n",
    "                    bordercolor=\"black\",\n",
    "                    textangle=0,\n",
    "                    xanchor=\"left\",\n",
    "                    yanchor=\"top\",\n",
    "                    ay=50 - offset_map[row[color_column]] * 20,\n",
    "                    ax=20,\n",
    "                )\n",
    "            )\n",
    "\n",
    "# Assign the created annotations to the plot layout\n",
    "fig[\"layout\"][\"annotations\"] = annotations\n",
    "\n",
    "# Display the plot\n",
    "fig.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}