{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "2b714898",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "from pina.problem import SpatialProblem, TimeDependentProblem\n",
    "from pina.operators import laplacian, grad\n",
    "from pina.geometry import CartesianDomain\n",
    "from pina.solvers import *\n",
    "from pina.trainer import Trainer\n",
    "from pina.equation import Equation\n",
    "from pina.equation.equation_factory import FixedValue\n",
    "from pina import Condition, Plotter\n",
    "from pytorch_lightning.loggers import TensorBoardLogger\n",
    "from pina.model import FeedForward\n",
    "from pina.callbacks import MetricTracker\n",
    "\n",
    "runs = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "6e9f489f",
   "metadata": {},
   "outputs": [],
   "source": [
    "class KleinGordon(TimeDependentProblem,SpatialProblem):\n",
    "    output_variables = ['u']\n",
    "    spatial_domain = CartesianDomain({'x':[-1,1],'y':[-1,1]})\n",
    "    temporal_domain = CartesianDomain({'t': [0, 10]})\n",
    "    def KG(input_,output_):\n",
    "        u = output_.extract(['u'])\n",
    "        x = input_.extract(['x'])\n",
    "        y = input_.extract(['y'])\n",
    "        t = input_.extract(['t'])\n",
    "        u_t = grad(output_, input_, components=['u'], d=['t'])\n",
    "        u_tt = grad(u_t, input_, components=['dudt'], d=['t'])\n",
    "        nabla_u = laplacian(output_, input_, components=['u'], d=['x', 'y'])\n",
    "        temp = (x + y) * torch.cos(2*t) + (x * y) * torch.sin(2*t)\n",
    "        f = temp**2 - 4*temp\n",
    "        \n",
    "        return u_tt - nabla_u + u**2 - f\n",
    "    def boundary_condition(input_,output_):\n",
    "        u = output_.extract(['u'])\n",
    "        x = input_.extract(['x'])\n",
    "        y = input_.extract(['y'])\n",
    "        t = input_.extract(['t'])\n",
    "        temp = (x + y) * torch.cos(2*t) + (x * y) * torch.sin(2*t)\n",
    "        return u - temp\n",
    "        \n",
    "        \n",
    "    def initial_condition(input_,output_):\n",
    "        x = input_.extract(['x'])\n",
    "        y = input_.extract(['y'])\n",
    "        u_exp = x+y\n",
    "        return output_.extract(['u']) - u_exp\n",
    "    \n",
    "    def kg_sol(self,pts):\n",
    "        x = pts.extract(['x'])\n",
    "        y = pts.extract(['y'])\n",
    "        t = pts.extract(['t'])\n",
    "        return (x + y) * torch.cos(2*t) + (x * y) * torch.sin(2*t)\n",
    "    truth_solution = kg_sol\n",
    "    \n",
    "    conditions = {\n",
    "        'gamma1': Condition(location=CartesianDomain({'x': [-1, 1], 'y':  1, 't': [0, 10]}), equation=Equation(boundary_condition)),\n",
    "        'gamma2': Condition(location=CartesianDomain({'x': [-1, 1], 'y': -1, 't': [0, 10]}), equation=Equation(boundary_condition)),\n",
    "        'gamma3': Condition(location=CartesianDomain({'x':  1, 'y': [-1, 1], 't': [0, 10]}), equation=Equation(boundary_condition)),\n",
    "        'gamma4': Condition(location=CartesianDomain({'x': -1, 'y': [-1, 1], 't': [0, 10]}), equation=Equation(boundary_condition)),\n",
    "        't0': Condition(location=CartesianDomain({'x': [-1, 1], 'y': [-1, 1], 't': 0}), equation=Equation(initial_condition)),\n",
    "        'D': Condition(location=CartesianDomain({'x': [-1, 1], 'y': [-1, 1], 't': [0, 10]}), equation=Equation(KG)),\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "d8eb6c25",
   "metadata": {},
   "outputs": [],
   "source": [
    "problem = KleinGordon()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "db9d4097",
   "metadata": {},
   "outputs": [],
   "source": [
    "problem.discretise_domain(64+64+64, 'random', locations=['D', 't0', 'gamma1', 'gamma2', 'gamma3', 'gamma4'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "a47dcaea",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = FeedForward(\n",
    "    layers=[64,64,64,64],\n",
    "    func=torch.nn.Tanh,\n",
    "    output_dimensions=len(problem.output_variables),\n",
    "    input_dimensions=len(problem.input_variables)\n",
    ")\n",
    "pinn = PINN(problem, model, optimizer_kwargs={'lr':0.001, 'weight_decay':1e-8})\n",
    "gpinn = GPINN(problem, model, optimizer_kwargs={'lr':0.001, 'weight_decay':1e-8})\n",
    "causalpinn = CausalPINN(problem, model, optimizer_kwargs={'lr':0.001, 'weight_decay':1e-8})\n",
    "sapinin = SAPINN(problem, model)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "6bd51d0c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "def error_giver(model):\n",
    "    x = np.linspace(-1,1,10)\n",
    "    y = np.linspace(-1,1,10)\n",
    "    t = np.linspace(0,10,20)\n",
    "    test_set = []\n",
    "    for X in x:\n",
    "        for Y in y:\n",
    "            for T in t:\n",
    "                test_set.append([X,Y,T])\n",
    "            \n",
    "    test_set = torch.Tensor(test_set)  \n",
    "\n",
    "    predicted_output= model(test_set).detach().numpy()\n",
    "    def kg_sol(pts):\n",
    "        x = pts[:,0]\n",
    "        y = pts[:,1]\n",
    "        t = pts[:,2]\n",
    "        return (x + y) * torch.cos(2*t) + (x * y) * torch.sin(2*t)\n",
    "    truth_output = kg_sol(test_set).detach().numpy().reshape((100*20,1))\n",
    "\n",
    "    error = np.linalg.norm(predicted_output-truth_output) / np.linalg.norm(truth_output)\n",
    "    return error"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "72a88e7b",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/vemuri/anaconda3/envs/pina/lib/python3.11/site-packages/lightning_fabric/plugins/environments/slurm.py:204: The `srun` command is available on your system but is not used. HINT: If your intention is to run Lightning on SLURM, prepend your python command with `srun` like so: srun python /home/vemuri/anaconda3/envs/pina/lib/python3.11/site ...\n",
      "GPU available: True (cuda), used: True\n",
      "TPU available: False, using: 0 TPU cores\n",
      "IPU available: False, using: 0 IPUs\n",
      "HPU available: False, using: 0 HPUs\n",
      "/home/vemuri/anaconda3/envs/pina/lib/python3.11/site-packages/pytorch_lightning/trainer/connectors/logger_connector/logger_connector.py:75: Starting from v1.9.0, `tensorboardX` has been removed as a dependency of the `pytorch_lightning` package, due to potential conflicts with other packages in the ML ecosystem. For this reason, `logger=True` will use `CSVLogger` as the default logger, unless the `tensorboard` or `tensorboardX` packages are found. Please `pip install lightning[extra]` or one of them to enable TensorBoard support by default\n",
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]\n",
      "/home/vemuri/anaconda3/envs/pina/lib/python3.11/site-packages/pytorch_lightning/loops/fit_loop.py:298: The number of training batches (1) is smaller than the logging interval Trainer(log_every_n_steps=50). Set a lower value for log_every_n_steps if you want to see logs for the training epoch.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0:   0%| | 0/1 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/vemuri/anaconda3/envs/pina/lib/python3.11/site-packages/torch/autograd/graph.py:744: UserWarning: Attempting to run cuBLAS, but there was no current CUDA context! Attempting to set the primary context... (Triggered internally at ../aten/src/ATen/cuda/CublasHandlePool.cpp:135.)\n",
      "  return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 29999: 100%|█| 1/1 [00:00<00:00"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "`Trainer.fit` stopped: `max_epochs=30000` reached.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 29999: 100%|█| 1/1 [00:00<00:00\n",
      "0.1748321\n"
     ]
    }
   ],
   "source": [
    "for i in range(runs):\n",
    "    trainer = Trainer(pinn, max_epochs=30000, callbacks=[MetricTracker()], accelerator='gpu',enable_model_summary=False,) \n",
    "    trainer.train()\n",
    "    print(error_giver(model))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "abc4ecf5",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True (cuda), used: True\n",
      "TPU available: False, using: 0 TPU cores\n",
      "IPU available: False, using: 0 IPUs\n",
      "HPU available: False, using: 0 HPUs\n",
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 29999: 100%|█| 1/1 [00:00<00:00"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "`Trainer.fit` stopped: `max_epochs=30000` reached.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 29999: 100%|█| 1/1 [00:00<00:00\n",
      "0.12129881\n"
     ]
    }
   ],
   "source": [
    "for i in range(runs):\n",
    "    trainer2 = Trainer(gpinn, max_epochs=30000, callbacks=[MetricTracker()], accelerator='gpu',enable_model_summary=False,) \n",
    "    trainer2.train()\n",
    "    print(error_giver(4model))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9fad6b67",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "fb99476b",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True (cuda), used: True\n",
      "TPU available: False, using: 0 TPU cores\n",
      "IPU available: False, using: 0 IPUs\n",
      "HPU available: False, using: 0 HPUs\n",
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 29999: 100%|█| 1/1 [00:00<00:00"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "`Trainer.fit` stopped: `max_epochs=30000` reached.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 29999: 100%|█| 1/1 [00:00<00:00\n",
      "0.18318304\n"
     ]
    }
   ],
   "source": [
    "for i in range(runs):\n",
    "    trainer4 = Trainer(sapinin, max_epochs=30000, callbacks=[MetricTracker()], accelerator='gpu',enable_model_summary=False,) \n",
    "    trainer4.train()\n",
    "    print(error_giver(model))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "f89cf5ef",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"\\nfor i in range(runs):\\n    trainer3 = Trainer(causalpinn, max_epochs=30000, callbacks=[MetricTracker()], accelerator='gpu',enable_model_summary=False,) \\n    trainer3.train()\\n    print(error_giver(model))\\n\""
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\"\"\"\n",
    "for i in range(runs):\n",
    "    trainer3 = Trainer(causalpinn, max_epochs=30000, callbacks=[MetricTracker()], accelerator='gpu',enable_model_summary=False,) \n",
    "    trainer3.train()\n",
    "    print(error_giver(model))\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "ef0cff78",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Plotting at t=0\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 69744 characters]",
      "text/plain": [
       "<Figure size 1600x600 with 6 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Plotting at t=5\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 70200 characters]",
      "text/plain": [
       "<Figure size 1600x600 with 6 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Plotting at t=10\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 64240 characters]",
      "text/plain": [
       "<Figure size 1600x600 with 6 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plotter = Plotter()\n",
    "\n",
    "# plotting at fixed time t = 0.0\n",
    "print('Plotting at t=0')\n",
    "plotter.plot(pinn, fixed_variables={'t': 0.0})\n",
    "\n",
    "# plotting at fixed time t = 0.5\n",
    "print('Plotting at t=5')\n",
    "plotter.plot(pinn, fixed_variables={'t':5})\n",
    "\n",
    "# plotting at fixed time t = 1.\n",
    "print('Plotting at t=10')\n",
    "plotter.plot(pinn, fixed_variables={'t': 10})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef9b0fcb",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44d9cde4",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff397a04",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0aaed6ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6308ab81",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "61ab8f88",
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'truth_output' 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[12], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m truth_output\u001b[38;5;241m.\u001b[39mshape\n",
      "\u001b[0;31mNameError\u001b[0m: name 'truth_output' is not defined"
     ]
    }
   ],
   "source": [
    "truth_output.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "127968ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "predicted_output.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8483904d",
   "metadata": {},
   "outputs": [],
   "source": [
    "error"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a6ab97e4",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8c09bc9",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pina",
   "language": "python",
   "name": "pina"
  },
  "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
