{
 "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 PINN\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"
   ]
  },
  {
   "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(500, 'random', locations=['D', 't0', 'gamma1', 'gamma2', 'gamma3', 'gamma4'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "a47dcaea",
   "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 1:   0%| | 0/1 [00:00<?, ?it/s, v_num=12, gamma1_loss=0.903, gamma2_loss=0.773, "
     ]
    },
    {
     "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, 19.88it/s, v_num=12, gamma1_loss=0.000187, gamm"
     ]
    },
    {
     "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,  7.27it/s, v_num=12, gamma1_loss=0.000187, gamm\n"
     ]
    }
   ],
   "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",
    "trainer = Trainer(pinn, max_epochs=30000, callbacks=[MetricTracker()], accelerator='gpu',enable_model_summary=False,) \n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "ef0cff78",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Plotting at t=0\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 72572 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, 65516 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, 58308 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": 14,
   "id": "ef9b0fcb",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "x = np.linspace(-1,1,100)\n",
    "y = np.linspace(-1,1,100)\n",
    "t = np.linspace(0,10,100)\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)           "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44d9cde4",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "ff397a04",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1000000, 3])"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test_set.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "0aaed6ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "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",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "6308ab81",
   "metadata": {},
   "outputs": [],
   "source": [
    "truth_output = kg_sol(test_set).detach().numpy().reshape((100**3,1))\n",
    "\n",
    "error = np.linalg.norm(predicted_output-truth_output) / np.linalg.norm(truth_output)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "61ab8f88",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1000000, 1)"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "truth_output.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "127968ec",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1000000, 1)"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "predicted_output.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "8483904d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.07557773"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "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
}
