Merge pull request #54 from sintefmath/dev

Test + fix for sparsity wrapper

Project.toml

@@ -1,7 +1,7 @@
 name = "Jutul"
 uuid = "2b460a1a-8a2b-45b2-b125-b5c536396eb9"
 authors = ["Olav Møyner <[email protected]>"]
-version = "0.2.18"
+version = "0.2.19"
 
 [deps]
 AlgebraicMultigrid = "2169fc97-5a83-5252-b627-83903c6c433c"

src/ad/sparsity.jl

@@ -35,19 +35,47 @@ Base.parent(A::SparsityTracingWrapper) = A.data
 Base.size(A::SparsityTracingWrapper) = size(A.data)
 Base.axes(A::SparsityTracingWrapper) = axes(A.data)
 
-function Base.getindex(A::SparsityTracingWrapper, i, j)
+function Base.getindex(A::SparsityTracingWrapper, i)
+    if i isa Colon
+        return A
+    else
+        return map(i -> A[i], i)
+    end
+end
+
+function Base.getindex(A::SparsityTracingWrapper{T, D, <:Any}, I, J) where {T, D}
+    @assert D > 1
+    if I isa Colon
+        I = axes(A, 1)
+    end
+    if J isa Colon
+        J = axes(A, 2)
+    end
+    Ts = Jutul.ST.ADval{T}
+    n = length(I)
+    m = length(J)
+    out = Matrix{Ts}(undef, n, m)
+    for (i, ix) in enumerate(I)
+        for (j, jx) in enumerate(J)
+            out[i, j] = A[ix, jx]
+        end
+    end
+    return out
+end
+
+function Base.getindex(A::SparsityTracingWrapper, i::Int, j::Int)
     return as_tracer(A.data[i, j], j)
 end
 
-function Base.getindex(A::SparsityTracingWrapper{T, 2, D}, ix) where {T, D}
+function Base.getindex(A::SparsityTracingWrapper{T, 2, D}, ix::Int) where {T, D}
     n, m = size(A)
     zero_ix = ix - 1
     i = (zero_ix ÷ m) + 1
     j = mod(zero_ix, m) + 1
     return as_tracer(A.data[i, j], j)
 end
 
-function Base.getindex(A::SparsityTracingWrapper{T, 1, D}, i) where {T, D}
+function Base.getindex(A::SparsityTracingWrapper{T, 1, D}, i::Int) where {T, D}
     return as_tracer(A.data[i], i)
 end
 

src/meshes/unstructured/types.jl

@@ -226,8 +226,8 @@ function UnstructuredMesh(
     int_neighbors = convert_neighborship(int_neighbors)
     @assert length(bnd_cells) == nb
     bnd_cells::AbstractVector
-    @assert maximum(bnd_cells) <= nc
-    @assert minimum(bnd_cells) > 0
+    @assert maximum(bnd_cells, init = 1) <= nc
+    @assert minimum(bnd_cells, init = nc) > 0
 
     @assert maximum(faces_to_nodes.vals, init = 0) <= nn "Too few nodes provided"
     return UnstructuredMesh(cells_to_faces, cells_to_bnd, faces_to_nodes, bnd_to_nodes, node_points, int_neighbors, bnd_cells; kwarg...)

src/partitioning.jl

@@ -294,7 +294,7 @@ end
 """
     partition_hypergraph(g, n::Int, partitioner = MetisPartitioner(); expand = true)
 
-Partition a hypergraph from [setup_partitioner_hypergraph](@ref) using a given
+Partition a hypergraph from [`setup_partitioner_hypergraph`](@ref) using a given
 partitioner. If the optional `expand` parameter is set to true the result will
 be expanded to the full graph (i.e. where groups are not condensed).
 """

src/simulator/print.jl

@@ -102,7 +102,7 @@ function final_simulation_message(simulator, p, rec, t_elapsed, reports, timeste
     if verbose && length(reports) > 0
         if aborted
             start_str = "Simulation aborted"
-            endstr = "$(stats.steps-1) of $(length(timesteps))"
+            endstr = "$(stats.steps) of $(length(timesteps))"
             str_c = :red
         else
             start_str = "Simulation complete"

test/runtests.jl

@@ -6,3 +6,4 @@ include("adjoints/basic_adjoint.jl")
 include("utils.jl")
 include("partitioning.jl")
 include("units.jl")
+include("sparsity.jl")

test/sparsity.jl

@@ -0,0 +1,51 @@
+using Jutul, Test
+
+@testset "SparsityTracingWrapper" begin
+    n = 10
+    m = 3
+    test_mat = zeros(m, n)
+    test_vec = zeros(n)
+
+    for i in 1:n
+        test_vec[i] = i
+        for j in 1:m
+            test_mat[j, i] = i + (j-1)*n
+        end
+    end
+
+    vec_st = Jutul.SparsityTracingWrapper(test_vec)
+    for i in 1:n
+        v = vec_st[i]
+        @test v isa Jutul.ST.ADval
+        @test v.derivnode.index == i
+        @test v.val == test_vec[i]
+    end
+
+    mat_st = Jutul.SparsityTracingWrapper(test_mat)
+    for i in 1:n
+        for j in 1:m
+            v = mat_st[j, i]
+            @test v isa Jutul.ST.ADval
+            @test v.derivnode.index == i
+            @test v.val == test_mat[j, i]
+        end
+    end
+
+    # Test ranges for vector
+    rng = 2:7
+    tmp = vec_st[rng]
+    @test size(tmp) == size(test_vec[rng])
+
+    for (i, ix) in enumerate(rng)
+        @test tmp[i] == vec_st[ix]
+    end
+
+    # Test subranges
+    tmp2 = mat_st[:, rng]
+    @test size(tmp2) == size(test_mat[:, rng])
+    for (i, ix) in enumerate(rng)
+        for j in 1:m
+            @test tmp2[j, i] == mat_st[j, ix]
+        end
+    end
+end