Use other solvers than Tsit5 in multiple_shoot (#521)

* Fix parameter name in docstring

* Fix error message

should say `< 1` and `> number of data points` instead of `<=` and `>=`.

* Relax types

* Make code solver agnostic

instead of relying on hardcoded `Tsit5()`.

* Remove preset grp_size from docstring

Preset grp_size might lead to DomainErrors if the dataset has less than 5 data points.

* Update tests to solver agnostic interface, add DomainError test

* Update docs to use solver agnostic multiple shooting

* Update Project.toml

* Update Project.toml

Co-authored-by: Christopher Rackauckas <[email protected]>

Author: adrhill

Project.toml

@@ -28,7 +28,7 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444"
 
 [compat]
-Adapt = "2, 3.0"
+Adapt = "3 - 3.2"
 ConsoleProgressMonitor = "0.1"
 DataInterpolations = "3.3"
 DiffEqBase = "6.41"

docs/src/examples/multiple_shooting.md

@@ -87,7 +87,7 @@ function loss_function_param(ode_data, pred):: Float32
 end
 
 function loss_neuralode(p)
-	return multiple_shoot(p, ode_data, tsteps, prob_param, loss_function_param, grp_size_param, loss_multiplier_param)
+	return multiple_shoot(p, ode_data, tsteps, prob_param, loss_function_param, Tsit5(), grp_size_param, loss_multiplier_param)
 end
 
 result_neuralode = DiffEqFlux.sciml_train(loss_neuralode, prob_neuralode.p,

src/multiple_shooting.jl

@@ -5,7 +5,7 @@ The default group size is 5 implying the whole dataset would be divided in group
 The default continuity term is 100 implying any losses arising from the non-continuity of 2 different groups will be scaled by 100.
 
 ```julia
-multiple_shoot(p,ode_data,tsteps,prob,loss_function,grp_size=5,continuity_strength=100)
+multiple_shoot(p,ode_data,tsteps,prob,loss_function,grp_size,continuity_strength=100)
 ```
 Arguments:
 - `p`: The parameters of the Neural Network to be trained.
@@ -14,40 +14,51 @@ Arguments:
 - `prob`: ODE problem that the Neural Network attempts to solve.
 - `loss_function`: Any arbitrary function to calculate loss.
 - `grp_size`: The group size achieved after splitting the ode_data into equal sizes.
-- `continuity_strength`: Multiplying factor to ensure continuity of predictions throughout different groups.
+- `continuity_term`: Multiplying factor to ensure continuity of predictions throughout different groups.
 
 !!!note
-The parameter 'continuity_strength' should be a relatively big number to enforce a large penalty whenever the last point of any group doesn't coincide with the first point of next group.
+The parameter 'continuity_term' should be a relatively big number to enforce a large penalty whenever the last point of any group doesn't coincide with the first point of next group.
 """
-function multiple_shoot(p :: Array, ode_data :: Array, tsteps, prob :: ODEProblem, loss_function ::Function, grp_size :: Integer = 5, continuity_term :: Integer = 100)
+function multiple_shoot(
+    p::AbstractArray,
+    ode_data::AbstractArray,
+    tsteps,
+    prob::ODEProblem,
+    loss_function::Function,
+    solver::DiffEqBase.AbstractODEAlgorithm,
+    grp_size::Integer,
+    continuity_term::Real=100
+)
 	datasize = length(ode_data[1,:])
 
-	@assert (grp_size >= 1 && grp_size <= datasize) "grp_size can't be <= 1 or >= number of data points"
+	if grp_size < 1 || grp_size > datasize
+        throw(DomainError(grp_size, "grp_size can't be < 1 or > number of data points"))
+    end
 
 	tot_loss = 0
 
 	if(grp_size == datasize)
-		grp_predictions = [solve(remake(prob, p = p, tspan = (tsteps[1],tsteps[datasize]), u0 = ode_data[:,1]), Tsit5(),saveat = tsteps)]
+		grp_predictions = [solve(remake(prob, p=p, tspan=(tsteps[1],tsteps[datasize]), u0=ode_data[:,1]), solver, saveat=tsteps)]
 		tot_loss = loss_function(ode_data, grp_predictions[1][:,1:grp_size])
 		return tot_loss, grp_predictions[1], grp_predictions
 	end
 
 	if(grp_size == 1)
 		# store individual single shooting predictions for each group
-		grp_predictions = [solve(remake(prob, p = p, tspan = (tsteps[i],tsteps[i+1]), u0 = ode_data[:,i]), Tsit5(),saveat = tsteps[i:i+1]) for i in 1:datasize-1]
+		grp_predictions = [solve(remake(prob, p=p, tspan=(tsteps[i],tsteps[i+1]), u0=ode_data[:,i]), solver, saveat=tsteps[i:i+1]) for i in 1:datasize-1]
 
 		# calculate multiple shooting loss from the single shoots done in above step
 		for i in 1:datasize-1
 		tot_loss += loss_function(ode_data[:,i:i+1], grp_predictions[i][:, 1:grp_size]) + (continuity_term * sum(abs,grp_predictions[i][:,2] - ode_data[:,i+1]))
 		end
 
 		# single shooting predictions from ode_data[:,1] (= u0)
-		pred = solve(remake(prob, p = p, tspan = (tsteps[1],tsteps[datasize]), u0 = ode_data[:,1]), Tsit5(),saveat = tsteps)
+		pred = solve(remake(prob, p=p, tspan=(tsteps[1],tsteps[datasize]), u0=ode_data[:,1]), solver, saveat=tsteps)
 		return tot_loss, pred, grp_predictions
 	end
 
 	# multiple shooting predictions
-	grp_predictions = [solve(remake(prob, p = p, tspan = (tsteps[i],tsteps[i+grp_size-1]), u0 = ode_data[:,i]), Tsit5(),saveat = tsteps[i:i+grp_size-1]) for i in 1:grp_size-1:datasize-grp_size]
+	grp_predictions = [solve(remake(prob, p=p, tspan=(tsteps[i],tsteps[i+grp_size-1]), u0=ode_data[:,i]), solver, saveat=tsteps[i:i+grp_size-1]) for i in 1:grp_size-1:datasize-grp_size]
 
 	# calculate multiple shooting loss
 	for i in 1:grp_size-1:datasize-grp_size
@@ -56,6 +67,6 @@ function multiple_shoot(p :: Array, ode_data :: Array, tsteps, prob :: ODEProble
 	end
 
 	# single shooting prediction
-	pred = solve(remake(prob, p = p, tspan = (tsteps[1],tsteps[datasize]), u0 = ode_data[:,1]), Tsit5(),saveat = tsteps)
+	pred = solve(remake(prob, p=p, tspan=(tsteps[1],tsteps[datasize]), u0=ode_data[:,1]), solver, saveat=tsteps)
 	return tot_loss, pred, grp_predictions
 end

test/multiple_shoot.jl

@@ -51,7 +51,7 @@ function loss_function_param(ode_data, pred):: Float32
 end
 
 function loss_neuralode_param(p)
-	return multiple_shoot(p, ode_data, tsteps, prob_param, loss_function_param, grp_size_param, loss_multiplier_param)
+	return multiple_shoot(p, ode_data, tsteps, prob_param, loss_function_param, Tsit5(), grp_size_param, loss_multiplier_param)
 end
 
 
@@ -76,3 +76,7 @@ multiple_shoot_loss_2 = general_loss_function(multiple_shoot_result_neuralode_2)
 println("multiple_shoot_loss_2: ",multiple_shoot_loss_2)
 
 @test multiple_shoot_loss_2 < single_shoot_loss
+
+# test for DomainErrors
+@test_throws DomainError multiple_shoot(prob_neuralode.p, ode_data, tsteps, prob_param, loss_function_param, Tsit5(), 0, loss_multiplier_param)
+@test_throws DomainError multiple_shoot(prob_neuralode.p, ode_data, tsteps, prob_param, loss_function_param, Tsit5(), datasize + 1, loss_multiplier_param)