{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "7e50db47",
   "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "8c9b7095",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Helmholtz(SpatialProblem):\n",
    "    output_variables = ['u']\n",
    "    spatial_domain = CartesianDomain({'x':[-1,1],'y':[-1,1],'z':[-1,1]})\n",
    "    def HH(input_,output_):\n",
    "        a1 = 4\n",
    "        a2 = 4\n",
    "        a3 = 3\n",
    "        u = output_.extract(['u'])\n",
    "        x = input_.extract(['x'])\n",
    "        y = input_.extract(['y'])\n",
    "        z = input_.extract(['z'])\n",
    "        nabla_u = laplacian(output_, input_, components=['u'], d=['x', 'y','z'])\n",
    "        q = -u * ((a1*torch.pi)**2 + (a2*torch.pi)**2 + (a3*torch.pi)**2 - 1)\n",
    "        return nabla_u + u - q\n",
    "  \n",
    "    def hh_sol(self,pts):\n",
    "        a1 = 4\n",
    "        a2 = 4\n",
    "        a3 = 3\n",
    "        x = pts.extract(['x'])\n",
    "        y = pts.extract(['y'])\n",
    "        z = pts.extract(['z'])\n",
    "        return torch.sin(a1*torch.pi*x)*torch.sin(a2*torch.pi*y)*torch.sin(a3*torch.pi*z)\n",
    "    truth_solution = hh_sol\n",
    "    \n",
    "    conditions = {\n",
    "        'D': Condition(location=CartesianDomain({'x': [-1, 1], 'y': [-1, 1], 'z': [-1, 1]}), equation=Equation(HH)),\n",
    "        'gamma1': Condition(location=CartesianDomain({'x': [-1, 1], 'y':  1, 'z': [-1, 1]}), equation=FixedValue(0.0)),\n",
    "        'gamma2': Condition(location=CartesianDomain({'x': [-1, 1], 'y': -1, 'z': [-1, 1]}), equation=FixedValue(0.0)),\n",
    "        'gamma3': Condition(location=CartesianDomain({'x':  1, 'y': [-1, 1], 'z': [-1, 1]}), equation=FixedValue(0.0)),\n",
    "        'gamma4': Condition(location=CartesianDomain({'x': -1, 'y': [-1, 1], 'z': [-1, 1]}), equation=FixedValue(0.0)),\n",
    "        'gamma5': Condition(location=CartesianDomain({'x': [-1, 1], 'y': [-1, 1], 'z': -1}), equation=FixedValue(0.0)),\n",
    "        'gamma6': Condition(location=CartesianDomain({'x': [-1, 1], 'y': [-1, 1], 'z': 1}), equation=FixedValue(0.0)),\n",
    "        \n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "acb5f2a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "problem = Helmholtz()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "24584d5a",
   "metadata": {},
   "outputs": [],
   "source": [
    "problem.discretise_domain(5000, 'random', locations=['D', 'gamma1', 'gamma2', 'gamma3', 'gamma4','gamma5','gamma6'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "84e514dd",
   "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": "92de2844",
   "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, 16.84it/s, v_num=2, gamma1_loss=2.34e-11, gamma2_loss=4.92e-11, gamma3_loss=2.91e-11, gamma4"
     ]
    },
    {
     "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,  5.90it/s, v_num=2, gamma1_loss=2.34e-11, gamma2_loss=4.92e-11, gamma3_loss=2.91e-11, gamma4\n"
     ]
    }
   ],
   "source": [
    "trainer = Trainer(pinn, max_epochs=30000, callbacks=[MetricTracker()], accelerator='gpu',enable_model_summary=False,) \n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "05a5913a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Plotting at z=0\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 54080 characters]",
      "text/plain": [
       "<Figure size 1600x600 with 6 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Plotting at z=0.5\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 79724 characters]",
      "text/plain": [
       "<Figure size 1600x600 with 6 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Plotting at z=0.5\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 82448 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 z=0')\n",
    "plotter.plot(pinn, fixed_variables={'z': 0.0})\n",
    "\n",
    "# plotting at fixed time t = 0.5\n",
    "print('Plotting at z=0.5')\n",
    "plotter.plot(pinn, fixed_variables={'z':0.5})\n",
    "\n",
    "# plotting at fixed time t = 1.\n",
    "print('Plotting at z=0.5')\n",
    "plotter.plot(pinn, fixed_variables={'z': 0.25})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3118333d",
   "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
}
