# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     name: python3
# ---

# + id="dIGIo9kpYFqs"
# If the repo was not cloned
from google.colab import drive
drive.mount('/content/drive')
# %cd /content/drive/MyDrive
# !git clone https://github.com/pschultzendorff/pf-pinn
# %cd /content/drive/MyDrive/pf-pinn

# + colab={"base_uri": "https://localhost:8080/"} id="ExFq7xsAYFqs" outputId="8e9778d3-656a-4053-c1aa-78b7554301a5"
# If the repo needs to be updated
from google.colab import drive
drive.mount('/content/drive')
# %cd /content/drive/MyDrive/pf-pinn
# !git pull origin main

# + id="0PPGV_MmYFqs"
# Run tensorboard
# %load_ext tensorboard
# %tensorboard --logdir runs

# + id="lLE47Y__YFqs"
# Speedup and reproducability
import torch
from src import utils_fast as utils
from google.colab import drive
drive.mount('/content/drive')
# %cd /content/drive/MyDrive/pf-pinn
# %env PYTHONOPTIMIZE = TRUE
torch.backends.cudnn.benchmark = True
utils.fix_seeds(0)

# + id="EXDshh0cSWqE"
# Imports
from src import geometry, models
from src import lbfgs_custom_nvb as lbfgs_custom
from src import utils_fast as utils
from src import train_fast as train
from src import visualization as vis

# + [markdown] id="1BghpXbOVMcc"
# ### Set parameters
#
# We use 25 Gauss points per element, since the lateral load example may be more complicated. The loading steps are taken from section 7.2 of *A physics-informed variational DeepONet for predicting the crack path in brittle materials*.

# + id="jJ2bhIgHVMFA"
from typing import List

G_C: float = 2.7
CRACK_WIDTH: float = 0.0125
LAME_LAMBDA: float = 121150.0
LAME_MU: float = 80770.0
NUM_POINTS: int = 16
LOADING_STEP: float = 0.0005

# increasements of 5x10^-4
loads: List[float] = [LOADING_STEP*i for i in range(26)]

# + [markdown] id="omRaKZTuS4Zc"
# ### Generate the geometry

# + colab={"base_uri": "https://localhost:8080/"} id="euQ6l4nGS_0c" outputId="d3c56615-8397-4c82-c8c1-fd67b2bdf8d7"
import numpy as np
import torch

device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')

x = np.linspace(0, 1, 17)
y = np.linspace(0, 1, 17)

y1_1 = np.linspace(0, y[5], 6)
y1_2 = np.linspace(y[11], 1, 6)

x2 = np.linspace(0, 1, 33)
y2_1 = np.linspace(y[5], y[7], 5)
y2_2 = np.linspace(y[9], y[11], 5)

x3 = np.linspace(0, 1, 65)
y3 = np.linspace(y[7], y[9], 9)

x1v, y1v_1 = np.meshgrid(x, y1_1)
x1v, y1v_2 = np.meshgrid(x, y1_2)

x2v, y2v_1 = np.meshgrid(x2, y2_1)
x2v, y2v_2 = np.meshgrid(x2, y2_2)

x3v, y3v = np.meshgrid(x3, y3)

nodes_l1_1 = np.stack((x1v, y1v_1), axis=2).tolist()
nodes_l1_2 = np.stack((x1v, y1v_2), axis=2).tolist()
nodes_l2_1 = np.stack((x2v, y2v_1), axis=2).tolist()
nodes_l2_2 = np.stack((x2v, y2v_2), axis=2).tolist()
nodes_l3 = np.stack((x3v, y3v), axis=2).tolist()

nodes_lists = [nodes_l1_1, nodes_l1_2, nodes_l2_1, nodes_l2_2, nodes_l3]
elements = []

for nodes in nodes_lists:
    for row_id, row in enumerate(nodes[:-1]):
        for node_id, node in enumerate(row[:-1]):
            elements.append([
                node[0],
                nodes[row_id+1][node_id+1][0],
                node[1],
                nodes[row_id+1][node_id+1][1],
                ])

points, weights = geometry.generate_points_weights(elements, NUM_POINTS, 2)

points_t: torch.Tensor = torch.tensor(points, dtype=torch.float).to(device)
weights_t: torch.Tensor = torch.tensor(weights, dtype=torch.float).to(device).unsqueeze(dim=-1)

print(sum(weights))

x = np.linspace(0, 1, 129)
y = np.flip(np.linspace(0, 1, 129))

xv, yv = np.meshgrid(x, y)

vis_points: np.ndarray = np.stack((xv, yv), axis=2)
vis_points_t: torch.Tensor = torch.tensor(vis_points, dtype=torch.float).to(device)


# + [markdown] id="6EsAuRgxVsvo"
# ### Initialize the history function

# + id="z6-hdgqXVvQk"
def H_init(point: torch.Tensor) -> float:
    """Inital history function.
    """
    x, y = point[0], point[1]
    if x <= 0.5:
        x_dist: float = 0.0
    else:
        x_dist: float = x-0.5
    dist = torch.sqrt(x_dist**2 +(y - 0.5)**2)
    B: float = 1000.0 if dist <= CRACK_WIDTH/2 else 1.0
    if dist <= CRACK_WIDTH/2:
        return B*G_C/(2*CRACK_WIDTH)*(1 - 2*dist/CRACK_WIDTH)
    else:
        return 0.0

H_init_t: torch.Tensor = torch.tensor([H_init(point) for point in points_t]).unsqueeze(dim=-1).to(device)
H_history: List[torch.Tensor] = [H_init_t]

H_init_vis_t: torch.Tensor = torch.tensor([H_init(point) for point in vis_points_t.flatten(start_dim=0, end_dim=1)]).unsqueeze(dim=-1).to(device)
H_history_vis: List[torch.Tensor] = [H_init_vis_t]


# + [markdown] id="rgj8r9A0XLyc"
# ### Train the model

# + id="WekKk3hzXN6Y"
import os
import torch.optim as optim

loads: List[float] = [LOADING_STEP*i for i in range(16)]
utils.fix_seeds(0)

for i, load in enumerate(loads):
    path: str = f'runs/5.2_lateral/history_function/load_{load}'
    def train_history():
        loss_cfg: utils.LossConfig = utils.LossConfig(
            G_c=G_C,
            crack_width=CRACK_WIDTH,
            lame_lambda=LAME_LAMBDA,
            lame_mu=LAME_MU,
            dis_field_dim=2,
            history=H_history,
            irreversibility='history_function',
            pen_param=0.5
            )
        vis_cfg: vis.VisConfig = vis.VisConfig.from_LossConfig(loss_cfg, history=H_history_vis)
        loss_func: utils.TotalEnergy = utils.TotalEnergy(loss_cfg, loss_split=True)
        pinnbc_cfg: models.PinnBCConfig = models.PinnBCConfig(
            depth=3,
            in_dim=2,
            hidden_dim=50,
            boundary_cond=models.forward_5_2_lateral,
            dis_field_dim=2,
            load=load
            )
        model: models.PinnBC = models.PinnBC(pinnbc_cfg).float().to(device)
        if i >= 1:
            model.load_state_dict(torch.load(f'runs/5.2_lateral/history_function/load_{loads[i-1]}/models/best_model.pt'))
            for fc in model.pinn.fcs[:-1]:
                for param in fc.parameters():
                    param.requires_grad = False
            adam_epochs: int = 2000
            lbfgs_epochs: int = 500
        else:
            adam_epochs: int = 5000
            lbfgs_epochs: int = 500
        adam_cfg: train.TrainConfig = train.TrainConfig(
            points=points_t,
            weights=weights_t,
            vis_points=vis_points_t,
            save_path=path,
            epochs=adam_epochs,
            patience=600,
            patience_lr=200,
            save_step=1000,
            vis_step=1000,
            num_points=NUM_POINTS
            )
        adam_trainer: train.Trainer = train.Trainer(adam_cfg, pinnbc_cfg, loss_cfg, vis_cfg, model=model, optimizer=optim.Adam(model.parameters()), loss_func=loss_func, jupyter=True)
        end_epoch: int = adam_trainer.train()
        model.load_state_dict(torch.load(os.path.join(path, 'models', 'best_model.pt')))
        lbfgs_cfg: train.TrainConfig = train.TrainConfig(
            points=points_t,
            weights=weights_t,
            vis_points=vis_points_t,
            save_path=path,
            epochs=lbfgs_epochs,
            patience=150,
            patience_lr=50,
            save_step=100,
            vis_step=100,
            start_epoch=end_epoch+1,
            num_points=NUM_POINTS,
            min_handling='adam'
            )
        lbfgs_trainer: train.Trainer = train.Trainer(lbfgs_cfg, pinnbc_cfg, loss_cfg, vis_cfg,
                                              model=model, optimizer=lbfgs_custom.LBFGSCustom(model.parameters(), line_search_fn='strong_wolfe', tolerance_change=1e-14), loss_func=loss_func, jupyter=True)
        lbfgs_trainer.train()
        lbfgs_trainer.save_vis()
        model.load_state_dict(torch.load(os.path.join(path, 'models', 'best_model.pt'))) 
        vis.update_H(model, points_t, loss_cfg, H_history)
        vis.update_H(model, vis_points_t, vis_cfg, H_history_vis)
        vis.save_losses([adam_trainer.losses, lbfgs_trainer.losses], ['Adam', 'L-BFGS'], os.path.join(path, 'images', 'losses.png'))

    train_history()

# + [markdown] id="2G3tIqS0MPqM"
# ### Finer visualization

# + id="LP7FbbvuMPqM"
x = np.linspace(0, 1, 513)
y = np.flip(np.linspace(0, 1, 513))

xv, yv = np.meshgrid(x, y)

vis_points: np.ndarray = np.stack((xv, yv), axis=2)
vis_points_t: torch.Tensor = torch.tensor(vis_points, dtype=torch.float).to(device)

# + id="DEZk0MaeMPqM" outputId="59962e56-9635-4721-ac3f-a35d236ea794"
for i, load in enumerate(loads):
  pinnbc_cfg: models.PinnBCConfig = models.PinnBCConfig(
      depth=3,
      in_dim=2,
      hidden_dim=50,
      boundary_cond=models.forward_5_2_lateral,
      dis_field_dim=2,
      load=load
      )
  vis_cfg: utils.LossConfig = utils.LossConfig(
    G_c=G_C,
    crack_width=CRACK_WIDTH,
    lame_lambda=LAME_LAMBDA,
    lame_mu=LAME_MU,
    dis_field_dim=2,
    history=[H_history_vis[i]],
    irreversibility='history_function',
    )
  model: models.PinnBC = models.PinnBC(pinnbc_cfg).float().to(device)
  model.load_state_dict(torch.load(f'runs/5.2_lateral/history_function/load_{load}/models/best_model.pt'))
  vis.visualize_model_2d(model, vis_points_t, vis_cfg, title=True)

# + [markdown] id="N-rhhzOAB4NU"
# ### Without transfer learning
#
# We observe that the loss for the subsequent time steps after the first one does not decrease as much. For these time steps the concept of transfer learning \cite{Goswami} is applied, where all layers of the PINN except the last ones are fixed. We try another run without fixing the layers but the model parameters from the last time step are still taken as initial parameters. It is not neccessary to repeat the first time step, hence we start at the second one.

# + id="NZ9DdwzlB0Ds"
import os
import torch.optim as optim

H_init_t: torch.Tensor = torch.tensor([H_init(point) for point in points_t]).unsqueeze(dim=-1).to(device)
H_history: List[torch.Tensor] = [H_init_t]

H_init_vis_t: torch.Tensor = torch.tensor([H_init(point) for point in vis_points_t.flatten(start_dim=0, end_dim=1)]).unsqueeze(dim=-1).to(device)
H_history_vis: List[torch.Tensor] = [H_init_vis_t]

loads: List[float] = [LOADING_STEP*i for i in range(16)]
utils.fix_seeds(0)

for i, load in enumerate(loads[1:]):
    path: str = f'runs/5.2/no_tl/load_{load}'
    def train_history():
        loss_cfg: utils.LossConfig = utils.LossConfig(
            G_c=G_C,
            crack_width=CRACK_WIDTH,
            lame_lambda=LAME_LAMBDA,
            lame_mu=LAME_MU,
            dis_field_dim=2,
            history=H_history,
            irreversibility='history_function',
            )
        vis_cfg: vis.VisConfig = vis.VisConfig.from_LossConfig(loss_cfg, history=H_history_vis)
        loss_func: utils.TotalEnergy = utils.TotalEnergy(loss_cfg, loss_split=True)
        pinnbc_cfg: models.PinnBCConfig = models.PinnBCConfig(
            depth=3,
            in_dim=2,
            hidden_dim=50,
            boundary_cond=models.forward_5_2_lateral,
            dis_field_dim=2,
            load=load
            )
        model: models.PinnBC = models.PinnBC(pinnbc_cfg).float().to(device)
        if i >= 2:
            model.load_state_dict(torch.load(f'runs/5.2/no_tl/load_{loads[i-1]}/models/best_model.pt'))
            adam_epochs: int = 2000
            lbfgs_epochs: int = 450
        else:
            adam_epochs: int = 5000
            lbfgs_epochs: int = 1000
        adam_cfg: train.TrainConfig = train.TrainConfig(
            points=points_t,
            weights=weights_t,
            vis_points=vis_points_t,
            save_path=path,
            epochs=adam_epochs,
            patience=600,
            patience_lr=200,
            save_step=1000,
            vis_step=1000,
            num_points=NUM_POINTS
            )
        adam_trainer: train.Trainer = train.Trainer(adam_cfg, pinnbc_cfg, loss_cfg, vis_cfg, model=model, optimizer=optim.Adam(model.parameters()), loss_func=loss_func, jupyter=True)
        end_epoch: int = adam_trainer.train()
        model.load_state_dict(torch.load(os.path.join(path, 'models', 'best_model.pt')))
        lbfgs_cfg: train.TrainConfig = train.TrainConfig(
            points=points_t,
            weights=weights_t,
            vis_points=vis_points_t,
            save_path=path,
            epochs=lbfgs_epochs,
            patience=150,
            patience_lr=50,
            save_step=100,
            vis_step=100,
            start_epoch=end_epoch+1,
            num_points=NUM_POINTS,
            min_handling='adam'
            )
        lbfgs_trainer: train.Trainer = train.Trainer(lbfgs_cfg, pinnbc_cfg, loss_cfg, vis_cfg,
                                              model=model, optimizer=lbfgs_custom.LBFGSCustom(model.parameters(), line_search_fn='strong_wolfe', tolerance_change=1e-14), loss_func=loss_func, jupyter=True)
        lbfgs_trainer.train()
        lbfgs_trainer.save_vis()
        model.load_state_dict(torch.load(os.path.join(path, 'models', 'best_model.pt'))) 
        vis.update_H(model, points_t, loss_cfg, H_history)
        vis.update_H(model, vis_points_t, vis_cfg, H_history_vis)
        vis.save_losses([adam_trainer.losses, lbfgs_trainer.losses], ['Adam', 'L-BFGS'], os.path.join(path, 'images', 'losses.png'))

    train_history()

StarcoderdataJupyterScriptsDedupFiltered