Add reshape support for DiffCache resizing

Project.toml

@@ -1,6 +1,6 @@
 name = "PreallocationTools"
 uuid = "d236fae5-4411-538c-8e31-a6e3d9e00b46"
-version = "1.2.2"
+version = "1.3.0"
 authors = ["Chris Rackauckas <accounts@chrisrackauckas.com>"]
 
 [deps]

docs/src/index.md

@@ -55,6 +55,20 @@ the `levels` keyword (`N` for each level computed using based on the size of the
 state vector) or by specifying `N` as an array of integers of chunk sizes, which
 enables full control of chunk sizes on all differentiation levels.
 
+For resizeable multidimensional workspaces, keep the backing `DiffCache`
+one-dimensional and reshape it when a multidimensional cache is needed:
+
+```julia
+storage = DiffCache(zeros(10), 2)
+cache = reshape(storage, 2, 5)
+
+resize!(storage, 12)
+cache = reshape(storage, 2, 6)
+```
+
+The reshaped cache shares the normal cache storage with `storage`, while the
+dual cache storage remains vector-backed for `get_tmp`.
+
 ### DiffCache Example 1: Direct Usage
 
 ```julia

src/PreallocationTools.jl

@@ -164,6 +164,28 @@ function get_tmp(dc::DiffCache, ::Type{T}) where {T <: Number}
     end
 end
 
+"""
+    reshape(dc::DiffCache, dims...)
+    reshape(dc::DiffCache, dims)
+
+Return a `DiffCache` whose normal cache has shape `dims`.
+
+This is useful for vector-backed caches that need to be resized. Resize the
+backing `DiffCache` with `resize!`, then call `reshape` again with the updated
+dimensions. The returned cache shares the normal cache storage with `dc`; the
+raw dual cache storage remains vector-backed so `get_tmp` can reinterpret it for
+the requested automatic differentiation element type.
+"""
+function Base.reshape(dc::DiffCache, dims::Tuple{Vararg{Integer}})
+    shape = map(Int, dims)
+    return DiffCache(_resizeable_reshape(dc.du, shape), dc.dual_du, dc.any_du, dc.warn_on_resize)
+end
+
+Base.reshape(dc::DiffCache, dims::Integer...) = reshape(dc, dims)
+
+_resizeable_reshape(a::AbstractVector, shape) = reshape(view(a, :), shape)
+_resizeable_reshape(a::AbstractArray, shape) = reshape(a, shape)
+
 get_tmp(dc, u) = dc
 
 """
@@ -178,9 +200,18 @@ function _restructure(normal_cache::Array, duals)
 end
 
 function _restructure(normal_cache::AbstractArray, duals)
+    if _has_vector_view_parent(normal_cache)
+        return reshape(duals, size(normal_cache)...)
+    end
     return ArrayInterface.restructure(normal_cache, duals)
 end
 
+function _has_vector_view_parent(a)
+    a isa SubArray{<:Any, 1, <:AbstractVector} && return true
+    p = applicable(parent, a) ? parent(a) : nothing
+    return p isa SubArray{<:Any, 1, <:AbstractVector}
+end
+
 """
     enlargediffcache!(dc::DiffCache, nelem::Integer)
 
@@ -305,15 +336,13 @@ Base.getindex(b::GeneralLazyBufferCache, u::T) where {T} = get_tmp(b, u)
 function Base.resize!(dc::DiffCache, n::Integer)
     # Only resize if the array is a vector
     if dc.du isa AbstractVector
+        dual_length = length(dc.du) == 0 ? n : cld(length(dc.dual_du), length(dc.du)) * n
         resize!(dc.du, n)
     else
         throw(ArgumentError("resize! is only supported for DiffCache with vector arrays, got $(typeof(dc.du))"))
     end
-    # dual_du is often pre-allocated for ForwardDiff dual numbers,
-    # and may need special handling based on chunk size
-    # Only resize if it's a vector
     if dc.dual_du isa AbstractVector
-        resize!(dc.dual_du, n)
+        resize!(dc.dual_du, dual_length)
     end
     # Always resize the any_du cache
     resize!(dc.any_du, n)

test/core_resizing.jl

@@ -1,4 +1,5 @@
-using Test, PreallocationTools, ForwardDiff, LinearAlgebra
+using Test, PreallocationTools, ForwardDiff, LinearAlgebra, LabelledArrays,
+    RecursiveArrayTools
 
 #test for downsizing cache
 randmat = rand(5, 3)
@@ -56,7 +57,7 @@ analyticalsolution = [3.0 0; 0 0]
 # Test resize! functionality for DiffCache
 @testset "resize! for DiffCache" begin
     u = rand(10)
-    dc = DiffCache(u)
+    dc = DiffCache(u, 2)
 
     # Initial size
     @test length(dc.du) == 10
@@ -65,15 +66,127 @@ analyticalsolution = [3.0 0; 0 0]
     # Resize to larger
     resize!(dc, 20)
     @test length(dc.du) == 20
+    @test length(dc.dual_du) == 20 * 3
 
     # Resize to smaller
     resize!(dc, 5)
     @test length(dc.du) == 5
+    @test length(dc.dual_du) == 5 * 3
 
     # Test that it returns the cache itself
     @test resize!(dc, 8) === dc
 end
 
+@testset "reshape standard DiffCache arrays" begin
+    dc = DiffCache(collect(1.0:10.0), 2)
+    cache = reshape(dc, 2, 5)
+
+    normal_tmp = get_tmp(cache, 1.0)
+    @test size(normal_tmp) == (2, 5)
+    @test vec(normal_tmp) == collect(1.0:10.0)
+
+    normal_tmp[2, 3] = 42.0
+    @test dc.du[6] == 42.0
+
+    dual = ForwardDiff.Dual{Nothing}(1.0, 1.0, 0.0)
+    dual_tmp = get_tmp(cache, dual)
+    @test size(dual_tmp) == (2, 5)
+    @test eltype(dual_tmp) <: ForwardDiff.Dual
+    dual_tmp[1] = ForwardDiff.Dual{Nothing}(2.0, 3.0, 4.0)
+    @test dc.dual_du[1:3] == [2.0, 3.0, 4.0]
+
+    dual_type_tmp = get_tmp(cache, typeof(dual))
+    @test size(dual_type_tmp) == (2, 5)
+    @test eltype(dual_type_tmp) == typeof(dual)
+    dual_type_tmp[2] = ForwardDiff.Dual{Nothing}(5.0, 6.0, 7.0)
+    @test dc.dual_du[4:6] == [5.0, 6.0, 7.0]
+
+    dual_array_tmp = get_tmp(cache, fill(dual, 2, 5))
+    @test size(dual_array_tmp) == (2, 5)
+    @test eltype(dual_array_tmp) == typeof(dual)
+    dual_array_tmp[3] = ForwardDiff.Dual{Nothing}(8.0, 9.0, 10.0)
+    @test dc.dual_du[7:9] == [8.0, 9.0, 10.0]
+
+    resize!(dc, 12)
+    @test length(dc.dual_du) == 36
+
+    resized_cache = reshape(dc, (2, 6))
+    @test size(get_tmp(resized_cache, 1.0)) == (2, 6)
+    @test size(get_tmp(resized_cache, dual)) == (2, 6)
+
+    matrix_dc = DiffCache(zeros(2, 5), 2)
+    matrix_cache = reshape(matrix_dc, 5, 2)
+    matrix_dual_tmp = get_tmp(matrix_cache, dual)
+    @test size(matrix_dual_tmp) == (5, 2)
+    matrix_dual_tmp[1] = ForwardDiff.Dual{Nothing}(5.0, 6.0, 7.0)
+    @test matrix_dc.dual_du[1:3] == [5.0, 6.0, 7.0]
+end
+
+@testset "reshape vector-backed non-standard DiffCache arrays" begin
+    dual = ForwardDiff.Dual{Nothing}(1.0, 1.0, 0.0)
+
+    parent = collect(1.0:10.0)
+    view_dc = DiffCache(view(parent, 1:10), 2)
+    view_cache = reshape(view_dc, 2, 5)
+    view_normal_tmp = get_tmp(view_cache, 1.0)
+    @test size(view_normal_tmp) == (2, 5)
+    view_normal_tmp[2, 4] = 21.0
+    @test view_dc.du[8] == 21.0
+
+    view_dual_tmp = get_tmp(view_cache, dual)
+    @test size(view_dual_tmp) == (2, 5)
+    @test eltype(view_dual_tmp) == typeof(dual)
+    view_dual_tmp[4] = ForwardDiff.Dual{Nothing}(11.0, 12.0, 13.0)
+    @test view_dc.dual_du[10:12] == [11.0, 12.0, 13.0]
+
+    lvector_dc = DiffCache(LVector(a = 1.0, b = 2.0, c = 3.0, d = 4.0), 2)
+    lvector_cache = reshape(lvector_dc, 2, 2)
+    lvector_normal_tmp = get_tmp(lvector_cache, 1.0)
+    @test size(lvector_normal_tmp) == (2, 2)
+    @test vec(lvector_normal_tmp) == [1.0, 2.0, 3.0, 4.0]
+    lvector_normal_tmp[2, 2] = 31.0
+    @test lvector_dc.du[4] == 31.0
+
+    lvector_dual_tmp = get_tmp(lvector_cache, dual)
+    @test size(lvector_dual_tmp) == (2, 2)
+    @test eltype(lvector_dual_tmp) == typeof(dual)
+    lvector_dual_tmp[2] = ForwardDiff.Dual{Nothing}(14.0, 15.0, 16.0)
+    @test lvector_dc.dual_du[4:6] == [14.0, 15.0, 16.0]
+
+    partition_dc = DiffCache(ArrayPartition([1.0, 2.0], [3.0, 4.0, 5.0]), 2)
+    partition_cache = reshape(partition_dc, 5)
+    partition_normal_tmp = get_tmp(partition_cache, 1.0)
+    @test size(partition_normal_tmp) == (5,)
+    @test collect(partition_normal_tmp) == [1.0, 2.0, 3.0, 4.0, 5.0]
+    partition_normal_tmp[4] = 41.0
+    @test partition_dc.du[4] == 41.0
+
+    partition_dual_tmp = get_tmp(partition_cache, dual)
+    @test size(partition_dual_tmp) == (5,)
+    @test eltype(partition_dual_tmp) == typeof(dual)
+    partition_dual_tmp[3] = ForwardDiff.Dual{Nothing}(17.0, 18.0, 19.0)
+    @test partition_dc.dual_du[7:9] == [17.0, 18.0, 19.0]
+end
+
+@testset "reshape non-vector non-standard DiffCache arrays" begin
+    dual = ForwardDiff.Dual{Nothing}(1.0, 1.0, 0.0)
+
+    larray_dc = DiffCache(LArray((2, 2); a = 1.0, b = 2.0, c = 3.0, d = 4.0), 2)
+    larray_cache = reshape(larray_dc, 4)
+    larray_normal_tmp = get_tmp(larray_cache, 1.0)
+    @test size(larray_normal_tmp) == (4,)
+    @test collect(larray_normal_tmp) == [1.0, 2.0, 3.0, 4.0]
+    larray_normal_tmp[3] = 51.0
+    @test larray_dc.du[3] == 51.0
+
+    larray_dual_tmp = get_tmp(larray_cache, dual)
+    @test size(larray_dual_tmp) == (4,)
+    @test larray_dual_tmp isa LArray
+    @test eltype(larray_dual_tmp) == typeof(dual)
+    larray_dual_tmp[1] = ForwardDiff.Dual{Nothing}(20.0, 21.0, 22.0)
+    @test larray_dual_tmp[1] == ForwardDiff.Dual{Nothing}(20.0, 21.0, 22.0)
+end
+
 # Test warn_on_resize option
 @testset "warn_on_resize option" begin
     # Default: warn_on_resize = true

test/qa/Project.toml

@@ -1,4 +1,5 @@
 [deps]
+AllocCheck = "9b6a8646-10ed-4001-bbdc-1d2f46dfbb1a"
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 PreallocationTools = "d236fae5-4411-538c-8e31-a6e3d9e00b46"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
@@ -9,6 +10,7 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 PreallocationTools = {path = "../.."}
 
 [compat]
+AllocCheck = "0.2"
 Aqua = "0.8.11"
 SafeTestsets = "0.1, 1"
 SciMLTesting = "1.6"

test/qa/allocation_tests.jl

@@ -0,0 +1,31 @@
+using AllocCheck, PreallocationTools, Test
+
+_qa_reshape_varargs(cache, rows, cols) = reshape(cache, rows, cols)
+_qa_reshape_tuple(cache, dims) = reshape(cache, dims)
+_qa_get_tmp_number(cache, x) = get_tmp(cache, x)
+_qa_get_tmp_type(cache, ::Type{T}) where {T} = get_tmp(cache, T)
+_qa_first_tmp_value(cache, x) = get_tmp(cache, x)[1]
+
+function _qa_test_no_allocs(f, argtypes)
+    allocs = check_allocs(f, argtypes)
+    @test isempty(allocs)
+    return allocs
+end
+
+@testset "reshaped DiffCache" begin
+    storage = DiffCache(collect(1.0:10.0), 2)
+    reshaped = reshape(storage, 2, 5)
+
+    parent = collect(1.0:10.0)
+    view_storage = DiffCache(view(parent, 1:10), 2)
+    view_reshaped = reshape(view_storage, 2, 5)
+
+    _qa_test_no_allocs(_qa_reshape_varargs, Tuple{typeof(storage), Int, Int})
+    _qa_test_no_allocs(_qa_reshape_tuple, Tuple{typeof(storage), Tuple{Int, Int}})
+    _qa_test_no_allocs(_qa_reshape_varargs, Tuple{typeof(view_storage), Int, Int})
+
+    _qa_test_no_allocs(_qa_get_tmp_number, Tuple{typeof(reshaped), Float64})
+    _qa_test_no_allocs(_qa_get_tmp_type, Tuple{typeof(reshaped), Type{Float64}})
+    _qa_test_no_allocs(_qa_first_tmp_value, Tuple{typeof(reshaped), Float64})
+    _qa_test_no_allocs(_qa_get_tmp_number, Tuple{typeof(view_reshaped), Float64})
+end

test/qa/qa.jl

@@ -45,3 +45,7 @@ run_qa(
         ),
     ),
 )
+
+@testset "AllocCheck" begin
+    include("allocation_tests.jl")
+end