From d5dbb9ad59f89aa1ad17cc05849d84bc786b33cf Mon Sep 17 00:00:00 2001
From: meggart <fabiangans@gmx.de>
Date: Tue, 18 Oct 2022 16:42:27 +0200
Subject: [PATCH] Add AggregatedDiskArray

---
 Project.toml          |  4 +--
 src/DiskArrayTools.jl | 63 ++++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 64 insertions(+), 3 deletions(-)

diff --git a/Project.toml b/Project.toml
index 8064dea..a9c90a5 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,7 +1,7 @@
 name = "DiskArrayTools"
 uuid = "fcd2136c-9f69-4db6-97e5-f31981721d63"
 authors = ["Fabian Gans <fgans@bgc-jena.mpg.de>"]
-version = "0.1.5"
+version = "0.1.6"
 
 [deps]
 DiskArrays = "3c3547ce-8d99-4f5e-a174-61eb10b00ae3"
@@ -10,7 +10,7 @@ IterTools = "c8e1da08-722c-5040-9ed9-7db0dc04731e"
 OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"
 
 [compat]
-DiskArrays = "0.3"
+DiskArrays = "0.3.7"
 Interpolations = "0.12, 0.13, 0.14"
 IterTools = "1"
 OffsetArrays = "1"
diff --git a/src/DiskArrayTools.jl b/src/DiskArrayTools.jl
index c7555b9..940c9f2 100644
--- a/src/DiskArrayTools.jl
+++ b/src/DiskArrayTools.jl
@@ -1,7 +1,7 @@
 module DiskArrayTools
 import DiskArrays: AbstractDiskArray, eachchunk, haschunks, Chunked,
 estimate_chunksize, GridChunks, findints, readblock!, writeblock!, 
-RegularChunks, IrregularChunks, ChunkType, approx_chunksize
+RegularChunks, IrregularChunks, ChunkType, approx_chunksize, chunktype_from_chunksizes
 using Interpolations
 using IterTools: imap
 using Base.Iterators: product
@@ -151,6 +151,67 @@ function fremap(xout,xin,newinds,ipmeth;bc = Flat())
   end
 end
 
+function aggregate_chunks(cs,agg)
+  lcs = cumsum(length.(cs))
+  lagg = cumsum(length.(agg))
+  inds = [searchsortedfirst(lagg,lc) for lc in lcs]
+  chunktype_from_chunksizes(diff([0;inds]))
+end
+interpret_aggsize(_, x::ChunkType) = x
+interpret_aggsize(sparent,x) = RegularChunks(x,0,sparent)
+
+
+struct AggregatedDiskArray{T,N,A<:AbstractArray{T,N},I,F,CS<:Union{Nothing,GridChunks{N}}} <: ResampledDiskArray{T,N}
+  a::A
+  newinds::I
+  aggfun::F
+  chunksize::CS
+  dimkw::Bool
+end
+Base.size(a::AggregatedDiskArray) = map(length,a.newinds)
+eachchunk(a::AggregatedDiskArray) = a.chunksize
+function AggregatedDiskArray(parent,aggsizes,aggfun;chunksizes=nothing,dimkw=true)
+  newinds = GridChunks(interpret_aggsize.(size(parent),aggsizes)...)
+  if chunksizes === nothing
+      chunksizes = GridChunks(aggregate_chunks.(eachchunk(parent).chunks,newinds.chunks)...)
+  end
+  AggregatedDiskArray(parent,newinds.chunks,aggfun,chunksizes,dimkw)
+end
+
+
+get_readinds(::AggregatedDiskArray, parentranges) = map(r->first(first(r)):last(last(r)),parentranges)
+function resample_disk(a::AggregatedDiskArray,aout,atemp2,parentranges)
+  if a.dimkw
+    resample_disk_dimkw(a,aout,atemp2,parentranges)
+  else
+    resample_disk_loop(a,aout,atemp2,parentranges)
+  end
+end
+
+#Generic fallback
+function resample_disk_loop(a::AggregatedDiskArray,aout,atemp2,parentranges)
+  allr = collect(Iterators.product(parentranges...))
+  for ic in eachindex(allr)
+      ii = allr[ic...]
+      aout[ic...] = a.aggfun(atemp2[ii...])
+  end
+end
+
+function resample_disk_dimkw(a::AggregatedDiskArray,aout,atemp2,parentranges)
+  allpure = map(parentranges) do pr
+    all(i->length(i)==1,pr) ? [Colon()] : pr
+  end
+  alloutinds = map(allpure) do pr
+    isa(pr[1],Colon) ? pr : 1:length(pr)
+  end
+  dimkw = (findall(i->i[1] !== Colon(), allpure)...,)
+  allr = Iterators.product(allpure...)
+  outr = Iterators.product(alloutinds...)
+  for (ii,iout) in zip(allr, outr)
+      aout[iout...] = a.aggfun(atemp2[ii...],dims=dimkw)
+  end
+end
+
 
 
 #Use of Sentinel missing value