{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "83d8670f-8576-4813-99cc-63ec28572db6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import json\n",
    "import os\n",
    "import matplotlib.pyplot as plt\n",
    "from ipywidgets import interact\n",
    "from skimage.transform import resize\n",
    "\n",
    "import sys\n",
    "# import band plotter\n",
    "sys.path.append('..')\n",
    "from src.band_plotters import plot, DATA_DIRECTORY, plot_from_bands_picture"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "e7b76ff8-37a0-407b-8a6e-8c3b2462e770",
   "metadata": {},
   "outputs": [],
   "source": [
    "THRESHOLD = 8 # eV\n",
    "MAX_NUMBER_OF_BANDS_WITHIN_THRESHOLD = 189 # for 8eV # set to false to work out again\n",
    "MAX_NUMBER_K_POINTS = 105"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "293b6b2f-61d3-4f97-a8c5-f33b93e909d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_bands_that_come_near_fermi_level(bands, threshold=THRESHOLD):\n",
    "    '''\n",
    "    bands - band energies with efermi subtracted\n",
    "    threshold - distance that bands must approach (at any point!) to be included\n",
    "    '''\n",
    "    distance_from_efermi = np.abs(bands)\n",
    "    mask = (distance_from_efermi < 8).any(axis=1)\n",
    "    return bands[mask]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "447089bb-93ff-462a-ad1a-aa52422b762c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c791e0930b6440708a7cac45be535acc",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "interactive(children=(IntSlider(value=258, description='x', max=774, min=-258), Output()), _dom_classes=('widg…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<function __main__.show(x)>"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def show(x):\n",
    "\n",
    "    example = DATA_DIRECTORY/f\"bands/2dm-{x}.json\"\n",
    "    bands_dict=json.load(open(example))\n",
    "    bands_minus_efermi = np.array(bands_dict[\"bands\"][\"1\"]) - bands_dict[\"efermi\"]\n",
    "    plt.imshow(bands_minus_efermi)\n",
    "    plt.show()\n",
    "    bands_minus_efermi_in_threshold = get_bands_that_come_near_fermi_level(bands_minus_efermi)\n",
    "    plt.imshow(bands_minus_efermi_in_threshold)\n",
    "\n",
    "interact(show, x=258)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "51325454-65f0-4236-80b2-884bc6f5b8d2",
   "metadata": {},
   "source": [
    "Find maximum number of bands that satify the criteria of being within 8ev of the fermi energy at any point"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "439412f5-f4a2-4d71-85f7-8ab4844f9084",
   "metadata": {},
   "outputs": [],
   "source": [
    "if not MAX_NUMBER_OF_BANDS_WITHIN_THRESHOLD:\n",
    "    df_template = pd.read_csv(\"../fingerprints/template.csv\", index_col=\"ID\")\n",
    "    max_bands_near_efermi = 0\n",
    "    for material_id in df_template.index:\n",
    "        file_name = DATA_DIRECTORY/f\"bands/{material_id}.json\"\n",
    "        bands_dict=json.load(open(file_name))\n",
    "        bands = np.array(bands_dict[\"bands\"][\"1\"])\n",
    "        distance_from_efermi = np.abs(bands - bands_dict[\"efermi\"])\n",
    "        mask = (distance_from_efermi < THRESHOLD).any(axis=1)\n",
    "    \n",
    "        if max_bands_near_efermi < mask.sum():\n",
    "            max_bands_near_efermi = mask.sum()\n",
    "            print(material_id,\":\",  max_bands_near_efermi)\n",
    "            print(distance_from_efermi.shape[1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "0ae6801c-fc1d-4d1b-a1b7-ec5bc62f2e70",
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'MAX_ENERGY' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[8], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[43mMAX_ENERGY\u001b[49m \u001b[38;5;129;01mand\u001b[39;00m MAX_ENERGY):\n\u001b[1;32m      2\u001b[0m     df_template \u001b[38;5;241m=\u001b[39m pd\u001b[38;5;241m.\u001b[39mread_csv(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m../fingerprints/template.csv\u001b[39m\u001b[38;5;124m\"\u001b[39m, index_col\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mID\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m      4\u001b[0m     max_energy \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m0\u001b[39m\n",
      "\u001b[0;31mNameError\u001b[0m: name 'MAX_ENERGY' is not defined"
     ]
    }
   ],
   "source": [
    "\n",
    "if not (MIN_ENERGY_MINUS_EFERMI and MIN_ENERGY_MINUS_EFERMI):\n",
    "    df_template = pd.read_csv(\"../fingerprints/template.csv\", index_col=\"ID\")\n",
    "\n",
    "    bands = np.array(bands_dict[\"bands\"][\"1\"])\n",
    "\n",
    "    for material_id in df_template.index:\n",
    "        file_name = DATA_DIRECTORY/f\"bands/{material_id}.json\"\n",
    "        bands_dict=json.load(open(file_name))\n",
    "        bands = np.array(bands_dict[\"bands\"][\"1\"])-"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3461a1a8-5142-41bc-849d-83b2cc0c8125",
   "metadata": {},
   "outputs": [],
   "source": [
    "x = \"2dm-258\"\n",
    "example = DATA_DIRECTORY/f\"bands/{x}.json\"\n",
    "bands_dict=json.load(open(example))\n",
    "bands_minus_efermi = np.array(bands_dict[\"bands\"][\"1\"]) - bands_dict[\"efermi\"]\n",
    "bands_minus_efermi_in_threshold = get_bands_that_come_near_fermi_level(bands_minus_efermi)\n",
    "\n",
    "plot_from_bands_picture(\"2dm-258\", bands_minus_efermi_in_threshold)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f36aaa8b-2c04-4250-86c0-609750e17d11",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot(\"2dm-258\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9fa75018-6259-4fc7-a217-a0faa3e9dcdd",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_from_bands_picture(\"2dm-258\", bands_minus_efermi_in_threshold[50:51])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ffd13c1-f2df-4399-9cbc-161e0b9ea6f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "bands_minus_efermi_in_threshold[50:51].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e95e4ba9-e07f-4eff-8583-b7b76ecf5575",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(bands_minus_efermi_in_threshold)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6293a30a-0224-4485-a5fd-b6d62f0d3129",
   "metadata": {},
   "outputs": [],
   "source": [
    "new =resize(bands_minus_efermi_in_threshold, (500, 105), preserve_range=True, mode=\"edge\", order=0)\n",
    "plt.imshow(new)\n",
    "plt.colorbar()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b3d380e2-6185-4fdc-b0e3-33f408f44cc4",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_from_bands_picture(\"2dm-258\", new)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6bcdc7e3-2414-4add-b48b-5938eefa60cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "from PIL import Image\n",
    "im = Image.fromarray(new)\n",
    "# im = im.convert(\"L\")\n",
    "im.save(\"your_file.tiff\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "88de319a-fbcf-47ad-bb78-bea821aa7e88",
   "metadata": {},
   "outputs": [],
   "source": [
    "im.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87873ce7-4361-4b7e-8247-d6ddc2ae676e",
   "metadata": {},
   "outputs": [],
   "source": [
    "I = np.asarray(Image.open(\"your_file.tiff\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "906bfc2f-ce14-497b-8d8b-279e9fce5023",
   "metadata": {},
   "outputs": [],
   "source": [
    "I.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "852a0daf-6767-402f-9c76-2df7d4bd91bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(I)\n",
    "plt.colorbar()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7415a3d9-ae6f-4b25-8809-66d4968386e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "for i, index in enumerate(df_template.index):\n",
    "    example = DATA_DIRECTORY/f\"bands/{index}.json\"\n",
    "    bands_dict=json.load(open(example))\n",
    "    bands_minus_efermi = np.array(bands_dict[\"bands\"][\"1\"]) - bands_dict[\"efermi\"]\n",
    "    bands_minus_efermi_in_threshold = get_bands_that_come_near_fermi_level(bands_minus_efermi)\n",
    "    bands_minus_efermi_in_threshold = get_bands_that_come_near_fermi_level(bands)\n",
    "\n",
    "    resized = resize(bands_minus_efermi_in_threshold, (224, bands_minus_efermi_in_threshold.shape[1]), preserve_range=True, mode=\"edge\", order=0)  \n",
    "    im = Image.fromarray(resized)\n",
    "    im.save(f\"../images/{index}.tiff\")\n",
    "\n",
    "    if i%200 == 0:\n",
    "        print(i, \"/\", len(df_template.index))\n",
    "        break"
   ]
  }
 ],
 "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}