# tfNeuralODE-Spirals

This vignette runs through a small example of how to use tfNeuralODE to learn differential equations. We start by importing all of the correct packages, and instantiating a few constants.

# initial checks for python
library(reticulate)
if(!py_available()){
install_python()
}

#checking for tensorflow and keras installation
#checking for tensorflow installation
if(!py_module_available("tensorflow")){
tensorflow::install_tensorflow()
}
library(tensorflow)
library(tfNeuralODE)
library(keras)
library(deSolve)

# constants
data_size = 1000
batch_time = 20 # this seems to works the best ...
niters = 200
batch_size = 16

Now we create the neural network that will define our system. We also instantiate a few more constants that define our system.

# ODE Model with time input

OdeModel(keras$Model) %py_class% { initialize <- function() { super$initialize()
self$block_1 <- layer_dense(units = 50, activation = 'tanh') self$block_2 <- layer_dense(units = 2, activation = 'linear')
}

call <- function(inputs) {
x<- inputs ^ 3
x <- self$block_1(x) self$block_2(x)
}
}
model<- OdeModel()
# more constants, time vectors
tsteps <- seq(0, 25, by = 25/data_size)
true_y0 = c(2., 0.)
true_A = rbind(c(-0.1, 2.0), c(-2.0, -0.1))

Now we solve the ODE and plot our results.

# solving a spiral ode
trueODEfunc<- function(du, u, p, t){
true_A = rbind(c(-0.1, 2.0), c(-2.0, -0.1))
du <- (u^3) %*% true_A
return(list(du))
}

# solved ode output
prob_trueode <- lsode(func = trueODEfunc, y = true_y0, times = tsteps)

Let’s start looking at how we’re going to train our model. We instantiate an optimizer and create a batching function.

#optimizer
optimizer = tf$keras$optimizers$legacy$Adam(learning_rate = 1e-3)

# batching function
get_batch<- function(prob_trueode, tsteps){
starts = sample(seq(1, data_size - batch_time), size = batch_size, replace = FALSE)
batch_y0 <- as.matrix(prob_trueode[starts,])
batch_yN <- as.matrix(prob_trueode[starts + batch_time,])
batch_y0 <- tf$cast((batch_y0), dtype = tf$float32)
batch_yN <- tf$cast((batch_yN), dtype = tf$float32)
return(list(batch_y0, batch_yN))
}

Now we can train our neural ODE, using naive backpropagation.

# Training Neural ODE
for(i in 1:niters){
#print(paste("Iteration", i, "out of", niters, "iterations."))
inp = get_batch(prob_trueode[,2:3], tsteps)
pred = forward(model, inputs = inp[[1]], tsteps = tsteps[1:batch_time])
with(tf$GradientTape() %as% tape, { tape$watch(pred)
loss = tf$reduce_mean(tf$abs(pred - inp[[2]]))
})
#print(paste("loss:", as.numeric(loss)))
dLoss = tape$gradient(loss, pred) list_w = backward(model, tsteps[1:batch_time], pred, output_gradients = dLoss) optimizer$apply_gradients(zip_lists(list_w[[3]], model$trainable_variables)) # graphing the Neural ODE if(i %% 200 == 0){ pred_y = forward(model = model, inputs = tf$cast(t(as.matrix(true_y0)), dtype = tf\$float32),
tsteps = tsteps, return_states = TRUE)
pred_y_c<- k_concatenate(pred_y[[2]], 1)
p_m<- as.matrix(pred_y_c)
plot(p_m, main = paste("iteration", i), type = "l", col = "red")
lines(prob_trueode[,2], prob_trueode[,3], col = "blue")
}
}
plot of the final result