coupled_higgs_pino.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "2ad2ddd4-26a5-4447-bb5e-7693b9b3185d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.init as init\n",
    "from torch.optim import AdamW, LBFGS\n",
    "from torch.autograd import Variable\n",
    "from tqdm import tqdm\n",
    "import numpy as np\n",
    "import os\n",
    "import pandas as pd\n",
    "from itertools import cycle \n",
    "\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\", category=UserWarning)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "84e9f9d8-cbe2-44cf-9b8a-e400cdae858a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def fourier_features(x, B):\n",
    "    x_transformed = torch.matmul(x, B)\n",
    "    return torch.cat([torch.sin(x_transformed), torch.cos(x_transformed)], dim=-1)\n",
    "\n",
    "def init_fixed_frequency_matrix(size, scale=1.0):\n",
    "    num_elements = size[0] * size[1]\n",
    "    lin_space = torch.linspace(-scale, scale, steps=num_elements)\n",
    "    B = lin_space.view(size).float()\n",
    "    return B\n",
    "\n",
    "class FourierFeatureNN(nn.Module):\n",
    "    def __init__(self, input_dim=1, shared_units=16, neuron_units=32, scale=1.0, \n",
    "                 activation=nn.Tanh, device='cpu'):\n",
    "        super(FourierFeatureNN, self).__init__()\n",
    "        self.Bx = init_fixed_frequency_matrix((input_dim, shared_units // 2), scale=scale).to(device)\n",
    "        self.Bt = init_fixed_frequency_matrix((input_dim, shared_units // 2), scale=scale).to(device)\n",
    "\n",
    "        # Define separate paths for x and t after Fourier transformation\n",
    "        self.path_x = nn.Sequential( \n",
    "            nn.Linear(shared_units, neuron_units),  # Adjusted from shared_units // 2 to shared_units\n",
    "            activation(),\n",
    "            nn.Linear(neuron_units, neuron_units),\n",
    "            activation(),\n",
    "            nn.Linear(neuron_units, neuron_units),\n",
    "            activation() )\n",
    "        self.path_t = nn.Sequential( \n",
    "            nn.Linear(shared_units, neuron_units),  # Same adjustment\n",
    "            activation(),\n",
    "            nn.Linear(neuron_units, neuron_units),\n",
    "            activation(),\n",
    "            nn.Linear(neuron_units, neuron_units),\n",
    "            activation() )\n",
    "\n",
    "        # Define separate FFN for u and v directly after the paths\n",
    "        self.ffn_u = nn.Sequential(\n",
    "            nn.Linear(neuron_units, neuron_units), activation(),\n",
    "            nn.Linear(neuron_units, neuron_units), activation(),\n",
    "            nn.Linear(neuron_units, 2)  # Outputs for u (real and imaginary parts)\n",
    "        )\n",
    "        self.ffn_v = nn.Sequential(\n",
    "            nn.Linear(neuron_units, neuron_units), activation(),\n",
    "            nn.Linear(neuron_units, neuron_units), activation(),\n",
    "            nn.Linear(neuron_units, 1)  # Output for v\n",
    "        )\n",
    "\n",
    "        self.apply(self.initialize_weights)\n",
    "\n",
    "    def forward(self, x, t):\n",
    "        # Apply Fourier feature transformations\n",
    "        x_fourier = fourier_features(x, self.Bx)\n",
    "        t_fourier = fourier_features(t, self.Bt)\n",
    "\n",
    "        # Pass through separate paths\n",
    "        x_path_output = self.path_x(x_fourier)\n",
    "        t_path_output = self.path_t(t_fourier)\n",
    "\n",
    "        # Pointwise multiplication of the separate path outputs\n",
    "        combined_features = x_path_output * t_path_output\n",
    "\n",
    "        # Directly pass through different FFNs for u and v\n",
    "        final_output_u = self.ffn_u(combined_features)\n",
    "        final_output_v = self.ffn_v(combined_features)\n",
    "\n",
    "        # Splitting the output for u into real and complex parts\n",
    "        output_1, output_2 = final_output_u.split(1, dim=-1)\n",
    "        output_3 = final_output_v\n",
    "        \n",
    "        return output_1, output_2, output_3\n",
    "\n",
    "    def initialize_weights(self, m):\n",
    "        if isinstance(m, nn.Linear):\n",
    "            init.xavier_uniform_(m.weight)\n",
    "            if m.bias is not None:\n",
    "                init.constant_(m.bias, 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "8bf33577-940a-4f3a-a754-858514cccf5c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def real_u1(x, t, k, omega, r):\n",
    "    complex_exp = torch.exp(1j * r * (omega * x + t))\n",
    "    tanh_val = torch.tanh((r * (k + x + omega * t)) / torch.sqrt(torch.tensor(2.0)))\n",
    "    result = torch.real(1j * r * complex_exp * torch.sqrt(torch.tensor(1) + omega**2) * tanh_val)\n",
    "    return result\n",
    "\n",
    "def imag_u1(x, t, k, omega, r):\n",
    "    complex_exp = torch.exp(1j * r * (omega * x + t))\n",
    "    tanh_val = torch.tanh((r * (k + x + omega * t)) / torch.sqrt(torch.tensor(2.0)))\n",
    "    result = torch.imag(1j * r * complex_exp * torch.sqrt(torch.tensor(1) + omega**2) * tanh_val)\n",
    "    return result\n",
    "\n",
    "def real_v1(x, t, k, omega, r):\n",
    "    result = (r * torch.tanh((r * (k + x + omega * t)) / torch.sqrt(torch.tensor(2.0))))**2\n",
    "    return result\n",
    "\n",
    "def compute_analytical_boundary_loss(model, x, t, mse_cost_function, k, omega, r):\n",
    "    pred_u_r, pred_u_i, pred_v = model(x, t)\n",
    "\n",
    "    real_u1_val = real_u1(x, t, k, omega, r)\n",
    "    imag_u1_val = imag_u1(x, t, k, omega, r)\n",
    "    real_v1_val = real_v1(x, t, k, omega, r)\n",
    " \n",
    "    boundary_loss_ur = mse_cost_function(pred_u_r, real_u1_val)\n",
    "    boundary_loss_ui = mse_cost_function(pred_u_i, imag_u1_val)\n",
    "    boundary_loss_v = mse_cost_function(pred_v, real_v1_val)\n",
    "    \n",
    "    return boundary_loss_ur, boundary_loss_ui, boundary_loss_v\n",
    "\n",
    "def cyclic_iterator(items):\n",
    "    return cycle(items)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "8ab5257a-6959-4595-bb23-5443087fe570",
   "metadata": {},
   "outputs": [],
   "source": [
    "def LBFGS_training(model, model_save_path, mse_cost_function, device, num_epochs, lr, num_samples, r, k, omega, gamma, beta, line_search_fn):\n",
    "    print('Starting LBFGS Fine Tuning')\n",
    "    optimizer = LBFGS(model.parameters(), lr=1, max_iter=20, max_eval=None, tolerance_grad=1e-07, tolerance_change=1e-09, history_size=100, line_search_fn=line_search_fn)\n",
    "    factor = -2\n",
    "\n",
    "    x_n = (torch.rand(num_samples, 1)*4 + factor ).to(device)  # x in range [-5, -3]\n",
    "    t_n = (torch.rand(num_samples, 1)).to(device)   \n",
    "    x_dom = (torch.rand(num_samples, 1)*4 + factor ).to(device)\n",
    "    t_dom = torch.rand(num_samples, 1).to(device) \n",
    "    x_bc_x0 = (torch.zeros(num_samples, 1)*4 + factor ).to(device)\n",
    "    t_bc_x0 = torch.rand(num_samples, 1).to(device)  # Uniformly distributed random values between 0 and 1\n",
    "    x_bc_x1 = (torch.zeros(num_samples, 1)*4 - factor ).to(device)\n",
    "    t_bc_x1 = torch.rand(num_samples, 1).to(device)  # Uniformly distributed random values between 0 and 1\n",
    "    x_bc_t0 = (torch.rand(num_samples, 1)*4 + factor ).to(device)  # Uniformly distributed random values between 0 and 1\n",
    "    t_bc_t0 = torch.zeros(num_samples, 1).to(device)\n",
    "\n",
    "    for epoch in tqdm(range(num_epochs),\n",
    "                  desc='Progress:',  \n",
    "                  leave=False,  \n",
    "                  ncols=75,\n",
    "                  mininterval=0.1,\n",
    "                  bar_format='{l_bar} {bar} | {remaining}',  # Only show the bar without any counters\n",
    "                  colour='blue'): \n",
    "        model.train()\n",
    "        \n",
    "        def closure():\n",
    "            optimizer.zero_grad()\n",
    "            x_dom = (torch.rand(num_samples, 1)*4 + factor ).to(device)\n",
    "            t_dom = torch.rand(num_samples, 1).to(device) \n",
    "            x_dom.requires_grad_(True)\n",
    "            t_dom.requires_grad_(True)\n",
    "\n",
    "            #physics_loss_ur, physics_loss_ui, physics_loss_v = compute_physics_loss(model, x_n, t_n, device, mse_cost_function) \n",
    "            u_real, u_imag, v = model(x_dom, t_dom)\n",
    "            u_abs = torch.square(u_real) + torch.square(u_imag)\n",
    "\n",
    "            # First order derivatives with retain_graph=True to reuse computational graph\n",
    "            u_real_x = torch.autograd.grad(u_real.sum(), x_dom, create_graph=True )[0]\n",
    "            u_real_t = torch.autograd.grad(u_real.sum(), t_dom, create_graph=True )[0]\n",
    "            u_imag_x = torch.autograd.grad(u_imag.sum(), x_dom, create_graph=True )[0]\n",
    "            u_imag_t = torch.autograd.grad(u_imag.sum(), t_dom, create_graph=True )[0]\n",
    "            v_x = torch.autograd.grad(v.sum(), x_dom, create_graph=True )[0]\n",
    "            v_t = torch.autograd.grad(v.sum(), t_dom, create_graph=True )[0]\n",
    "    \n",
    "            # Second order derivatives\n",
    "            u_real_xx = torch.autograd.grad(u_real_x.sum(), x_dom, create_graph=True )[0]\n",
    "            u_real_tt = torch.autograd.grad(u_real_t.sum(), t_dom, create_graph=True )[0]\n",
    "            u_imag_xx = torch.autograd.grad(u_imag_x.sum(), x_dom, create_graph=True )[0]\n",
    "            u_imag_tt = torch.autograd.grad(u_imag_t.sum(), t_dom, create_graph=True )[0]\n",
    "            v_xx = torch.autograd.grad(v_x.sum(), x_dom, create_graph=True )[0]\n",
    "            v_tt = torch.autograd.grad(v_t.sum(), t_dom, create_graph=True )[0]\n",
    "\n",
    "            # Compute u_abs_xx with retain_graph if further gradients need to be calculated\n",
    "            u_abs_x = torch.autograd.grad(u_abs.sum(), x_dom, create_graph=True )[0]\n",
    "            u_abs_xx = torch.autograd.grad(u_abs_x.sum(), x_dom, create_graph=True )[0]\n",
    "\n",
    "            # Define du_r, du_i, dv according to given formulas\n",
    "            du_r = u_real_tt - u_real_xx + u_abs * u_real - 2 * u_real * v\n",
    "            du_i = u_imag_tt - u_imag_xx + u_abs * u_imag - 2 * u_imag * v\n",
    "            dv = v_tt + v_xx - u_abs_xx\n",
    "\n",
    "            zero_target = torch.zeros_like(du_r)  # Assuming du_r, du_i, dv have the same shape\n",
    "            physics_loss_ur = mse_cost_function(du_r, zero_target)\n",
    "            physics_loss_ui = mse_cost_function(du_i, zero_target)\n",
    "            physics_loss_v = mse_cost_function(dv, zero_target)\n",
    "            #print(physics_loss_ur)\n",
    "            #print(physics_loss_ur)\n",
    "            #print(physics_loss_v)\n",
    "\n",
    "            boundary_loss_ur_x0, boundary_loss_ui_x0, boundary_loss_v_x0 = compute_analytical_boundary_loss(model, x_bc_x0, t_bc_x0, mse_cost_function, k, omega, r)\n",
    "            boundary_loss_ur_x1, boundary_loss_ui_x1, boundary_loss_v_x1 = compute_analytical_boundary_loss(model, x_bc_x1, t_bc_x1, mse_cost_function, k, omega, r)\n",
    "            boundary_loss_ur_t0, boundary_loss_ui_t0, boundary_loss_v_t0 = compute_analytical_boundary_loss(model, x_bc_t0, t_bc_t0, mse_cost_function, k, omega, r)\n",
    "            # boundary_loss_ur_t1, boundary_loss_ui_t1, boundary_loss_v_t1 = compute_analytical_boundary_loss(model, x_bc_t1, t_bc_t1, mse_cost_function, k, omega, r)\n",
    "            domain_loss_ur_t, domain_loss_ui_t, domain_loss_v_t = compute_analytical_boundary_loss(model, x_dom, t_dom, mse_cost_function, k, omega, r)\n",
    "            \n",
    "            # Total loss \n",
    "            loss_ur = gamma*(physics_loss_ur) + beta*( boundary_loss_ur_x0 + boundary_loss_ur_t0 + domain_loss_ur_t)\n",
    "            loss_ui = gamma*(physics_loss_ui) + beta*( boundary_loss_ui_x0 + boundary_loss_ui_t0 + domain_loss_ui_t)\n",
    "            loss_v = gamma*(physics_loss_v) + beta*( boundary_loss_v_x0 + boundary_loss_v_t0 + domain_loss_v_t )\n",
    "            total_loss = loss_ur + loss_ui + loss_v\n",
    "            total_loss.backward()\n",
    "\n",
    "            return total_loss \n",
    "    \n",
    "        optimizer.step(closure)\n",
    "        if epoch % 10 == 0:\n",
    "            current_loss = closure()  # Optionally recompute to print\n",
    "            print(f' Epoch {epoch}, Loss: {current_loss.item()}') \n",
    "            model_filename = os.path.join(model_save_path, f'C_HIGGS_second_training_epoch_{epoch}.pth')\n",
    "            torch.save(model.state_dict(), model_filename)\n",
    "            plot_model_results(epoch, model, device, k, omega, r, sigma=1, cmap='viridis', image_save_path='results') \n",
    "            \n",
    "    model_filename = os.path.join(model_save_path, f'C_HIGGS_second_training.pth')\n",
    "    torch.save(model.state_dict(), model_filename)\n",
    "    print('TRAINING COMPLETED')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "eb18d7cc-4719-4832-8506-8d3174de888e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from mpl_toolkits.mplot3d import Axes3D\n",
    "from scipy.ndimage import gaussian_filter\n",
    "\n",
    "def plot_model_results(epoch, model, device, k, omega, r, sigma=1, cmap='viridis', image_save_path='results'):\n",
    "    x = torch.linspace(-1.8, 1.8, 400)\n",
    "    t = torch.linspace(0.2, 0.8, 400)\n",
    "    X, T = torch.meshgrid(x, t)  # Create a 2D grid of x and t\n",
    "    X_flat = X.flatten().unsqueeze(-1).to(device)\n",
    "    T_flat = T.flatten().unsqueeze(-1).to(device)\n",
    "    \n",
    "    model_save_path = 'model_weights' \n",
    "    model_state = torch.load(os.path.join(model_save_path, f'C_HIGGS_second_training_epoch_{epoch}.pth'), map_location=device)\n",
    "    model.load_state_dict(model_state)\n",
    "    model.eval()\n",
    "\n",
    "    # Get predictions from the trained models\n",
    "    with torch.no_grad():\n",
    "        pred_u_r, pred_u_i, pred_v = model(X_flat, T_flat) \n",
    "\n",
    "    pred_u_r = pred_u_r.cpu().reshape(X.shape).numpy()\n",
    "    pred_u_i = pred_u_i.cpu().reshape(X.shape).numpy()\n",
    "    pred_v = pred_v.cpu().reshape(X.shape).numpy()\n",
    "\n",
    "    real_u1_analytical = real_u1(X_flat, T_flat, k, omega, r).cpu().reshape(X.shape).numpy()\n",
    "    imag_u1_analytical = imag_u1(X_flat, T_flat, k, omega, r).cpu().reshape(X.shape).numpy()\n",
    "    real_v1_analytical = real_v1(X_flat, T_flat, k, omega, r).cpu().reshape(X.shape).numpy()\n",
    "\n",
    "    pred_v_smooth = gaussian_filter(pred_v, sigma=sigma)\n",
    "\n",
    "    shrink = 0.3\n",
    "    aspect = 50\n",
    "\n",
    "    # Plotting predictions\n",
    "    fig = plt.figure(figsize=(24, 16))\n",
    "\n",
    "    ax1 = fig.add_subplot(231, projection='3d')\n",
    "    ax1.plot_surface(X.numpy(), T.numpy(), pred_u_r, cmap=cmap)\n",
    "    ax1.set_title('Predicted Real Part of $u_1(x, t)$')\n",
    "    ax1.set_xlabel('x')\n",
    "    ax1.set_ylabel('t')\n",
    "    ax1.set_zlabel('Real part of $u_1$')\n",
    "\n",
    "    ax2 = fig.add_subplot(232, projection='3d')\n",
    "    ax2.plot_surface(X.numpy(), T.numpy(), pred_u_i, cmap=cmap)\n",
    "    ax2.set_title('Predicted Imaginary Part of $u_1(x, t)$')\n",
    "    ax2.set_xlabel('x')\n",
    "    ax2.set_ylabel('t')\n",
    "    ax2.set_zlabel('Imag part of $u_1$')\n",
    "\n",
    "    ax3 = fig.add_subplot(233, projection='3d')\n",
    "    ax3.plot_surface(X.numpy(), T.numpy(), pred_v_smooth, cmap=cmap)\n",
    "    ax3.set_title('Predicted Real Part of $v_1(x, t)$')\n",
    "    ax3.set_xlabel('x')\n",
    "    ax3.set_ylabel('t')\n",
    "    ax3.set_zlabel('Real part of $v_1$')\n",
    "\n",
    "    ax4 = fig.add_subplot(234, projection='3d')\n",
    "    ax4.plot_surface(X.numpy(), T.numpy(), real_u1_analytical, cmap=cmap)\n",
    "    ax4.set_title('Analytical Real Part of $u_1(x, t)$')\n",
    "    ax4.set_xlabel('x')\n",
    "    ax4.set_ylabel('t')\n",
    "    ax4.set_zlabel('Real part of $u_1$')\n",
    "\n",
    "    ax5 = fig.add_subplot(235, projection='3d')\n",
    "    ax5.plot_surface(X.numpy(), T.numpy(), imag_u1_analytical, cmap=cmap)\n",
    "    ax5.set_title('Analytical Imaginary Part of $u_1(x, t)$')\n",
    "    ax5.set_xlabel('x')\n",
    "    ax5.set_ylabel('t')\n",
    "    ax5.set_zlabel('Imag part of $u_1$')\n",
    "\n",
    "    ax6 = fig.add_subplot(236, projection='3d')\n",
    "    ax6.plot_surface(X.numpy(), T.numpy(), real_v1_analytical, cmap=cmap)\n",
    "    ax6.set_title('Analytical Real Part of $v_1(x, t)$')\n",
    "    ax6.set_xlabel('x')\n",
    "    ax6.set_ylabel('t')\n",
    "    ax6.set_zlabel('Real part of $v_1$')\n",
    "\n",
    "    plt.tight_layout()\n",
    "    plt.savefig(os.path.join(image_save_path, f'chiggs_model_comparison_3d_epoch_{epoch}.png'))\n",
    "    plt.close(fig)  # Close the figure to free memory\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "89e64a9b-ec97-4750-b565-6d916c706d43",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CUDA is available! Training on GPU.\n",
      "FourierFeatureNN(\n",
      "  (path_x): Sequential(\n",
      "    (0): Linear(in_features=16, out_features=32, bias=True)\n",
      "    (1): Tanh()\n",
      "    (2): Linear(in_features=32, out_features=32, bias=True)\n",
      "    (3): Tanh()\n",
      "    (4): Linear(in_features=32, out_features=32, bias=True)\n",
      "    (5): Tanh()\n",
      "  )\n",
      "  (path_t): Sequential(\n",
      "    (0): Linear(in_features=16, out_features=32, bias=True)\n",
      "    (1): Tanh()\n",
      "    (2): Linear(in_features=32, out_features=32, bias=True)\n",
      "    (3): Tanh()\n",
      "    (4): Linear(in_features=32, out_features=32, bias=True)\n",
      "    (5): Tanh()\n",
      "  )\n",
      "  (ffn_u): Sequential(\n",
      "    (0): Linear(in_features=32, out_features=32, bias=True)\n",
      "    (1): Tanh()\n",
      "    (2): Linear(in_features=32, out_features=32, bias=True)\n",
      "    (3): Tanh()\n",
      "    (4): Linear(in_features=32, out_features=2, bias=True)\n",
      "  )\n",
      "  (ffn_v): Sequential(\n",
      "    (0): Linear(in_features=32, out_features=32, bias=True)\n",
      "    (1): Tanh()\n",
      "    (2): Linear(in_features=32, out_features=32, bias=True)\n",
      "    (3): Tanh()\n",
      "    (4): Linear(in_features=32, out_features=1, bias=True)\n",
      "  )\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "\n",
    "if torch.cuda.is_available():\n",
    "    print(\"CUDA is available! Training on GPU.\")\n",
    "else:\n",
    "    print(\"CUDA is not available. Training on CPU.\")\n",
    "\n",
    "model = FourierFeatureNN(device=device).to(device)\n",
    "\n",
    "print(model)\n",
    "num_epochs_lbfgs = 500  # Number of training epochs\n",
    "num_samples_lbfgs = 1000*3 # Number of samples for training\n",
    "num_epochs_sq = 36000\n",
    "num_samples_sq = 1000\n",
    "lr_sq = 1e-4 \n",
    "lr_lbfgs = 1e-3\n",
    "r = 1.1\n",
    "omega = 5 \n",
    "k = 0.5\n",
    "gamma = 0\n",
    "beta = 1\n",
    "model_save_path = 'model_weights' \n",
    "mse_cost_function = torch.nn.MSELoss()\n",
    "os.makedirs(model_save_path, exist_ok=True)\n",
    "losses = []\n",
    "line_search_fn = \"strong_wolfe\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "8663dca4-1f75-4867-8640-ff323f7ef06f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Starting LBFGS Fine Tuning\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Progress::   0%| \u001b[34m                                                      \u001b[0m | ?\u001b[0m"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Epoch 0, Loss: 40.32649612426758\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Progress::   2%| \u001b[34m█                                                 \u001b[0m | 15:13\u001b[0m"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Epoch 10, Loss: 11.141274452209473\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Progress::   4%| \u001b[34m██                                                \u001b[0m | 16:05\u001b[0m"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Epoch 20, Loss: 0.36170750856399536\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Progress::   6%| \u001b[34m███                                               \u001b[0m | 11:55\u001b[0m"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Epoch 30, Loss: 0.2836231589317322\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Progress::   8%| \u001b[34m████                                              \u001b[0m | 13:01\u001b[0m"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Epoch 40, Loss: 0.12260878086090088\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Progress::  10%| \u001b[34m█████                                             \u001b[0m | 11:34\u001b[0m"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Epoch 50, Loss: 0.11815682798624039\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Progress::  12%| \u001b[34m██████                                            \u001b[0m | 09:44\u001b[0m"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Epoch 60, Loss: 0.1038997620344162\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Progress::  14%| \u001b[34m███████                                           \u001b[0m | 10:38\u001b[0m"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Epoch 70, Loss: 0.09809969365596771\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Progress::  16%| \u001b[34m████████                                          \u001b[0m | 10:10\u001b[0m"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Epoch 80, Loss: 0.07776270806789398\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Progress::  18%| \u001b[34m█████████                                         \u001b[0m | 13:05\u001b[0m"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Epoch 90, Loss: 0.07059863954782486\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Progress::  20%| \u001b[34m██████████                                        \u001b[0m | 16:16\u001b[0m"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Epoch 100, Loss: 0.05988970398902893\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Progress::  22%| \u001b[34m███████████                                       \u001b[0m | 07:44\u001b[0m"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Epoch 110, Loss: 0.057288337498903275\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Progress::  24%| \u001b[34m████████████                                      \u001b[0m | 08:42\u001b[0m"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Epoch 120, Loss: 0.05308811366558075\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Progress::  26%| \u001b[34m█████████████                                     \u001b[0m | 12:48\u001b[0m"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Epoch 130, Loss: 0.05034566670656204\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Progress::  28%| \u001b[34m██████████████                                    \u001b[0m | 09:53\u001b[0m"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Epoch 140, Loss: 0.04841305688023567\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Progress::  30%| \u001b[34m███████████████                                   \u001b[0m | 05:49\u001b[0m"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Epoch 150, Loss: 0.045653000473976135\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Progress::  32%| \u001b[34m████████████████                                  \u001b[0m | 06:57\u001b[0m"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Epoch 160, Loss: 0.04532000795006752\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Progress::  34%| \u001b[34m█████████████████                                 \u001b[0m | 04:54\u001b[0m"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Epoch 170, Loss: 0.043359436094760895\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Progress::  36%| \u001b[34m██████████████████                                \u001b[0m | 09:21\u001b[0m"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Epoch 180, Loss: 0.041150398552417755\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Progress::  38%| \u001b[34m███████████████████                               \u001b[0m | 09:42\u001b[0m"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Epoch 190, Loss: 0.039673756808042526\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Progress::  40%| \u001b[34m████████████████████                              \u001b[0m | 07:32\u001b[0m"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Epoch 200, Loss: 0.037515297532081604\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Progress::  42%| \u001b[34m█████████████████████                             \u001b[0m | 06:02\u001b[0m"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Epoch 210, Loss: 0.037325166165828705\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Progress::  44%| \u001b[34m██████████████████████                            \u001b[0m | 07:17\u001b[0m"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Epoch 220, Loss: 0.035366810858249664\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Progress::  46%| \u001b[34m███████████████████████                           \u001b[0m | 08:18\u001b[0m"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Epoch 230, Loss: 0.03297027572989464\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Progress::  48%| \u001b[34m████████████████████████                          \u001b[0m | 07:23\u001b[0m"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Epoch 240, Loss: 0.030817843973636627\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Progress::  50%| \u001b[34m█████████████████████████                         \u001b[0m | 09:24\u001b[0m"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Epoch 250, Loss: 0.029109612107276917\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Progress::  52%| \u001b[34m██████████████████████████                        \u001b[0m | 07:40\u001b[0m"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Epoch 260, Loss: 0.02840748056769371\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Progress::  54%| \u001b[34m███████████████████████████                       \u001b[0m | 08:24\u001b[0m"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Epoch 270, Loss: 0.027436252683401108\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Progress::  56%| \u001b[34m████████████████████████████                      \u001b[0m | 06:43\u001b[0m"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Epoch 280, Loss: 0.026334404945373535\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Progress::  58%| \u001b[34m████████████████████████████▉                     \u001b[0m | 07:59\u001b[0m"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Epoch 290, Loss: 0.025803817436099052\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Progress::  60%| \u001b[34m██████████████████████████████                    \u001b[0m | 04:25\u001b[0m"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Epoch 300, Loss: 0.024911588057875633\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Progress::  62%| \u001b[34m███████████████████████████████                   \u001b[0m | 05:31\u001b[0m"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Epoch 310, Loss: 0.023505810648202896\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Progress::  64%| \u001b[34m████████████████████████████████                  \u001b[0m | 06:49\u001b[0m"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Epoch 320, Loss: 0.023122970014810562\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Progress::  66%| \u001b[34m█████████████████████████████████                 \u001b[0m | 05:07\u001b[0m"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Epoch 330, Loss: 0.021974485367536545\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Progress::  68%| \u001b[34m██████████████████████████████████                \u001b[0m | 04:51\u001b[0m"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Epoch 340, Loss: 0.020499899983406067\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Progress::  70%| \u001b[34m███████████████████████████████████               \u001b[0m | 03:39\u001b[0m"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Epoch 350, Loss: 0.020918721333146095\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Progress::  72%| \u001b[34m████████████████████████████████████              \u001b[0m | 04:46\u001b[0m"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Epoch 360, Loss: 0.020186275243759155\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Progress::  74%| \u001b[34m█████████████████████████████████████             \u001b[0m | 02:44\u001b[0m"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Epoch 370, Loss: 0.019499406218528748\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Progress::  76%| \u001b[34m██████████████████████████████████████            \u001b[0m | 04:05\u001b[0m"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Epoch 380, Loss: 0.019053664058446884\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Progress::  78%| \u001b[34m███████████████████████████████████████           \u001b[0m | 03:04\u001b[0m"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Epoch 390, Loss: 0.018598293885588646\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Progress::  80%| \u001b[34m████████████████████████████████████████          \u001b[0m | 02:28\u001b[0m"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Epoch 400, Loss: 0.017988935112953186\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Progress::  82%| \u001b[34m█████████████████████████████████████████         \u001b[0m | 03:05\u001b[0m"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Epoch 410, Loss: 0.01772921159863472\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Progress::  84%| \u001b[34m██████████████████████████████████████████        \u001b[0m | 02:26\u001b[0m"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Epoch 420, Loss: 0.016262587159872055\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Progress::  86%| \u001b[34m███████████████████████████████████████████       \u001b[0m | 01:51\u001b[0m"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Epoch 430, Loss: 0.01592869684100151\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Progress::  88%| \u001b[34m████████████████████████████████████████████      \u001b[0m | 02:12\u001b[0m"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Epoch 440, Loss: 0.015365242958068848\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Progress::  90%| \u001b[34m█████████████████████████████████████████████     \u001b[0m | 01:41\u001b[0m"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Epoch 450, Loss: 0.014659320935606956\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Progress::  92%| \u001b[34m██████████████████████████████████████████████    \u001b[0m | 00:46\u001b[0m"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Epoch 460, Loss: 0.014546466991305351\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Progress::  94%| \u001b[34m███████████████████████████████████████████████   \u001b[0m | 01:15\u001b[0m"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Epoch 470, Loss: 0.01357647217810154\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Progress::  96%| \u001b[34m████████████████████████████████████████████████  \u001b[0m | 00:22\u001b[0m"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Epoch 480, Loss: 0.012881951406598091\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Progress::  98%| \u001b[34m█████████████████████████████████████████████████ \u001b[0m | 00:13\u001b[0m"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Epoch 490, Loss: 0.01287191640585661\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                           m | 00:00\u001b[0m"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TRAINING COMPLETED\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r"
     ]
    }
   ],
   "source": [
    "LBFGS_training(model, model_save_path, mse_cost_function, device, num_epochs_lbfgs, lr_lbfgs, num_samples_lbfgs, r, k, omega, gamma, beta, line_search_fn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "46caacfd-eab4-41c1-a417-054e5742b7b5",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "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.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}