Wed, 22 Dec 2021 11:13:38 +0200
Planar finite elements, simple linear solvers for fixed dimensions
src/AlgTools.jl | file | annotate | diff | comparison | revisions | |
src/DifferentiableFN.jl | file | annotate | diff | comparison | revisions | |
src/FunctionalProgramming.jl | file | annotate | diff | comparison | revisions | |
src/LinSolve.jl | file | annotate | diff | comparison | revisions | |
src/Loops.jl | file | annotate | diff | comparison | revisions | |
src/Metaprogramming.jl | file | annotate | diff | comparison | revisions | |
src/PlanarFE.jl | file | annotate | diff | comparison | revisions | |
src/ZipArrays.jl | file | annotate | diff | comparison | revisions |
--- a/src/AlgTools.jl Wed Dec 15 01:09:09 2021 +0200 +++ b/src/AlgTools.jl Wed Dec 22 11:13:38 2021 +0200 @@ -9,7 +9,7 @@ For further documentation, see the submodules `FunctionalProgramming`, `StructTools`, `LinkedLists`, `Logger`, `Iterate`, `VectorMath`, `Util`, `ThreadUtil`, `Comms`, `LinOps`, `DifferentiableFN`, -`Metaprogramming`, `Loops`, and `ZipArrays`. +`Metaprogramming`, `Loops`, `ZipArrays`, `LinSolve`, and `PlanarFE`. """ module AlgTools @@ -27,5 +27,7 @@ include("Metaprogramming.jl") include("Loops.jl") include("ZipArrays.jl") +include("LinSolve.jl") +include("PlanarFE.jl") end
--- a/src/DifferentiableFN.jl Wed Dec 15 01:09:09 2021 +0200 +++ b/src/DifferentiableFN.jl Wed Dec 22 11:13:38 2021 +0200 @@ -69,8 +69,8 @@ # Squared norm of a differentiable function as a differentiable function -ℝ = Float64 -ℝⁿ = Vector{Float64} +const ℝ = Float64 +const ℝⁿ = Vector{Float64} struct Squared <: DiffF{ℝⁿ, ℝ, MatrixOp{ℝ}} r :: DiffF{ℝⁿ, ℝⁿ, MatrixOp{ℝ}}
--- a/src/FunctionalProgramming.jl Wed Dec 15 01:09:09 2021 +0200 +++ b/src/FunctionalProgramming.jl Wed Dec 22 11:13:38 2021 +0200 @@ -17,7 +17,8 @@ ############## export curry, - curryflip + curryflip, + maybe ########### # Currying @@ -57,4 +58,12 @@ return x ->f(x, y...; kwargs...) end +""" +`maybe(f, x)` + +Returns `nothing` if `x` is `nothing, otherwise `f(x)`. +""" +maybe( :: Function, :: Nothing) = nothing +maybe(f :: Function, x :: T) where T = f(x) + end # module
--- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/src/LinSolve.jl Wed Dec 22 11:13:38 2021 +0200 @@ -0,0 +1,190 @@ +""" +Linear solvers for small problems. +""" +module LinSolve + +using ..Metaprogramming + +export linsolve, + TupleMatrix + +const TupleMatrix{M,N} = NTuple{M, NTuple{N, Float64}} + +""" +`linsolve(AB :: TupleMatrix{M,N}, :: Type{TupleMatrix{M, K}}) :: TupleMatrix{M, K}` + +where + +`TupleMatrix{M, N} = NTuple{M, NTuple{N, Float64}}` + +“Immutable” Gaussian elimination on tuples: solve AX=B for X, +Both A and B are stored in AB. The second type parameter indicates the size of B. +""" +@polly function linsolve₀(AB :: TupleMatrix{M,N}, :: Val{K}) :: TupleMatrix{M, K} where {N,M,K} + @assert(M == N - K) + + k = 0 + + # Convert to row-echelon form + for h = 1:(M-1) + # Find pivotable column (has some non-zero entries in rows ≥ h) + v = 0.0 + î = h + while k ≤ N-1 && v == 0 + k = k + 1 + v = abs(AB[h][k]) + # Find row ≥ h of maximum absolute value in this column + for i=(h+1):M + local v′ = abs(AB[i][k]) + if v′ > v + î = i + v = v′ + end + end + end + + if v > 0 + AB = ( + AB[1:(h-1)]..., + AB[î], + (let ĩ = (i==î ? h : i), h̃ = î, f = AB[ĩ][k] / AB[h̃][k] + ((0.0 for _ = 1:k)..., (AB[ĩ][j]-AB[h̃][j]*f for j = (k+1):N)...,) + end for i = (h+1):M)... + ) + end + end + + # Solve UAX=UB for X where UA with U presenting the transformations above an + # upper triangular matrix. + X = () + for i=M:-1:1 + r = .+(AB[i][M+1:end], (-AB[i][j].*X[j-i] for j=i+1:M)...)./AB[i][i] + X = (r, X...) + end + return X +end + +@inline function linsolve₀(AB :: TupleMatrix{M,N}) :: NTuple{M,Float64} where {N,M} + X = linsolve₀(AB, Val(1)) + return ((x for (x,) ∈ X)...,) +end + +#@generated function linsolve₁(AB :: TupleMatrix{M,N}, :: Type{TupleMatrix{M, K}}) :: TupleMatrix{M, K} where {N,M,K} +function generate_linsolve(AB :: Symbol, M :: Int, N :: Int, K :: Int) + @assert(M == N - K) + # The variables of ABN collect the stepwise stages of the transformed matrix. + # Initial i-1 entries of ABN[i] are never used, as the previous steps have already + # finalised the corresponding rows, but are included as “missing” (at compile time) + # for the sake of indexing clarity. The M-1:th row-wise step finalises the row-echelon + # form, so ABN has M-1 rows itself. + step_ABN(step) = ((missing for i=1:step-1)..., + (gensym("ABN_$(step)_$(i)") for i ∈ step:M)...,) + ABN = ((step_ABN(step) for step=1:M-1)...,) + # UAB “diagonally” refers to ABN to collate the final rows of the transformed matrix UAB. + # Since the M-1:th row-wise step finalises the row-echelon form, the last row comes already + # from the M-1:th step. + UAB = ((ABN[i][i] for i=1:M-1)..., ABN[M-1][M]) + # The variables of X collect the rows of the solution to AX=B. + X = ((gensym("X_$(i)") for i ∈ 1:M)...,) + + # Convert to row-echelon form. On each step we strip leading zeroes from ABN. + # In the end UAB should be upper-triangular. + convert_row(ABNout, h, ABNin) = quote + # Find pivotable column (has some non-zero entries in rows ≥ h) + $(ABNout[h]) = $(ABNin[h]) + local v = abs($(ABNout[h])[1]) + # Find row ≥ h of maximum absolute value in this column + $(sequence_exprs(:( + let v′ = abs($(ABNin[i])[1]) + if v′ > v + $(ABNout[h]), $(ABNin[i]) = $(ABNin[i]), $(ABNout[h]) + v = v′ + end + end + ) for i=(h+1):M)) + + $(lift_exprs(ABNout[h+1:M])) = v > 0 ? ( $(lift_exprs( :( + # Transform + $(ABNin[i])[2:$N-$h+1] .- $(ABNout[h])[2:$N-$h+1].*( $(ABNin[i])[1] / $(ABNout[h])[1]) + ) for i=h+1:M)) ) : ( $(lift_exprs( :( + # Strip leading zeroes + $(ABNin[i])[2:$N-$h+1] + ) for i=h+1:M)) ) + end + + # Solve UAX=UB for X where UA with U presenting the transformations above an + # upper triangular matrix. + solve_row(UAB, i) = :( + $(X[i]) = $(lift_exprs( :( + +($(UAB[i])[$M-$i+1+$k], $(( :( -$(UAB[i])[$j-$i+1]*$(X[j])[$k] ) for j=i+1:M)...)) + ) for k=1:K )) ./ $(UAB[i])[1] + ) + + return X, quote + $(lift_exprs(ABN[1][i] for i=1:M)) = $(lift_exprs( :( $AB[$i] ) for i=1:M)) + $(convert_row(ABN[1], 1, ABN[1])) + $((convert_row(ABN[h], h, ABN[h-1]) for h = 2:(M-1))...) + $((solve_row(UAB, i) for i=M:-1:1)...) + end +end + +@inline @generated function linsolve₁(AB :: TupleMatrix{M,N}, :: Val{K}) :: TupleMatrix{M, K} where {N,M,K} + X, solver = generate_linsolve(:AB, M, N, K) + return quote + $solver + return $(lift_exprs( X[i] for i=1:M )) + end +end + +@inline @generated function linsolve₁(AB :: TupleMatrix{M,N}) :: NTuple{M,Float64} where {N,M} + X, solver = generate_linsolve(:AB, M, N, 1) + return quote + $solver + return $(lift_exprs( :( $(X[i])[1] ) for i=1:M )) + end +end + +const linsolve = linsolve₁ + +function tuplify(M, N, A, b) + ((((A[i][j] for j=1:N)..., b[j]) for i=1:M)...,) +end + + +function compare(; dim=5, n_matrices=10000, n_testvectors=100) + testmatrices=[] + testvectors=[] + while length(testmatrices)<n_matrices + A=randn(dim, dim) + push!(testmatrices, A) + end + test_vectors=[randn(dim) for _ = 1:n_testvectors] + + + function evaluate(fn) + for A ∈ testmatrices + for b ∈ testvectors + fn(A, b) + end + end + end + + function evaluate_and_report(fn, name) + printstyled("Evaluating $name…\n", color=:cyan) + @time evaluate(fn) + end + + evaluate_and_report("tuple-linsolve, ungenerated") do A, b + linsolve₀(tuplify(A, b)) + end + + evaluate_and_report("tuple-linsolve, generated") do A, b + linsolve₁(tuplify(A, b)) + end + + evaluate_and_report("backslash") do A, b + A \ b + end +end + +end # module
--- a/src/Loops.jl Wed Dec 15 01:09:09 2021 +0200 +++ b/src/Loops.jl Wed Dec 22 11:13:38 2021 +0200 @@ -100,7 +100,7 @@ # Looping over indices and transformation ########################################## -Box{N, T} = NTuple{N, Tuple{T, T}} +const Box{N, T} = NTuple{N, Tuple{T, T}} struct Grid{R} end """
--- a/src/Metaprogramming.jl Wed Dec 15 01:09:09 2021 +0200 +++ b/src/Metaprogramming.jl Wed Dec 22 11:13:38 2021 +0200 @@ -35,7 +35,7 @@ Does `foldl` on the `Expr`s in `a` with the `f`, a `Function` or another `Expr` representing a function. Returns new `Expr` with the expanded `foldl` operation. """ -function foldl_exprs(f :: Expr, a) +function foldl_exprs(f :: Union{Expr,Symbol}, a) return foldl((b, c) -> :($f($b, $c)), a) end
--- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/src/PlanarFE.jl Wed Dec 22 11:13:38 2021 +0200 @@ -0,0 +1,937 @@ +""" +Simple planar (2D) finite element discretisation on a box. +""" +module PlanarFE + +import ..Loops: Box +import ..LinSolve: linsolve +using ..Metaprogramming + +export FEModel, + FEUniform, + UniformP0, + PlanarP0, + LineP0, + UniformP1, + PlanarP1, + LineP1, + UniformP2, + PlanarP2, + LineP2, + simplices, + nodes, + interpolate, + differential, + minimise + +const Location{N} = NTuple{N, Float64} +abstract type FEModel{N} end +abstract type FEUniform{N} <: FEModel{N} end + +const Index = Int +const CaseIndex = Int + +const NodeId{N} = NTuple{N, Index} + +struct Node{N} + id :: NodeId{N} + x :: Location{N} +end + +@inline function Base.getindex(a :: Array{T,N}, n :: Node{N}) where {T,N} + return getindex(a, n.id...) +end + +@inline function Base.setindex!(a :: Array{T,N}, x :: T, n :: Node{N}) where {T,N} + return setindex!(a, x, n.id...) +end + +############################################## +# Simplices +############################################## + +# (D-1)-dimensional simplex in ℝ^N. D is the node count. +struct Simplex{N,D,Id} + id :: Id + nodes :: NTuple{D, Node{N}} + @inline Simplex{Id}(id :: Id, nodes :: Vararg{Node{N}, D}) where {N,D,Id} = new{N,D,Id}(id, nodes) + @inline Simplex{N,D,Id}(id :: Id, nodes :: Vararg{Node{N}, D}) where {N,D,Id} = new{N,D,Id}(id, nodes) +end + +# Uniformly indexed simplices +const UCubeId{N} = NTuple{N, Int} +const USimplexId{N} = Tuple{UCubeId{N}, CaseIndex} +const USimplex{N,D} = Simplex{N,D,USimplexId{N}} +const PlanarSimplex = USimplex{2,3} +const RealInterval = USimplex{1,2} + +@inline function Base.getindex(a :: Array{T,Nplus1}, s :: USimplex{N,D}) where {T,Nplus1,D,N} + idx = (s.id[1]..., s.id[2]) :: NTuple{Nplus1,Int} + return getindex(a, idx...) +end + +@inline function Base.setindex!(a :: Array{T,Nplus1}, x :: T, s :: USimplex{N,D}) where {T,Nplus1,D,N} + idx = (s.id[1]..., s.id[2]) :: NTuple{Nplus1,Int} + return setindex!(a, x, idx...) +end + +# We do this by invididual cases to allow typing return values +# (N+1 is not calculable at compile time) +@inline function nodes(s :: Simplex{N,D,Id}) where {N,D,Id} + return s.nodes +end + +@inline @generated function dot(x :: Location{N}, y :: Location{N}) where N + return foldl_exprs( :( + ), :( x[$i]*y[$i] ) for i=1:N) +end + +@inline function orthog((x, y) :: Location{2}) + return (y, -x) +end + +@inline function in_simplex(t :: RealInterval, x :: Location{1}) + (x₀, x₁) = nodes(t) + return ((x .- x₀)*(x₁ .- x₀) ≥ 0 && + (x .- x₁)*(x₀ .- x₁) ≥ 0) +end + +@inline function in_simplex(t :: PlanarSimplex, x :: Location{2}) + (x₀, x₁, x₁) = nodes(t) + return (dot(x .- x₀, orthog(x₁ .- x₀)) ≥ 0 && + dot(x .- x₁, orthog(x₂ .- x₁)) ≥ 0 && + dot(x .- x₂, orthog(x₀ .- x₂)) ≥ 0) +end + + +############################################## +# Simple planar grid and triangulation +############################################## + +@inline function cellwidth(box :: Box{N, Float64}, gridsize :: Dims{N}) where N + return (((box[i][2]-box[i][1])/(gridsize[i]-1) for i=1:N)...,) +end + +# Cannot calculate N+1 compile-time +struct UniformGrid{N,Nplus1} + box :: Box{N, Float64} + size :: Dims{N} + cellwidth :: NTuple{N, Float64} + + function UniformGrid{N,Nplus1}(box :: Box{N, Float64}, gridsize :: Dims{N}) where {N,Nplus1} + # TODO: Nested tuple access is inefficient + @assert(N+1 == Nplus1) + @assert(all(box[i][2]>box[i][1] for i=1:N)) + @assert(all(gridsize .≥ 2)) + return new{N,Nplus1}(box, gridsize, cellwidth(box, gridsize)) + end +end + +const LineGrid = UniformGrid{1,2} +const PlanarGrid = UniformGrid{2,3} + +# Number of simplices in a cube of dimension N. +n_simplex_cases( :: LineGrid) = 1 :: CaseIndex +n_simplex_cases( :: PlanarGrid) = 2 :: CaseIndex +max_cube_id(grid :: UniformGrid{N, Nplus1}) where {N, Nplus1} = grid.size .- 1 +max_simplex_id(grid :: UniformGrid{N, Nplus1}) where {N, Nplus1} = + (max_cube_id(grid)..., n_simplex_cases(grid)) + +∈(x :: Location{N}, b :: Box{N}) where N = all(b[i][1] ≤ x[i] ≤ b[i][2] for i=1:N) +∉(x :: Location{N}, b :: Box{N}) where N = !(x ∈ b) +∈(x :: Location{N}, grid :: UniformGrid{N,Nplus1}) where {N,Nplus1} = x ∈ grid.box +∉(x :: Location{N}, grid :: UniformGrid{N,Nplus1}) where {N,Nplus1} = x ∉ grid.box + +@inline function width(grid :: UniformGrid{N,Nplus1}) where {N,Nplus1} + return ((grid.box[i][2]-grid.box[i][1] for i=1:N)...,) +end + +@inline function shift(grid :: UniformGrid{N,Nplus1}) where {N,Nplus1} + return ((grid.box[i][1] for i=1:N)...,) +end + +function construct_node(grid :: UniformGrid{N,Nplus1}, idx :: NodeId{N}) :: Node{N} where {N,Nplus1} + loc = grid.cellwidth.*(idx.-1).+shift(grid) + return Node{N}(idx, loc) +end + +@inline function construct_simplex(grid :: LineGrid, id :: USimplexId{1}) :: RealInterval + # TODO: optimise case away + (idx, _) = id + return RealInterval(construct_node(grid, idx), construct_node(grid, idx.+(1,))) +end + + +# In 2D we alternate ⧄ and ⧅. The “base” 90° corner is always listed first. +# It has importance for P0 elements. Otherwise arbitrary we proceed clockwise. +# Julia is real GARBAGE. It does tons of memory alloc trying to access this in +# any form, as Array or Tuple. Therefore things are in code below. +# offset₂ = [[ +# [(1, 0), (0, 0), (1, 1)], # ◪ +# [(0, 1), (1, 1), (0, 0)], # ◩ +# ],[ +# [(0, 0), (1, 0), (0, 1)], # ⬕ +# [(1, 1), (1, 0), (0, 1)], # ⬔ +# ]] :: Array{Array{Array{Tuple{Int,Int},1},1},1} + + +@inline function sqform₂(idx :: UCubeId{2}) + return 1+(idx[1]+idx[2])%2 +end + +# 2D simplices are identified by their “leading” node, i.e., the lower-left corner of the grid square +# that contains them, and the “case” number 1 or 2. +function construct_simplex(grid :: PlanarGrid, simplexid :: USimplexId{2}, form :: Index) :: PlanarSimplex + (idx, case) = simplexid + @assert(all(idx .< grid.size)) + @assert(form==1 || form==2) + @assert(case==1 || case==2) + (off1 :: NodeId{2}, off2 :: NodeId{2}, off3 :: NodeId{2}) = (form==1 ? + (case==1 ? + ((1, 0), (0, 0), (1, 1)) # ◪ + : + ((0, 1), (1, 1), (0, 0)) # ◩ + ) : (case==1 ? + ((0, 0), (1, 0), (0, 1)) # ⬕ + : + ((1, 1), (1, 0), (0, 1)) # ⬔ + )) + + return PlanarSimplex(simplexid, + construct_node(grid, idx .+ off1), + construct_node(grid, idx .+ off2), + construct_node(grid, idx .+ off3)) +end + +function construct_simplex(grid :: PlanarGrid, simplexid :: USimplexId{2}) :: PlanarSimplex + (idx, _) = simplexid + return construct_simplex(grid, simplexid, sqform₂(idx)) +end + + +@inline function get_cubeid(grid :: UniformGrid{N,Nplus1}, x :: Location{N}) :: UCubeId{N} where {N,Nplus1} + (( + # Without the :: Int typeassert this generates lots of memory allocs, although a typeof() + # check here confirms that the type is Int. Julia is just as garbage as everything. + max(min(ceil(Index, (x[i]-grid.box[i][1])/grid.cellwidth[i]), grid.size[i] - 1), 1) :: Int + for i=1:N)...,) +end + + +@inline function get_offset(grid :: UniformGrid{N,Nplus1}, cubeid :: UCubeId{N}, x :: Location{N}) :: Location{N} where {N,Nplus1} + (( + (x[i] - (cubeid[i] - 1) * grid.cellwidth[i]) + for i=1:N)...,) +end + +@inline function simplex_at(grid :: LineGrid, x :: Location{1}) :: RealInterval + if x ∉ grid.box + throw("Coordinate $x out of domain $(grid.box)") + else + return construct_simplex(grid, get_cubeid(grid, x)) + end +end + +@inline function simplex_at(grid :: PlanarGrid, x :: Location{2}) :: PlanarSimplex + if x ∉ grid.box + throw("Coordinate $x out of domain $(grid.box)") + else + # Lower left node is is cube id. Therefore maximum cubeid is grid size - 1 + # in each dimension. + cubeid = get_cubeid(grid, x) + off = get_offset(grid, cubeid, x) + form = sqform₂(cubeid) + if form==0 + # ⧄ → ◪ or ◩ + case = (off[1]>off[2] ? 1 : 2 ) :: CaseIndex + else + # ⧅ → ⬕ or ⬔ + case = (off[1]+off[2]<1 ? 1 : 2) :: CaseIndex + end + return construct_simplex(grid, (cubeid, case), form) + end +end + +############################################## +# Edges for P2 elements +############################################## + +const Edge{N,Id} = Simplex{N,2,Id} +const UEdgeId{N} = USimplexId{N} +const UEdge{N} = Edge{N,UEdgeId{N}} + +@inline midpoint(e :: Edge{N,Id}) where {N,Id} = (e.nodes[1].x .+ e.nodes[2].x) ./ 2.0 + +# For our purposes, an edge is something useful for putting an interpolation midpoint on. +# This in 1D there's one “edge”, which is just the original simplex. +n_cube_edges( :: LineGrid) = 1 +n_cube_edges( :: PlanarGrid) = 3 +max_edge_id(grid :: UniformGrid{N, Nplus1}) where {N, Nplus1} = (grid.size..., n_cube_edges(grid)) + +@inline function edge_id(node₁ :: NodeId{1}, node₂ :: NodeId{1}) :: UEdgeId{1} + return (min(node₁[1], node₂[1]),) +end + +# Edges are identified by the lower left node of the “least-numbered” cube that +# contains them, along with 1=horizontal, 2=vertical, 3=diagonal. +# Both of the alternating cases of diagonal edges have the same number. +# Since the least-numbered cube identifies edges, only both diagonal the left and +# bottom border of the cube are counted for that cube, so there are 3 possible +# edges counted towards a cube. +function edge_id(idx₁ :: NodeId{2}, idx₂ :: NodeId{2}) :: UEdgeId{2} + # For performance this does not verify that idx₁ and idx₂ are in the + # same simplex, it only uses their ordering! + if idx₂[1] > idx₁[1] # one of ↗︎, →, ↘︎ + if idx₂[2] > idx₁[2] + return (idx₁, 3) # diagonal ↗︎ in cube(idx₁) + elseif idx₂[2] == idx₁[2] + return (idx₁, 1) # lower side → in cube(idx₁) + else # idx₂[2] < idx₁[2] + return ((idx₁[1], idx₂[2]), 3) # ↘︎ in cube(prev(idx₁,idx₂)) + end + elseif idx₂[1] == idx₁[1] # one of ↑, ↓ + if idx₂[2] > idx₁[2] + return (idx₁, 2) # ↑ in cube(idx₁) + elseif idx₂[2] == idx₁[2] + throw("Edge cannot end in starting idx.") + else # idx₂[2] < idx₁[2] + return (idx₂, 2) # ↓ cube(idx₂) + end + else # idx₂[1] < idx₁[1] # one of ↙︎, ←, ↖︎ + if idx₂[2] > idx₁[2] + return ((idx₁[1], idx₂[2]), 3) # diagonal ↖︎ in cube(prev(idx₁,idx₂)) + elseif idx₂[2] == idx₁[2] + return (idx₂, 1) # lower side ← in cube(idx₂) + else # idx₂[2] < idx₁[2] + return (idx₂, 3) # ↙︎ in cube(idx₂) + end + end +end + +function construct_edge(n₁ :: Node{N}, n₂ :: Node{N}) where N + return UEdge{N}(edge_id(n₁.id, n₂.id), n₁, n₂) +end + +############################################## +# Iteration over simplices and nodes +############################################## + +struct Iter{G,W} + grid :: G +end + +# @inline function nodes(s :: PlanarSimplex{N}) where N +# return Iter{PlanarSimplex{N}, :nodes}(s) +# end + +# @inline function Base.iterate(sn :: Iter{PlanarSimplex{N}, :nodes}) where N +# return (sn.p.base, Index(0)) +# end + +# @inline function Base.iterate(sn :: Iter{PlanarSimplex{N}, :nodes}, idx :: Index) where N +# if idx==length(sn.p.tips) +# return nothing +# else +# idx=idx+1 +# return (sn.p.tips[idx], idx) +# end +# end + +@inline function nodes(p :: M) where {N, M <: FEUniform{N}} + return nodes(p.grid) +end + +@inline function nodes(grid :: UniformGrid{N,Nplus1}) where {N,Nplus1} + return Iter{UniformGrid{N,Nplus1}, Node{N}}(grid) +end + +@inline function first_nodeid( :: UniformGrid{N,Nplus1}) :: NodeId{N} where {N,Nplus1} + return (( 1 :: Int for i=1:N)...,) +end + +# Need to hard-code this for each N to avoid Julia munching memory. +# Neither @generated nor plain code works. +@inline function next_nodeid(grid :: LineGrid, idx :: NodeId{1}) :: Union{Nothing, NodeId{1}} + if idx[1] < grid.size[1] + return (idx[1]+1,) + else + return nothing + end +end + +@inline function next_nodeid(grid :: PlanarGrid, idx :: NodeId{2}) :: Union{Nothing, NodeId{2}} + if idx[1] < grid.size[1] + return (idx[1]+1, idx[2]) + elseif idx[2] < grid.size[2] + return (1 :: Int, idx[2]+1) + else + return nothing + end +end + +function Base.iterate(iter :: Iter{UniformGrid{N,Nplus1}, Node{N}}) :: + Union{Nothing, Tuple{Node{N}, NodeId{N}}} where {N,Nplus1} + idx = first_nodeid(iter.grid) + return construct_node(iter.grid, idx), idx +end + +function Base.iterate(iter :: Iter{UniformGrid{N,Nplus1}, Node{N}}, idx :: NodeId{N}) :: + Union{Nothing, Tuple{Node{N}, NodeId{N}}} where {N,Nplus1} + nidx = next_nodeid(iter.grid, idx) + if isnothing(nidx) + return nothing + else + return construct_node(iter.grid, nidx), nidx + end +end + +# @generated function Base.iterate(iter :: Iter{N, Node{N}}, idx :: NodeId{N}) where N +# # This generation reduces memory allocations +# do_coord(i) = :( +# if idx[$i] < iter.grid.size[$i] +# nidx = $(lift_exprs(( (:( 1 :: Int ) for _=1:i-1)..., :( idx[$i]+1 ), ( :( idx[$j] ) for j=i+1:N)...,))) :: NTuple{N, Int} +# return construct_node(iter.grid, nidx), nidx +# end +# ) + +# return quote +# $(sequence_exprs( do_coord(i) for i=1:N )) +# return nothing +# end +# end + +# function Base.iterate(iter :: Iter{N, Node{N}}, idx :: NodeId{N}) where N +# for i=1:N +# if idx[i]<iter.grid.size[i] +# nidx = (( 1 :: Int for _=1:i-1)..., idx[i]+1, idx[i+1:N]...) +# return construct_node(iter.grid, nidx), nidx +# end +# end + +# return nothing +# end + +@inline function simplices(p :: M) where {N, M <: FEUniform{N}} + return simplices(p.grid) +end + +@inline function simplices(grid :: UniformGrid{N,Nplus1}) where {N,Nplus1} + return Iter{UniformGrid{N,Nplus1}, USimplex{N,D}}(grid) +end + +@inline function Base.iterate(iter :: Iter{UniformGrid{N,Nplus1}, Simplex{N,Nplus1}}) :: + Union{Nothing, Tuple{USimplex{N,Nplus1}, USimplexId{N}}} where {N,Nplus1} + cubeid = ((Index(1) for i=1:N)...,) + simplexid = (cubeid, Index(1)) + return (construct_simplex(iter.grid, simplexid), simplexid) +end + +@inline function Base.iterate(iter :: Iter{UniformGrid{N,Nplus1}, Simplex{N,Nplus1}}, + (cubeid, case) :: USimplexId{N}) :: + Union{Nothing, Tuple{USimplex{N,Nplus1}, USimplexId{N}}} where {N,Nplus1} + if case < n_simplex_cases(iter.grid) + simplexid = (cubeid, case+1) + return (construct_simplex(iter.grid, simplexid), simplexid) + else + for i=1:N + if cubeid[i]+1<iter.grid.size[i] + new_cubeid=((Index(1) for _=1:i-1)..., cubeid[i]+1, cubeid[i+1:N]...) + simplexid = (new_cubeid, Index(1)) + return (construct_simplex(iter.grid, simplexid), simplexid) + end + end + end + + return nothing +end + +# TODO: wrong for far-side boundary? Need to do by cumbe numbers +@inline function edges(s :: RealInterval) :: Tuple{UEdge{1}} + return (construct_edge(s.nodes...),) +end + +# TODO: Edge -> Simplex +@inline function edges(s :: PlanarSimplex) :: NTuple{3,UEdge{2}} + return (construct_edge(s.nodes[1], s.nodes[2]), + construct_edge(s.nodes[1], s.nodes[3]), + construct_edge(s.nodes[2], s.nodes[3])) +end + +@inline function edges(p :: M) where {N, M <: FEUniform{N}} + return edges(p.grid) +end + +@inline function edges(grid :: LineGrid) + return Iter{LineGrid, UEdge{1}}(grid) +end + +@inline function edges(grid :: PlanarGrid) + return Iter{PlanarGrid, UEdge{2}}(grid) +end + +@inline function UEdge(grid :: LineGrid, edgeid :: UEdgeId{1}) :: UEdge{1} + (cubeid, case) = edgeid + @assert(case==1) + return UEdge(edgeid, construct_node(grid, cubeid), construct_node(grid, cubeid.+(1,))) +end + +@inline function next_edge_case(grid :: LineGrid, (cubeid, case) :: UEdgeId{1}) :: CaseIndex + return -1 +end + +@inline function first_edge_case(grid :: LineGrid, cubeid :: UCubeId{1}) :: CaseIndex + return cubeid[1] < grid.size[1] ? 1 : -1 +end + +@inline function UEdge(grid :: PlanarGrid, edgeid :: UEdgeId{2}) :: UEdge{2} + (cubeid, case) = edgeid + e(idx1, idx2) = UEdge{2}(edgeid, construct_node(grid, idx1), construct_node(grid, idx2)) + if case==1 + return e(cubeid, cubeid.+(1,0)) + elseif case==2 + return e(cubeid, cubeid.+(0,1)) + elseif case==3 + if sqform₂(cubeid)==1 + return e(cubeid, cubeid.+(1,1)) + else + return e(cubeid.+(1,0), cubeid.+(0,1)) + end + else + throw("Invalid edge") + end +end + +@inline function first_edge_case(grid :: PlanarGrid, cubeid :: UCubeId{2}) :: CaseIndex + if cubeid[1]<grid.size[1] + return 1 + elseif cubeid[2]<grid.size[2] + return 2 + else + return -1 + end +end + + +@inline function next_edge_case(grid :: PlanarGrid, (cubeid, case) :: UEdgeId{2}) :: CaseIndex + if case<3 && cubeid[1]<grid.size[1] && cubeid[2]<grid.size[2] + return case + 1 + else + return -1 + end +end + + +@inline function Base.iterate(iter :: Iter{UniformGrid{N,Nplus1}, UEdge{N}}) :: + Union{Nothing, Tuple{UEdge{N}, UEdgeId{N}}} where {N,Nplus1} + cubeid = first_nodeid(iter.grid) + case = first_edge_case(iter.grid, cubeid) + @assert(case ≥ 1) + edgeid = (cubeid, case) + return (UEdge(iter.grid, edgeid), edgeid) +end + +@inline function Base.iterate(iter :: Iter{UniformGrid{N,Nplus1}, UEdge{N}}, (cubeid, case) :: UEdgeId{N}) :: + Union{Nothing, Tuple{UEdge{N}, UEdgeId{N}}} where {N,Nplus1} + case = next_edge_case(iter.grid, (cubeid, case)) + if case ≥ 1 + edgeid = (cubeid, case) + return (UEdge(iter.grid, edgeid), edgeid) + else + while true + cubeid = next_nodeid(iter.grid, cubeid) + if isnothing(cubeid) + return nothing + end + case = first_edge_case(iter.grid, cubeid) :: Index + if case ≥ 1 + edgeid = (cubeid, case) :: UEdgeId{N} + return (UEdge(iter.grid, edgeid), edgeid) + end + end + end + + return nothing +end + + +############################################## +# Evaluation, differentiation, minimisation +############################################## + +function interpolate(p :: M, x :: Location{N}) where {N, M <: FEModel{N}} + return interpolate(p, simplex_at(p.grid, x), x) +end + +# Apparent Julia bug / incomplete implementation: +# https://stackoverflow.com/questions/46063872/parametric-functors-in-julia +# How to do individually. +# function (p :: M)(x :: Location{N}) where {N, M <: FEModel{N}} +# return interpolate(p, x) +# end + +function differential(p :: M, x :: Location{N}) where {N, M <: FEModel{N}} + differential(p, simplex_at(p.grid, x), x) +end + +function minimise(p :: M) where {N, M <: FEModel{N}} + minloc = nothing + minv = nothing + for simplex ∈ simplices(p) + thisminloc, thisminv = minimise(p, simplex) + if isnothing(minv) || minv < thisminv + minv=thisminv + minloc=thisminloc + end + end + return minv, minloc +end + +############################################## +# P0 elements +############################################## + +struct UniformP0{N,Nplus1} <: FEUniform{N} + simplex_values :: Array{Float64, Nplus1} + grid :: UniformGrid{N,Nplus1} + + function UniformP0{N,Nplus1}(simplex_values :: Array{Float64, Nplus1}, + box :: Box{N, Float64}, gridsize :: Dims{N}) where {N,Nplus1} + grid = UniformGrid{N,Nplus1}(box, gridsize) + return new{N,Nplus1}(simplex_values, grid) + end + + function UniformP0{N,Nplus1}(box :: Box{N, Float64}, gridsize :: Dims{N}) where {N,Nplus1} + grid = UniformGrid{N,Nplus1}(box, gridsize) + a = Array{Float64, Nplus1}(undef, max_simplex_id(grid)) + return new{N,Nplus1}(simplex_values, grid) + end + +end + +const LineP0 = UniformP0{1,1} +const PlanarP0 = UniformP0{2,3} + +# f is not specified a type to allow callable objects and functions +function UniformP0{N,Nplus1}(f, box :: Box{N, Float64}, gridsize :: Dims{N}) where {N,Nplus1} + p = UniformP0{N,Nplus1}(box, gridsize) + discretise!(f, p) + return p +end + +function (p :: UniformP0{N,Nplus1})(x :: Location{N}) where {N,Nplus1} + return interpolate(p, x) +end + +@inline function interpolate(p :: UniformP0{N,Nplus1}, s :: Simplex{N,Nplus1}, x :: Location{N}) where {N,Nplus1} + return p.simplex_values[s] +end + +function differential(p :: UniformP0{N,Nplus1}, s :: Simplex{N,Nplus1}, x :: Location{N}) where {N,Nplus1} + return ((0.0 for i=1:N)...,) +end + +@inline function minimise(p :: UniformP0{N,Nplus1}, s :: Simplex{N,Nplus1}) where {N,Nplus1} + return s.nodes[1].x, p.simplex_values[s] +end + +function discretise!(f, p :: UniformP0{N,Nplus1}) where {N,Nplus1} + for s ∈ simplices(p) + p.simplex_values[s]=f(s.nodes[1].x) + end +end + +############################################## +# P1 elements +############################################## + +struct UniformP1{N,Nplus1} <: FEUniform{N} + node_values :: Array{Float64, N} + grid :: UniformGrid{N,Nplus1} + + function UniformP1{N,Nplus1}(node_values :: Array{Float64, N}, + box :: Box{N, Float64}, gridsize :: Dims{N}) where {N,Nplus1} + grid = UniformGrid{N,Nplus1}(box, gridsize) + return new{N,Nplus1}(node_values, grid) + end + + function UniformP1{N,Nplus1}(box :: Box{N, Float64}, gridsize :: Dims{N}) where {N,Nplus1} + return UniformP1{N,Nplus1}(Array{Float64, N}(undef, gridsize), box, gridsize) + end +end + +const LineP1 = UniformP1{1,1} +const PlanarP1 = UniformP1{2,3} + +# f is not specified a type to allow callable objects and functions +function UniformP1{N,Nplus1}(f, box :: Box{N, Float64}, gridsize :: Dims{N}) where {N,Nplus1} + p = UniformP1{N,Nplus1}(box, gridsize) + discretise!(f, p) + return p +end + +const P1NodalData{N} = Tuple{Location{N}, Float64} + +# 1D model functions + +@inline function p1_simplicialf(((x₀,), v₀) :: P1NodalData{1}, + ((x₁,), v₁) :: P1NodalData{1}) :: NTuple{2,Float64} + a₁ = (v₁-v₀)/(x₁-x₀) + a₀ = v₀ - a₁*x₀ + return a₀, a₁ +end + +function p1_modelf(p :: UniformP1{1}, s :: RealInterval) + n₀, n₁ = nodes(s) + return p1_simplicialf((n₀.x, p.node_values[n₀]), + (n₁.x, p.node_values[n₁])) +end + + +# 2D model functions + +@inline function p1_simplicialf((x₀, v₀) :: P1NodalData{2}, + (x₁, v₁) :: P1NodalData{2}, + (x₂, v₂) :: P1NodalData{2}) :: NTuple{3,Float64} + d = x₁ .- x₀ + q = x₂ .- x₀ + r = q[1]*d[2] - d[1]*q[2] + a₁ = ((v₂-v₀)*d[2] - (v₁-v₀)*q[2])/r + a₂ = ((v₁-v₀)*q[1] - (v₂-v₀)*d[1])/r + a₀ = v₀-a₁*x₀[1]-a₂*x₀[2] + return a₀, a₁, a₂ +end + +@inline function p1_modelf(p :: UniformP1{N,Nplus1}, s :: Simplex{N,Nplus1}) where {N,Nplus1} + return p1_simplicialf(((n.x, p.node_values[n]) for n ∈ nodes(s))...) +end + +# @generated function p1_modelf(p :: UniformP1{N}, s :: PlanarSimplex{N}) where N +# return quote +# ns = nodes(s) +# p1_simplicialf($(lift_exprs( :((ns[$i].x, p.node_values[ns[$i].id...])) for i=1:3 ))...) +# end +# end + +#@inline function g(a, (i,)) +# return a[i] +#end + +# @inline function g(a, (i, j)) +# return a[i,j] +# end + +# function p1_modelf(p :: UniformP1{2}, s :: PlanarSimplex{2}) +# n₀, n₁, n₂ = nodes(s) +# a, b = n₀.id[1], n₀.id[2] +# v₀ = p.node_values[a, b] +# a, b = n₁.id[1], n₁.id[2] +# v₁ = p.node_values[a, b] +# a, b = n₂.id[1], n₂.id[2] +# v₂ = p.node_values[a, b] +# return p1_simplicialf(n₀.x, v₀, +# n₁.x, v₁, +# n₂.x, v₂) +# end + + +# Common + +function (p :: UniformP1{N,Nplus1})(x :: Location{N}) where {N,Nplus1} + return interpolate(p, x) +end + +function interpolate(p :: UniformP1{N,Nplus1}, s :: Simplex{N,Nplus1}, x :: Location{N}) where {N,Nplus1} + a = p1_modelf(p, s) + return +((a.*(1, x...))...) +end + +@inline function differential(p :: UniformP1{N,Nplus1}, s :: Simplex{N,Nplus1}, :: Location{N}) :: Gradient{N} where {N,Nplus1} + a = p1_modelf(p, s) + return a[2:end] +end + +@inline function minimise(p :: UniformP1{N,Nplus1}, t :: Simplex{N,Nplus1}) where {N,Nplus1} + bestnode = s.nodes[1] + bestvalue = p.node_values[bestnode] + # The minima is found at a node, so only need to go through them + for node ∈ s.nodes[2:end] + value = p.node_values[node] + if value < bestvalue + bestvalue = value + bestnode = node + end + end + return bestnode.x, bestvalue +end + +function discretise!(f, p :: UniformP1{N,Nplus1}) where {N,Nplus1} + for n ∈ nodes(p) + p.node_values[n]=f(n.x) + end +end + +############################################## +# P2 elements +############################################## + +const Gradient{N} = NTuple{N, Float64} +const P2NodalData{N} = Tuple{Location{N}, Float64, Gradient{N}} +const LocationAndGradient{N} = Tuple{Location{N}, Gradient{N}} + +struct UniformP2{N,Nplus1} <: FEUniform{N} + node_values :: Array{Float64, N} + edge_values :: Array{Float64, Nplus1} + grid :: UniformGrid{N,Nplus1} + + function UniformP2{N,Nplus1}( + (node_values, edge_values) :: Tuple{Array{Float64, N}, Array{Float64, Nplus1}}, + box :: Box{N, Float64}, gridsize :: Dims{N}) where {N, Nplus1} + return new{N,Nplus1}(node_values, edge_values, UniformGrid{N,Nplus1}(box, gridsize)) + end + + function UniformP2{N,Nplus1}(box :: Box{N, Float64}, gridsize :: Dims{N}) where {N,Nplus1} + grid = UniformGrid{N,Nplus1}(box, gridsize) + na = Array{Float64, N}(undef, gridsize) + ne = Array{Float64, Nplus1}(undef, max_edge_id(grid)) + return new{N,Nplus1}(na, ne, grid) + end +end + +const LineP2 = UniformP2{1,2} +const PlanarP2 = UniformP2{2,3} + +# f is not specified a type to allow callable objects and functions +function UniformP2{N,Nplus1}(f, box :: Box{N, Float64}, gridsize :: Dims{N}) where {N,Nplus1} + p = UniformP2{N,Nplus1}(box, gridsize) + discretise!(f, p) + return p +end + +# 1D model function + +@inline function p2powers(x :: Location{1}) :: NTuple{3, Float64} + return (1, x[1], x[1]*x[1]) +end + +@inline function p2powers_diff(x :: Location{1}) :: NTuple{3, Float64} + return ((0, 1, 2*x[1]),) +end + +@inline function p2_simplicialf(xv :: Vararg{P1NodalData{1},3}) :: NTuple{3, Float64} + Ab = (((p2powers(xv[i][1])..., xv[i][2]) for i = 1:3)...,) + return linsolve(Ab) +end + +# 2D model function + +@inline function p2powers(x :: Location{2}) #:: NTuple{6, Float64} + return (1.0, x[1], x[2], x[1]*x[1], x[1]*x[2], x[2]*x[2]) +end + +function p2powers_diff(x :: Location{2}) :: NTuple{6, Float64} + return ((0.0, 1.0, 0.0, 2*x[1], 2*x[2], 0.0), + (0.0, 0.0, 1.0, 0.0, 2*x[1], 2*x[2])) +end + +function p2_simplicialf(xv :: Vararg{P1NodalData{2},6}) :: NTuple{6, Float64} + Ab = (((p2powers(xv[i][1])..., xv[i][2]) for i = 1:6)...,) + return linsolve(Ab) +end + +# Common + +@noinline function p2_modelf(p :: UniformP2{N,Nplus1}, s :: Simplex{N,Nplus1}) where {N,Nplus1} + a = (((n.x, p.node_values[n]) for n ∈ nodes(s))...,) + b = (((midpoint(e), p.edge_values[e]) for e ∈ edges(s))...,) + return p2_simplicialf(a..., + b...) +end + +@inline function interpolate(p :: UniformP2{N,Nplus1}, s :: Simplex{N,Nplus1}, x :: Location{N}) where {N,Nplus1} + a = p2_modelf(p, s) + p = p2powers(x) + return +((a.*p)...) +end + +function (p :: UniformP2{N,Nplus1})(x :: Location{N}) where {N,Nplus1} + return interpolate(p, x) +end + +@inline function differential(p :: UniformP2{N,Nplus1}, s :: Simplex{N,Nplus1}, + x :: Location{N}) :: Gradient{N} where {N,Nplus1} + a = p2_modelf(p, s) + d = p2powers_diff(x) + return (+((a.*di)...) for di ∈ d) +end + +@inline function minimise_p2_1d(a, (x₀, x₁)) + (_, a₁, a₁₁) = a + # We do this in cases, first trying for an interior solution, then edges. + # For interior solution, first check determinant; no point trying if non-positive + if a₁₁ > 0 + # An interior solution x[1] has to satisfy + # 2a₁₁*x[1] + a₁ =0 + # This gives + x = (-a₁/(2*a₁₁),) + if in_simplex(t, x) + return x, +(a.*p2powers(x)) + end + end + + (x₀, x₁) = nodes(t) + v₀ = +((a.*p2powers(x₀))...) + v₁ = +((a.*p2powers(x₁))...) + return v₀ > v₁ ? (x₀, v₀) : (x₁, v₁) +end + +function minimise(p :: LineP2, t :: RealInterval) + minimise_p2_1d(p2_modelf(p, s), nodes(t)) +end + +@inline function minimise_p2_edge((a₀, a₁, a₂, a₁₁, a₁₂, a₂₂), x₀, x₁) + # Minimise the 2D model on the edge {x₀ + t(x₁ - x₀) | t ∈ [0, 1] } + d = x₁ .- x₀ + b₀ = a₀ + a₁*x₀[1] + a₂*x₀[2] + a₁₁*x₀[1]*x₀[1] + a₁₂*x₀[1]*x₀[2] + a₂₂*x₀[2]*x₀[2] + b₁ = a₁*d[1] + a₂*d[2] + 2*a₁₁*d[1]*x₀[1] + a₁₂*(d[1]*x₀[2] + d[2]*x₀[1]) + 2*a₂₂*d[2]*x₀[2] + b₁₁ = a₁₁*d[1]*d[1] + a₁₂*d[1]*d[2] + a₂₂*d[2]*d[2] + t, v = minimise_p2_1d((b₀, b₁, b₁₁), (0, 1)) + return @.(x₀ + t*d), v +end + +@inline function minimise(p :: PlanarP2, t :: PlanarSimplex) + a = (_, a₁, a₂, a₁₁, a₁₂, a₂₂) = p2_modelf(p, s) + # We do this in cases, first trying for an interior solution, then edges. + + # For interior solution, first check determinant; no point trying if non-positive + r = 2*(a₁₁*a₂₂-a₁₂*a₁₂) + if r > 0 + # An interior solution (x[1], x[2]) has to satisfy + # 2a₁₁*x[1] + 2a₁₂*x[2]+a₁ =0 and 2a₂₂*x[1] + 2a₁₂*x[1]+a₂=0 + # This gives + x = ((a₂₂*a₁-a₁₂*a₂)/r, (a₁₂*a₁-a₁₁*a₂)/r) + if in_simplex(t, x) + return x, +(a.*p2powers(x)) + end + end + + (x₀, x₁, x₁) = nodes(t) + + x₀₁, v₀₁ = minimise_p2_edge(a, x₀, x₁) + x₁₂, v₁₂ = minimise_p2_edge(a, x₁, x₂) + x₂₀, v₂₀ = minimise_p2_edge(a, x₂, x₀) + if v₀₁ > v₁₂ + return v₀₁ > v₂₀ ? (x₀₁, v₀₁) : (x₂₀, v₂₀) + else + return v₁₂ > v₂₀ ? (x₁₂, v₁₂) : (x₂₀, v₂₀) + end +end + +function discretise!(f, p :: UniformP2{N,Nplus1}) where {N,Nplus1} + for n ∈ nodes(p) + p.node_values[n] = f(n.x) + end + for e ∈ edges(p) + p.edge_values[e] = f(midpoint(e)) + end +end + +end # module \ No newline at end of file
--- a/src/ZipArrays.jl Wed Dec 15 01:09:09 2021 +0200 +++ b/src/ZipArrays.jl Wed Dec 22 11:13:38 2021 +0200 @@ -36,21 +36,21 @@ A way to refer to a `ZipArray` without specifying the container types. """ -AbstractZipArray{S,T,N} = ZipArray{S,T,N,A,B} where {A <: AbstractArray{S,N}, B <: AbstractArray{T,N}} +const AbstractZipArray{S,T,N} = ZipArray{S,T,N,A,B} where {A <: AbstractArray{S,N}, B <: AbstractArray{T,N}} """ `ZipVector{S,T,A,B}` One-dimensional `ZipArray`. """ -ZipVector{S,T,A,B} = ZipArray{S,T,1,A,B} +const ZipVector{S,T,A,B} = ZipArray{S,T,1,A,B} """ `AbstractZipVector{S,T}` One-dimensional `AbstractZipArray`. """ -AbstractZipVector{S,T} = AbstractZipArray{S,T,1} +const AbstractZipVector{S,T} = AbstractZipArray{S,T,1} function Base.IndexStyle( :: Type{<:ZipArray{S,T,N,A,B}}) where {S,T,N,A,B}