PINNs.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# A Hands-on Introduction to Physics-Informed Neural Networks\n",
    "\n",
    "**By Atharva Hans and Ilias Bilionis (2021)**  \n",
    "*Publication Link:* [A Hands-on Introduction to Physics-Informed Neural Networks](https://nanohub.org/resources/handsonpinns)  \n",
    "*DOI:* 10.21981/FN6Y-NC12\n",
    "\n",
    "---\n",
    "\n",
    "### Authors:\n",
    "\n",
    "**Atharva Hans**\n",
    "- Role: Graduate Research Assistant\n",
    "- Affiliation: [Predictive Science Lab](https://www.predictivesciencelab.org/), School of Mechanical Engineering, Purdue University\n",
    "- Contact: [hans1@purdue.edu](mailto:hans1@purdue.edu)\n",
    "\n",
    "**Ilias Bilionis**\n",
    "- Role: Associate Professor\n",
    "- Affiliation: [Predictive Science Lab](https://www.predictivesciencelab.org/), School of Mechanical Engineering, Purdue University\n",
    "- Contact: [ibilion@purdue.edu](mailto:ibilion@purdue.edu)\n",
    "\n",
    "---\n",
    "\n",
    "### Learning Objectives:\n",
    "\n",
    "1. Understand how to leverage neural networks for solving ODEs.\n",
    "2. Grasp the concept of using neural networks for simulator-free forward model evaluations, illustrated through a solid mechanics example.\n",
    "3. Gain insights into how the structure of a network influences its convergence.\n",
    "4. Delve into the challenge of spectral bias in Deep Neural Networks (DNNs) and explore potential solutions.\n",
    "\n",
    "---\n",
    "\n",
    "### Key References:\n",
    "\n",
    "- [Artificial Neural Networks for Solving Ordinary and Partial Differential Equations](https://arxiv.org/pdf/physics/9705023.pdf)\n",
    "- [Deep Residual Learning for Image Recognition](https://arxiv.org/pdf/1512.03385.pdf)\n",
    "- [On the Spectral Bias of Neural Networks](http://proceedings.mlr.press/v97/rahaman19a/rahaman19a.pdf)\n",
    "- [Fourier Features Let Networks Learn High Frequency Functions in Low Dimensional Domains](https://arxiv.org/pdf/2006.10739.pdf)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import necessary libraries\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import numpy as np\n",
    "import os\n",
    "import requests\n",
    "from time import perf_counter\n",
    "from PIL import Image\n",
    "import matplotlib.pyplot as plt\n",
    "from functools import partial\n",
    "from IPython.display import display\n",
    "\n",
    "# Determine computation device: Use GPU if available, else CPU\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "def download_file_from_url(url, local_filename=None):\n",
    "    \"\"\"\n",
    "    Downloads a file from the provided URL and saves it to the specified path or current working directory.\n",
    "    \n",
    "    Parameters:\n",
    "    - url (str): The URL of the file to be downloaded.\n",
    "    - local_filename (str, optional): The path where the downloaded file should be saved.\n",
    "    \n",
    "    Returns:\n",
    "    None\n",
    "    \"\"\"\n",
    "    response = requests.get(url)\n",
    "    \n",
    "    # If no filename specified, use the filename from the URL\n",
    "    if local_filename is None:\n",
    "        local_filename = os.path.basename(url)\n",
    "        \n",
    "    with open(local_filename, 'wb') as file:\n",
    "        file.write(response.content)\n",
    "\n",
    "# Download an image of ResNet architecture\n",
    "download_file_from_url(\"https://www.dropbox.com/s/2xpxzko3d03onwq/ResNet.png?dl=1\", 'ResNet.png')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "TLMUNkY-Ss9L"
   },
   "source": [
    "## Example 1: Solving a Single Ordinary Differential Equation (ODE)\n",
    "\n",
    "Consider the ode:\n",
    "$$\n",
    "\\frac{d\\Psi}{dx} = f(x, \\Psi),\n",
    "$$\n",
    "with $x \\in [0,1]$ and initial conditions (IC):\n",
    "$$\n",
    "\\Psi(0) = A.\n",
    "$$\n",
    "We write the trial solution by:\n",
    "$$\n",
    "\\hat{\\Psi}(x; \\theta) = A + x N(x; \\theta),\n",
    "$$\n",
    "where $N(x; \\theta)$ is a neural network (NN).\n",
    "The solution is $\\hat{\\Psi}(x;\\theta)$ automatically satisfied the initial conditions.\n",
    "The loss function we would like to minimize to train the NN is:\n",
    "$$\n",
    "L(\\theta) = \\int_0^1 \\left[\\frac{d\\hat{\\Psi}(x;\\theta)}{dx} - f(x,\\hat{\\Psi}(x;\\theta))\\right]^2dx\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the neural network architecture (N) as described by Lagaris et al. 1997\n",
    "class NeuralNetwork(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(NeuralNetwork, self).__init__()\n",
    "        self.layer1 = nn.Linear(1, 50)\n",
    "        self.layer2 = nn.Linear(50, 1, bias=False)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        x = self.layer1(x)\n",
    "        x = nn.Sigmoid()(x)\n",
    "        x = self.layer2(x)\n",
    "        return x\n",
    "\n",
    "# Initialize the neural network\n",
    "N = NeuralNetwork()\n",
    "\n",
    "# Given initial condition\n",
    "A = 0.\n",
    "\n",
    "# Trial solution function using the neural network\n",
    "Psi_t = lambda x: A + x * N(x)\n",
    "\n",
    "# Right-hand side of the differential equation\n",
    "f = lambda x, Psi: torch.exp(-x / 5.0) * torch.cos(x) - Psi / 5.0\n",
    "\n",
    "# Loss function for training the neural network\n",
    "def loss(x):\n",
    "    \"\"\"Compute the loss based on the discrepancy between the ODE solution and its approximation.\"\"\"\n",
    "    x.requires_grad = True\n",
    "    outputs = Psi_t(x)\n",
    "    Psi_t_x = torch.autograd.grad(outputs, x, grad_outputs=torch.ones_like(outputs), create_graph=True)[0]\n",
    "    return torch.mean((Psi_t_x - f(x, outputs)) ** 2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "3wrjeMPnDVhX"
   },
   "source": [
    "First, we will use the method that we find in [Lagaris et al](https://arxiv.org/pdf/physics/9705023.pdf).\n",
    "Instead of using stochastic optimization, they use a lot of points to estimate the loss integral (We will use 100) and then they just do gradient-based optimization (We will do BFGS)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define optimizer for the neural network (Algorithm as mentioned by Lagaris)\n",
    "optimizer = torch.optim.LBFGS(N.parameters())\n",
    "\n",
    "# Define the collocation points (as used by Lagaris)\n",
    "x = torch.Tensor(np.linspace(0, 2, 100)[:, None])\n",
    "\n",
    "# Define the closure function for optimization\n",
    "def closure():\n",
    "    \"\"\"Compute the loss and perform backpropagation.\"\"\"\n",
    "    optimizer.zero_grad()\n",
    "    l = loss(x)\n",
    "    l.backward()\n",
    "    return l\n",
    "\n",
    "# Train the neural network using the optimizer\n",
    "for i in range(10):\n",
    "    optimizer.step(closure)\n",
    "\n",
    "# Compare the neural network's approximation to the true solution\n",
    "xx = np.linspace(0, 2, 100)[:, None]\n",
    "\n",
    "with torch.no_grad():\n",
    "    yy = Psi_t(torch.Tensor(xx)).numpy()\n",
    "\n",
    "yt = np.exp(-xx / 5.0) * np.sin(xx)\n",
    "\n",
    "# Plotting the results\n",
    "fig, ax = plt.subplots(dpi=100)\n",
    "ax.plot(xx, yt, label='True Solution')\n",
    "ax.plot(xx, yy, '--', label='Neural Network Approximation')\n",
    "ax.set_title('Comparison of True Solution vs Neural Network Approximation')\n",
    "ax.set_xlabel('$x$')\n",
    "ax.set_ylabel('$\\Psi(x)$')\n",
    "ax.legend(loc='best')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "N-T0lLLhSs9M"
   },
   "source": [
    "Now, let's use stochastic gradient descent to minimize the loss integral."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Reinitialize the neural network (N) for fresh training\n",
    "N = nn.Sequential(\n",
    "    nn.Linear(1, 50),\n",
    "    nn.Sigmoid(),\n",
    "    nn.Linear(50, 1, bias=False)\n",
    ")\n",
    "\n",
    "# Define the stochastic optimizer (Adam in this case)\n",
    "adam_optimizer = torch.optim.Adam(N.parameters(), lr=0.01)\n",
    "\n",
    "# Define the batch size (number of points to use per iteration) and maximum iterations\n",
    "n_batch = 5\n",
    "max_iterations = 1000\n",
    "\n",
    "print(\"Training the neural network...\")\n",
    "\n",
    "# Training loop\n",
    "for iteration in range(max_iterations):\n",
    "    # Randomly sample n_batch data points from the interval [0, 2]\n",
    "    x_batch = 2 * torch.rand(n_batch, 1)\n",
    "    \n",
    "    # Reset the gradients\n",
    "    adam_optimizer.zero_grad()\n",
    "    \n",
    "    # Compute the loss for the sampled batch\n",
    "    current_loss = loss(x_batch)\n",
    "    \n",
    "    # Perform backpropagation to compute gradients\n",
    "    current_loss.backward()\n",
    "    \n",
    "    # Update the neural network parameters\n",
    "    adam_optimizer.step()\n",
    "    \n",
    "    # Print the progress at regular intervals\n",
    "    if iteration % 100 == 99:\n",
    "        print(f\"Iteration {iteration + 1} completed.\")\n",
    "\n",
    "print(\"Training completed.\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate a range of x-values for evaluating the true solution and its neural network approximation\n",
    "x_values = np.linspace(0, 2, 100)[:, None]\n",
    "\n",
    "# Evaluate the neural network's approximation over the generated x-values\n",
    "with torch.no_grad():\n",
    "    nn_approximation = Psi_t(torch.Tensor(x_values)).numpy()\n",
    "\n",
    "# Calculate the true solution over the x-values\n",
    "true_solution = np.exp(-x_values / 5.0) * np.sin(x_values)\n",
    "\n",
    "# Plotting the results for comparison\n",
    "fig, ax = plt.subplots(dpi=100)\n",
    "\n",
    "# Plot the true solution\n",
    "ax.plot(x_values, true_solution, label='True Solution')\n",
    "\n",
    "# Plot the neural network approximation\n",
    "ax.plot(x_values, nn_approximation, '--', label='Neural Network Approximation')\n",
    "\n",
    "# Set axis labels and legend\n",
    "ax.set_xlabel('$x$')\n",
    "ax.set_ylabel('$\\Psi(x)$')\n",
    "ax.legend(loc='best')\n",
    "plt.title('Comparison of True Solution vs Neural Network Approximation')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5bDFdGfBSs9N"
   },
   "source": [
    "### Questions\n",
    "\n",
    "+ Change `n_batch` to just 1. Does the algorithm work? Did you have to increase the iterations to achieve the same accuracy?\n",
    "+ Modify the code, so that you now solve the problem for $x$ between 0 and 5. Play with the `n_batch` and `max_it` until you get a good solution."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "8d4_hy-VSs9O"
   },
   "source": [
    "## Example 2: Deformation of a compressible neo-Hookean material\n",
    "\n",
    "+ Consider a neo-Hookean square body defined on $(x,y) \\in [0,1]^2$. Let $u(x,y)$ describe the displacement field for this body. This body is subjected to the following displacement boundary conditions:\n",
    "$$\n",
    "u_x(0,y) = 0,\n",
    "$$\n",
    "$$\n",
    "u_y(0,y) = 0,\n",
    "$$\n",
    "$$\n",
    "u_x(1,y) = \\delta,\n",
    "$$\n",
    "$$\n",
    "u_y(1,y) = 0,\n",
    "$$\n",
    "with $\\delta$ referring to the applied displacement along the x-direction.\n",
    "\n",
    "\n",
    "+ For this hyperelastic material, the stored energy $E_b$ in the body can be expressed in as:\n",
    "$$\n",
    "E_b(u(x,y)) = \\int_{[0,1]^2}\\left\\{\\frac{1}{2}(\\sum_{i}\\sum_{j}{F_{ij}^2} - 2)- \\ln(\\det(F)) + 50\\ln(\\det(F))^2\\right\\} dxdy,\n",
    "$$\n",
    "with,\n",
    "$$\n",
    "F = I + \\nabla u,\n",
    "$$\n",
    "where $I$ is an identity matrix.\n",
    "\n",
    "\n",
    "+ The final orientation of this body is described by a displacement field that minimizes the stored energy $E_b$.\n",
    "\n",
    "\n",
    "+ Let's describe the horizontal and the vertical components of displacement field ($u_x$ and $u_y$ respectively) for this body as:\n",
    "$$\n",
    "u_x(x,y) = \\delta - \\delta(1-x) + x(1-x)N_x(x,y;\\theta),\n",
    "$$\n",
    "and,\n",
    "$$\n",
    "u_y(x,y) = x(1-x)N_y(x,y;\\theta)\n",
    "$$\n",
    "where $N_x(x,y;\\theta)$ and $N_y(x,y;\\theta)$ are neural networks. Notice $u_x(x,y)$ and $u_y(x,y)$ automatically satisfies the boundary conditions. \n",
    "\n",
    "\n",
    "+ The loss function that we would like to minimize to train the NN is:\n",
    "$$\n",
    "L(\\theta) = E_b(u(x,y))\n",
    "$$\n",
    "We will Monte Carlo for estimating the integral.\n",
    "\n",
    "Next, let's define a few helper functions which will be used throughout this activity. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4u5fFz42Ss9O"
   },
   "source": [
    "### Helper Functions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "78X_7HCaKrMz"
   },
   "source": [
    "##### Gradient Function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "7OvqXYQwKmVp"
   },
   "outputs": [],
   "source": [
    "def grad(outputs, inputs):\n",
    "    \"\"\"\n",
    "    Compute the gradient of 'outputs' with respect to 'inputs'.\n",
    "    \n",
    "    Parameters:\n",
    "    - outputs (torch.Tensor): Tensor for which the gradient needs to be computed.\n",
    "    - inputs (torch.Tensor): Tensor with respect to which the gradient is computed.\n",
    "    \n",
    "    Returns:\n",
    "    - torch.Tensor: Gradient of 'outputs' with respect to 'inputs'.\n",
    "    \"\"\"\n",
    "    return torch.autograd.grad(\n",
    "        outputs, \n",
    "        inputs, \n",
    "        grad_outputs=torch.ones_like(outputs), \n",
    "        create_graph=True\n",
    "    )[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "xTRk41JOSs9O"
   },
   "source": [
    "##### Displacement Function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ybYzV1TeSs9O"
   },
   "outputs": [],
   "source": [
    "def u_r(x, delta, network1, network2, scale_network_out):\n",
    "    \"\"\"\n",
    "    Compute the displacement field \\( u_r \\) given the input tensor `x`, boundary conditions `delta`, and neural networks `network1` and `network2`.\n",
    "    \n",
    "    Parameters:\n",
    "    - x (torch.Tensor): Input tensor containing the X and Y coordinates. Shape: (batch size, 2).\n",
    "    - delta (torch.Tensor): Direchlet boundary conditions. Typically something like: torch.tensor(([right_edge_displacement, 0.])).\n",
    "    - network1 (torch.nn.Module): Neural network to compute the horizontal component of the displacement.\n",
    "    - network2 (torch.nn.Module): Neural network to compute the vertical component of the displacement.\n",
    "    - scale_network_out (float): Scaling factor for the network output.\n",
    "    \n",
    "    Returns:\n",
    "    - torch.Tensor: Displacement field \\( u_r \\) combining both horizontal (u_x) and vertical (u_y) displacements.\n",
    "    \"\"\"\n",
    "    \n",
    "    # Calculate horizontal displacement component\n",
    "    u_x = delta[0] - delta[0]*(1. - x[:, 0][:, None]) + x[:, 0][:, None] * (1. - x[:, 0][:, None]) * network1(x) / scale_network_out\n",
    "    \n",
    "    # Calculate vertical displacement component\n",
    "    u_y = x[:, 0][:, None] * (1. - x[:, 0][:, None]) * network2(x) / scale_network_out\n",
    "    \n",
    "    # Combine horizontal and vertical components\n",
    "    u_r_combined = torch.cat((u_x, u_y), dim=1)\n",
    "    \n",
    "    return u_r_combined"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "eCzUk-TDSs9P"
   },
   "source": [
    "##### Loss Function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "dZ1iUA2ZSs9P"
   },
   "outputs": [],
   "source": [
    "def loss(X, u_r):\n",
    "    \"\"\"\n",
    "    Compute the loss for given input data and displacement function.\n",
    "    \n",
    "    Parameters:\n",
    "    - X (torch.Tensor): Input data tensor. Shape: (batch_size, 2).\n",
    "    - u_r (function): Displacement function that accepts X as input and returns the displacement.\n",
    "    \n",
    "    Returns:\n",
    "    - torch.Tensor: Computed loss value.\n",
    "    \"\"\"\n",
    "    \n",
    "    # Enable gradient computation for the input tensor\n",
    "    X.requires_grad = True\n",
    "    \n",
    "    # Calculate the displacement using provided function\n",
    "    displacement = u_r(X)\n",
    "    \n",
    "    # Compute gradient of displacement components with respect to input X\n",
    "    gradient_u_x = grad(displacement[:, 0], X)[:, None, :]\n",
    "    gradient_u_y = grad(displacement[:, 1], X)[:, None, :]\n",
    "    \n",
    "    # Construct the Jacobian matrix\n",
    "    J = torch.cat([gradient_u_x, gradient_u_y], 1)\n",
    "    \n",
    "    # Define the identity matrix\n",
    "    I = torch.eye(2).repeat(X.shape[0], 1, 1)\n",
    "    \n",
    "    # Calculate the deformation gradient tensor\n",
    "    F = I + J\n",
    "    \n",
    "    # Compute the log determinant of the deformation gradient tensor\n",
    "    log_det_F = torch.log(F[:, 0, 0] * F[:, 1, 1] - F[:, 0, 1] * F[:, 1, 0])\n",
    "    \n",
    "    # Compute the residual for the loss\n",
    "    residual = (0.5 * (torch.einsum('ijk,ijk->i', F, F) - 2) - log_det_F + 50 * log_det_F ** 2)\n",
    "    \n",
    "    # Return the mean of the residuals as the final loss\n",
    "    return torch.mean(residual)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "e8tnnbU2Ss9P"
   },
   "source": [
    "##### Training Loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "4wnfZCsESs9P"
   },
   "outputs": [],
   "source": [
    "def train_MultiFieldFracture_seperate_net(net_u, net_v, right_edge_displacement, batch_size, max_iterations, print_frequency, scaling_factor):\n",
    "    \"\"\"\n",
    "    Trains the provided neural networks for multi-field fracture problems with separate networks for u_x and u_y.\n",
    "\n",
    "    Parameters:\n",
    "    - net_u (torch.nn.Module): Neural network for horizontal displacement.\n",
    "    - net_v (torch.nn.Module): Neural network for vertical displacement.\n",
    "    - right_edge_displacement (float): Boundary condition on the right edge.\n",
    "    - batch_size (int): Number of samples to use for each iteration.\n",
    "    - max_iterations (int): Maximum number of training iterations.\n",
    "    - print_frequency (int): Frequency at which to print the training progress.\n",
    "    - scaling_factor (float): Scaling factor for the network outputs.\n",
    "\n",
    "    Returns:\n",
    "    - list: List of loss values.\n",
    "    - torch.nn.Module: Trained network for horizontal displacement.\n",
    "    - torch.nn.Module: Trained network for vertical displacement.\n",
    "    \"\"\"\n",
    "    \n",
    "    print('\\nStarting training with separate networks for u_x and u_y...\\n')\n",
    "    print(f'Right Edge Displacement: {right_edge_displacement:.3f}', \n",
    "          f'\\tBatch Size: {batch_size}', \n",
    "          f'\\tMax Iterations: {max_iterations}\\n')\n",
    "\n",
    "    net_u = net_u.to(device)\n",
    "    net_v = net_v.to(device)\n",
    "\n",
    "    # Define the boundary condition\n",
    "    delta = torch.tensor([right_edge_displacement, 0.], device=device)\n",
    "\n",
    "    # Combine parameters of both networks\n",
    "    parameters = list(net_u.parameters()) + list(net_v.parameters())\n",
    "\n",
    "    # Initialize the optimizer\n",
    "    optimizer = torch.optim.Adam(parameters, lr=1e-4)\n",
    "    \n",
    "    start_time = perf_counter()\n",
    "    loss_values = []\n",
    "    accumulated_loss = 0.0\n",
    "\n",
    "    for iteration in range(max_iterations):\n",
    "        # Sample input data uniformly\n",
    "        X = torch.distributions.Uniform(0, 1).sample((batch_size, 2)).to(device)\n",
    "\n",
    "        # Reset gradients\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        # Evaluate the loss\n",
    "        current_loss = loss(X, partial(u_r, delta=delta, network1=net_u, network2=net_v, scale_network_out=scaling_factor))\n",
    "        accumulated_loss += current_loss.item()\n",
    "\n",
    "        # Compute gradients\n",
    "        current_loss.backward()\n",
    "\n",
    "        # Update network parameters\n",
    "        optimizer.step()\n",
    "\n",
    "        # Print progress at specified frequency\n",
    "        if (iteration + 1) % print_frequency == 0:\n",
    "            elapsed_time = perf_counter() - start_time\n",
    "            avg_loss = accumulated_loss / print_frequency\n",
    "            print(f'[Iteration: {iteration + 1}] Elapsed Time: {elapsed_time:.2f} secs, Loss: {avg_loss:.4f}')\n",
    "            loss_values.append(avg_loss)\n",
    "            accumulated_loss = 0.0\n",
    "\n",
    "    return loss_values, net_u, net_v"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_ygUtzIvSs9P"
   },
   "source": [
    "##### Model Capacity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "aUMnbGrNSs9Q"
   },
   "outputs": [],
   "source": [
    "def model_capacity(net):\n",
    "    \"\"\"\n",
    "    Determine and print the capacity of the neural network.\n",
    "\n",
    "    Parameters:\n",
    "    - net (torch.nn.Module): Neural network for which capacity is to be determined.\n",
    "\n",
    "    Returns:\n",
    "    - tuple: Number of layers and number of learnable parameters in the model.\n",
    "    \"\"\"\n",
    "    \n",
    "    # Compute the number of learnable parameters\n",
    "    number_of_learnable_params = sum(p.numel() for p in net.parameters() if p.requires_grad)\n",
    "    \n",
    "    # Determine the number of layers in the model\n",
    "    num_layers = len(list(net.parameters()))\n",
    "    \n",
    "    # Print the results\n",
    "    print(f\"\\nThe number of layers in the model: {num_layers}\")\n",
    "    print(f\"The number of learnable parameters in the model: {number_of_learnable_params}\\n\")\n",
    "    \n",
    "    return num_layers, number_of_learnable_params"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "W8yh0IfKSs9Q"
   },
   "source": [
    "##### Plots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "us0TL2KjSs9Q"
   },
   "outputs": [],
   "source": [
    "def plot_loss(loss_values, label_name):\n",
    "    \"\"\"\n",
    "    Plot the loss values over iterations.\n",
    "\n",
    "    Parameters:\n",
    "    - loss_values (list): List containing the loss values over iterations.\n",
    "    - label_name (str): Label for the plot representing the loss values.\n",
    "    \"\"\"\n",
    "    \n",
    "    ax.plot(100 * np.arange(len(loss_values)), loss_values, label=label_name)\n",
    "    ax.set_xlabel('Iterations')\n",
    "    ax.set_ylabel('Loss')\n",
    "    ax.legend(loc='best')\n",
    "    \n",
    "    \n",
    "def plot_displacement(net_u, net_v, right_edge_displacement, scaling_factor, num_samples=200, marker_size=4):\n",
    "    \"\"\"\n",
    "    Visualize the displacement field using the given neural networks.\n",
    "\n",
    "    Parameters:\n",
    "    - net_u (torch.nn.Module): Neural network for horizontal displacement.\n",
    "    - net_v (torch.nn.Module): Neural network for vertical displacement.\n",
    "    - right_edge_displacement (float): Boundary condition on the right edge.\n",
    "    - scaling_factor (float): Scaling factor for the network outputs.\n",
    "    - num_samples (int, optional): Number of grid points for visualization. Defaults to 200.\n",
    "    - marker_size (int, optional): Size of the scatter plot markers. Defaults to 4.\n",
    "    \"\"\"\n",
    "    \n",
    "    # Define the boundary condition\n",
    "    delta = torch.tensor([right_edge_displacement, 0.])\n",
    "    \n",
    "    fig, axes = plt.subplots(1, 2, figsize=(20, 5))\n",
    "    \n",
    "    # Create a mesh grid for visualization\n",
    "    x_vals = torch.linspace(0, 1, num_samples)\n",
    "    xx, yy = torch.meshgrid(x_vals, x_vals)\n",
    "    grid_points = torch.cat((xx.reshape(-1)[:, None], yy.reshape(-1)[:, None]), axis=1)\n",
    "\n",
    "    # Compute displacements using the neural networks\n",
    "    displacements = u_r(grid_points, delta, net_u, net_v, scaling_factor).detach().numpy()\n",
    "    \n",
    "    # Calculate final positions after displacements\n",
    "    final_positions = grid_points + displacements \n",
    "    final_positions = final_positions.numpy()\n",
    "    \n",
    "    # Plot horizontal displacements\n",
    "    sc1 = axes[0].scatter(final_positions[:, 0], final_positions[:, 1], s=marker_size, c=displacements[:, 0], cmap='copper')\n",
    "    fig.colorbar(sc1, ax=axes[0])\n",
    "    axes[0].set_title('Horizontal Displacement ($u_x$)')\n",
    "    \n",
    "    # Plot vertical displacements\n",
    "    sc2 = axes[1].scatter(final_positions[:, 0], final_positions[:, 1], s=marker_size, c=displacements[:, 1], cmap='copper')\n",
    "    fig.colorbar(sc2, ax=axes[1])\n",
    "    axes[1].set_title('Vertical Displacement ($u_y$)')\n",
    "    \n",
    "    plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "7W4wGF4lSs9R"
   },
   "source": [
    "Let's start with displacement ($\\delta$) of 0.5, so\n",
    "$$\n",
    "u_x(1,y) = 0.5.\n",
    "$$\n",
    "\n",
    "<b>Note: Solving this problem for displacement ($\\delta$) of 0.5 in FEM requires us to break down displacement in smaller chunks. For example, we would first solve for the displacement field for $\\delta=0.1$ then using that displacement field as baseline, we would further solve it for $\\delta=0.1$ and so on...(until we reach $\\delta=0.5$).\n",
    "</b> "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "VQxmkxI7Ss9R"
   },
   "outputs": [],
   "source": [
    "# Configuration for the boundary condition and network output scaling.\n",
    "\n",
    "# Boundary condition (delta) on the right edge.\n",
    "# Possible values to try include 0.05.\n",
    "right_edge_displacement = 0.5\n",
    "\n",
    "# Scaling factor for the network output. The output is scaled down to ensure \n",
    "# the Jacobian's determinant is positive, preserving physical meaning. \n",
    "# For instance, a negative determinant would imply a negative volume, \n",
    "# which isn't physically meaningful.\n",
    "scale_network_out = 5"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "DJwlOKMcSs9R"
   },
   "source": [
    "### Part A: Simple Fully Connected Neural Network\n",
    "\n",
    "Let's begin with a simplest case where we have simple fully connected neural networks for both the horizontal displacement $u_x$ and the vertical displacement $u_y$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "JsfGHBpeSs9R"
   },
   "source": [
    "##### Define Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def initialize_simple_net():\n",
    "    \"\"\"\n",
    "    Initialize a simple neural network with specific layers and activations.\n",
    "    \n",
    "    Returns:\n",
    "    - torch.nn.Module: Initialized neural network.\n",
    "    \"\"\"\n",
    "    return nn.Sequential(\n",
    "        nn.Linear(2, 50),\n",
    "        nn.Sigmoid(),\n",
    "        nn.Linear(50, 50),\n",
    "        nn.Sigmoid(),\n",
    "        nn.Linear(50, 50),\n",
    "        nn.Sigmoid(),\n",
    "        nn.Linear(50, 50),\n",
    "        nn.Sigmoid(),\n",
    "        nn.Linear(50, 50),\n",
    "        nn.Sigmoid(),\n",
    "        nn.Linear(50, 50),\n",
    "        nn.Sigmoid(),\n",
    "        nn.Linear(50, 50),\n",
    "        nn.Sigmoid(),\n",
    "        nn.Linear(50, 1)\n",
    "    )\n",
    "\n",
    "# Initialize networks for u_x and u_y\n",
    "simple_net1 = initialize_simple_net()\n",
    "simple_net2 = initialize_simple_net()\n",
    "\n",
    "# Display model capacities for both networks\n",
    "print(\"Model capacity for simple_net1 (u_x):\")\n",
    "model_capacity(simple_net1)\n",
    "\n",
    "print(\"Model capacity for simple_net2 (u_y):\")\n",
    "model_capacity(simple_net2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "OJVkYaOiSs9S"
   },
   "source": [
    "Here is just a sanity check that we start from a candidate solution that indeed satisfies the boundary conditions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Nqaeba48Ss9S"
   },
   "outputs": [],
   "source": [
    "plot_displacement(simple_net1, simple_net2,\n",
    "                  right_edge_displacement, scale_network_out)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "RpYGZ2qSSs9S"
   },
   "source": [
    "##### Train Model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "DQkJRv44Ss9S"
   },
   "source": [
    "Let's train the network for 5000 iterations and batch size of 10. Training should take around 5-7 minutes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "AeOp2-AdSs9S",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "loss_list_simple_net, simple_net1, simple_net2 = train_MultiFieldFracture_seperate_net(\n",
    "    net_u=simple_net1, \n",
    "    net_v=simple_net2, \n",
    "    right_edge_displacement=right_edge_displacement, \n",
    "    batch_size=10, \n",
    "    max_iterations=5000, \n",
    "    print_frequency=100, \n",
    "    scaling_factor=scale_network_out\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5uON5owvSs9S"
   },
   "source": [
    "##### Analyze Training Loss and Material Displacement"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ev0hcYjBSs9T"
   },
   "outputs": [],
   "source": [
    "# Let's visualize the training loss\n",
    "figure, ax = plt.subplots(dpi=100)\n",
    "plot_loss(loss_list_simple_net, 'Simple Net')\n",
    "\n",
    "if right_edge_displacement==0.05:\n",
    "    ax.plot(100*np.arange(len(loss_list_simple_net)), 0.006*np.ones(len(loss_list_simple_net)), label='FEM Energy')\n",
    "elif right_edge_displacement==0.5:\n",
    "    ax.plot(100*np.arange(len(loss_list_simple_net)), 0.43*np.ones(len(loss_list_simple_net)), label='FEM Energy')\n",
    "plt.legend(loc='best'); "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "VvqVuvuCSs9T"
   },
   "outputs": [],
   "source": [
    "# Let's visualize the body post-training\n",
    "plot_displacement(simple_net1, simple_net2, right_edge_displacement, scale_network_out)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Ch29lC-eSs9T"
   },
   "source": [
    "The above loss might or might not converge (to the FEM solution) depending on the number of iterations we trained the network and on the boundary displacement that we choose. If we run it long enough (~20,000 iterations), it might converge. \n",
    "\n",
    "Unlike the first example, where we were trying to solve a simple ODE (and the neural network wasn't very deep), here we have relatively deeper networks which may pose two issues:\n",
    "+ [Vanishing gradients](https://arxiv.org/pdf/1512.03385.pdf)\n",
    "+ [Bias towards lower frequency functions](http://proceedings.mlr.press/v97/rahaman19a/rahaman19a.pdf)\n",
    "\n",
    "Is there anything we can do to speed up the convergence?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "fFTXbQvqSs9T"
   },
   "source": [
    "### Part B: ResNet Architecture \n",
    "\n",
    "In this part let's introduce a ResNet architecture for our neural networks and investigate if we get better convergence. \n",
    "ResNet architecture includes skip-connections between the layers which has been [shown](https://arxiv.org/pdf/1512.03385.pdf) to overcome the vanishing gradient problem (and help with convergence).\n",
    "\n",
    "Figure 1a below shows a schematic of a deep ResNet with a dense input layer, followed by K residual blocks and an output layer. \n",
    "\n",
    "Figure 1b shows s single residual block with L layers. Pay attention to how we have a skip-connection from the input $z^{(i-1,0)}$ to the output $z^{(i,0)}$. The skip connections lead to better backpropagation of the loss and helps in better convergence. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "VwRe5pJ2Ss9T"
   },
   "outputs": [],
   "source": [
    "# Load and display the schematic of a deep ResNet with a caption\n",
    "resnet_image = Image.open('ResNet.png')\n",
    "\n",
    "# Display the image with a caption\n",
    "display(resnet_image)\n",
    "print(\"Figure: Schematic of a deep ResNet.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "zpypkSOqSs9T"
   },
   "source": [
    "Here is a `DenseResNet` class:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DenseResNet(nn.Module):\n",
    "    \"\"\"\n",
    "    Dense ResNet with options for Fourier feature mappings and tunable beta parameters.\n",
    "\n",
    "    Parameters:\n",
    "    ----------\n",
    "    dim_in : int, default=2\n",
    "        Network's input dimension.\n",
    "    dim_out : int, default=1\n",
    "        Network's output dimension.\n",
    "    num_resnet_blocks : int, default=3\n",
    "        Number of ResNet blocks.\n",
    "    num_layers_per_block : int, default=2\n",
    "        Number of layers per ResNet block.\n",
    "    num_neurons : int, default=50\n",
    "        Number of neurons in each layer.\n",
    "    activation : torch.nn.Module, default=nn.Sigmoid()\n",
    "        Non-linear activations function.\n",
    "    fourier_features : bool, default=False\n",
    "        Whether to pass the inputs through Fourier mapping.\n",
    "    m_freqs : int, default=100\n",
    "        Number of frequencies for Fourier mapping.\n",
    "    sigma : float, default=10\n",
    "        Controls the spectrum of frequencies.\n",
    "    tune_beta : bool, default=False\n",
    "        Whether to consider the beta parameter in the activation function.\n",
    "\n",
    "    Methods:\n",
    "    --------\n",
    "    forward(x) : torch.Tensor\n",
    "        Forward pass of the network.\n",
    "    model_capacity() -> Tuple[int, int]\n",
    "        Returns the number of layers and parameters in the network.\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self, dim_in=2, dim_out=1, num_resnet_blocks=3, \n",
    "                 num_layers_per_block=2, num_neurons=50, activation=nn.Sigmoid(), \n",
    "                 fourier_features=False, m_freqs=100, sigma=10, tune_beta=False):\n",
    "        super(DenseResNet, self).__init__()\n",
    "\n",
    "        self.num_resnet_blocks = num_resnet_blocks\n",
    "        self.num_layers_per_block = num_layers_per_block\n",
    "        self.fourier_features = fourier_features\n",
    "        self.activation = activation\n",
    "        self.tune_beta = tune_beta\n",
    "\n",
    "        # Beta parameters\n",
    "        if tune_beta:\n",
    "            self.beta0 = nn.Parameter(torch.ones(1, 1))\n",
    "            self.beta = nn.Parameter(torch.ones(num_resnet_blocks, num_layers_per_block))\n",
    "        else: \n",
    "            self.beta0 = torch.ones(1, 1)\n",
    "            self.beta = torch.ones(num_resnet_blocks, num_layers_per_block)\n",
    "\n",
    "        # First layer\n",
    "        input_dim = 2 * m_freqs if fourier_features else dim_in\n",
    "        self.first = nn.Linear(input_dim, num_neurons)\n",
    "        nn.init.xavier_uniform_(self.first.weight)\n",
    "        nn.init.zeros_(self.first.bias)\n",
    "\n",
    "        # ResNet blocks\n",
    "        self.resblocks = nn.ModuleList([\n",
    "            nn.ModuleList([nn.Linear(num_neurons, num_neurons) for _ in range(num_layers_per_block)]) \n",
    "            for _ in range(num_resnet_blocks)])\n",
    "        for block in self.resblocks:\n",
    "            for layer in block:\n",
    "                nn.init.xavier_uniform_(layer.weight)\n",
    "                nn.init.zeros_(layer.bias)\n",
    "\n",
    "        # Last layer\n",
    "        self.last = nn.Linear(num_neurons, dim_out)\n",
    "        nn.init.xavier_uniform_(self.last.weight)\n",
    "        nn.init.zeros_(self.last.bias)\n",
    "\n",
    "        # Fourier features\n",
    "        if fourier_features:\n",
    "            self.B = nn.Parameter(sigma * torch.randn(dim_in, m_freqs))\n",
    "\n",
    "    def forward(self, x):\n",
    "        # Fourier feature transformation\n",
    "        if self.fourier_features:\n",
    "            x = torch.cat((torch.cos(x @ self.B), torch.sin(x @ self.B)), dim=1)\n",
    "        x = self.activation(self.beta0 * self.first(x))\n",
    "\n",
    "        # ResNet blocks\n",
    "        for i, block in enumerate(self.resblocks):\n",
    "            z = x\n",
    "            for j, layer in enumerate(block):\n",
    "                z = self.activation(self.beta[i][j] * layer(z))\n",
    "            x = z + x\n",
    "\n",
    "        return self.last(x)\n",
    "\n",
    "    def model_capacity(self):\n",
    "        \"\"\"Returns the number of layers and parameters in the network.\"\"\"\n",
    "        num_params = sum(p.numel() for p in self.parameters() if p.requires_grad)\n",
    "        num_layers = len(list(self.parameters()))\n",
    "        return num_layers, num_params"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Nm0XbkyTSs9U"
   },
   "source": [
    "##### Define Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def initialize_dense_resnet():\n",
    "    \"\"\"\n",
    "    Initialize a DenseResNet with the given parameters.\n",
    "    \n",
    "    Returns:\n",
    "    - torch.nn.Module: Initialized DenseResNet model.\n",
    "    \"\"\"\n",
    "    return DenseResNet(dim_in=2, dim_out=1, num_resnet_blocks=3, \n",
    "                       num_layers_per_block=2, num_neurons=50, activation=nn.Sigmoid(), \n",
    "                       fourier_features=False, m_freqs=100, sigma=10, tune_beta=False)\n",
    "\n",
    "# Initialize two DenseResNet models\n",
    "dense_net1 = initialize_dense_resnet()\n",
    "dense_net2 = initialize_dense_resnet()\n",
    "\n",
    "# Print the model capacities\n",
    "print(\"Model capacity for dense_net1:\")\n",
    "model_capacity(dense_net1)\n",
    "\n",
    "print(\"Model capacity for dense_net2:\")\n",
    "model_capacity(dense_net2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "kWwq38GjSs9U"
   },
   "source": [
    "<b>Pay attention: We have the same number of parameters as in Part A</b>. Only difference now is that the network has a ResNet architecture (allows better back propagation of loss). Let's train this model and see if this helps in convergence.\n",
    "\n",
    "Here is just a sanity check that we start from a candidate solution that indeed satisfies the boundary conditions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ntZ2OYmsSs9U"
   },
   "outputs": [],
   "source": [
    "# Let's visualize the body pre-training and make sure the boundary conditions are satified\n",
    "plot_displacement(dense_net1, dense_net2, right_edge_displacement, scale_network_out)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "EzofPKumSs9U"
   },
   "source": [
    "##### Train Model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "C8foxKzMSs9U"
   },
   "source": [
    "Let's train the network for 5000 iterations and batch size of 10. Training should take around 5-7 minutes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "xSuE3K1nSs9U",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Training configuration and execution\n",
    "training_params = {\n",
    "    'net_u': dense_net1,\n",
    "    'net_v': dense_net2,\n",
    "    'right_edge_displacement': right_edge_displacement,\n",
    "    'batch_size': 10,\n",
    "    'max_iterations': 5000,\n",
    "    'print_frequency': 100,\n",
    "    'scaling_factor': scale_network_out\n",
    "}\n",
    "\n",
    "loss_list_dense, trained_net1, trained_net2 = train_MultiFieldFracture_seperate_net(**training_params)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "2xlWsV1FSs9U"
   },
   "source": [
    "##### Analyze Training Loss and Material Displacement"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "fRk0V6CjSs9U"
   },
   "outputs": [],
   "source": [
    "# Let's visualize the training loss\n",
    "figure, ax = plt.subplots(dpi=100)\n",
    "plot_loss(loss_list_simple_net, 'Simple Net')\n",
    "plot_loss(loss_list_dense, 'ResNet')\n",
    "\n",
    "if right_edge_displacement==0.05:\n",
    "    ax.plot(100*np.arange(len(loss_list_dense)), 0.006*np.ones(len(loss_list_dense)), label='FEM Energy')\n",
    "elif right_edge_displacement==0.5:\n",
    "    ax.plot(100*np.arange(len(loss_list_dense)), 0.43*np.ones(len(loss_list_dense)), label='FEM Energy')\n",
    "plt.legend(loc='best');"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "VObRb_qoSs9V"
   },
   "outputs": [],
   "source": [
    "# Let's visualize the body post-training\n",
    "plot_displacement(dense_net1, dense_net2, right_edge_displacement, scale_network_out)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "lcRbm_gGSs9V"
   },
   "source": [
    "Using ResNet architecture seems to have improved the convergence as compared to the case of the simple network in part A. \n",
    "\n",
    "One reason for this is as we make the network deeper, it can hurt the ability to train the network to do well on the training set. ResNet overcomes this issue by introducing skip connections between layers; this helps better backpropagation of loss and better convergence.\n",
    "\n",
    "Can we further improve this?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "qa2hLUO-Ss9V"
   },
   "source": [
    "### Part C: Fourier Features"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "05E69vSqSs9V"
   },
   "source": [
    "In this part let's not only use a ResNet architecture for our neural networks but also introduce a Fourier feature mapping for the input data. \n",
    "It has been observed that deeper networks are biased towards low frequency functions ([\"Over-parameterized networks prioritize learning simple patterns that generalize\n",
    "across data samples\"](http://proceedings.mlr.press/v97/rahaman19a/rahaman19a.pdf)). \n",
    "\n",
    "To solve this bias towards the low frequency function, [Tancik et al.](https://arxiv.org/pdf/2006.10739.pdf) show that passing the neural network inputs through a simple Fourier feature mapping not only helps neural network learn high-frequency function, but also, helps achieve faster convergence for high frequency components.  \n",
    "\n",
    "The create a neural network with Fourier feature, just set the parameter `fourier_features` as `True` in the `DenseResNet` class I defined above:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "QvTNp-1WSs9V"
   },
   "source": [
    "##### Define Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_dense_net(fourier_features=False):\n",
    "    \"\"\"\n",
    "    Helper function to create a DenseResNet instance with specified parameters.\n",
    "    \"\"\"\n",
    "    return DenseResNet(\n",
    "        dim_in=2, \n",
    "        dim_out=1, \n",
    "        num_resnet_blocks=3,\n",
    "        num_layers_per_block=2, \n",
    "        num_neurons=50, \n",
    "        activation=nn.Sigmoid(),\n",
    "        fourier_features=fourier_features, \n",
    "        m_freqs=100, \n",
    "        sigma=10, \n",
    "        tune_beta=False\n",
    "    )\n",
    "\n",
    "Fourier_dense_net1 = create_dense_net(fourier_features=True)\n",
    "Fourier_dense_net2 = create_dense_net(fourier_features=True)\n",
    "\n",
    "# Display model capacity\n",
    "def display_model_capacity(net, name):\n",
    "    \"\"\"\n",
    "    Helper function to display model capacity with network name.\n",
    "    \"\"\"\n",
    "    print(f\"Model capacity for {name}:\")\n",
    "    model_capacity(net)\n",
    "    print(\"------\\n\")\n",
    "\n",
    "display_model_capacity(Fourier_dense_net1, \"Fourier_dense_net1\")\n",
    "display_model_capacity(Fourier_dense_net2, \"Fourier_dense_net2\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "FjL3XlHaSs9V"
   },
   "source": [
    "Here is just a sanity check that we start from a candidate solution that indeed satisfies the boundary conditions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "KlC45M3vSs9V"
   },
   "outputs": [],
   "source": [
    "# Let's visualize the body pre-training and make sure the boundary conditions are satified\n",
    "plot_displacement(Fourier_dense_net1, Fourier_dense_net2, right_edge_displacement, scale_network_out)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "nfKddC6uSs9W"
   },
   "source": [
    "##### Train Model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "nzKJP-BYSs9W"
   },
   "source": [
    "Let's train the network for 5000 iterations and batch size of 10. Training should take around 5-7 minutes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "GA6RfvSbSs9W",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Training configuration and execution using the refactored approach\n",
    "training_params_fourier = {\n",
    "    'net_u': Fourier_dense_net1,\n",
    "    'net_v': Fourier_dense_net2,\n",
    "    'right_edge_displacement': right_edge_displacement,\n",
    "    'batch_size': 10,\n",
    "    'max_iterations': 5000,\n",
    "    'print_frequency': 100,\n",
    "    'scaling_factor': scale_network_out\n",
    "}\n",
    "\n",
    "loss_list_dense_fourier, trained_net1_fourier, trained_net2_fourier = train_MultiFieldFracture_seperate_net(**training_params_fourier)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "iAkTT3-bSs9W"
   },
   "source": [
    "##### Analyze Training Loss and Material Displacement"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Vm6Ua-p3Ss9W"
   },
   "outputs": [],
   "source": [
    "# Let's visualize the training loss\n",
    "figure, ax = plt.subplots(dpi=100)\n",
    "plot_loss(loss_list_simple_net, 'Simple Net')\n",
    "plot_loss(loss_list_dense, 'ResNet')\n",
    "plot_loss(loss_list_dense_fourier, 'Fourier-ResNet')\n",
    "\n",
    "if right_edge_displacement==0.05:\n",
    "    ax.plot(100*np.arange(len(loss_list_dense)), 0.006*np.ones(len(loss_list_dense)), label='FEM Energy')\n",
    "elif right_edge_displacement==0.5:\n",
    "    ax.plot(100*np.arange(len(loss_list_dense)), 0.43*np.ones(len(loss_list_dense)), label='FEM Energy')\n",
    "\n",
    "plt.legend(loc='best');"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "HLrZW3mMSs9W"
   },
   "outputs": [],
   "source": [
    "# Let's visualize the body post-training\n",
    "plot_displacement(Fourier_dense_net1, Fourier_dense_net2, right_edge_displacement, scale_network_out)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "nDG4YJlhSs9W"
   },
   "source": [
    "This look great! The loss seems to have converged to the FEM solution pretty quickly. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "iXJFbu60Ss9W"
   },
   "source": [
    "### Questions\n",
    "\n",
    "+ For part A, rerun the code for 20,000 iterations. Does the simple network in part A converge? How does the prediction look?\n",
    "+ How does batch size affect the convergence? Try different batch sizes (20, 40,...) by changing `batch_size_X` in the `train_MultiFieldFracture_seperate_net` function. Which one works the best? \n",
    "+ Change the activation function in the `DenseResNet` class from `nn.Sigmoid()` to `nn.ReLU()` and rerun the code. How does the convergence change?\n",
    "+ In the `DenseResNet` class, read the docstring to see what `tune_beta=True` does. Set `tune_beta=True` and rerun the code. Do you observe better convergence?"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [
    "iXJFbu60Ss9W"
   ],
   "name": "PINNs.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}