{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "880b887f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import numpy as np\n",
    "from contextualized.modules import MLP\n",
    "from contextualized.functions import identity_link, identity\n",
    "\n",
    "\n",
    "n, x_dim, y_dim = 1000, 100, 20\n",
    "X = torch.rand((n, x_dim)) * 2 - 1\n",
    "Y = torch.rand((n, y_dim)) * 2 - 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "ed6d68dd",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "time: 0.13069415092468262\n"
     ]
    }
   ],
   "source": [
    "mse_loss = lambda y_true, y_pred: ((y_true - y_pred)**2).mean()\n",
    "\n",
    "def train(model):\n",
    "    opt = torch.optim.Adam(model.parameters(), lr=1e-3)\n",
    "    start = time.time()\n",
    "    for _ in range(100):\n",
    "        opt.zero_grad()\n",
    "        loss = mse_loss(Y, model(X))\n",
    "        loss.backward()\n",
    "        opt.step()\n",
    "    end = time.time()\n",
    "    print(f\"time: {end - start}\")\n",
    "    \n",
    "train(MLP(x_dim, y_dim, 25, 2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "947d7a00",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "time: 3.776171922683716\n"
     ]
    }
   ],
   "source": [
    "# Current implementation\n",
    "class NGAM(nn.Module):\n",
    "    \"\"\"\n",
    "    Neural generalized additive model\n",
    "    \"\"\"\n",
    "    def __init__(self, input_dim, output_dim, width, layers, activation=nn.ReLU, link_fn=identity_link):\n",
    "        super(NGAM, self).__init__()\n",
    "        self.intput_dim = input_dim\n",
    "        self.output_dim = output_dim\n",
    "        self.nams = nn.ModuleList([MLP(1, output_dim, width, layers, activation=activation, link_fn=identity_link) for _ in range(input_dim)])\n",
    "        self.link_fn = link_fn\n",
    "\n",
    "    def forward(self, x):\n",
    "        batch_size = x.shape[0]\n",
    "        ret = torch.zeros((batch_size, self.output_dim))\n",
    "        for i, nam in enumerate(self.nams):\n",
    "            ret += nam(x[:, i].unsqueeze(-1))\n",
    "        return self.link_fn(ret)\n",
    "    \n",
    "train(NGAM(x_dim, y_dim, 25, 1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "81dd045a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "time: 5.214995861053467\n"
     ]
    }
   ],
   "source": [
    "# Current implementation\n",
    "class NGAM(nn.Module):\n",
    "    \"\"\"\n",
    "    Neural generalized additive model\n",
    "    \"\"\"\n",
    "    def __init__(self, input_dim, output_dim, width, layers, activation=nn.ReLU, link_fn=identity_link):\n",
    "        super(NGAM, self).__init__()\n",
    "        self.intput_dim = input_dim\n",
    "        self.output_dim = output_dim\n",
    "        self.nams = nn.ModuleList([MLP(1, output_dim, width, layers, activation=activation, link_fn=identity_link) for _ in range(input_dim)])\n",
    "        self.link_fn = link_fn\n",
    "        self.register_buffer('ret_buffer', torch.zeros(input_dim, 1000, output_dim))\n",
    "\n",
    "    def forward(self, x):\n",
    "        batch_size = x.shape[0]\n",
    "        ret = [] # self.nams[0](x[:, 0].unsqueeze(-1))\n",
    "        for i, nam in enumerate(self.nams):\n",
    "            ret += [nam(x[:, i].unsqueeze(-1)).unsqueeze(0)]\n",
    "        return self.link_fn(torch.cat(ret, dim=0).sum(dim=0))\n",
    "    \n",
    "train(NGAM(x_dim, y_dim, 25, 1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "cba2ce8d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "time: 3.8774161338806152\n"
     ]
    }
   ],
   "source": [
    "class FastNGAM(nn.Module):\n",
    "    \"\"\"\n",
    "    Fast training neural generalized additive model: requires extra memory\n",
    "    \"\"\"\n",
    "    def __init__(self, input_dim, output_dim, width, layers, activation=nn.ReLU, link_fn=identity_link):\n",
    "        super(NGAM, self).__init__()\n",
    "        self.intput_dim = input_dim\n",
    "        self.output_dim = output_dim\n",
    "        self.weights1 = nn.parameter.Parameter(torch.rand() * 2e-2 - 1e-2, requires_grad=True)\n",
    "        self.nams = nn.ModuleList([MLP(1, output_dim, width, layers, activation=activation, link_fn=identity_link) for _ in range(input_dim)])\n",
    "        self.link_fn = link_fn\n",
    "        self.register_buffer('ret_buffer', torch.zeros(input_dim, 1000, output_dim))\n",
    "\n",
    "    def forward(self, x):\n",
    "        ret = self.nams[0](x[:, 0].unsqueeze(-1))\n",
    "        for i, nam in enumerate(self.nams[1:]):\n",
    "            ret += nam(x[:, i].unsqueeze(-1))\n",
    "        return self.link_fn(ret)\n",
    "    \n",
    "train(NGAM(x_dim, y_dim, 25, 1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b30fc21d",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:contextualized]",
   "language": "python",
   "name": "conda-env-contextualized-py"
  },
  "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
