examples/notebooks/Wren.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "from torch import __version__ as TORCH_VERSION\n",
    "\n",
    "print(f\"{TORCH_VERSION=}\")\n",
    "\n",
    "!pip install -U git+https://github.com/CompRhys/aviary.git  # install aviary\n",
    "!wget -O taata.json.gz https://figshare.com/ndownloader/files/34423997"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import gzip\n",
    "import json\n",
    "\n",
    "import pandas as pd\n",
    "import torch\n",
    "from pymatgen.core import Structure\n",
    "from sklearn.model_selection import train_test_split as split\n",
    "\n",
    "from aviary.utils import results_multitask, train_ensemble\n",
    "from aviary.wren.data import WyckoffData\n",
    "from aviary.wren.data import collate_batch as wren_cb\n",
    "from aviary.wren.model import Wren\n",
    "from aviary.wren.utils import (\n",
    "    count_wyckoff_positions,\n",
    "    get_protostructure_label_from_spglib,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {},
   "outputs": [],
   "source": [
    "with gzip.open(\"taata.json.gz\", \"r\") as fin:\n",
    "    json_bytes = fin.read()\n",
    "\n",
    "json_str = json_bytes.decode(\"utf-8\")\n",
    "data = json.loads(json_str)\n",
    "\n",
    "df = pd.DataFrame(data[\"data\"], columns=data[\"columns\"])\n",
    "\n",
    "df[\"final_structure\"] = [Structure.from_dict(x) for x in df.final_structure]\n",
    "\n",
    "df[\"composition\"] = [x.composition.reduced_formula for x in df.final_structure]\n",
    "df[\"volume_per_atom\"] = [x.volume / len(x) for x in df.final_structure]\n",
    "df[\"wyckoff\"] = df[\"final_structure\"].map(get_protostructure_label_from_spglib)\n",
    "\n",
    "df = df[df.wyckoff.map(count_wyckoff_positions) < 16]\n",
    "df[\"n_sites\"] = df.final_structure.map(len)\n",
    "df = df[df.n_sites < 64]\n",
    "df = df[df.volume_per_atom < 500]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {},
   "outputs": [],
   "source": [
    "resume = False\n",
    "fine_tune = None\n",
    "transfer = None\n",
    "\n",
    "optim = \"AdamW\"\n",
    "learning_rate = 3e-4\n",
    "momentum = 0.9\n",
    "weight_decay = 1e-6\n",
    "batch_size = 128\n",
    "workers = 0\n",
    "device = torch.device(\"cuda\") if torch.cuda.is_available() else torch.device(\"cpu\")\n",
    "\n",
    "targets = [\"E_vasp_per_atom\"]\n",
    "tasks = [\"regression\"]\n",
    "losses = [\"L1\"]\n",
    "robust = True\n",
    "\n",
    "data_seed = 42\n",
    "test_size = 0.2\n",
    "sample = 1\n",
    "\n",
    "ensemble = 1\n",
    "run_id = 1\n",
    "epochs = 100\n",
    "log = False\n",
    "\n",
    "# NOTE setting workers to zero means that the data is loaded in the main\n",
    "# process and enables caching\n",
    "\n",
    "data_params = {\n",
    "    \"batch_size\": batch_size,\n",
    "    \"num_workers\": workers,\n",
    "    \"pin_memory\": False,\n",
    "    \"shuffle\": True,\n",
    "}\n",
    "\n",
    "setup_params = {\n",
    "    \"optim\": optim,\n",
    "    \"learning_rate\": learning_rate,\n",
    "    \"weight_decay\": weight_decay,\n",
    "    \"momentum\": momentum,\n",
    "    \"device\": device,\n",
    "}\n",
    "\n",
    "restart_params = {\n",
    "    \"resume\": resume,\n",
    "    \"fine_tune\": fine_tune,\n",
    "    \"transfer\": transfer,\n",
    "}\n",
    "\n",
    "task_dict = dict(zip(targets, tasks))\n",
    "loss_dict = dict(zip(targets, losses))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.manual_seed(0)  # ensure reproducible results\n",
    "\n",
    "elem_embedding = \"matscholar200\"\n",
    "sym_emb = \"bra-alg-off\"\n",
    "model_name = \"wren-reg-test\"\n",
    "\n",
    "data_params[\"collate_fn\"] = wren_cb\n",
    "data_params[\"shuffle\"] = True\n",
    "\n",
    "dataset = WyckoffData(\n",
    "    df=df, elem_embedding=elem_embedding, sym_emb=sym_emb, task_dict=task_dict\n",
    ")\n",
    "n_targets = dataset.n_targets\n",
    "elem_emb_len = dataset.elem_emb_len\n",
    "sym_emb_len = dataset.sym_emb_len\n",
    "\n",
    "train_idx = list(range(len(dataset)))\n",
    "\n",
    "print(f\"using {test_size} of training set as test set\")\n",
    "train_idx, test_idx = split(train_idx, random_state=data_seed, test_size=test_size)\n",
    "test_set = torch.utils.data.Subset(dataset, test_idx)\n",
    "\n",
    "print(\"No validation set used, using test set for evaluation purposes\")\n",
    "# NOTE that when using this option care must be taken not to\n",
    "# peak at the test-set. The only valid model to use is the one\n",
    "# obtained after the final epoch where the epoch count is\n",
    "# decided in advance of the experiment.\n",
    "val_set = test_set\n",
    "\n",
    "train_set = torch.utils.data.Subset(dataset, train_idx[0::sample])\n",
    "\n",
    "model_params = {\n",
    "    \"task_dict\": task_dict,\n",
    "    \"robust\": robust,\n",
    "    \"n_targets\": n_targets,\n",
    "    \"elem_emb_len\": elem_emb_len,\n",
    "    \"elem_fea_len\": 32,\n",
    "    \"sym_emb_len\": sym_emb_len,\n",
    "    \"sym_fea_len\": 32,\n",
    "    \"n_graph\": 3,\n",
    "    \"elem_heads\": 1,\n",
    "    \"elem_gate\": [256],\n",
    "    \"elem_msg\": [256],\n",
    "    \"cry_heads\": 1,\n",
    "    \"cry_gate\": [256],\n",
    "    \"cry_msg\": [256],\n",
    "    \"trunk_hidden\": [128, 128],\n",
    "    \"out_hidden\": [64, 64],\n",
    "}\n",
    "\n",
    "train_ensemble(\n",
    "    model_class=Wren,\n",
    "    model_name=model_name,\n",
    "    run_id=run_id,\n",
    "    ensemble_folds=ensemble,\n",
    "    epochs=epochs,\n",
    "    train_set=train_set,\n",
    "    val_set=val_set,\n",
    "    log=log,\n",
    "    data_params=data_params,\n",
    "    setup_params=setup_params,\n",
    "    restart_params=restart_params,\n",
    "    model_params=model_params,\n",
    "    loss_dict=loss_dict,\n",
    ")\n",
    "\n",
    "data_params[\"shuffle\"] = False  # need fixed data order due to ensembling\n",
    "\n",
    "roost_results_dict = results_multitask(\n",
    "    model_class=Wren,\n",
    "    model_name=model_name,\n",
    "    run_id=run_id,\n",
    "    ensemble_folds=ensemble,\n",
    "    test_set=test_set,\n",
    "    data_params=data_params,\n",
    "    robust=robust,\n",
    "    task_dict=task_dict,\n",
    "    device=device,\n",
    "    eval_type=\"checkpoint\",\n",
    "    save_results=False,\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10.8 ('py310')",
   "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.10.8"
  },
  "vscode": {
   "interpreter": {
    "hash": "8022b3e932e045c760cb4633b91dd1cb8bc60d104ca9808334cbd1645adbe837"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}