Improve DiskArrayEngine integration (#578)

* diskarrayengine updates

* bump DAE version

* Allow xmap computations on disconnected graphs

* custom loopranges

Author: meggart

Project.toml

@@ -1,7 +1,7 @@
 name = "YAXArrays"
 uuid = "c21b50f5-aa40-41ea-b809-c0f5e47bfa5c"
 authors = ["Fabian Gans <fgans@bgc-jena.mpg.de>"]
-version = "0.7.2"
+version = "0.7.3"
 
 [deps]
 CFTime = "179af706-886a-5703-950a-314cd64e0468"
@@ -36,7 +36,7 @@ YAXArrayBase = "90b8fcef-0c2d-428d-9c56-5f86629e9d14"
 CFTime = "0.0, 0.1, 0.2"
 DataStructures = "0.17, 0.18, 0.19"
 DimensionalData = "0.27, 0.28, 0.29"
-DiskArrayEngine = "0.3"
+DiskArrayEngine = "0.3.2"
 DiskArrayTools = "0.1.12"
 DiskArrays = "0.4.18"
 DocStringExtensions = "0.8, 0.9"

src/DAT/xmap.jl

@@ -181,6 +181,7 @@ intervals = MovingIntervals(:open, :closed; left=left_bounds, width=width)
 
 # Access the first interval
 first_interval = intervals[1]  # Interval(1, 3)
+```
 """
 struct MovingIntervals{T,O,L,R} <: DD.AbstractBins
     i1::T
@@ -618,7 +619,7 @@ Computes the YAXArrays dataset `ods` and saves it to a Zarr dataset at `path`.
 - `max_cache`: The maximum amount of data to cache in memory while computing the dataset.
 - `overwrite`: Whether to overwrite the dataset at `path` if it already exists.
 """
-function compute_to_zarr(ods, path; max_cache=5e8,overwrite=false)
+function compute_to_zarr(ods, path; max_cache=5e8, custom_loopranges=nothing, overwrite=false)
     if !isa(ods,Dataset)
         throw(ArgumentError("Direct saving of YAXArrays is not supported. Please wrap your array `a` into a Dataset by calling `Dataset(layer=a)`"))
     end
@@ -629,35 +630,49 @@ function compute_to_zarr(ods, path; max_cache=5e8,overwrite=false)
     end
     rpd = DAE.remove_aliases!(g)
     DAE.fuse_graph!(g)
-    op = DAE.gmwop_from_conn(only(g.connections), g.nodes)
-    lr = DAE.optimize_loopranges(op,max_cache)
-
-    newcubes = map(collect(keys(outnodes))) do k
-        looprange = lr.lr.members
-        lw = op.outspecs[1].lw
-        mylr = DAE.mysub(lw,looprange)
-        newcs = map(mylr,lw.windows.members) do mlr, w
-            map(mlr) do lr
-                DAE.windowmin(w[first(lr)]):DAE.windowmax(w[last(lr)]) |> length
-            end |> chunktype_from_chunksizes
-        end |> GridChunks
-        k=>YAXArrays.Datasets.setchunks(ods.cubes[k],newcs)
-    end
-    newds = Dataset(;newcubes...)
-
-    emptyds = savedataset(newds,path=path,skeleton=true, overwrite=true)
-    outars = Array{Any}(undef,length(op.outspecs))
-    fill!(outars,nothing)
-    for (k,v) in outnodes
-        v = get(rpd, v, v)
-        outars[v] = emptyds.cubes[k].data
+    opinfo = map(g.connections) do conn
+        op = DAE.gmwop_from_conn(conn, g.nodes)
+        lr = if custom_loopranges === nothing
+            DAE.optimize_loopranges(op, max_cache)
+        else
+            DAE.custom_loopranges(op, custom_loopranges)
+        end
+
+        newcubes = map(conn.outputids, op.outspecs) do oid, ospec
+            looprange = lr.lr.members
+            lw = ospec.lw
+            mylr = DAE.mysub(lw, looprange)
+            newcs = map(mylr, lw.windows.members) do mlr, w
+                map(mlr) do lr
+                    DAE.windowmin(w[first(lr)]):DAE.windowmax(w[last(lr)]) |> length
+                end |> chunktype_from_chunksizes
+            end |> GridChunks
+            k = findfirst(==(oid), outnodes)
+            if k === nothing
+                l = findall(==(oid), rpd)
+                k = findfirst(in(l), outnodes)
+            end
+            k => YAXArrays.Datasets.setchunks(ods.cubes[k], newcs)
+        end
+        op, lr, newcubes
     end
-    runner = if DAE.Distributed.nworkers() > 1
-        DAE.DaggerRunner(op,lr,outars)
-    else
-        DAE.LocalRunner(op,lr,outars)
+
+    newds = Dataset(; reduce(vcat, last.(opinfo))...)
+
+    emptyds = savedataset(newds, path=path, skeleton=true, overwrite=overwrite)
+
+    for (op, lr, newcubes) in opinfo
+
+        outars = map(newcubes) do (k, _)
+            emptyds.cubes[k].data
+        end
+        runner = if DAE.Distributed.nworkers() > 1
+            DAE.DaggerRunner(op, lr, outars)
+        else
+            DAE.LocalRunner(op, lr, outars)
+        end
+        run(runner)
     end
-    run(runner)
     emptyds
 end
 

test/DAT/xmap.jl

@@ -19,6 +19,20 @@
         #end
         @test r.data == a2.data .+ reshape(a3.data,(4,1,5))
     end
+    @testset "Disconnected Graphs" begin
+        using Test
+        using YAXArrays
+        a1 = YAXArray((Dim{:d1}(1:10),), 1:10)
+        a2 = YAXArray((Dim{:d2}(1:20),), 1:20)
+        outds = Dataset(a=a1 .* 2, b=a2 .+ 5)
+        @test outds.a.data isa YAXArrays.Xmap.DAE.GMWOPResult
+        @test outds.b.data isa YAXArrays.Xmap.DAE.GMWOPResult
+        outpath = tempname()
+        isfile(outpath)
+        dsdisk = compute_to_zarr(outds, outpath)
+        @test dsdisk.a[:] == 2:2:20
+        @test dsdisk.b[:] == 6:25
+    end
 end
 @testitem "set outtype" begin
     using YAXArrays