demo/dj_integration.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MEI Demo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import datajoint as dj\n",
    "\n",
    "dj.config[\"enable_python_native_blobs\"] = True\n",
    "dj.config[\"schema_name\"] = \"nnfabrik_tutorial\"\n",
    "\n",
    "schema = dj.schema(\"nnfabrik_tutorial\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "from matplotlib import pyplot as plt\n",
    "from torch import load\n",
    "\n",
    "from mei.main import TrainedEnsembleModelTemplate, CSRFV1SelectorTemplate, MEISeed, MEIMethod, MEITemplate\n",
    "from nnfabrik.template import TrainedModelBase\n",
    "from nnfabrik.main import Dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Define Tables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@schema\n",
    "class TrainedModel(TrainedModelBase):\n",
    "    pass\n",
    "\n",
    "\n",
    "@schema\n",
    "class TrainedEnsembleModel(TrainedEnsembleModelTemplate):\n",
    "    dataset_table = Dataset\n",
    "    trained_model_table = TrainedModel\n",
    "\n",
    "\n",
    "@schema\n",
    "class CSRFV1Selector(CSRFV1SelectorTemplate):\n",
    "    dataset_table = Dataset\n",
    "\n",
    "\n",
    "@schema\n",
    "class MEI(MEITemplate):\n",
    "    trained_model_table = TrainedEnsembleModel\n",
    "    selector_table = CSRFV1Selector"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Reset Tables For Demo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "CSRFV1Selector().drop()\n",
    "TrainedEnsembleModel().drop()\n",
    "MEIMethod().drop()\n",
    "MEISeed().drop()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Create Ensemble Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "TrainedEnsembleModel().create_ensemble(dict(dataset_fn=\"csrf_v1\", dataset_hash=\"3d94500a46b792bbb480aedfc30f9753\"), comment=\"Happy little ensemble\")\n",
    "TrainedEnsembleModel()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Populate Selector Table"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "CSRFV1Selector.populate()\n",
    "CSRFV1Selector()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Specify MEI Method Parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "method_fn = \"mei.methods.gradient_ascent\"\n",
    "method_config = dict(\n",
    "    initial=dict(path=\"mei.initial.RandomNormal\"),\n",
    "    optimizer=dict(path=\"torch.optim.SGD\", kwargs=dict(lr=0.1)),\n",
    "    stopper=dict(path=\"mei.stoppers.NumIterations\", kwargs=dict(num_iterations=1000)),\n",
    "    objectives=[dict(path=\"mei.objectives.EvaluationObjective\", kwargs=dict(interval=10))],\n",
    "    device=\"cuda\",\n",
    ")\n",
    "MEIMethod().add_method(method_fn, method_config, comment=\"My MEI method\")\n",
    "MEIMethod()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Add Seed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "MEISeed().insert1(dict(mei_seed=42))\n",
    "MEISeed()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. Generate MEIs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "MEI().populate(dict(neuron_id=188), display_progress=True)\n",
    "MEI()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 8. Look at MEIs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "neuron_id = 188\n",
    "mei_path = (MEI() & dict(ensemble_hash=0, neuron_id=neuron_id)).fetch1(\"mei\")\n",
    "plt.imshow(load(mei_path).squeeze(), cmap=\"gray\")\n",
    "plt.gca().axis(\"off\")\n",
    "os.remove(mei_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 9. Plot Evaluations Across Time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_path = (MEI() & dict(ensemble_id=0, neuron_id=neuron_id)).fetch1(\"output\")\n",
    "output = load(output_path)\n",
    "os.remove(output_path)\n",
    "\n",
    "plt.plot(\n",
    "    output[\"mei.objectives.EvaluationObjective\"][\"times\"],\n",
    "    output[\"mei.objectives.EvaluationObjective\"][\"values\"],\n",
    ")\n",
    "plt.gca().set_xlabel(\"# iteration\")\n",
    "plt.gca().set_ylabel(\"evaluation\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}