Stiff 2D Transient Diffusion (Fick's 2nd Law) for UHPC and RC (Diffusion coefficient of 1e-12 m2/s) and positivity constraints

[jimenezastorga1990-arch]

Dear community, I am Luis, and I am currently working on my MSc thesis regarding UHPC and RC durability, and as a part of it, I need to calculate the resulting concentration field on a concrete section in 2D produced by concentration on the boundaries, provided by Dirichlet conditions. For this I am using PINNs to model the 2D concentration field, but I have run into some issues. I wonder if someone could give me some insights on this, I think I am facing "stiffness" in the PINN and convergence issues:

I have explored two different approaches in DeepXDE, but neither has fully resolved the physics. Some common code cells are here. I am working on deepxde in Python:

import torch
import torch.nn as nn
import torch.optim as optim
import matplotlib.pyplot as plt
import deepxde as dde
import numpy as np
from scipy.io import savemat
# ============================================================
# Geometry
# ============================================================

geom = dde.geometry.Rectangle([0,0],[1,1])
timedomain = dde.geometry.TimeDomain(0,1)
geomtime = dde.geometry.GeometryXTime(geom, timedomain)
# ============================================================
# Definition of the pde loss:
# ============================================================
D = 1e-12 #* (torch.exp(5 * x[:,2:3]))                               # m2/s concrete has a very low diffusion coefficient
time = 100*365*24*3600                                                # s I am doing this to normalize the time
#section dimensions
b = 1
h = 1

Lx = b
Ly = h

def pde(x, y):                                           #x is input, y is output (Concentration in this case)
    alphax = D * time / (Lx*Lx)                 # the value calculateed to normalize the PINN in x
    alphay = D * time / (Ly*Ly)                 # the value calculateed to normalize the PINN in y 
   
    dy_t = dde.grad.jacobian(y, x, i=0, j=2)   #dc/dt
    dy_xx = dde.grad.hessian(y, x, i=0, j=0)   #d2c/dx2
    dy_yy = dde.grad.hessian(y, x, i=0, j=1)   #d2c/dy2
    
    return alphax * dy_xx + alphay *dy_yy  - dy_t     #the equation is solved like this to avoid large residual values, since alpha is low

1.- Hard PINNs: I have the following code for HPINNs. For the BCs and ICs, I have used an Ansatz like the following:

# ============================================================
# Data of model:
# ============================================================
data = dde.data.TimePDE(
    geomtime,
    pde,
    [],            # Only IC, BC is handled by output transform
    num_domain=10000,     # Points inside domain (as recommended)
    num_boundary=0,       # in HPINNs these are encoded inside the ansatz  
    num_initial=0,        # in HPINNs these are encoded inside the ansatz
    train_distribution="LHS"      #for better distribution of points (as suggested)
)

# ============================================================
# Ansatz:
# ============================================================

def A(x):
    f0y = 0                           #function at the left edge in f(y)
    f1y = 1* x[:,2:3]                #function at the rigth edge in f(y)
    g0x = x[:,0:1] * x[:,2:3]         #function at the bottom edge in f(x)
    g00 = 0                          #function at the bottom edge in f(x=0)  
    g01 = x[:,2:3]                    #function at the bottom edge in f(x=1)
    g1x = x[:,0:1] * x[:,2:3]         #function at the top edge in f(x)
    g10 = 0                           #function at the top edge in f(x=0)
    g11 = x[:,2:3]                    #function at the top edge in f(x=1)
    return ( 1 - x [:,0:1] ) * f0y + x[:,0:1] * f1y + ( 1 - x[:,1:2] ) * ( g0x - ( ( 1 - x[:,0:1] ) * g00 + x[:,0:1] * g01 ) ) + x[:,1:2] * ( g1x - ( ( 1 - x[:,0:1] ) *g10 + x[:,0:1] * g11) ) 

def ansatz(x):
    return x[:,2:3] * x[:,0:1] * (1 - x[:, 0:1]) * x[:,1:2] * ( 1 - x[:,1:2])   #to force the NN to be 0 at the Boundaries

def output_transform(x, y):
    return A(x) + ansatz(x) * y
    
# ============================================================
# Definition of the FNN:
# ============================================================
net = dde.nn.FNN([3,50,50,50,50,50,1], "tanh", "Glorot uniform")   #NN
net = net.to(torch.float64)                   # so smaller values can be used
net.apply_output_transform(output_transform)  # Hard BC applied with the ansatz
model = dde.Model(data, net)

# The iterations start with Adam
model.compile("adam", lr=1e-3)

# The model is trained for 9000 iterations using Adam
losshistory, train_state = model.train(iterations=500)    #as recommended

# Then the model is run for 6000 iterations using LBFGS
dde.optimizers.config.set_LBFGS_options(maxiter=500)      #I also added a second phase using LBFGS, but the results obtained are very similar from both
model.compile("L-BFGS")
model.train()

# The loss and train states are saved
dde.saveplot(losshistory, train_state, issave=True, isplot=True)

model.save("Model_4_export_matlab")

And so when I use this code I find the BCs perfect, but inside the domain there is a localized valley that goes to a negative value at the bottom (always close to the boundaries of the BCs).

The main issue with this code is that I need positivity, and I cant apply softplus (because if I apply it to the complete output transform, the model does not converge, especialy with LBFGS). And if I apply it only to y, the limits go from 0 to 1, to 0.62 and 1.32 coming from softplus, and this keeps the valley I mentioned, it moves from the negative value to 0.25).

The BCs are test functions that we believe could be used as design parameters in concrete durability, so initially they were to be kept like this, but the results are so stiff that it is not viable.

2.- After finding the previous issues, I have switched to soft PINNs (to make use of strict positivity with softplus), and for this I have developed the following code. Since in this case I can use softplus, I have defined that: instead of 0 I am using a small epsilon (razonably inside the domains of Softplus). I also implemented adaptive weights and resampling of points, to help the training:

# ============================================================
# Adaptive weights (I applied a class for this)
# ============================================================

class AdaptiveWeights(dde.callbacks.Callback):
    def __init__(self, period=2000, alpha=0.95):
        super().__init__()
        self.period = period
        self.alpha = alpha

    def on_epoch_end(self):
  
        if self.model.train_state.step % self.period != 0 or self.model.train_state.step == 0:
            return

        curr_weights = self.model.loss_weights
        losses = self.model.train_state.loss_train
        
        losses = [float(l) for l in losses]
        total = sum(losses)
        
        if total == 0: return

        new_weights_list = []
        for i, loss in enumerate(losses):
            if loss < 1e-12:
                new_weights_list.append(curr_weights[i])
                continue
            target_weight = loss / (total / len(losses))
            #target_weight = total / (len(losses) * loss) ## I was using this before
            
            old_w = float(curr_weights[i])
            updated_w = (self.alpha * old_w) + (1 - self.alpha) * target_weight

            updated_w = max(min(updated_w, 100.0), 1e-3)
            if i == 0:  # Keep PDE loss important
                updated_w = max(updated_w, 1.0)
            
            new_weights_list.append(updated_w)

        # The new weights are:
        self.model.loss_weights = new_weights_list
        
        print(f"\nStep {self.model.train_state.step}: Weights updated")
        print([f"{w:.3e}" for w in new_weights_list])

# ============================================================
# Definition of the BCs and ICs:
# ============================================================
# Boundary Locators
def boundary_left(x, on_boundary):
    return on_boundary and dde.utils.isclose(x[0], 0)

def boundary_right(x, on_boundary):
    return on_boundary and dde.utils.isclose(x[0], 1)

def boundary_bottom(x, on_boundary):
    return on_boundary and dde.utils.isclose(x[1], 0)

def boundary_top(x, on_boundary):
    return on_boundary and dde.utils.isclose(x[1], 1)

def is_initial_time(x, on_initial):
    return on_initial

# The BCs with the epsilon are:
epsilon = 1e-5

# Left: Conc = epsilon
bc_l = dde.icbc.DirichletBC(geomtime, lambda x: epsilon, boundary_left)

# Right: Conc = 1 * t (Note: x[:, 2:3] is time)
bc_r = dde.icbc.DirichletBC(geomtime, lambda x: x[:, 2:3] + epsilon, boundary_right)

# Bottom: Conc = x * t + epsilon
bc_b = dde.icbc.DirichletBC(geomtime, lambda x: x[:, 0:1] * x[:, 2:3] + epsilon, boundary_bottom)

# Top: Conc = x * t + epsilon
bc_t = dde.icbc.DirichletBC(geomtime, lambda x: x[:, 0:1] * x[:, 2:3] + epsilon, boundary_top)

# Initial Condition: at t=0, starts at epsilon
ic = dde.icbc.IC(geomtime, lambda x: epsilon, is_initial_time)

# I collect all constrains to input into the model
bcs = [bc_l, bc_r, bc_b, bc_t, ic]

# ============================================================
# Data of model
# ============================================================

data = dde.data.TimePDE(
    geomtime,
    pde,
    bcs,                  # [bc_l, bc_r, bc_b, bc_t, ic]
    num_domain=4000,
    num_boundary=4000,
    num_initial=4000,
    train_distribution="LHS",
)
# ============================================================
# FNN with softplus:
# ============================================================
net = dde.nn.FNN([3, 40, 40, 40, 40, 40, 1], "tanh", "Glorot uniform")
net.apply_output_transform(lambda x, y: torch.nn.functional.softplus(y))
net = net.to(torch.float64)

model = dde.Model(data, net)

# ============================================================
# Callback using adaptive weights
# Order:
# [PDE, bc_l, bc_r, bc_b, bc_t, ic]
# ============================================================

loss_weights = [
    torch.tensor(1.0),
    torch.tensor(1.0),
    torch.tensor(1.0),
    torch.tensor(1.0),
    torch.tensor(1.0),
    torch.tensor(1.0),
]
# ============================================================
# Compliation and training
# ============================================================

# Initial weights as floats or standard list
initial_weights = [1.0] * 6 

model.compile(
    "adam",
    lr=1e-3,
    loss_weights=initial_weights,
)


# ============================================================
# Adaptive sampling
# ============================================================

resampler = dde.callbacks.PDEPointResampler(
    period=2000,
    pde_points=True,
    bc_points=False,
)


# ============================================================
# Adaptive weighting
# ============================================================

adaptive_weights = AdaptiveWeights(period=200, alpha=0.99)


# ============================================================
# Model train
# ============================================================

losshistory, train_state = model.train(
    iterations=12000,
    callbacks=[resampler, adaptive_weights],
)

In this case convergence is extremely slow. Despite normalization and adaptive techniques, the model struggles to "push" the concentration front into the domain (especially at the edges). The PDE residual remains orders of magnitude higher than the BC losses, and I don't know what else could I do to improve this. Possitivity is ensured, but the ICs and BCs are not within a satisfactory distance of the specified value.

For now I have moved on with my thesis to the next steps. I understand that the problem of having values in a corner can be hard for PINNs to learn, and that is why I have put these functions so the values at each corner match with the functions on each edge.

I am open ot any suggestion of anyone, I would be very grateful. Whether if the BCs are wrong, or any suggestion about how to implement positivity or any other are welcome!.

[vl-dud]

Hi @jimenezastorga1990-arch,
you can try reformulating the hard boundary constraints as follows.

import os

os.environ["DDEBACKEND"] = "pytorch"

import torch
import deepxde as dde
import numpy as np


geom = dde.geometry.Rectangle([0, 0], [1, 1])
timedomain = dde.geometry.TimeDomain(0, 1)
geomtime = dde.geometry.GeometryXTime(geom, timedomain)

D = 1e-12
time = 100 * 365 * 24 * 3600
b = 1
h = 1
Lx = b
Ly = h


def pde(x, y):
    alpha = D * time / (Lx * Lx)
    dy_t = dde.grad.jacobian(y, x, i=0, j=2)
    dy_xx = dde.grad.hessian(y, x, i=0, j=0)
    dy_yy = dde.grad.hessian(y, x, i=0, j=1)
    residual = alpha * (dy_xx + dy_yy) - dy_t
    return residual


data = dde.data.TimePDE(
    geomtime,
    pde,
    [],
    num_domain=10000,
    num_boundary=0,
    num_initial=0, 
    train_distribution="LHS" 
)


def A(x):
    x_coord = x[:, 0:1]
    t_coord = x[:, 2:3]
    return x_coord * t_coord


def B(x):
    x_coord = x[:, 0:1]
    y_coord = x[:, 1:2]
    t_coord = x[:, 2:3]

    return (
        x_coord * (1.0 - x_coord)
        * y_coord * (1.0 - y_coord)
        * t_coord
    )


def output_transform(x, y):
    return A(x) * torch.exp(B(x) * y)


net = dde.nn.FNN([3, 50, 50, 50, 50, 50, 1], "LAAF-10 silu", "Glorot uniform")

net.apply_output_transform(output_transform)
model = dde.Model(data, net)

model.compile("adam", lr=5e-5)

losshistory, train_state = model.train(
    iterations=50000,
    display_every=1000,
    callbacks=[
        dde.callbacks.EarlyStopping(
            min_delta=1e-5,
            patience=5000,
            start_from_epoch=10000
        )
    ]
)

The idea is to use a boundary-satisfying base function and multiply the trainable correction by a factor that vanishes on the boundary, so the network can modify only the interior solution without breaking the hard constraints.

I would also recommend using layer-wise locally adaptive activation functions (LAAF-10 silu here), it may improve convergence.

[jimenezastorga1990-arch]

Hi @vl-dud!

Thanks a lot for your help! It actually helped! :D.

If I could just ask one question, I am working now on the implementation of the same problem, but to a 3D shell (that will represent bridge decks and slabs). I have tried Soft PINNs, where the results present differences with the BCs and so I was looking for how to approach it now.

After checking your code, my question is do you think that for 3D shells, using HPINNs with the same LAAL would help? I will start working on it right away, just checking if there might be a better activation function or way for 3D problems. Thanks again!

[vl-dud]

I’m glad to help. Yes, I think using LAAF is worth trying for 3D shells. It should help, or at least it should not make the results worse compared to standard activation functions.

I would also suggest starting with a few short trial runs with a limited number of iterations, testing different activation functions and learning rates.

[jimenezastorga1990-arch]

Thank you for your help, but if it was not too much, I would like to ask you something else.

I was asked that instead of making the boundary conditions increase linearly with time, they should increase either by the square root of the time or even exponentially:

def A(x):
x_coord = x[:, 0:1]
t_coord = x[:, 2:3]
return x_coord * torch.sqrt(t_coord)

When I try to run this with the code you proposed, the concentration near the boundary at time equal to 0.1 (roughly ten years) presents a weird shape at the bottom right corner. And also I noticed that at somepoints inside the domain, the concentration is not close to 0 for some time steps (there are small hills like 0.02, but noticeable in a plot).

If you could, I wanted to ask you if the following ideas I was thinking of couldd hep here:

1. Adding sample points in the domain where the PDE loss is high.
2. I think the hills are appearing due to the low learning rate, probably the problem is overfitting.

I am already testing what you suggested, but if you might have another idea please let me know. I would be very grateful. :D

[vl-dud]

I would try using s=sqrt(t) as the time coordinate and rewriting the PDE in terms of s.