{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "9436f296-9ab5-41f0-b5d7-2ffeb8b1dbac",
   "metadata": {},
   "source": [
    "# Create data splits for AI4Spec"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "6fcaa87a-ff10-40b9-9a9a-0d1445b9799a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current\n",
      "                                 Dload  Upload   Total   Spent    Left  Speed\n",
      "100  496M  100  496M    0     0  53.6M      0  0:00:09  0:00:09 --:--:-- 58.0M0:00:02  0:00:09 43.3M\n",
      "x 132k_16_opt_eV/\n",
      "x 132k_16_opt_eV/energies.txt\n",
      "x 132k_16_opt_eV/coulomb.npz\n",
      "x 132k_16_opt_eV/free.xyz\n",
      "x 132k_16_opt_eV/data.xyz\n",
      "x 132k_16_opt_eV/spectra.npz\n"
     ]
    }
   ],
   "source": [
    "# download data\n",
    "!curl -o data.tar.gz \"https://zenodo.org/records/3386508/files/132k_16_opt_eV.tar.gz?download=1\"\n",
    "!tar -xvzf data.tar.gz"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "e63d4e98-5c30-4307-96dc-0e2c86c1986c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "coulomb.npz  data.xyz     energies.txt free.xyz     spectra.npz\n"
     ]
    }
   ],
   "source": [
    "!ls 132k_16_opt_eV"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "abbbcc33-382e-45bb-bf21-0f3eb3ba8798",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "coulomb = np.load(\"132k_16_opt_eV/coulomb.npz\")['coulomb']\n",
    "energies = np.loadtxt(\"132k_16_opt_eV/energies.txt\")\n",
    "spectra = np.load(\"132k_16_opt_eV/spectra.npz\")['spectra']\n",
    "N_total = len(energies)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "791d36f2-807d-4c2a-b969-2e72a32cd2f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# we will select first 7000 molecules\n",
    "np.savez(\"coulomb_7000.npz\", coulomb[:7000])\n",
    "np.savez(\"spectra_7000.npz\", spectra[:7000])\n",
    "np.savez(\"energies_7000.npz\", energies[:7000])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "ec5da9ed-ea7c-43e4-9d01-92e465348d7c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-510.77096 -2.53181\n"
     ]
    }
   ],
   "source": [
    "print(min(energies.flatten()), max(energies.flatten()))"
   ]
  }
 ],
 "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
