Compatibility update

POLLU/POLLU.jl

@@ -10,6 +10,11 @@ using Flux.Optimise: update!
 using Flux.Losses: mae, mse
 using BSON: @save, @load
 
+dir         = @__DIR__
+dir         = dir*"/"
+cd(dir)
+mkpath(dir*"figs")
+mkpath(dir*"checkpoint")
 
 is_restart = false
 n_epoch = 25000;
@@ -117,7 +122,7 @@ t_end = tspan[2];
 k = [0.04, 3e7, 1e4];
 tsteps = range(0, t_end, length=ntotal);
 prob_rober = ODEProblem(pollu!, u0, tspan, k);
-sol_rober = solve(prob_rober, ode_solver, saveat=tsteps, atol=1.f-6, rtol=1e-12);
+sol_rober = solve(prob_rober, ode_solver, saveat=tsteps, abstol=1.f-6, reltol=1e-12);
 normdata = Array(sol_rober)
 
 
@@ -142,13 +147,13 @@ function dudt!(du, u, p, t)
 end
 
 prob = ODEProblem(dudt!, u0[i_slow], tspan)
-sense = BacksolveAdjoint(checkpointing=true; autojacvec=ZygoteVJP());
+sense = DiffEqSensitivity.BacksolveAdjoint(checkpointing=true; autojacvec=DiffEqSensitivity.ZygoteVJP());
 
 
 function predict_n_ode(p, sample)
     global rep = re(p)
     _prob = remake(prob, p=p, tspan=[0, tsteps[sample]])
-    pred = Array(solve(_prob, ode_solver, saveat=tsteps[1:sample], atol=lb, sensalg=sense))
+    pred = Array(solve(_prob, ode_solver, saveat=tsteps[1:sample], abstol=lb, sensealg=sense))
 end
 pred = predict_n_ode(p, ntotal)
 
@@ -193,7 +198,7 @@ cb = function (p, loss_mean, g_norm)
         xlabel!(plt_grad, "Epoch")
         ylabel!(plt_loss, "Loss")
         ylabel!(plt_grad, "Gradient Norm")
-        plt_all = plot([plt_loss, plt_grad]..., framestyle=:box, size=(1000, 400))
+        plt_all = plot([plt_loss, plt_grad]..., framestyle=:box, layout=(1,2))#, size=(1000, 400))
         plot!(plt_all, xtickfontsize=10, ytickfontsize=10, xguidefontsize=12, yguidefontsize=12)
         png(plt_all, "figs/loss_grad")
 

README.md

@@ -1,5 +1,37 @@
-# StiffNeuralODE
+# Introduction
 
-Code associated with:
+This repository contains the source codes of the demonstration examples with Robertson and Pollution models in [this paper](https://doi.org/10.1063/5.0060697).
 
-> Kim, Suyong, Weiqi Ji, Sili Deng, and Christopher Rackauckas. "Stiff neural ordinary differential equations." arXiv preprint arXiv:2103.15341 (2021). https://arxiv.org/abs/2103.15341
+# Update logs
+
+- July 25, 2021: First commit
+
+- June 6, 2023: The source code has been updated for compatibility with the updated dependencies such as DiffEqFlux.jl and DiffEqSensitivity.jl.
+
+# Citation
+
+If you use the code in your research or if you find our paper useful, please cite [this paper](https://doi.org/10.1063/5.0060697):
+
+```
+@Article{kim2021stiff,
+    author = {Kim, Suyong and Ji, Weiqi and Deng, Sili and Ma, Yingbo and Rackauckas, Christopher},
+    title = "{Stiff neural ordinary differential equations}",
+    journal = {Chaos: An Interdisciplinary Journal of Nonlinear Science},
+    volume = {31},
+    number = {9},
+    year = {2021},
+    month = {09},
+    issn = {1054-1500},
+    doi = {10.1063/5.0060697},
+}
+```
+
+# Contact
+
+If you have any questions or if the code does not work due to deprecated functions in the dependencies, please contact one of the authors.
+
+Suyong Kim: suyong@mit.edu
+
+Dr. Weiqi Ji: weiqiji@mit.edu
+
+Dr. Sili Deng: silideng@mit.edu

ROBER/ROBER.jl

@@ -10,6 +10,13 @@ using Flux.Optimise: update!
 using Flux.Losses: mae, mse
 using BSON: @save, @load
 
+
+dir         = @__DIR__
+dir         = dir*"/"
+cd(dir)
+mkpath(dir*"figs")
+mkpath(dir*"checkpoint")
+
 is_restart  = false;
 n_epoch     = 50000;
 ntotal      = 50
@@ -34,7 +41,7 @@ t_end       = tspan[2];
 k           = [0.04, 3e7, 1e4];
 tsteps      = 10 .^ (range(-5, log10(tspan[2]), length=ntotal));
 prob_rober  = ODEProblem(rober!, u0, tspan, k);
-sol_rober   = solve(prob_rober, ode_solver, saveat=tsteps, atol=1.f-8);
+sol_rober   = solve(prob_rober, ode_solver, saveat=tsteps, abstol=1.f-8);
 normdata    = Array(sol_rober)
 
 yscale      = maximum(normdata, dims = 2)  #scale for each species
@@ -60,11 +67,11 @@ function dudt!(du, u, p, t)
 end
 
 prob        = ODEProblem(dudt!, u0[i_slow], tspan)
-sense       = BacksolveAdjoint(checkpointing=true; autojacvec=ZygoteVJP());
+sense       = DiffEqSensitivity.BacksolveAdjoint(checkpointing=true; autojacvec=DiffEqSensitivity.ZygoteVJP());
 function predict_n_ode(p, sample)
     global rep  = re(p)
     _prob       = remake(prob, p=p, tspan=[0, tsteps[sample]])
-    pred        = Array(solve(_prob, ode_solver, saveat=tsteps[1:sample], atol=lb, sensalg=sense))
+    pred        = Array(solve(_prob, ode_solver, saveat=tsteps[1:sample], abstol=lb, sensealg=sense))
 end
 pred        = predict_n_ode(p, ntotal)
 
@@ -110,7 +117,7 @@ cb = function (p, loss_mean, g_norm)
         xlabel!(plt_grad, "Epoch")
         ylabel!(plt_loss, "Loss")
         ylabel!(plt_grad, "Gradient Norm")
-        plt_all = plot([plt_loss, plt_grad]..., framestyle=:box, size = (1000, 400))
+        plt_all = plot([plt_loss, plt_grad]..., framestyle=:box, layout=(1,2))#, size = (1000, 400))
         plot!(plt_all, xtickfontsize=10, ytickfontsize=10, xguidefontsize=12, yguidefontsize=12)
         png(plt_all, "figs/loss_grad")