Finished first pass at generate_bc_rules

Author: xtalax

src/MOL_utils.jl

@@ -74,3 +74,5 @@ function unitindices(N::Int) #create unit CartesianIndex for each dimension
 end
 
 half_range(x) = -div(x,2):div(x,2)
+
+

src/MethodOfLines.jl

@@ -11,10 +11,10 @@ module MethodOfLines
 
     include("MOL_utils.jl")
     include("discretization/fornberg.jl")
+    include("bcs/generate_bc_eqs.jl")
     include("discretization/discretize_vars.jl")
     include("discretization/differential_discretizer.jl")
     include("discretization/generate_rules.jl")
-    include("bcs/generate_bc_eqs.jl")
 
     include("discretization/MOL_discretization.jl")
 

src/bcs/generate_bc_eqs.jl

@@ -2,13 +2,17 @@
 
 abstract type AbstractBoundary end
 
-struct LowerBoundary <: AbstractBoundary
+abstract type AbstractTruncatingBoundary <: AbstractBoundary end
+
+abstract type AbstractExtendingBoundary <: AbstractBoundary end
+
+struct LowerBoundary <: AbstractTruncatingBoundary
 end
 
-struct UpperBoundary<: AbstractBoundary
+struct UpperBoundary<: AbstractTruncatingBoundary
 end
 
-struct CompleteBoundary <: AbstractBoundary
+struct CompleteBoundary <: AbstractTruncatingBoundary
 end#
 
 struct PeriodicBoundary <: AbstractBoundary
@@ -20,7 +24,7 @@ end
 
 """
 Mutates bceqs and u0 by finding relevant equations and discretizing them.
-TODO: return a handler for use with generate_finite_difference_rules
+TODO: return a handler for use with generate_finite_difference_rules and pull out initial condition
 """
 function BoundaryHandler!!(u0, bceqs, bcs, s, depvar_ops, tspan, derivweights)
 
@@ -63,7 +67,7 @@ function BoundaryHandler!!(u0, bceqs, bcs, s, depvar_ops, tspan, derivweights)
                 for term in terms, r in lower_boundary_rules
                     if r(term) !== nothing
                         u_, x_ = (term, r(term))
-                        boundary = :lower
+                        boundary = LowerBoundary()
                         for term_ in setdiff(terms, term)
                             for r in upper_boundary_rules
                                 if r(term_) !== nothing
@@ -77,7 +81,7 @@ function BoundaryHandler!!(u0, bceqs, bcs, s, depvar_ops, tspan, derivweights)
                 for term in terms, r in upper_boundary_rules
                     if r(term) !== nothing
                         u_, x_ = (term, r(term))
-                        boundary = :upper
+                        boundary = UpperBoundary()
                         break
                     end
                 end
@@ -86,7 +90,7 @@ function BoundaryHandler!!(u0, bceqs, bcs, s, depvar_ops, tspan, derivweights)
 
                 push!(bceqs, vec(map(s.Iedge[x_][boundary]) do II
                     rules = generate_bc_rules(II, s, bc, u_, boundary, derivweights)
-                    rules = vcat(rules, generate_finite_difference_rules(II, s, bc, derivweights))
+                    rules = vcat(rules, )
 
                     substitute(bc.lhs, rules) ~ substitute(bc.rhs, rules)
                 end))
@@ -96,17 +100,45 @@ function BoundaryHandler!!(u0, bceqs, bcs, s, depvar_ops, tspan, derivweights)
     end
 end
 
-function get_active_variable(bc, s, depvar_ops)
-    bcdepvar = first(get_depvars(bc.lhs, depvar_ops))
-    out = Array{typeof(operation(s.vars[1]))}()
+function generate_bc_rules(II, dim, s, bc, u_, x_, boundary::AbstractTruncatingBoundary, derivweights, G::CenterAlignedGrid) 
+    # depvarbcmaps will dictate what to replace the variable terms with in the bcs
+    # replace u(t,0) with u₁, etc
+    ufunc(v, I, x) = s.discvars[v][I]
+
     for u in s.vars
-        if u isa Sym && isequal(operation(u), operation(bcdepvar))
-            push!(out, u)   
+        if isequal(operation(u), operation(u_))
+            # What to replace derivatives at the boundary with
+            depvarderivbcmaps = [(Differential(x)^d)(u_) => central_difference(derivweights.map[Differential(x_)^d], II, s, (s.x2i[x_],x_), u, ufunc) for d in derivweights.orders[x_]]
+            # ? Does this need to be done for all variables at the boundary?
+            depvarbcmaps = [u_ => s.discvars[u][II]]
         end
     end
-    return out
-end
 
+    fd_rules = generate_finite_difference_rules(II, s, bc, derivweights)
+    varrules = axiesvals(s, II)
+    
+    return vcat(depvarderivbcmaps, depvarbcmaps, fd_rules, varrules)
+end
 
+function generate_bc_rules(II, dim, s, bc, u_, x_, boundary::AbstractTruncatingBoundary, derivweights, G::EdgeAlignedGrid) 
+    
+    boundaryoffset(::LowerBoundary) = 1/2
+    boundaryoffset(::UpperBoundary) = -1/2
+    ufunc(v, I, x) = s.discvars[v][I]
 
-#function generate_u0_and_bceqs_with_rules!!(u0, bceqs, bcs, t, s, depvar_ops)
+    # depvarbcmaps will dictate what to replace the variable terms with in the bcs
+    # replace u(t,0) with u₁, etc
+    for u in s.vars
+        if isequal(operation(u), operation(u_))
+            depvarderivbcmaps = [(Differential(x)^d)(u_) => half_offset_centered_difference(derivweights.halfoffsetmap[Differential(x_)^d], II, s, offset(boundary), (s.x2i[x_],x_), u, ufunc) for d in derivweights.orders[x_]]
+    
+            depvarbcmaps = [u_ => half_offset_centered_difference(derivweights.interpmap[x_], II, s, offset(boundary), (s.x2i[x_],x_), u, ufunc)]
+        end
+    end
+    
+    fd_rules = generate_finite_difference_rules(II, s, bc, derivweights)
+    varrules = axiesvals(s, II)
+    valr = valrules(s, II)
+    
+    return vcat(depvarderivbcmaps, depvarbcmaps, fd_rules, varrules)
+end

src/discretization/MOL_discretization.jl

@@ -1,21 +1,29 @@
 # Method of lines discretization scheme
+abstract type AbstractGrid end
 
-@enum GridAlign center_align edge_align
+struct CenterAlignedGrid <: AbstractGrid 
+end
+
+struct EdgeAlignedGrid <: AbstractGrid 
+end
+
+const center_align=CenterAlignedGrid()
+const edge_align=EdgeAlignedGrid()
 
-struct MOLFiniteDifference{T,T2} <: DiffEqBase.AbstractDiscretization
-    dxs::T
-    time::T2
+struct MOLFiniteDifference{G} <: DiffEqBase.AbstractDiscretization
+    dxs
+    time
     approx_order::Int
-    grid_align::GridAlign
+    grid_align::G
 end
 
 # Constructors. If no order is specified, both upwind and centered differences will be 2nd order
-function MOLFiniteDifference(dxs, time=nothing; upwind_order = 1, centered_order = 2, grid_align=center_align)
+function MOLFiniteDifference(dxs, time=nothing; upwind_order = 1, centered_order = 2, grid_align=CenterAlignedGrid())
 
     if centered_order % 2 != 0
         warn("Discretization centered_order must be even, rounding up to $(centered_order+1)")
     end
-    return MOLFiniteDifference(dxs, time, centered_order, grid_align)
+    return MOLFiniteDifference{typeof(grid_align)}(dxs, time, centered_order, grid_align)
 end
 
 function SciMLBase.symbolic_discretize(pdesys::PDESystem, discretization::MethodOfLines.MOLFiniteDifference)
@@ -81,7 +89,7 @@ function SciMLBase.symbolic_discretize(pdesys::PDESystem, discretization::Method
 
             # Discretize the equation on the interior
             pdeeqs = vec(map(interior) do II
-                rules = generate_finite_difference_rules(II, s, pde, derivweights)
+                rules = vcat(generate_finite_difference_rules(II, s, pde, derivweights), valrules(s, II))
                 substitute(pde.lhs,rules) ~ substitute(pde.rhs,rules)
             end)
 

src/discretization/differential_discretizer.jl

@@ -9,6 +9,7 @@ end
 
 function DifferentialDiscretizer(pde, s, discretization)
     approx_order = discretization.approx_order
+    # TODO: Include bcs in this calculation
     d_orders(x) = reverse(sort(collect(union(differential_order(pde.rhs, x), differential_order(pde.lhs, x)))))
 
     # central_deriv_rules = [(Differential(s)^2)(u) => central_deriv(2,II,j,k) for (j,s) in enumerate(s.nottime), (k,u) in enumerate(s.vars)]
@@ -23,16 +24,18 @@ function DifferentialDiscretizer(pde, s, discretization)
         # TODO: Only generate weights for derivatives that are actually used and avoid redundant calculations
         rs = [(Differential(x)^d) => CompleteCenteredDifference(d, approx_order, s.dxs[x] ) for d in last(orders).second]
 
+        nonlinlap = vcat(nonlinlap, [Differential(x)^d => CompleteHalfCenteredDifference(d, approx_order, s.dxs[x]) for d in last(orders).second])
         differentialmap = vcat(differentialmap, rs)
-        push!(nonlinlap, x => CompleteHalfCenteredDifference(0, approx_order, s.dxs[x])
-        push!(interp, x => CompleteHalfCenteredDifference(1, approx_order, s.dxs[x]))
+        push!(interp, x => CompleteHalfCenteredDifference(0, approx_order, s.dxs[x])
     end
 
     return DifferentialDiscretizer{eltype(orders), typeof(Dict(differentialmap))}(approx_order, Dict(differentialmap), Dict(nonlinlap), Dict(interp), Dict(orders))
 end
 
-
-# ufunc is a function that returns the correct discretization indexed at Itap, it is designed this way to allow for central differences of arbitrary expressions which may be needed in some schemes 
+"""
+Performs a centered difference in `x` centered at index `II` of `u`
+ufunc is a function that returns the correct discretization indexed at Itap, it is designed this way to allow for central differences of arbitrary expressions which may be needed in some schemes 
+"""
 function central_difference(D, II, s, jx, u, ufunc)
     j, x = jx
     # unit index in direction of the derivative
@@ -88,10 +91,9 @@ TODO: consider refactoring this to harmonize with centered difference
 function get_half_offset_weights_and_stencil(D, II, s, offset, jx)
     j, x = jx
     I1 = unitindices(nparams(s))[j]
-
     # Shift the current index to the correct offset
-    II_prime = II + offset*I1
-    
+    II_prime = II + Int(offset-0.5)*I1
+    @assert all(i-> 0<i<=length(s,x), II_prime) "Index out of bounds"
     return _get_weights_and_stencil(D, II_prime, I1, s, offset, j, x)
 end
 

src/discretization/discretize_vars.jl

@@ -3,7 +3,7 @@
 Return a map of of variables to the gridpoints at the edge of the domain
 """
 @inline function get_edgevals(params, axies, i)
-    return [params[i] => first(axies[i]), params[i] => last(axies[i])]
+    return [params[i] => first(axies[params[i]]), params[i] => last(axies[params[i]])]
 end
 
 """
@@ -42,17 +42,28 @@ params(s::DiscreteSpace{N,M}) where {N,M}= s.params
 grid_idxs(s::DiscreteSpace) = CartesianIndices(((axes(g)[1] for g in s.grid)...,))
 edge_idxs(s::DiscreteSpace{N}) where {N} = reduce(vcat, [[vcat([Colon() for j = 1:i-1], 1, [Colon() for j = i+1:N]), vcat([Colon() for j = 1:i-1], length(s.axies[i]), [Colon() for j = i+1:N])] for i = 1:N])
 
-axiesvals(s::DiscreteSpace{N}) where {N} = map(y -> [s.nottime[i] => s.axies[s.nottime[i]][y.I[i]] for i = 1:N], s.Iaxies)
-gridvals(s::DiscreteSpace{N}) where {N} = map(y -> [s.nottime[i] => s.grid[s.nottime[i]][y.I[i]] for i = 1:N], s.Igrid)
+axiesvals(s::DiscreteSpace{N}, I) where {N} = [x => s.axies[x][I[j]] for (j,x) in enumerate(s.nottime)]
+gridvals(s::DiscreteSpace{N}, I) where {N} = [x => s.grid[x][I[j]] for (j,x) in enumerate(s.nottime)]
 
 ## Boundary methods ##
 edgevals(s::DiscreteSpace{N}) where {N} = reduce(vcat, [get_edgevals(s.nottime, s.axies, i) for i = 1:N])
-edgevars(s::DiscreteSpace) = [[d[e...] for e in s.Iedge] for d in s.discvars]
+edgevars(s::DiscreteSpace, I) = [u => s.discvars[u][I] for u in s.vars]
 
-@inline function edgemaps(s::DiscreteSpace)
-    bclocs(s::DiscreteSpace) = map(e -> substitute.(s.params, e), edgevals(s))
-    return Dict(bclocs(s) .=> [axiesvals(s)[e...] for e in s.Iedge])
+"""
+Generate a map of variables to the gridpoints at the edge of the domain
+"""
+@inline function edgemaps(s::DiscreteSpace, ::LowerBoundary)
+    return [x => first(s.axies[x]) for x in s.nottime]
 end
+@inline function edgemaps(s::DiscreteSpace, ::LowerBoundary)
+    return [x => last(s.axies[x]) for x in s.nottime]
+end
+
+varmaps(s::DiscreteSpace, II) = [u => s.discvars[u][II] for u in s.vars]
+
+valmaps(s::DiscreteSpace, II) = vcat(varmaps(s,II), gridvals(s,II))
+
+
 
 Iinterior(s::DiscreteSpace) = s.Igrid[[2:(length(s, x)-1) for x in s.nottime]...]
 
@@ -109,7 +120,7 @@ function DiscreteSpace(domain, depvars, indvars, nottime, discretization)
     end
 
     # Build symbolic maps for boundaries
-    Iedge = vcat((vcat(vec(selectdim(Igrid, dim, 1)), vec(selectdim(Igrid, dim, length(grid[dim].second)))) for dim in 1:nspace)...)
+    Iedge = Dict([x => Dict([LowerBoundary() => vec(selectdim(Igrid, dim, 1)), UpperBoundary() => vec(selectdim(Igrid, dim, length(grid[dim].second)))]) for (dim, x) in enumerate(s.nottime)])
 
     nottime2dim = [nottime[i] => i for i in 1:nspace]
     dim2nottime = [i => nottime[i] for i in 1:nspace]

src/discretization/generate_rules.jl renamed to src/discretization/generate_finite_difference_rules.jl

@@ -5,7 +5,6 @@ Interpolate gridpoints by taking the average of the values of the discrete point
 """
 function interpolate_discrete_param(II, s, itap, j, x)
     # * This will need to be updated to dispatch on grid type when grids become more general
-    offset = itap+1/2
     if (II[j]+itap) < 1
         return s.grid[x][1]+s.dxs[x]*offset
     elseif (II[j]+itap) > (length(x) -  1)
@@ -47,13 +46,14 @@ function cartesian_nonlinear_laplacian(expr, II, derivweights, s, x, u)
     # See scheme 1, namely the term without the 1/r dependence. See also #354 and #371 in DiffEqOperators, the previous home of this package.
 
     jx = j, x = (s.x2i(x), x)
-    D_inner = derivweights.halfoffsetmap[x]
+    D_inner = derivweights.halfoffsetmap[Differential(x)]
     inner_interpolater = derivweights.interpmap[x]
 
-    # Get the outer weights and stencil to generate the required 
-    outerweights, outerstencil = get_half_offset_weights_and_stencil(D_inner, II, s, 0, jx)
+    # Get the outer weights and stencil. 
+    # ! The offset should eliminate the need for bounds checking, the inner function ensuring that the taps lie inbounds.
+    outerweights, outerstencil = _get_weights_and_stencil(D_inner, II, s, 1/2, jx)
     # Index offsets of each stencil in the inner finite difference to get the correct stencil for each needed half grid point, 0 corresopnds to x+1/2
-    itaps = getindex.(outerstencil, (j,))
+    itaps = getindex.(outerstencil .- II, (j,)) .+ 0.5
 
     # Get the correct weights and stencils for this II
     inner_deriv_weights_and_stencil = [get_half_offset_weights_and_stencil(D_inner, II, s, itap, jx) for itap in itaps]
@@ -128,6 +128,7 @@ function spherical_diffusion(innerexpr, II, derivweights, s, r, u)
     return exprhere*(D_1_u/substitute(r, _rsubs(r, II)) + cartesian_nonlinear_laplacian(innerexpr, II, derivweights, s, r, u))
 end
 
+
 """
 `generate_finite_difference_rules`
 
@@ -150,8 +151,6 @@ There are of course more specific schemes that are used to improve stability/spe
 Please submit an issue if you know of any special cases that are not implemented, with links to papers and/or code that demonstrates the special case.
 """
 function generate_finite_difference_rules(II, s, pde, derivweights)
-    valrules = vcat([u => s.discvars[u][II] for u in s.vars],
-                    [x => s.grid[x][II[j]] for (j,x) in enumerate(s.nottime)])
     # central_deriv_rules = [(Differential(s)^2)(u) => central_deriv(2,II,j,k) for (j,s) in enumerate(s.nottime), (k,u) in enumerate(s.vars)]
 
     central_ufunc(u, I, x) = s.discvars[u][I]
@@ -177,9 +176,9 @@ function generate_finite_difference_rules(II, s, pde, derivweights)
     cartesian_deriv_rules = vcat(vec(cartesian_deriv_rules),vec(
                             [@rule ($(Differential(x))($(Differential(x))(u)/~a)) => cartesian_nonlinear_laplacian(1/~a, II, derivweights, s, x, u) for x in s.nottime, u in s.vars]))
 
-    spherical_deriv_rules = [@rule *(~~a, (r^-2), ($(Differential(r))(*(~~c, (r^2), ~~d, $(Differential(r))(u), ~~e))), ~~b) =>
+    spherical_deriv_rules = vec([@rule *(~~a, (r^-2), ($(Differential(r))(*(~~c, (r^2), ~~d, $(Differential(r))(u), ~~e))), ~~b) =>
             *(~a..., spherical_diffusion(*(~c..., ~d..., ~e...), II, derivweights, s, r, u), ~b...)
-            for r in s.nottime, u in s.vars]
+            for r in s.nottime, u in s.vars])
 
     rhs_arg = istree(pde.rhs) && (SymbolicUtils.operation(pde.rhs) == +) ? SymbolicUtils.arguments(pde.rhs) : [pde.rhs]
     lhs_arg = istree(pde.lhs) && (SymbolicUtils.operation(pde.lhs) == +) ? SymbolicUtils.arguments(pde.lhs) : [pde.lhs]
@@ -200,24 +199,7 @@ function generate_finite_difference_rules(II, s, pde, derivweights)
         end
     end
     rules = vcat(vec(nonlinlap_rules),
-                vec(central_deriv_rules_cartesian),
-                valrules)
+                vec(central_deriv_rules_cartesian))
     return rules
 end
 
-function generate_bc_rules(II, dim, s, bc, edgemaps)
-    # ! Recognise which dim a BC is on, and use that to get the axiesvals at the boundary
-    # ! loop through the Iedge at that boundary
-    # ! replace the symbolic variables at either end of the boundary with the appropriate discvars, interpolated if nessecary
-    # ! eventually move to multi dim interpolations to improve validity
-
-    # depvarbcmaps will dictate what to replace the variable terms with in the bcs
-    # replace u(t,0) with u₁, etc
-    if grid_align == center_align
-        depvarbcmaps = reduce(vcat,[substitute(depvar, edgevals(s)) .=> edgevar for (depvar, edgevar) in zip(s.vars, edgevars(s, II))])
-    elseif grid_align == edge_align
-        
-    end
-
-    varrules = edgemaps(s)
-    rules = vcat(depvarbcmaps, edgemaps)