Add discrete callbacks

.gitignore

@@ -16,3 +16,4 @@ dev/*
 # Build artifacts for creating documentation generated by the Documenter package
 docs/build/
 docs/site/
+.DS_Store

Project.toml

@@ -31,7 +31,7 @@ Interpolations = "0.14"
 Latexify = "0.15, 0.16"
 ModelingToolkit = "8"
 OrdinaryDiffEq = "6"
-PDEBase = "0.1.4"
+PDEBase = "0.1.6"
 PrecompileTools = "1"
 RuntimeGeneratedFunctions = "0.5.9"
 SafeTestsets = "0.0.1"

src/MethodOfLines.jl

@@ -50,6 +50,7 @@ import Base.ndims
 # Interface
 include("interface/grid_types.jl")
 include("interface/scheme_types.jl")
+include("interface/callbacks.jl")
 include("interface/disc_strategy_types.jl")
 include("interface/MOLFiniteDifference.jl")
 
@@ -72,6 +73,7 @@ include("discretization/schemes/upwind_difference/upwind_diff_weights.jl")
 include("discretization/schemes/half_offset_weights.jl")
 include("discretization/schemes/extrapolation_weights.jl")
 include("discretization/differential_discretizer.jl")
+include("discretization/schemes/callbacks/callback_rules.jl")
 
 # System Parsing
 include("system_parsing/pde_system_transformation.jl")
@@ -104,6 +106,6 @@ include("precompile.jl")
 
 # Export
 export MOLFiniteDifference, discretize, symbolic_discretize, ODEFunctionExpr, generate_code, grid_align, edge_align, center_align, get_discrete, chebyspace
-export UpwindScheme, WENOScheme, FunctionalScheme
+export UpwindScheme, WENOScheme, FunctionalScheme, MOLDiscCallback
 
 end

src/discretization/differential_discretizer.jl

@@ -8,13 +8,15 @@ struct DifferentialDiscretizer{T, D1, S} <: AbstractDifferentialDiscretizer
     interpmap::Dict{Num,DerivativeOperator}
     orders::Dict{Num,Vector{Int}}
     boundary::Dict{Num,DerivativeOperator}
+    callbacks
 end
 
 function PDEBase.construct_differential_discretizer(pdesys, s::DiscreteSpace, discretization::MOLFiniteDifference, orders)
     pdeeqs = pdesys.eqs isa Vector ? pdesys.eqs : [pdesys.eqs]
     bcs = pdesys.bcs isa Vector ? pdesys.bcs : [pdesys.bcs]
     approx_order = discretization.approx_order
     advection_scheme = discretization.advection_scheme
+    callbacks = discretization.callbacks
     upwind_order = advection_scheme isa UpwindScheme ? advection_scheme.order : 1
 
     differentialmap = Array{Pair{Num,DerivativeOperator},1}()
@@ -63,5 +65,5 @@ function PDEBase.construct_differential_discretizer(pdesys, s::DiscreteSpace, di
         push!(boundary, x => BoundaryInterpolatorExtrapolator(max(6, approx_order), dx))
     end
 
-    return DifferentialDiscretizer{eltype(orders),typeof(Dict(differentialmap)),typeof(advection_scheme)}(approx_order, advection_scheme, Dict(differentialmap), (Dict(nonlinlap_inner), Dict(nonlinlap_outer)), (Dict(windpos), Dict(windneg)), Dict(interp), Dict(orders), Dict(boundary))
+    return DifferentialDiscretizer{eltype(orders),typeof(Dict(differentialmap)),typeof(advection_scheme)}(approx_order, advection_scheme, Dict(differentialmap), (Dict(nonlinlap_inner), Dict(nonlinlap_outer)), (Dict(windpos), Dict(windneg)), Dict(interp), Dict(orders), Dict(boundary), callbacks)
 end

src/discretization/generate_finite_difference_rules.jl

@@ -80,6 +80,8 @@ function generate_finite_difference_rules(II::CartesianIndex, s::DiscreteSpace,
         integration_rules = []
     end
 
+    cb_rules = generate_cb_rules(II, s, depvars, derivweights, bmap, indexmap, terms)
+
     integration_rules = vcat(integration_rules, vec(generate_whole_domain_integration_rules(II, s, depvars, indexmap, terms)))
-    return vcat(vec(spherical_diffusion_rules), vec(nonlinlap_rules), vec(mixed_deriv_rules_cartesian), vec(central_deriv_rules_cartesian), vec(advection_rules), integration_rules)
+    return vcat(cb_rules, vec(spherical_diffusion_rules), vec(nonlinlap_rules), vec(mixed_deriv_rules_cartesian), vec(central_deriv_rules_cartesian), vec(advection_rules), integration_rules)
 end

src/discretization/schemes/callbacks/callback_rules.jl

@@ -0,0 +1,5 @@
+function generate_cb_rules(II, s, depvars, derivweights, bmap, indexmap, terms)
+    cbs = derivweights.callbacks
+    cb_rules = [cb.sym => cb(s) for cb in cbs]
+    return cb_rules
+end

src/interface/MOLFiniteDifference.jl

@@ -33,11 +33,12 @@ struct MOLFiniteDifference{G,D} <: AbstractEquationSystemDiscretization
     use_ODAE::Bool
     disc_strategy::D
     useIR::Bool
+    callbacks
     kwargs
 end
 
 # Constructors. If no order is specified, both upwind and centered differences will be 2nd order
-function MOLFiniteDifference(dxs, time=nothing; approx_order = 2, advection_scheme = UpwindScheme(), grid_align=CenterAlignedGrid(), discretization_strategy = ScalarizedDiscretization(), upwind_order = nothing, should_transform = true, use_ODAE = false, useIR = true, kwargs...)
+function MOLFiniteDifference(dxs, time=nothing; approx_order = 2, advection_scheme = UpwindScheme(), grid_align=CenterAlignedGrid(), discretization_strategy = ScalarizedDiscretization(), upwind_order = nothing, should_transform = true, use_ODAE = false, useIR = true, callbacks = [], kwargs...)
     if upwind_order !== nothing
         @warn "`upwind_order` no longer does anything, and will be removed in a future release. See the docs for the current interface."
     end
@@ -50,7 +51,7 @@ function MOLFiniteDifference(dxs, time=nothing; approx_order = 2, advection_sche
 
     dxs = dxs isa Dict ? dxs : Dict(dxs)
 
-    return MOLFiniteDifference{typeof(grid_align), typeof(discretization_strategy)}(dxs, time, approx_order, advection_scheme, grid_align, should_transform, use_ODAE, discretization_strategy, useIR, kwargs)
+    return MOLFiniteDifference{typeof(grid_align), typeof(discretization_strategy)}(dxs, time, approx_order, advection_scheme, grid_align, should_transform, use_ODAE, discretization_strategy, useIR, callbacks, kwargs)
 end
 
 PDEBase.get_time(disc::MOLFiniteDifference) = disc.time

src/interface/callbacks.jl

@@ -0,0 +1,13 @@
+struct MOLDiscCallback
+    f::Function
+    p
+    sym
+end
+
+function MOLDiscCallback(f, disc_ps)
+    symh = hash(f)
+    sym = unwrap(first(@parameters(Symbol("MOLDiscCallback_$symh"))))
+    return MOLDiscCallback(f, disc_ps, sym)
+end
+
+(M::MOLDiscCallback)(s) = M.f(s, M.p)

src/interface/solution/timeindep.jl

@@ -11,7 +11,7 @@ function SciMLBase.PDENoTimeSolution(sol::SciMLBase.NonlinearSolution{T}, metada
     umap = mapreduce(vcat, dvs) do u
         let discu = discretespace.discvars[u]
             solu = map(CartesianIndices(discu)) do I
-                i = sym_to_index(discu[I], odesys.states)
+                i = sym_to_index(discu[I], get_states(odesys))
                 # Handle Observed
                 if i !== nothing
                     sol.u[i]

src/system_parsing/pde_system_transformation.jl

@@ -35,6 +35,9 @@
 end
 
 function PDEBase.should_transform(pdesys::PDESystem, disc::MOLFiniteDifference, boundarymap)
+    if !disc.should_transform
+        return false
+    end
     if has_interfaces(boundarymap)
         @warn "The system contains interface boundaries, which are not compatible with system transformation. The system will not be transformed. Please post an issue if you need this feature."
         return false

test/components/callbacks.jl

@@ -0,0 +1,38 @@
+using ModelingToolkit, MethodOfLines, DomainSets, Test, Symbolics, SymbolicUtils, OrdinaryDiffEq
+
+@testset "Discrete callback" begin
+    @parameters x, t
+    @variables u(..)
+    Dt = Differential(t)
+    Dxx = Differential(x)^2
+    t_min = 0.0 
+    t_max = 2.0
+    x_min = 0.0
+    x_max = 2.0
+
+    a = 0.1*pi
+
+    eq = Dt(u(t, x)) ~ a*Dxx(u(t, x))
+
+    bcs = [u(t_min, x) ~ sinpi(x), u(t, x_min) ~ sinpi(t+x), u(t, x_max) ~ sinpi(t+x)]
+
+    domains = [t ∈ Interval(t_min, t_max), x ∈ Interval(x_min, x_max)]
+
+    @named pdesys = PDESystem(eq, bcs, domains, [t, x], [u(t, x)])
+
+    cb = MOLDiscCallback((s, p) -> prod(p), [0.1, pi])
+    a_cb = cb.sym
+    cbeq = Dt(u(t, x)) ~ a_cb*Dxx(u(t, x))
+
+    @named cbpdesys = PDESystem(cbeq, bcs, domains, [t, x], [u(t, x)])
+
+    disc = MOLFiniteDifference([x => 0.1], t; approx_order=2, callbacks = [cb])
+
+    prob = discretize(pdesys, disc)
+    cbprob = discretize(cbpdesys, disc)
+
+    sol1 = solve(prob, Tsit5(), saveat=0.1)
+    sol2 = solve(cbprob, Tsit5(), saveat=0.1) 
+
+    @test sol1[u(t, x)] ≈ sol2[u(t, x)]
+end

test/runtests.jl

@@ -71,6 +71,9 @@ const is_TRAVIS = haskey(ENV, "TRAVIS")
         @time @safetestset "ODEFunction" begin
             include("components/ODEFunction_test.jl")
         end
+        @time @safetestset "Discrete Callbacks" begin
+            include("components/callbacks.jl")
+        end
         #@time @safetestset "Finite Difference Schemes" begin include("components/finite_diff_schemes.jl") end
     end