Sort method and better docs

Author: DanielVandH

Project.toml

@@ -1,7 +1,7 @@
 name = "EpithelialDynamics1D"
 uuid = "ace8a2d7-7779-48a6-a8a4-cf6831a7e55b"
 authors = ["Daniel VandenHeuvel <danj.vandenheuvel@gmail.com>"]
-version = "1.4.0"
+version = "1.5.0"
 
 [deps]
 CommonSolve = "38540f10-b2f7-11e9-35d8-d573e4eb0ff2"

src/problem.jl

@@ -21,6 +21,14 @@ Struct for representing a cell simulation problem.
 You can solve a `CellProblem` like you would solve a problem from DifferentialEquations.jl, e.g. with 
 
     solve(prob, Tsit5(), saveat = 0.1)
+
+The other kwargs for `solve` that we provide are:
+
+- `proliferation::Bool = false`: Whether to include proliferation.
+- `rng::AbstractRNG = Random.default_rng()`: The random number generator to use for proliferation.
+- `sort::Bool = false`: Whether to sort the cells after each step. This is useful for preventing the cells from becoming unsorted, which can happen due to numerical errors or force laws like `F(q) = k/q^(n+1)`, `n ≥ 0`.
+
+For proliferation problems, `EnsembleProblem` can be useful - see the docs.
 """
 Base.@kwdef struct CellProblem{F,Fp,G,Gp}
     force_law::F

src/solve.jl

@@ -7,7 +7,7 @@ function default_prob_func(prob, proliferation, rng)
         end
     else
         (ode_problem, i, repeat) -> (ode_problem)
-    end 
+    end
     return prob_func
 end
 
@@ -16,30 +16,45 @@ function SciMLBase.ODEProblem(prob::CellProblem;
     jac_prototype=jacobian_sparsity(prob),
     proliferation=false,
     rng=Random.default_rng(),
+    sort=false,
     kwargs...) where {S}
     initial_time = prob.initial_time
     final_time = prob.final_time
     time_span = (initial_time, final_time)
     initial_condition = prob.initial_condition
+    sort_callback = DiscreteCallback(
+        (u, t, integrator) -> true,
+        integrator -> sort!(integrator.u),
+        save_positions=(false, false),
+        initialize=(c, u, t, integrator) -> c.affect!(integrator),
+        finalize=(c, u, t, integrator) -> c.affect!(integrator))
     if proliferation
         f = ODEFunction{true,S}(cell_odes!)
         callback = build_proliferation_callback(prob, rng)
+        if sort
+            callback = CallbackSet(sort_callback, callback)
+        end
         ode_problem = ODEProblem{true,S}(f, initial_condition, time_span, prob; callback=callback, kwargs...)
     else
         f = ODEFunction{true,S}(cell_odes!; jac_prototype=jac_prototype)
-        ode_problem = ODEProblem{true,S}(f, initial_condition, time_span, prob; kwargs...)
+        if sort
+            ode_problem = ODEProblem{true,S}(f, initial_condition, time_span, prob; callback=sort_callback, kwargs...)
+        else
+            ode_problem = ODEProblem{true,S}(f, initial_condition, time_span, prob; kwargs...)
+        end
     end
+    return ode_problem
 end
-function SciMLBase.NonlinearProblem(prob::SteadyCellProblem; proliferation=false, rng=Random.default_rng(), kwargs...)
+function SciMLBase.NonlinearProblem(prob::SteadyCellProblem; proliferation=false, rng=Random.default_rng(), sort=false, kwargs...)
     proliferation && @warn "Cannot compute the steady state of problems with proliferation. Ignoring the proliferation kwarg."
-    ode_prob = ODEProblem(prob.prob; rng, kwargs...)
+    ode_prob = ODEProblem(prob.prob; rng, sort, kwargs...)
     nl_prob = NonlinearProblem{true}(ode_prob.f, ode_prob.u0, ode_prob.p; kwargs...)
     return nl_prob
 end
-function SciMLBase.EnsembleProblem(prob::CellProblem; proliferation=true, rng=Random.default_rng(),
+function SciMLBase.EnsembleProblem(prob::CellProblem; proliferation=true, rng=Random.default_rng(), sort=false,
     prob_func=default_prob_func(prob, proliferation, rng), kwargs...)
-    return EnsembleProblem(ODEProblem(deepcopy(prob); proliferation, rng, kwargs...); prob_func, kwargs...)
+    return EnsembleProblem(ODEProblem(deepcopy(prob); proliferation, rng, sort, kwargs...); prob_func, kwargs...)
 end
 
-CommonSolve.init(prob::CellProblem, alg; proliferation=false, rng=Random.default_rng(), kwargs...) = CommonSolve.init(ODEProblem(prob; proliferation, rng, kwargs...), alg; kwargs...)
-CommonSolve.solve(prob::SteadyCellProblem, alg; proliferation=false, rng=Random.default_rng(), kwargs...) = CommonSolve.solve(NonlinearProblem(prob; proliferation, rng, kwargs...), alg; kwargs...)
+CommonSolve.init(prob::CellProblem, alg; proliferation=false, rng=Random.default_rng(), sort=false, kwargs...) = CommonSolve.init(ODEProblem(prob; proliferation, rng, sort, kwargs...), alg; kwargs...)
+CommonSolve.solve(prob::SteadyCellProblem, alg; proliferation=false, rng=Random.default_rng(), sort=false, kwargs...) = CommonSolve.solve(NonlinearProblem(prob; proliferation, rng, sort, kwargs...), alg; kwargs...)

test/runtests.jl

@@ -17,4 +17,4 @@ end
 end
 @safetestset "Step Function" begin
     include("step_function.jl")
-end
+end