Fixed interface, rules, bcs

Author: xtalax

src/discretization/MOL_discretization.jl renamed to src/MOL_discretization.jl

@@ -54,7 +54,7 @@ function SciMLBase.symbolic_discretize(pdesys::PDESystem, discretization::Method
             s = DiscreteSpace(domain, depvars, indvars, x̄, discretization)
 
             # Get the finite difference weights
-            derivweights = DifferentialDiscretizer(pde, s, discretization)
+            derivweights = DifferentialDiscretizer(pde, pdesys.bcs, s, discretization)
 
             # Get the boundary conditions
             BoundaryHandler!!(u0, bceqs, pdesys.bcs, s, depvar_ops, tspan, derivweights)

src/MOL_utils.jl

@@ -81,10 +81,10 @@ function split_additive_terms(eq)
     return vcat(lhs_arg,rhs_arg)
 end
 
-function clip(I::CartesianIndex, s::DiscreteSpace{N}, j::Int) where N
-    # Clip the index by 1 at each end of the dimension
+@inline function clip(I::CartesianIndex, s::DiscreteSpace{N}, j::Int, bpc::Int) where N
     I1 = unitindices(N)[j]
-    return I[j] > length(s, j) ? I-I1 : I+I1
+    return I-I1
+ 
 end
 
 subsmatch(expr, rule) = isequal(substitute(expr, rule), expr) ? false : true

src/MethodOfLines.jl

@@ -9,21 +9,21 @@ module MethodOfLines
     using SymbolicUtils: operation, arguments
     import DomainSets
 
-    include("MOL_utils.jl")
     include("discretization/fornberg.jl")
 
     include("grid_types.jl")
     include("MOLFiniteDifference.jl")
     include("bcs/boundary_types.jl")
-
+    
     include("discretization/discretize_vars.jl")
+    include("MOL_utils.jl")
 
     include("discretization/differential_discretizer.jl")
     include("discretization/generate_finite_difference_rules.jl")
 
     include("bcs/generate_bc_eqs.jl")
 
-    include("discretization/MOL_discretization.jl")
+    include("MOL_discretization.jl")
 
     export MOLFiniteDifference, discretize, symbolic_discretize, grid_align, edge_align, center_align
 end

src/bcs/generate_bc_eqs.jl

@@ -4,9 +4,9 @@ function _boundary_rules(s, orders, x, val)
 
     args = substitute.(args, (x=>val,))
 
-    rules = [operation(u)(args...) => (operation(u)(args...), x) for u in ū]
+    rules = [operation(u)(args...) => (operation(u)(args...), x) for u in s.ū]
 
-    return vcat(rules, vec([(Differential(x)^d)(operation(u)(args...)) => (operation(u)(args...), x) for d in orders[x], u in ū]))
+    return vcat(rules, vec([(Differential(x)^d)(operation(u)(args...)) => (operation(u)(args...), x) for d in orders[x], u in s.ū]))
 end
 
 function generate_boundary_matching_rules(s, orders)
@@ -31,9 +31,9 @@ function BoundaryHandler!!(u0, bceqs, bcs, s::DiscreteSpace, depvar_ops, tspan,
     t=s.time
 
     if t === nothing
-        initmaps = ū
+        initmaps = s.ū
     else
-        initmaps = substitute.(ū,[t=>tspan[1]])
+        initmaps = substitute.(s.ū,[t=>tspan[1]])
     end
 
     # Create some rules to match which bundary/variable a bc concerns
@@ -52,13 +52,13 @@ function BoundaryHandler!!(u0, bceqs, bcs, s::DiscreteSpace, depvar_ops, tspan,
         # * Assume in the form `u(...) ~ ...` for now
         bcdepvar = first(get_depvars(bc.lhs, depvar_ops))
 
-        if any(u -> isequal(operation(u), operation(bcdepvar)), ū)
+        if any(u -> isequal(operation(u), operation(bcdepvar)), s.ū)
             if t !== nothing && operation(bc.lhs) isa Sym && !any(x -> isequal(x, t.val), arguments(bc.lhs))
                 # initial condition
                 # * Assume that the initial condition is not in terms of the initial derivative
                 initindex = findfirst(isequal(bc.lhs), initmaps) 
                 if initindex !== nothing
-                    push!(u0,vec(s.discvars[ū[initindex]] .=> substitute.((bc.rhs,),gridvals(s))))
+                    push!(u0,vec(s.discvars[s.ū[initindex]] .=> substitute.((bc.rhs,),gridvals(s))))
                 end
             else
                 # Split out additive terms
@@ -94,7 +94,7 @@ function BoundaryHandler!!(u0, bceqs, bcs, s::DiscreteSpace, depvar_ops, tspan,
                 end
 
                 @assert boundary !== nothing "Boundary condition $bc is not on a boundary of the domain, or is not a valid boundary condition"
-                
+                @show bc
                 push!(bceqs, vec(map(s.Iedge[x_][idx(boundary)]) do II
                     rules = generate_bc_rules(II, derivweights, s, bc, u_, x_, boundary)
 
@@ -112,9 +112,9 @@ function generate_bc_rules(II, derivweights, s::DiscreteSpace{N,M,G}, bc, u_, x_
 
     depvarderivbcmaps = []
     depvarbcmaps = []
-
+    
     # * Assume that the BC is in terms of an explicit expression, not containing references to variables other than u_ at the boundary
-    for u in ū
+    for u in s.ū
         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_]]
@@ -123,7 +123,7 @@ function generate_bc_rules(II, derivweights, s::DiscreteSpace{N,M,G}, bc, u_, x_
             break
         end
     end
-    
+    @show depvarderivbcmaps
     fd_rules = generate_finite_difference_rules(II, s, bc, derivweights)
     varrules = axiesvals(s, x_, II)
 
@@ -143,7 +143,8 @@ function generate_bc_rules(II, derivweights, s::DiscreteSpace{N,M,G}, bc, u_, x_
     # depvarbcmaps will dictate what to replace the variable terms with in the bcs
     # replace u(t,0) with u₁, etc
     # * Assume that the BC is in terms of an explicit expression, not containing references to variables other than u_ at the boundary
-    for u in ū
+    
+    for u in s.ū
         if isequal(operation(u), operation(u_))
             depvarderivbcmaps = [(Differential(x_)^d)(u_) => half_offset_centered_difference(derivweights.halfoffsetmap[Differential(x_)^d], II, s, (s.x2i[x_],x_), u, ufunc) for d in derivweights.orders[x_]]
 

src/discretization/differential_discretizer.jl

@@ -7,24 +7,28 @@ struct DifferentialDiscretizer{T, D1}
     orders::Dict{Num, Vector{Int}}
 end
 
-function DifferentialDiscretizer(pde, s, discretization)
+function DifferentialDiscretizer(pde, bcs, 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)))))
+    d_orders(x) = reverse(sort(collect(union(differential_order(pde.rhs, x), differential_order(pde.lhs, x), (differential_order(bc.rhs, x) for bc in bcs)..., (differential_order(bc.lhs, x) for bc in bcs)...))))
 
-    # central_deriv_rules = [(Differential(s)^2)(u) => central_deriv(2,II,j,k) for (j,s) in enumerate(s.x̄), (k,u) in enumerate(ū)]
+    # central_deriv_rules = [(Differential(s)^2)(u) => central_deriv(2,II,j,k) for (j,s) in enumerate(s.x̄), (k,u) in enumerate(s.ū)]
     differentialmap = Array{Pair{Num,DiffEqOperators.DerivativeOperator},1}()
     nonlinlap = []
     interp = []
     orders = []
     # Hardcoded to centered difference, generate weights for each differential
     # TODO: Add handling for upwinding
+
     for x in s.x̄
-        push!(orders, x => d_orders(x))
+        orders_ = d_orders(x)
+        push!(orders, x => orders_)
+        _orders = Set(vcat(orders_, [1,2]))
         # 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]
+        rs = [(Differential(x)^d) => CompleteCenteredDifference(d, approx_order, s.dxs[x] ) for d in _orders]
+        
 
-        nonlinlap = vcat(nonlinlap, [Differential(x)^d => CompleteHalfCenteredDifference(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 _orders])
         differentialmap = vcat(differentialmap, rs)
         # A 0th order derivative off the grid is an interpolation
         push!(interp, x => CompleteHalfCenteredDifference(0, approx_order, s.dxs[x]))
@@ -92,7 +96,7 @@ end
 
 # i is the index of the offset, assuming that there is one precalculated set of weights for each offset required for a first order finite difference
 function half_offset_centered_difference(D, II, s, jx, u, ufunc)
-
-    (weights, Itap) = get_half_offset_weights_and_stencil(D, II, s, jx)
+    j, x = jx
+    weights, Itap = get_half_offset_weights_and_stencil(D, II, s, jx)
     return dot(weights, ufunc(u, Itap, x))
 end

src/discretization/discretize_vars.jl

@@ -17,7 +17,7 @@ end
     return substitute.((depvar,), edge_vals)
 end
 struct DiscreteSpace{N,M,G}
-    vars
+    ū
     discvars
     time
     x̄ # Note that these aren't necessarily @parameters
@@ -35,6 +35,7 @@ nparams(::DiscreteSpace{N,M}) where {N,M} = N
 nvars(::DiscreteSpace{N,M}) where {N,M} = M
 
 Base.length(s::DiscreteSpace, x) = length(s.grid[x])
+Base.length(s::DiscreteSpace, j::Int) = length(s.grid[s.x̄[j]])
 Base.size(s::DiscreteSpace) = Tuple(length(s.grid[z]) for z in s.x̄)
 
 params(s::DiscreteSpace{N,M}) where {N,M}= s.params
@@ -57,7 +58,7 @@ gridvals(s::DiscreteSpace{N}, I::CartesianIndex) where N = [Pair(x, s.grid[x][I[
 
 ## Boundary methods ##
 edgevals(s::DiscreteSpace{N}) where {N} = reduce(vcat, [get_edgevals(s.x̄, s.axies, i) for i = 1:N])
-edgevars(s::DiscreteSpace, I) = [u => s.discvars[u][I] for u in ū]
+edgevars(s::DiscreteSpace, I) = [u => s.discvars[u][I] for u in s.ū]
 
 """
 Generate a map of variables to the gridpoints at the edge of the domain
@@ -69,15 +70,15 @@ end
     return [x => last(s.axies[x]) for x in s.x̄]
 end
 
-varmaps(s::DiscreteSpace, II) = [u => s.discvars[u][II] for u in ū]
+varmaps(s::DiscreteSpace, II) = [u => s.discvars[u][II] for u in s.ū]
 
 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.x̄]...]
 
-map_symbolic_to_discrete(II::CartesianIndex, s::DiscreteSpace{N,M}) where {N,M} = vcat([ū[k] => s.discvars[k][II] for k = 1:M], [s.x̄[j] => s.grid[j][II[j]] for j = 1:N])
+map_symbolic_to_discrete(II::CartesianIndex, s::DiscreteSpace{N,M}) where {N,M} = vcat([s.ū[k] => s.discvars[k][II] for k = 1:M], [s.x̄[j] => s.grid[j][II[j]] for j = 1:N])
 
 # ? How rude is this? Makes Iedge work
 
@@ -141,5 +142,5 @@ end
 
 
 
-#varmap(s::DiscreteSpace{N,M}) where {N,M} = [ū[i] => i for i = 1:M]
+#varmap(s::DiscreteSpace{N,M}) where {N,M} = [s.ū[i] => i for i = 1:M]
 

src/discretization/generate_finite_difference_rules.jl

@@ -3,15 +3,8 @@
 
 Interpolate gridpoints by taking the average of the values of the discrete points, or if the offset is outside the grid, extrapolate the value with dx.
 """
-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
-    if (II[j]+itap) < 1
-        return s.grid[x][1]+s.dxs[x]*offset
-    elseif (II[j]+itap) > (length(x) -  1)
-        return s.grid[x][length(x)]+s.dxs[x]*offset
-    else
-        return (s.grid[x][II[j]+itap]+s.grid[x][II[j]+itap+1])/2
-    end
+@inline function interpolate_discrete_param(i, s, itap, x, bpc)
+    return s.grid[x][i+itap]+s.dxs[x]*.5
 end
 
 """
@@ -44,35 +37,38 @@ And so on.
 function cartesian_nonlinear_laplacian(expr, II, derivweights, s, x, u)
     # Based on the paper https://web.mit.edu/braatzgroup/analysis_of_finite_difference_discretization_schemes_for_diffusion_in_spheres_with_variable_diffusivity.pdf 
     # 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)
+    @assert II[j] != 1 "The nonlinear laplacian is only defined on the interior of the grid, it is unsupported in boundary conditions."
+    @assert II[j] != length(s, x) "The nonlinear laplacian is only defined on the interior of the grid, it is unsupported in boundary conditions."
 
-    jx = j, x = (s.x2i(x), x)
     D_inner = derivweights.halfoffsetmap[Differential(x)]
     inner_interpolater = derivweights.interpmap[x]
 
-    # Get the outer weights and stencil. 
-    outerweights, outerstencil = _get_weights_and_stencil(D_inner, clip(II, s, j), s, jx)
+    # Get the outer weights and stencil. clip() essentially removes a point from either end of the grid, for this reason this function is only defined on the interior, not in bcs
+    outerweights, outerstencil = get_half_offset_weights_and_stencil(D_inner, clip(II, s, j, D_inner.boundary_point_count), s, jx)
 
     # Get the correct weights and stencils for this II
     inner_deriv_weights_and_stencil = [get_half_offset_weights_and_stencil(D_inner, I, s, jx) for I in outerstencil]
     interp_weights_and_stencil = [get_half_offset_weights_and_stencil(inner_interpolater, I, s, jx) for I in outerstencil]
 
     # map variables to symbolically inerpolated/extrapolated expressions
-    map_vars_to_interpolated(stencil, weights) = [v => dot(weights, s.discvars[v][stencil]) for v in ū]
+    map_vars_to_interpolated(stencil, weights) = [v => dot(weights, s.discvars[v][stencil]) for v in s.ū]
 
     # Map parameters to interpolated values. Using simplistic extrapolation/interpolation for now as grids are uniform
-    map_params_to_interpolated(itap) = [z => interpolate_discrete_param(II, s, itap, i, z) for (i,z) in enumerate(s.x̄) ]
+    map_params_to_interpolated(itap) = x => interpolate_discrete_param(II[j], s, itap[j]-II[j], x, D_inner.boundary_point_count)
 
     # Take the inner finite difference
     inner_difference = [dot(inner_weights, s.discvars[u][inner_stencil]) for (inner_weights, inner_stencil) in inner_deriv_weights_and_stencil]
 
     # Symbolically interpolate the multiplying expression
-    interpolated_expr = [Num(substitute(substitute(expr, map_vars_to_interpolated(interpstencil, interpweights)), map_params_to_interpolated(itap))) 
-                        for (itap,(interpweights, interpstencil)) in zip(itaps, interp_weights_and_stencil)]
+    interpolated_expr = [Num(substitute(substitute(expr, map_vars_to_interpolated(interpstencil, interpweights)), map_params_to_interpolated.(interpstencil))) 
+                        for (interpweights, interpstencil) in interp_weights_and_stencil]
 
     # multiply the inner finite difference by the interpolated expression, and finally take the outer finite difference
     return dot(outerweights, inner_difference .* interpolated_expr)
 end
-ū
+
 """
 `cartesian_nonlinear_laplacian`
 
@@ -106,16 +102,16 @@ function spherical_diffusion(innerexpr, II, derivweights, s, r, u)
     D_2 = derivweights.map[Differential(r)^2]
 
     # What to replace parameter x with given I
-    _rsubs(x, I) = x => s.grid[x][I[s.x2i(x)]]
+    _rsubs(x, I) = x => s.grid[x][I[s.x2i[x]]]
     # Full rules for substituting parameters in the inner expression
-    rsubs(I) = vcat([v => s.discvars[v][I] for v in ū], [_rsubs(x, I) for x in params(s)])
+    rsubs(I) = vcat([v => s.discvars[v][I] for v in s.ū], [_rsubs(x, I) for x in s.x̄])
     # Discretization func for u
     ufunc_u(v, I, x) = s.discvars[v][I]
 
     # 2nd order finite difference in u
     exprhere = Num(substitute(innerexpr, rsubs(II)))
     # Catch the r ≈ 0 case
-    if substitute(r, _rsubs(r, II)) ≈ 0
+    if Symbolics.unwrap(substitute(r, _rsubs(r, II))) ≈ 0
         D_2_u = central_difference(D_2, II, s, (s.x2i(r), r), u, ufunc_u)
         return 3exprhere*D_2_u # See appendix B of the paper
     end
@@ -127,24 +123,24 @@ end
 
 @inline function generate_cartesian_rules(II, s, derivweights, terms)
     central_ufunc(u, I, x) = s.discvars[u][I]
-    return reduce(vcat, [[(Differential(x)^d)(u) => central_difference(derivweights.map[Differential(x)^d], II, s, (j,x), u, central_ufunc) for d in derivweights.orders[x], u in ū] for (j,x) in enumerate(s.x̄)])
+    return reduce(vcat, [[(Differential(x)^d)(u) => central_difference(derivweights.map[Differential(x)^d], II, s, (j,x), u, central_ufunc) for d in derivweights.orders[x], u in s.ū] for (j,x) in enumerate(s.x̄)])
 end
 
 @inline function generate_upwinding_rules(II, s, derivweights, terms)    
     #forward_weights(II,j) = calculate_weights(discretization.upwind_order, 0.0, s.grid[j][[II[j],II[j]+1]])
     #reverse_weights(II,j) = calculate_weights(discretization.upwind_order, 0.0, s.grid[j][[II[j]-1,II[j]]])
     # upwinding_rules = [@rule(*(~~a, $(Differential(s.x̄[j]))(u),~~b) => IfElse.ifelse(*(~~a..., ~~b...,)>0,
-    #                         *(~~a..., ~~b..., dot(reverse_weights(II,j),ū[k][central_neighbor_idxs(II,j)[1:2]])),
-    #                         *(~~a..., ~~b..., dot(forward_weights(II,j),ū[k][central_neighbor_idxs(II,j)[2:3]]))))
-    #                         for j in 1:nparams(s), k in 1:length(pdesys.ū)]
+    #                         *(~~a..., ~~b..., dot(reverse_weights(II,j),s.ū[k][central_neighbor_idxs(II,j)[1:2]])),
+    #                         *(~~a..., ~~b..., dot(forward_weights(II,j),s.ū[k][central_neighbor_idxs(II,j)[2:3]]))))
+    #                         for j in 1:nparams(s), k in 1:length(pdesys.s.ū)]
 
 end
 
 @inline function generate_nonlinlap_rules(II, s, derivweights, terms)
-    cartesian_deriv_rules = [@rule ($(Differential(x))(*(~~a, $(Differential(x))(u), ~~b))) => cartesian_nonlinear_laplacian(*(a..., b...), II, derivweights, s, x, u) for x in s.x̄, u in ū]
+    cartesian_deriv_rules = [@rule *(~~c, $(Differential(x))(*(~~a, $(Differential(x))(u), ~~b)), ~~d) => cartesian_nonlinear_laplacian(*(a..., b...), II, derivweights, s, x, u) for x in s.x̄, u in s.ū]
 
     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.x̄, u in ū]))
+                            [@rule ($(Differential(x))($(Differential(x))(u)/~a)) => cartesian_nonlinear_laplacian(1/~a, II, derivweights, s, x, u) for x in s.x̄, u in s.ū]))
 
     nonlinlap_rules = []
     for t in terms
@@ -158,12 +154,14 @@ end
 end
 
 @inline function generate_spherical_diffusion_rules(II, s, derivweights, terms)
-    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.x̄, u in ū])
+    rules = vec([@rule *(~~a, 1/(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.x̄, u in s.ū])
+    rules = vcat(rules, vec([@rule /(($(Differential(r))(*(~~c, (r^2), ~~d, $(Differential(r))(u), ~~e))), (r^2)) => spherical_diffusion(*(~c..., ~d..., ~e...), II, derivweights, s, r, u)
+    for r in s.x̄, u in s.ū]))
 
     spherical_diffusion_rules = []
     for t in terms
-        for r in spherical_deriv_rules
+        for r in rules
             if r(t) !== nothing
                 push!(spherical_diffusion_rules, t => r(t))
             end