Adaptive Lambas as a custom loss function (PyTorch backend)

[Vancuum]

Hi, I am trying to implement adaptive lambdas as described in this lecture. This requires the definition of a custom loss function. I've implemented it as shown below. This seems to do ...something. I'm quite new to ML in general, so I'm not quite sure the best way to test if this is working. Does this implementation make sense? Does anyone have ideas for tests I can run?

def adaptive_loss(y_true, y_pred):
    global losses_weight
    losses = torch.square(y_true-y_pred)
    # Update weights manually
    with torch.no_grad():
        losses_weight = losses_weight + lr * losses.detach()
        losses_weight = losses_weight / losses_weight.max()
    # Stack weighted losses across components
    weighted_losses = torch.stack([losses_weight[i] * losses[:, i] for i in range(losses.shape[1])], dim=1)

    # Reduce mean over all elements
    return torch.mean(weighted_losses)
 
'''NN Initalization Boilerplate goes here'''

model.compile("adam", lr=lr, loss = adaptive_loss) 

[vl-dud]

Hi @Vancuum

I doubt it will work as you expect. And it might be better to add your weights inside the function _compile_pytorch

[Vancuum]

Hi @vl-dud thanks for replying. Could you let me know what you think is wrong with this implementation?

My concern with adding the adaptive weights in _compile_pytorch is that the losses returned from losses_fn seem to be each term in the loss function (i.e. each different summation), but I need to extract the loss at each collocation point, and I'm not sure how I would access that in _compile_pytorch.

Thanks for the help!

[vl-dud]

Could you let me know what you think is wrong with this implementation?

I'm not sure about your loss stacking. Could you show all the code along with the physical task?

[Vancuum]

Sure. I am trying to do forward pulse propagation according to the linear part of the non-linear Schrodinger equation (see page 24-27) with $\beta_2 = 1$ and $\alpha = 0$. I have some doubts about the ICs and BCs we're using, but it is not important to whether or not adaptive lambdas are working. They are:

$$E_x(z,0) = 0 = E_x(z, 10)$$ $$E_x(0, t) = e^{-10(t-1.5)^2}$$

on the rectangular domain $t\in[0,10]$ and $x\in[0, 1]$.
Note that these are technically inconsistent, but $E_x(0,0)\approx 0$ and $E_x(0,10)\approx 0$ so the PINN should still be able to learn these.

I am also combining this with hard boundary constraints and Fourier features

Here is the full code:

# adapted from https://github.com/lululxvi/deepxde/blob/master/examples/pinn_forward/Schrodinger.ipynb

# Arithmetic libraries
import numpy as np
from scipy.interpolate import griddata
from scipy.special import erf

#ML Libraries
import deepxde as dde
import torch #Using backend pytorch
from deepxde import nn
from deepxde import callbacks
from deepxde import backend as bkd

#Visualization libraries
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation

# Set up geometry of the problem 
L = 1
x_lower = 0
x_upper = L
num_points_space = 1000 #This is for simulations purposes only
x = np.linspace(x_lower, x_upper, num_points_space)

a = 10
t0 = 1.5
tf = 10

t_lower = 0
t_upper = tf
num_points_time = 500
t = np.linspace(t_lower, t_upper, num_points_time)

#Define the geometry
space_domain = dde.geometry.Interval(-1, 1)
time_domain  = dde.geometry.TimeDomain(-1, 1)
geomtime = dde.geometry.GeometryXTime(space_domain, time_domain)

# Normalize x and t to both be in [-1,1] and extract them into an array
x_norm = 2*(x-x_lower)/(x_upper-x_lower) - 1
t_norm = 2*(t-t_lower)/(t_upper-t_lower) - 1
X,T = np.meshgrid(x_norm, t_norm)
domain = np.hstack((X.flatten()[:, None], T.flatten()[:, None]))

beta = 1
alpha = 0.2

def pde(x,y):
    """
    INPUTS:
        x: x[:,0] is x-coordinate
           x[:,1] is t-coordinate
        y: Network output, in this case:
            y[:,0] is u(x,t) the real part
            y[:,1] is v(x,t) the imaginary part
    OUTPUT:
        The pde in standard form i.e. something = zero
    """
    u = y[:, 0]
    v = y[:, 1]

    u_x = dde.grad.jacobian(y, x, i=0, j=0)
    v_x = dde.grad.jacobian(y, x, i=1, j=0)

    u_tt = dde.grad.hessian(y, x, component = 0, i = 1, j = 1)
    v_tt = dde.grad.hessian(y, x, component = 1, i = 1, j = 1)

    pot_u = (alpha/2) * u
    pot_v = (alpha/2) * v

    #with torch.no_grad():
     #   potential_u = (alpha/2) * u
      #  potential_v = (alpha/2) * v

    f_u = u_x - (beta/2)*v_tt - pot_u
    f_v = v_x + (beta/2)*u_tt - pot_v

    return [f_u, f_v]

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

#net = dde.nn.FNN([2]+[500]*5+[2], "tanh", "Glorot normal")

class FourierFeatures(torch.nn.Module):
    def __init__(self, in_features, num_frequencies, mu=0.0, sigma=1.0):
        super().__init__()
        self.in_features = in_features
        self.num_frequencies = num_frequencies

        # Gamma ~ N(mu, sigma^2)
        self.gamma = torch.nn.Parameter(
            torch.normal(mu, sigma, size=(num_frequencies, in_features)),
            requires_grad=False  # Gamma is fixed, not trainable
        )

    def forward(self, x):
        # x: shape (N, in_features)
        gamma = self.gamma.to(x.device)
        x_proj = 2 * np.pi * torch.matmul(x, gamma.T)  # (N, num_frequencies)
        return torch.cat([torch.cos(x_proj), torch.sin(x_proj)], dim=-1)  # (N, 2 * num_frequencies)

def create_pinn_model(in_dim, out_dim, num_frequencies, mu=0.0, sigma=1.0):
    # Define the Fourier feature transformation
    fourier_layer = FourierFeatures(in_dim, num_frequencies, mu, sigma)

    # Define the network using DeepXDE's FNN with feature transformation
    net = dde.maps.FNN([in_dim*num_frequencies] + [200]*2 + [out_dim], "tanh", "Glorot normal")
    net.apply_feature_transform(fourier_layer)
    return net

net = create_pinn_model(2, 2, 30, mu = 0, sigma = 2)

losses_weight = torch.ones((num_domain, 1), dtype=torch.float32)
lr = 1e-3

def adaptive_loss(y_true, y_pred):
    global losses_weight
    losses = torch.square(y_true-y_pred)
    # Update weights manually
    with torch.no_grad():
        losses_weight = losses_weight + lr * losses.detach()
        losses_weight = losses_weight / losses_weight.max()
    # Stack weighted losses across components
    weighted_losses = torch.stack([losses_weight[i] * losses[:, i] for i in range(losses.shape[1])], dim=1)

    # Reduce mean over all elements
    return torch.mean(weighted_losses)

# This function applies hard boundary and initial constraints
def output_transform(domain, outputs):
    x_norm = domain[:, 0:1]
    t_norm = domain[:, 1:2]

    x = (x_upper-x_lower)*(x_norm+1)/2 + x_lower
    t = (t_upper-t_lower)*(t_norm+1)/2 + t_lower

    u = outputs[:, 0:1]
    v = outputs[:, 1:2]

    bc_u = x*u + torch.exp(-a*((t-t0)**2)) 
    bc_v = x*v

    return torch.cat([bc_u,bc_v], dim = 1)

net.apply_output_transform(output_transform)

model = dde.Model(data, net)

model.compile("adam", lr=lr, loss = adaptive_loss)

def no_test_loss(self, best=False):
    self.train_state.loss_test = None
    return None

model.test = no_test_loss.__get__(model, dde.Model)  # Bind method to the model instance

# Save only the best model
checkpoint_cb = dde.callbacks.ModelCheckpoint(
    "./models/nlse-opti-best",  
    save_better_only=True, 
    period=1           
)

losshistory, train_state = model.train(iterations = 30000, display_every = 1, callbacks = [checkpoint_cb])
#model.restore('./models/nlse-opti-best.pt')

prediction = model.predict(domain, operator = None)

# Extract the predicted data
real = griddata(domain, prediction[:, 0], (X,T), method='nearest')
imag = griddata(domain, prediction[:, 1], (X,T), method='nearest')
mod2 = real**2 + imag**2

'''Reference Solution (probably incorrect)'''
def coeff(m):
    pi = np.pi
    i = 1j  # imaginary unit
    
    # Common terms
    sin_term = np.sin(pi * m * t0 / tf)
    cos_term = np.cos(pi * m * t0 / tf)
    exp_term = np.exp(-pi**2 * m**2 / (4 * a * tf**2))
    sqrt_a = np.sqrt(a)
    inv_sqrt_a = 1 / sqrt_a
    conj_sqrt_a = np.conj(sqrt_a)
    conj_inv_sqrt_a = np.conj(inv_sqrt_a)

    # First group of error function arguments
    term1 = erf((2 * a * tf * (tf-t0) + i * pi * m) / (2 * sqrt_a * tf))
    term2 = erf((2 * a * tf * (tf-t0) - i * pi * m) / (2 * sqrt_a * tf))
    term3 = erf((2 * conj_sqrt_a * tf * (tf-t0) + i * pi * conj_inv_sqrt_a * m) / (2 * tf))
    term4 = erf((2 * conj_sqrt_a * tf * (tf-t0) - i * pi * conj_inv_sqrt_a * m) / (2 * tf))

    # Second group of error function arguments
    term5 = erf((2 * a * t0 * tf + i * pi * m) / (2 * sqrt_a * tf))
    term6 = erf((2 * a * t0 * tf - i * pi * m) / (2 * sqrt_a * tf))
    term7 = erf((2 * conj_sqrt_a * t0 * tf + i * pi * conj_inv_sqrt_a * m) / (2 * tf))
    term8 = erf((2 * conj_sqrt_a * t0 * tf - i * pi * conj_inv_sqrt_a * m) / (2 * tf))

    # Putting it all together  
    part1 = ((sin_term + i * cos_term) * term1 +
             (sin_term - i * cos_term) * term2 +
             (sin_term + i * cos_term) * term3 +
             (sin_term - i * cos_term) * term4 +
             (sin_term - i * cos_term) * term5 +
             (sin_term + i * cos_term) * term6 +
             (sin_term - i * cos_term) * term7 +
             (sin_term + i * cos_term) * term8) 
    '''
    part1 = ((sin_term - i * cos_term) * term1 +
             (sin_term + i * cos_term) * term2 +
             (sin_term - i * cos_term) * term3 +
             (sin_term + i * cos_term) * term4)

    part2 = ((sin_term - i * cos_term) * term5 +
             (sin_term + i * cos_term) * term6 +
             (sin_term - i * cos_term) * term7 +
             (sin_term + i * cos_term) * term8) 
    '''
    result = (1 * np.sqrt(pi) * exp_term * part1) / ( 4 * tf * sqrt_a) 
    return result

coeff_table = {m: coeff(m) for m in range(1, 20 + 1)} #Precompute fourier coefficients for improved performance

def lambda_m(m):
    return 1j*beta*(m**2)*(np.pi**2) / (2*(tf ** 2)) - alpha/2

ref = np.array([[np.sum( [coeff_table[m] * np.exp(lambda_m(m)* k)*np.sin(m*np.pi*l/tf) for m in range (1, 20)]) for k in x] for l in t])

ref_real = ref.real
ref_imag = ref.imag

def animate_model(pred, ref):
    if pred.shape != ref.shape:
        raise ValueError("Predicted and reference solutions do not have the same dimensions")
    
    # Find out the scale of the animation
    max_intensity = np.max([np.max(ref[i]) for i in range(len(ref))])
    
    # Create the plot
    fig, ax = plt.subplots(figsize=(10, 4))
    line, = ax.plot([], [], lw=2)
    line_ref, = ax.plot([],[], lw=2, color = 'green')
    ax.set_xlim(x.min(), x.max())
    ax.set_ylim(-max_intensity, max_intensity*1.1)
    ax.set_xlabel("Space")
    ax.set_ylabel("Energy")

    # Initialize the animation
    def init():
        line.set_data([], [])
        line_ref.set_data([], [])
        return line, line_ref

    # Update function for each frame
    def update(frame):
        line.set_data(x, pred[frame])
        line_ref.set_data(x, ref[frame])
        return line, line_ref

    # Create the animation
    ani = FuncAnimation(fig, update, frames=len(pred)-1, init_func=init,
                        blit=True, interval=50)

    """ Uncomment this to save the animation """
    #ani.save("./Figures/best.mp4", "ffmpeg", 60)
    plt.show()

animate_model(mod2, ref)
animate_model(real, ref_real)
animate_model(imag, ref_imag)

[vl-dud]

Please check your PDEs. May be you should use

    u = y[:, 0:1]
    v = y[:, 1:2]

instead of

    u = y[:, 0]
    v = y[:, 1]

Then you can use just losses_weight * losses instead of stacking in adaptive_loss