Deepxde glycolysis model

[HenrikGundestrup]

Hi everyone,

I'm working on comparing different models' abilities to handle biological systems. As part of this, I'm using DeepXDE to fit a Physics-Informed Neural Network (PINN) model to synthetic data and learn the system's parameters. My implementation is based on the Lorenz system example.

Here’s my code:

import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import odeint
from sklearn.metrics import r2_score

def gen_traindata():
    data = np.load("CellData.npz")
    return data["t"], data["y"]

def Cell(y, t, p):
    dn = p[0] - p[1] * (1 + (y[1] / p[2])**p[3]) * y[0]
    dc = p[1] * (1 + (y[1] / p[2])**p[3]) * y[0] - p[4] * y[1]
    return [dn, dc]

# Simulation parameters
times = np.arange(0, 15, 0.01)
v, k1, K, n, k2 = 15, 1, 2, 3, 4

# Generate synthetic data
y = odeint(Cell, t=times, y0=[0, 0], args=((v, k1, K, n, k2),), rtol=1e-8)


# Save the data in .npz format
np.savez("CellData.npz", t=times, y=y)

print("Data saved as CellData.npz")

C1 = dde.Variable(1.0)
C2 = dde.Variable(1.0)
C3 = dde.Variable(1.0)
C4 = dde.Variable(1.0)
C5 = dde.Variable(1.0)

# glycolysis_ode
def glycolysis_ode(x, y):
    '''glycolysis_ode
    dn = v - k1 * (1 + (y2 / K) ** n) * y1
    dc = k1 * (1 + (y2 / K) ** n) * y1 - k2 * y2
    '''
    y1, y2 = y[:, 0:1], y[:, 1:2]
    dy1_x = dde.grad.jacobian(y, x, i=0)
    dy2_x = dde.grad.jacobian(y, x, i=1)
    
    return [
        dy1_x - C1 + C2 * (1 + (y2 / C3)**C4) * y1,
        dy2_x - C2 * (1 + (y2 / C3)**C4) * y1 + C5 * y2 ,
    ]



def boundary(_, on_initial):
    return on_initial


geom = dde.geometry.TimeDomain(0, 15)

# Initial conditions
ic1 = dde.icbc.IC(geom, lambda X: 0, boundary, component=0)
ic2 = dde.icbc.IC(geom, lambda X: 0, boundary, component=1)


# Get the train data
observe_t, ob_y = gen_traindata()
observe_y0 = dde.icbc.PointSetBC(observe_t, ob_y[:, 0:1], component=0)
observe_y1 = dde.icbc.PointSetBC(observe_t, ob_y[:, 1:2], component=1)

observe_t = np.vstack(observe_t)


data = dde.data.PDE(
    geom,
    glycolysis_ode,
    [ic1, ic2, observe_y0, observe_y1],
    num_domain=400,
    num_boundary=2,
    anchors=observe_t,
)

net = dde.nn.FNN([1] + [40] * 3 + [2], "tanh", "Glorot uniform")
model = dde.Model(data, net)

external_trainable_variables = [C1, C2, C3, C4, C5]
variable = dde.callbacks.VariableValue(
    external_trainable_variables, period=600, filename="variables.dat"
)

# train adam
model.compile(
    "adam", lr=0.001, external_trainable_variables=external_trainable_variables
)
losshistory, train_state = model.train(iterations=20000, callbacks=[variable])

# train lbfgs 
model.compile("L-BFGS", external_trainable_variables=external_trainable_variables)
losshistory, train_state = model.train(callbacks=[variable])

dde.saveplot(losshistory, train_state, issave=True, isplot=True)


y_pred = model.predict(observe_t)

plt.plot(observe_t, ob_y[:, 0:1], label='obs',linewidth=4.0)
plt.plot(observe_t, ob_y[:, 1:2], label='obs',linewidth=4.0)

plt.plot(observe_t, y_pred, label='predict')
plt.legend()
plt.show()


print(C1, C2, C3, C4, C5)

However, I’m running into the following error in dde.data.PDE:

ValueError: all the input array dimensions except for the concatenation axis must match exactly, but along dimension 1, the array at index 0 has size 1 and the array at index 2 has size 1500

Has anyone encountered this before or can spot where the mismatch is happening?
Any help would be greatly appreciated! Thanks in advance.