How to predict the new system from the pre-trained weight in NeuralPDE

I used the code below to solve the ODE system and save trained weight.

using NeuralPDE, Lux, ModelingToolkit, Optimization, OptimizationOptimJL
import ModelingToolkit: Interval, infimum, supremum

@parameters t
@variables u1(..), u2(..)
D = Differential(t)
eq = [D(u1(t)) ~ u2(t),
    D(u2(t)) ~ 5 - 10*sin(u1(t))-1.7*u2(t)];

import ModelingToolkit: Interval
bcs = [u1(0) ~ -1, u2(0) ~ 7]
domains = [t ∈ Interval(0.0,10.0)]
dt = 0.01

input_ = length(domains) # number of dimensions
n = 16
chain =[Lux.Chain(Dense(input_,n,Lux.σ),Dense(n,1)) for _ in 1:2]

@named pde_system = PDESystem(eq,bcs,domains,[t],[u1(t),u2(t)])

strategy = NeuralPDE.GridTraining(dt)
discretization = PhysicsInformedNN(chain, strategy)
sym_prob = NeuralPDE.symbolic_discretize(pde_system, discretization)

pde_loss_functions = sym_prob.loss_functions.pde_loss_functions
bc_loss_functions = sym_prob.loss_functions.bc_loss_functions

callback = function (p, l)
    println("loss: ", l)
    return false
end

loss_functions =  [pde_loss_functions;bc_loss_functions]

function loss_function(θ,p)
    sum(map(l->l(θ) ,loss_functions))
end

f_ = OptimizationFunction(loss_function, Optimization.AutoZygote())
prob = Optimization.OptimizationProblem(f_, sym_prob.flat_init_params)

res = Optimization.solve(prob,OptimizationOptimJL.BFGS(); callback = callback, maxiters = 10000)
phi = discretization.phi

#Saving weight.
using JLD2
save("SIMB.jld2",Dict("params" => res.u.depvar))

#Loading pre-trained weight.
loadata = load("SIMB.jld2")

After that, I created the new discretization with the initial parameters from the pre-trained weight as follows:

discretization = PhysicsInformedNN(chain, strategy, init_params = loadata["params"])

However, I do not know how to use the new discretization to predict the solution for the ODE system with a new initial condition.
Because in the tutorial, I need to solve the ODE by res = Optimization. solve() then using phi() to get a trial solution:

ts = 0.0:0.01:10.0
minimizers_ = [res.u.depvar[sym_prob.depvars[i]] for i in 1:2]
u_predict  = [[phi[i]([t],minimizers_[i])[1] for t in ts] for i in 1:2]

And during the solving phase, the weight had been changed and it take time to solve similarly solving for the first time. Please help me on how could I predict the new values without resolving the ODE system.
Thank you all.

The PINN needs new weights to solve a different ODE. The PINN’s solution are the weights for a specific ODE system.

Let’s take a step back, what exactly are you trying to do?

I want to train the first network and then use it to predict the ODE system when changing the initial condition. For the ODE above, the time needed to solve the ODEs system of NeuralPDE is longer than the classical one (Tsit5()). So I think it could use the pre-trained weight as transfer learning for the ODE with changing conditions to reduce the predicting time.