aimm-server-data-access.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67684c22",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tiled.client import from_uri, logout_all\n",
    "from getpass import getpass\n",
    "\n",
    "import functools\n",
    "import ipywidgets as widgets\n",
    "import numpy as np\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9f792d12",
   "metadata": {},
   "source": [
    "# Access the online data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "393ea8c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "logout_all() # For demo purposes, ensure logout first"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cdd9f7b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "client = from_uri(\"https://aimm.lbl.gov/api\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c52c82c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "client"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e3672187",
   "metadata": {},
   "source": [
    "# Browse through all the datasets by code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "04e8bf17",
   "metadata": {},
   "outputs": [],
   "source": [
    "client[\"element\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e47b822",
   "metadata": {},
   "outputs": [],
   "source": [
    "client[\"element\"][\"Cr\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ac7b694b",
   "metadata": {},
   "outputs": [],
   "source": [
    "client[\"element\"][\"Cr\"]['edge']['K']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8beee3ce",
   "metadata": {},
   "source": [
    "# Browse through all the datasets with widgets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a582a4b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "class DropDownGroup:\n",
    "    \n",
    "    def __init__(self, root_node):\n",
    "        self.root_node = root_node\n",
    "        root_list = list(self.root_node)\n",
    "        root_list.insert(0, \"-\")\n",
    "        \n",
    "        dd_root = widgets.Dropdown(\n",
    "            options=root_list,\n",
    "            value=root_list[0],\n",
    "            description='Node_0:',\n",
    "            disabled=False,\n",
    "        )\n",
    "        \n",
    "        self.dd_box = widgets.VBox([dd_root])\n",
    "        self.selected_node = None\n",
    "           \n",
    "    def find_child_index(self, description):\n",
    "        counter = 0\n",
    "        for child in self.dd_box.children:\n",
    "            if child.description == description:\n",
    "                break\n",
    "            counter += 1\n",
    "        return counter\n",
    "    \n",
    "    def detect_change(self, change, node=client):\n",
    "        if change['new'] == '-':\n",
    "            index = self.find_child_index(change['owner'].description)\n",
    "            self.dd_box.children = self.dd_box.children[:index+1]\n",
    "        else:\n",
    "            child_id = int(change['owner'].description.split('_')[-1][:-1])\n",
    "            if child_id < len(self.dd_box.children)-1:\n",
    "                self.dd_box.children = self.dd_box.children[:child_id+1]\n",
    "\n",
    "            if node[change['new']].item['attributes']['structure_family'] == 'node':\n",
    "                next_level = list(node[change['new']])\n",
    "                next_level.insert(0, \"-\")\n",
    "                dd_next_level = widgets.Dropdown(\n",
    "                    options=next_level,\n",
    "                    value=next_level[0],\n",
    "                    description=f'Node_{child_id+1}:',\n",
    "                    disabled=False,\n",
    "                )\n",
    "                self.dd_box.children += (dd_next_level,)\n",
    "                dd_next_level.observe(functools.partial(self.detect_change, node=node[change['new']]), names='value')\n",
    "            else:\n",
    "                df_data = widgets.HTML(\n",
    "                    value=node[change['new']].read().head().to_html(),\n",
    "                    placeholder='Dataframe',\n",
    "                    description='Dataframe:',\n",
    "                )\n",
    "                self.dd_box.children += (df_data,)\n",
    "                self.selected_node = node[change['new']]\n",
    "                \n",
    "    def display_box(self):\n",
    "        self.dd_box.children[0].observe(self.detect_change, names='value')\n",
    "        display(self.dd_box)\n",
    "        \n",
    "    def get_selected_node(self):\n",
    "        return self.selected_node"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87a34827",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run dynamic dropdown menus\n",
    "dd_group = DropDownGroup(client)\n",
    "dd_group.display_box()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb95991b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Use this method to use the dataset that was selected from the dropdown menus previously\n",
    "dd_node = dd_group.get_selected_node()\n",
    "dd_node[[\"energy\", 'i0']]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "47c897c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "dd_node.metadata"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9cdd5383",
   "metadata": {},
   "source": [
    "# Filtering rows with sliders"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93c1d790",
   "metadata": {},
   "outputs": [],
   "source": [
    "# save the data in our local environment\n",
    "df = dd_node.read()\n",
    "\n",
    "preslit_data = df['Preslit']\n",
    "energy_data = df['energy']\n",
    "\n",
    "print(type(df))\n",
    "print(df.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68dfc8ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "preslit_step = (max(preslit_data) - min(preslit_data))/(len(preslit_data)*2)\n",
    "preslit_value = (max(preslit_data) - min(preslit_data)) * np.random.random() + min(preslit_data)\n",
    "slider = widgets.FloatSlider(\n",
    "    value=preslit_value,\n",
    "    min=min(preslit_data),\n",
    "    max=max(preslit_data),\n",
    "    step=preslit_step,\n",
    "    description='Test:',\n",
    "    disabled=False,\n",
    "    continuous_update=False,\n",
    "    orientation='horizontal',\n",
    "    readout=True,\n",
    "    readout_format='.1f',\n",
    ")\n",
    "\n",
    "display(slider)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e78cdd2c",
   "metadata": {},
   "outputs": [],
   "source": [
    "energy_step = (max(energy_data) - min(energy_data))/(len(energy_data)*2)\n",
    "energy_value_min = round(((max(energy_data) - min(energy_data))/2) * np.random.random() + min(energy_data),5)\n",
    "energy_value_max = round(max(energy_data) - ((max(energy_data) - min(energy_data))/2) * np.random.random(),5)\n",
    "range_slider = widgets.FloatRangeSlider(\n",
    "    value=[energy_value_min, energy_value_max],\n",
    "    min=min(energy_data),\n",
    "    max=max(energy_data),\n",
    "    step=energy_step,\n",
    "    description='Range:',\n",
    "    disabled=False,\n",
    "    continuous_update=False,\n",
    "    orientation='horizontal',\n",
    "    readout=False,\n",
    "    readout_format='.5f',\n",
    ")\n",
    "\n",
    "min_limit = widgets.Text(\n",
    "    value=\"0.0, 0.0\",\n",
    "    placeholder='Limits',\n",
    "    description='Limits:',\n",
    "    disabled=False\n",
    ")\n",
    "\n",
    "link = widgets.jslink((range_slider, 'value'), (min_limit, 'value'))\n",
    "\n",
    "display(range_slider, min_limit)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1fdfb652",
   "metadata": {},
   "outputs": [],
   "source": [
    "# A range slider can be used to define the max and min limits the filter applied to the dataframe\n",
    "filtered_df = df[df['energy'].between(range_slider.value[0], range_slider.value[1])]\n",
    "print(filtered_df.shape)\n",
    "filtered_df"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bf8ddace",
   "metadata": {},
   "source": [
    "# Filtering columns with a selection box"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29251337",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_column_names = list(df.columns)\n",
    "df_selection = widgets.SelectMultiple(\n",
    "    options=df_column_names,\n",
    "    value=[df_column_names[0]],\n",
    "    #rows=10,\n",
    "    description='dataframe',\n",
    "    disabled=False\n",
    ")\n",
    "\n",
    "display(df_selection)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8312e485",
   "metadata": {},
   "outputs": [],
   "source": [
    "xdi_df = df[list(df_selection.value)]\n",
    "xdi_df"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "141b5b8e",
   "metadata": {},
   "source": [
    "# Creating simple plots with selection boxes and dataframes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c7f3c9d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "column_list = list(df.columns)\n",
    "x_axis = widgets.Dropdown(\n",
    "    options=column_list,\n",
    "    value=column_list[0],\n",
    "    description='X axis:',\n",
    "    disabled=False,\n",
    ")\n",
    "\n",
    "y_axis = widgets.Dropdown(\n",
    "    options=column_list,\n",
    "    value=column_list[0],\n",
    "    description='Y axis:',\n",
    "    disabled=False,\n",
    ")\n",
    "\n",
    "display(x_axis, y_axis)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "013991a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "df.plot(x_axis.value, y_axis.value)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "24436477",
   "metadata": {},
   "source": [
    "# More plots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e40e54e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from aimmdb.tree import RawMongoQuery as RawMongo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4c75a9b",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "ax.set_title(client[\"dataset\"][\"aimm_ncm\"][\"sample\"][\"h6bjpUSB4Qc\"].metadata[\"sample\"][\"name\"])\n",
    "\n",
    "for k, v in client[\"dataset\"][\"aimm_ncm\"][\"sample\"][\"h6bjpUSB4Qc\"][\"uid\"].search(RawMongo({\"metadata.charge.cycle\" : 0})).items():\n",
    "    df = v.read()\n",
    "    element = v.metadata[\"element\"][\"symbol\"]\n",
    "    edge = v.metadata[\"element\"][\"edge\"]\n",
    "    if \"mutrans\" in df.columns:\n",
    "        df.plot(\"energy\", \"mutrans\", ax=ax, label=f\"{element}-{edge}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "64e7daef",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "for k, v in client[\"dataset\"][\"aimm_ncm\"][\"sample\"][\"h6bjpUSB4Qc\"][\"element\"][\"Ni\"][\"edge\"][\"L3\"][\"uid\"].items():\n",
    "    charge = v.metadata[\"charge\"]\n",
    "    cycle, voltage, state = [charge[k] for k in [\"cycle\", \"voltage\", \"state\"]]\n",
    "    df = v.read()\n",
    "    df.plot(\"energy\", \"mu_tey\", ax=ax, label=f\"cycle {cycle} {voltage}V {state}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b17bd33f",
   "metadata": {},
   "source": [
    "# Search and queries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d0892fff",
   "metadata": {},
   "outputs": [],
   "source": [
    "# register queries\n",
    "\n",
    "from tiled.query_registration import QueryTranslationRegistry, register\n",
    "from dataclasses import dataclass\n",
    "import collections\n",
    "import json\n",
    "\n",
    "@register(name=\"raw_mongo\", overwrite=True)\n",
    "@dataclass\n",
    "class RawMongo:\n",
    "    \"\"\"\n",
    "    Run a MongoDB query against a given collection.\n",
    "    \"\"\"\n",
    "\n",
    "    query: str  # We cannot put a dict in a URL, so this a JSON str.\n",
    "\n",
    "    def __init__(self, query):\n",
    "        if isinstance(query, collections.abc.Mapping):\n",
    "            query = json.dumps(query)\n",
    "        self.query = query"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9015472f",
   "metadata": {},
   "outputs": [],
   "source": [
    "client[\"uid\"].search(RawMongo({\"metadata.sample.dataset\" : \"aimm_ncm\", \"metadata.element.edge\" : \"K\", \"metadata.element.symbol\" : \"Ni\"}))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "84800290",
   "metadata": {},
   "outputs": [],
   "source": [
    "client[\"dataset\"][\"aimm_ncm\"][\"sample\"][\"h6bjpUSB4Qc\"][\"element\"][\"Ni\"][\"edge\"][\"L3\"][\"uid\"].search(RawMongo({\"metadata.charge.cycle\" : 0})).values_indexer[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50e6bbb1",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}