Merge pull request #83 from jipolanco/similar_dims

`similar`: always return a `PencilArray`

CHANGELOG.md

@@ -0,0 +1,15 @@
+# Changelog
+
+The format is based on [Keep a Changelog] and [Common Changelog].
+
+## [0.19.0] - 2023-07-14
+
+### Changed
+
+-   **Breaking:** change behaviour of `similar(u::PencilArray, [T], dims)` ([#83])
+
+    When the `dims` argument is passed, we now try to return a new `PencilArray` instead of another (non-distributed) array type. Since this is only possible when `dims` matches the array size, an error is now thrown if that is not the case. This allows things to play nicely with other packages such as [StructArrays.jl](https://github.com/JuliaArrays/StructArrays.jl), which in some cases end up calling `similar` with the `dims` argument.
+
+  [Keep a Changelog]: https://keepachangelog.com/en/1.1.0/
+  [Common Changelog]: https://common-changelog.org/
+  [#83]: https://github.com/jipolanco/PencilArrays.jl/pull/83

Project.toml

@@ -1,7 +1,7 @@
 name = "PencilArrays"
 uuid = "0e08944d-e94e-41b1-9406-dcf66b6a9d2e"
 authors = ["Juan Ignacio Polanco <juan-ignacio.polanco@cnrs.fr> and contributors"]
-version = "0.18.1"
+version = "0.19.0"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

docs/make.jl

@@ -17,7 +17,7 @@ doctest(PencilArrays; fix = false)
 function main()
     makedocs(;
         modules = [PencilArrays],
-        authors = "Juan Ignacio Polanco <jipolanc@gmail.com> and contributors",
+        authors = "Juan Ignacio Polanco <juan-ignacio.polanco@cnrs.fr>",
         repo = "https://github.com/jipolanco/PencilArrays.jl/blob/{commit}{path}#L{line}",
         sitename = "PencilArrays.jl",
         format = Documenter.HTML(;

src/arrays.jl

@@ -205,17 +205,14 @@ end
 Base.axes(x::PencilArray) = permutation(x) \ axes(parent(x))
 
 """
-    similar(x::PencilArray, [element_type=eltype(x)], [dims])
+    similar(x::PencilArray, [element_type=eltype(x)], [dims]) -> PencilArray
 
-Returns an array similar to `x`.
+Returns a `PencilArray` similar to `x`.
 
-The actual type of the returned array depends on whether `dims` is passed:
-
-- if `dims` is *not* passed, then a `PencilArray` of same dimensions of `x` is
-  returned.
-
-- otherwise, an array similar to that wrapped by `x` (typically a regular
-  `Array`) is returned, with the chosen dimensions.
+In particular, the new array shares the same parallel decomposition (the same `Pencil`) than
+`x`. This means that the dimensions of the new array must be the same as those of `x`. Note
+that the optional `dims` argument is allowed for the sole reason of making things work
+nicely with other packages (such as StructArrays.jl), but things will fail if `dims ≠ size(x)`.
 
 # Examples
 
@@ -227,20 +224,23 @@ julia> u = PencilArray{Float64}(undef, pen);
 julia> similar(u) |> summary
 "20×10×12 PencilArray{Float64, 3}(::Pencil{3, 2, NoPermutation, Array})"
 
+julia> similar(u, size(u)) |> summary
+"20×10×12 PencilArray{Float64, 3}(::Pencil{3, 2, NoPermutation, Array})"
+
 julia> similar(u, ComplexF32) |> summary
 "20×10×12 PencilArray{ComplexF32, 3}(::Pencil{3, 2, NoPermutation, Array})"
 
-julia> similar(u, (4, 3, 8)) |> summary
-"4×3×8 Array{Float64, 3}"
+julia> similar(u, (4, 3, 8))
+ERROR: DimensionMismatch: cannot construct a similar PencilArray with different size
 
 julia> similar(u, (4, 3)) |> summary
-"4×3 Matrix{Float64}"
+ERROR: DimensionMismatch: cannot construct a similar PencilArray with different size
 
 julia> similar(u, ComplexF32) |> summary
 "20×10×12 PencilArray{ComplexF32, 3}(::Pencil{3, 2, NoPermutation, Array})"
 
-julia> similar(u, ComplexF32, (4, 3)) |> summary
-"4×3 Matrix{ComplexF32}"
+julia> similar(u, ComplexF32, (4, 3))
+ERROR: DimensionMismatch: cannot construct a similar PencilArray with different size
 ```
 
 ---
@@ -289,8 +289,11 @@ function Base.similar(x::PencilArray, ::Type{S}) where {S}
     PencilArray(x.pencil, similar(parent(x), S, dims_perm))
 end
 
-Base.similar(x::PencilArray, ::Type{S}, dims::Dims) where {S} =
-    similar(parent(x), S, dims)
+function Base.similar(x::PencilArray, ::Type{S}, dims::Dims) where {S}
+    dims == size(x) ||
+        throw(DimensionMismatch("cannot construct a similar PencilArray with different size"))
+    similar(x, S)
+end
 
 function Base.similar(x::PencilArray, ::Type{S}, p::Pencil) where {S}
     dims_mem = (size_local(p, MemoryOrder())..., extra_dims(x)...)

test/Project.toml

@@ -14,6 +14,8 @@ OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd"
 StaticArrayInterface = "0d7ed370-da01-4f52-bd93-41d350b8b718"
+StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
+StructArrays = "09ab397b-f2b6-538f-b94a-2f83cf4a842a"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [compat]

test/ode.jl

@@ -6,8 +6,16 @@ using MPI
 using PencilArrays
 using OrdinaryDiffEq
 using RecursiveArrayTools: ArrayPartition
+using StructArrays: StructArray
+using StaticArrays: SVector
 using Test
 
+function to_structarray(us::NTuple{N, A}) where {N, A <: AbstractArray}
+    T = eltype(A)
+    Vec = SVector{N, T}
+    StructArray{Vec}(us)
+end
+
 MPI.Init()
 
 comm = MPI.COMM_WORLD
@@ -25,14 +33,10 @@ u0 = PencilArray{Float64}(undef, pen)
 @. u0 = grid.x * grid.y + grid.z
 
 function rhs!(du, u, p, t)
-    @. du = 0.1 * u
+    @. du = -0.1 * u
     du
 end
 
-@static if !isdefined(Base, :allequal)
-    allequal(xs) = all(==(first(xs)), xs)
-end
-
 @testset "DiffEqBase" begin
     unorm = DiffEqBase.ODE_DEFAULT_NORM(u0, 0.0)
     unorms = MPI.Allgather(unorm, comm)
@@ -80,4 +84,21 @@ end
             adaptive = true, save_everystep = false,
         )
     end
+
+    # Solve the equation for a 2D vector field represented by a StructArray.
+    @testset "StructArray" begin
+        v0 = to_structarray((u0, 2u0))
+        @assert eltype(v0) <: SVector{2}
+        tspan = (0.0, 1.0)
+        prob = @inferred ODEProblem{true}(rhs!, v0, tspan, params)
+        integrator = init(
+            prob, Tsit5();
+            adaptive = true, save_everystep = false,
+        )
+        @test integrator.u == v0
+        for _ ∈ 1:10
+            step!(integrator)
+        end
+        @test integrator.u ≠ v0
+    end
 end

test/pencils.jl

@@ -107,17 +107,13 @@ function test_array_wrappers(p::Pencil, ::Type{T} = Float64) where {T}
             @test pencil(v) === pencil(u)
         end
 
-        let v = @inferred similar(u, (3, 4))
-            @test v isa Matrix
-            @test size(v) == (3, 4)
-            @test eltype(v) === eltype(u)
-        end
+        @test (@inferred similar(u, size(u))) isa PencilArray{T}
+        @test (@inferred similar(u, Int, size(u))) isa PencilArray{Int}
 
-        let v = @inferred similar(u, Int, (3, 4))
-            @test v isa Matrix
-            @test size(v) == (3, 4)
-            @test eltype(v) === Int
-        end
+        @test_throws DimensionMismatch similar(u, 12)
+        @test_throws DimensionMismatch similar(u, Int, 12)
+        @test_throws DimensionMismatch similar(u, size(u) .+ 1)
+        @test_throws DimensionMismatch similar(u, Int, size(u) .+ 1)
 
         let A = DummyArray{Int}(undef, size_local(p, MemoryOrder()))
             pdummy = Pencil(DummyArray, p)
@@ -127,14 +123,6 @@ function test_array_wrappers(p::Pencil, ::Type{T} = Float64) where {T}
             v = @inferred similar(u)
             @test typeof(v) === typeof(u)
             @test size(v) == size(u)
-
-            w = @inferred similar(u, (4, 2))
-            @test w isa DummyArray{Int,2}
-            @test size(w) == (4, 2)
-
-            z = @inferred similar(u, Float32, (3, 4, 2))
-            @test z isa DummyArray{Float32,3}
-            @test size(z) == (3, 4, 2)
         end
 
         # Test similar(u, [T], q::Pencil)