How to use deepxde to train an ODE model (time-dependent) using two datasets?

[sbalasbas3]

I don't think the below code is correct, and I don't even know if deepxde can already do what I want. I wanted to train a model. I have two datasets and both are assumed to come from the same underlying model, but the initial conditions differ between dataset1 and dataset2. How can I fix this code to do what I want to do?

What I did is to basically just concatenate datasets 1 and 2 and I also tried to define the initial conditions through this:

# Define the initial condition function for both datasets
def initial_condition(X):
    # Apply different initial conditions based on time
    return np.where(X[:, 0] < 5, [1.0, 0.0, 0.0], [10.0, 1.0, 0.0])  # Initial conditions for Dataset 1 and 2

# Define the boundary function
def boundary(_, on_initial):
    return on_initial

# Define the initial conditions using dde.icbc.IC
ic1 = dde.icbc.IC(geom, initial_condition, boundary, component=0)  # Component 0 is A
ic2 = dde.icbc.IC(geom, initial_condition, boundary, component=1)  # Component 1 is B
ic3 = dde.icbc.IC(geom, initial_condition, boundary, component=2)  # Component 2 is C

My entire script would look like this, but I don't know if I am doing this properly???

import numpy as np
import deepxde as dde

from scipy.integrate import solve_ivp
import matplotlib.pyplot as plt


# Function to simulate the reaction network ODEs
def reaction_ode_system(t, y, k1, k2, KmA, KmB, k3):
    A, B, C = y
    dA_dt = (-k1 * A / (1+ A/KmA)) + k2 * B / (1+B/KmB)
    dB_dt = k1 * A - k2 * B
    dC_dt = k2 * B - k3 * C
    return [dA_dt, dB_dt, dC_dt]

# Generating synthetic data for training with realistic noise
def gen_traindata(noise_std=0.02):  # noise_std controls the standard deviation of the noise
    # Time range for simulation
    t_span = (0, 480)  # Start and end time
    t_eval = np.array([0, 5, 10, 15, 20, 25, 30, 60, 90, 120, 240, 480])  # Time points for evaluation
    
    # Initial concentrations of A, B, and C
    y0 = [1.0, 0.0, 0.0]  # A(0) = 1, B(0) = 0, C(0) = 0
    
    # Known rate constants for synthetic data generation
    k1_true = 0.55  # Rate constant for A -> B
    k2_true = 0.3   # Rate constant for B -> A
    KmA_true = 0.05
    KmB_true = 0.0001
    k3_true = 0.001  # Rate constant for B -> C
    
    # Solving the ODEs to generate synthetic data
    sol = solve_ivp(reaction_ode_system,
                    t_span,
                    y0,
                    t_eval=t_eval,
                    args=(k1_true, k2_true, KmA_true, KmB_true, k3_true))
    
    # Extract the time and solution (concentrations of A, B, and C)
    t = sol.t[:, None]  # Time points (column vector)
    y = sol.y.T  # Transpose the solution to match the format [A(t), B(t), C(t)] at each time point
    
    # Add normally distributed noise to the synthetic data
    noise = np.random.normal(0, noise_std, y.shape)
    
    # Add noise to the data, ensuring that concentrations remain realistic (non-negative)
    y_noisy = y + noise
    
    # Clip the values to ensure no negative concentrations
    y_noisy = np.clip(y_noisy, 0, None)
    
    return t, y_noisy

# Generating synthetic data for training with realistic noise
def gen_traindata2(noise_std=0.02):  # noise_std controls the standard deviation of the noise
    # Time range for simulation
    t_span = (0, 480)  # Start and end time
    t_eval = np.array([0, 5, 10, 15, 20, 25, 30, 60, 90, 120, 240, 480])  # Time points for evaluation
    
    # Initial concentrations of A, B, and C
    y0 = [10.0, 1.0, 0.0]  # A(0) = 1, B(0) = 0, C(0) = 0
    
    # Known rate constants for synthetic data generation
    k1_true = 0.55  # Rate constant for A -> B
    k2_true = 0.3   # Rate constant for B -> A
    KmA_true = 0.05
    KmB_true = 0.0001
    k3_true = 0.001  # Rate constant for B -> C
    
    # Solving the ODEs to generate synthetic data
    sol = solve_ivp(reaction_ode_system,
                    t_span,
                    y0,
                    t_eval=t_eval,
                    args=(k1_true, k2_true, KmA_true, KmB_true, k3_true))
    
    # Extract the time and solution (concentrations of A, B, and C)
    t = sol.t[:, None]  # Time points (column vector)
    y = sol.y.T  # Transpose the solution to match the format [A(t), B(t), C(t)] at each time point
    
    # Add normally distributed noise to the synthetic data
    noise = np.random.normal(0, noise_std, y.shape)
    
    # Add noise to the data, ensuring that concentrations remain realistic (non-negative)
    y_noisy = y + noise
    
    # Clip the values to ensure no negative concentrations
    y_noisy = np.clip(y_noisy, 0, None)
    
    return t, y_noisy

# Trainable parameters (rate constants)
k1 = dde.Variable(25.7)
k2 = dde.Variable(13.6)
k3 = dde.Variable(5.86)

# Reaction network ODE system
def reaction_network_ode(t, y):
    """
    Reaction network ODEs for a two-species reversible reaction:
    A <-> B with rate constants k1 and k2.
    
    dy1/dt = -k1 * [A] + k2 * [B]
    dy2/dt = k1 * [A] - k2 * [B]
    """
    A, B, C = y[:, 0:1], y[:, 1:2], y[:, 2:3]
    dA_dt = dde.grad.jacobian(y, t, i=0)
    dB_dt = dde.grad.jacobian(y, t, i=1)
    dC_dt = dde.grad.jacobian(y, t, i=2)

    
    return [
        dA_dt + k1 * A - k2 * B,
        dB_dt - k1 * A + k2 * B,
        dC_dt - k2 * B + k3 * C
    ]


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

# Get the training data (from dataset or synthetic data)
observe_t, ob_y = gen_traindata(noise_std=0.0005)

# Get the training data (from dataset or synthetic data)
observe_t2, ob_y2 = gen_traindata2(noise_std=0.0005)

# Concatenate the time values and observations
observe_t = np.vstack([observe_t, observe_t2])  # Concatenate time points
ob_y = np.vstack([ob_y, ob_y2])  # Concatenate observation data

# Define the geometry (time domain) for the combined dataset
geom = dde.geometry.TimeDomain(0, 480)  # Time domain from 0 to 480

# Define the initial condition function for both datasets
def initial_condition(X):
    # Apply different initial conditions based on time
    return np.where(X[:, 0] < 5, [1.0, 0.0, 0.0], [10.0, 1.0, 0.0])  # Initial conditions for Dataset 1 and 2

# Define the boundary function
def boundary(_, on_initial):
    return on_initial

# Define the initial conditions using dde.icbc.IC
ic1 = dde.icbc.IC(geom, initial_condition, boundary, component=0)  # Component 0 is A
ic2 = dde.icbc.IC(geom, initial_condition, boundary, component=1)  # Component 1 is B
ic3 = dde.icbc.IC(geom, initial_condition, boundary, component=2)  # Component 2 is C

# Create the neural network model with the concatenated dataset
observe_y0 = dde.icbc.PointSetBC(observe_t, ob_y[:, 0:1], component=0)  # Observations for A
observe_y1 = dde.icbc.PointSetBC(observe_t, ob_y[:, 1:2], component=1)  # Observations for B
observe_y2 = dde.icbc.PointSetBC(observe_t, ob_y[:, 2:3], component=2)  # Observations for C


# Define the data object
data = dde.data.PDE(
    geom,
    reaction_network_ode,
    [ic1, ic2, ic3, observe_y0, observe_y1, observe_y2], # No observe_C!!!
    num_domain=24,
    num_boundary=2,
    anchors=observe_t,
)

# Neural network model
net = dde.nn.FNN([1] + [30] * 6 + [3], "tanh", "Glorot uniform")
model = dde.Model(data, net)

external_trainable_variables = [k1, k2, k3]
variable = dde.callbacks.VariableValue(
    external_trainable_variables, period=800, filename="variables.dat"
)

# Train with Adam optimizer
model.compile(
    "adam", lr=0.001, external_trainable_variables=external_trainable_variables
)
losshistory, train_state = model.train(iterations=50000, callbacks=[variable])

# Train with L-BFGS optimizer (if supported by the backend)
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)


# Generate predictions from the trained model
t_test = observe_t  # Use the same time points as in the synthetic data
y_pred = model.predict(t_test)  # Predict concentrations

# Extract A, B and C from synthetic data
A_true = ob_y[:, 0]  # True A(t) values from the synthetic data
B_true = ob_y[:, 1]  # True B(t) values from the synthetic data
C_true = ob_y[:, 2]

# Extract predicted A, B and C from the model
A_pred = y_pred[:, 0]  # Predicted A(t) values from the model
B_pred = y_pred[:, 1]  # Predicted B(t) values from the model
C_pred = y_pred[:, 2]

# Plotting the results
plt.figure(figsize=(15, 6))

# Plot A(t)
plt.subplot(3, 1, 1)
plt.plot(t_test, A_true, 'b-o', label='True A(t)', linewidth=2)
plt.plot(t_test, A_pred, 'r--', label='Predicted A(t)', linewidth=2)
plt.xlabel('Time (t)')
plt.ylabel('[A]')
plt.legend()
plt.grid("on")
plt.title('Comparison of True and Predicted A(t)')

# Plot B(t)
plt.subplot(3, 1, 2)
plt.plot(t_test, B_true, 'b-o', label='True B(t)', linewidth=2)
plt.plot(t_test, B_pred, 'r--', label='Predicted B(t)', linewidth=2)
plt.xlabel('Time (t)')
plt.ylabel('[B]')
plt.legend()
plt.grid("on")
plt.title('Comparison of True and Predicted B(t)')

# Plot C(t)
plt.subplot(3, 1, 3)
plt.plot(t_test, C_true, 'b-o', label='True C(t)', linewidth=2)
plt.plot(t_test, C_pred, 'r--', label='Predicted C(t)', linewidth=2)
plt.xlabel('Time (t)')
plt.ylabel('[C]')
plt.legend()
plt.grid("on")
plt.title('Comparison of True and Predicted C(t)')
# Display the plots
plt.tight_layout()
plt.show()

[lululxvi]

Use DeepONet or PI-DeepONet.