diff --git a/src/DAT/xmap.jl b/src/DAT/xmap.jl index d0c93d3b..9438b8a2 100644 --- a/src/DAT/xmap.jl +++ b/src/DAT/xmap.jl @@ -17,6 +17,39 @@ function DD._group_indices(dim::DD.Dimension, ::Whole; labels=nothing) look = DD.lookup(DD.format(DD.rebuild(dim,[first(dim) .. last(dim)]))) DD.rebuild(dim,look), [1:length(dim)] end + +""" + ⊘(a, b) + +A convenience operator, used as `a ⊘ b`, for creating a windowed view of a `DimArrayOrStack` where one or more dimensions are treated as a single, complete window. This is equivalent to calling `windows(a, b => Whole())`. + +This operator is particularly useful with `xmap`, where it specifies that a function should be applied over an entire dimension (or multiple dimensions) at once. For example, `xmap(mean, a ⊘ :time)` will compute the mean over the entire `:time` dimension. + +# How to type + +The `⊘` symbol can be typed in the Julia REPL or compatible editors by typing `\\oslash` and then pressing the Tab key. + +# Arguments +- `a`: A `YAXArray` or other `DimArrayOrStack`. +- `b`: A `Symbol` or a `Tuple` of `Symbol`s representing the dimension(s) to be treated as a single window. + +# Returns +A `DimWindowArray` which can be passed to `xmap` for windowed processing. + +# Examples +```julia +using YAXArrays, Dates, Statistics +using YAXArrays: YAXArrays as YAX + +a = YAXArray((YAX.time(1:5), lon(1:3)), rand(5,3)) + +# Create a windowed view where the `:time` dimension is a single window +w = a ⊘ :time + +# Use this view with xmap to calculate the mean over the time dimension. +time_mean = xmap(mean, w, inplace=false) +``` +""" ⊘(a,b::Tuple) = windows(a,map(Base.Fix2(Pair,Whole()),map(Symbol,b))...) ⊘(a,b) = ⊘(a,(b,)) windows(A::DimArrayOrStack) = DimWindowArray(A,DD.dims(A),map(d->1:length(d),DD.dims(A)),DD.dims(A)) @@ -306,6 +339,72 @@ function xmap(f, ars::Union{YAXArrays.Cubes.YAXArray,DimWindowArray}...; end end +""" + xmap(f, arrays::Union{YAXArray,DimWindowArray}...; output=XOutput(), inplace=default_inplace(f), + function_args=(), function_kwargs=()) + +Apply a function `f` across multiple YAXArrays or DimWindowArrays, with support for windowed operations and +dimension-aware broadcasting. + +# Arguments +- `f`: Function to apply. Can be a regular function or an `XFunction` +- `arrays`: One or more YAXArrays or DimWindowArrays to operate on + +# Keywords +- `output`: Single `XOutput` or tuple of `XOutput` specifying output dimensions and properties +- `inplace`: Whether the function `f` modifies its first argument (default depends on `f`) +- `function_args`: Additional positional arguments to pass to `f` +- `function_kwargs`: Additional keyword arguments to pass to `f` + +# Returns +- A single YAXArray if `output` is a single `XOutput` +- A tuple of YAXArrays if `output` is a tuple of `XOutput`s + +# Examples +```julia +using YAXArrays, Dates +using YAXArrays: YAXArrays as YAX + +# Create example arrays +axlist = ( + YAX.time(Date("2022-01-01"):Day(1):Date("2022-01-05")), + lon(1:3), + lat(1:2) +) +array1 = YAXArray(axlist, rand(5,3,2)) +array2 = YAXArray(axlist, rand(5,3,2)) + +# Element-wise arithmetic using broadcasting +result = array1 .+ array2 + +# Custom function with multiple outputs +function myfunc(xout1, xout2, x, y) + xout1 = x .+ y # Broadcasting with dot operator + xout2 = x .* y # Broadcasting with dot operator + return nothing +end + +out1 = XOutput(array1.time) +out2 = XOutput(array1.time) +sum_arr, prod_arr = xmap(myfunc, array1⊘:time, array2⊘:time, + output=(out1, out2), inplace=true) + +# Access computed values +computed_sum = sum_arr[:,:,:] # Now works correctly +computed_prod = prod_arr[:,:,:] # Now works correctly + +# Time series operation +using Statistics +time_mean = xmap(Statistics.mean, array1⊘:time, + output=XOutput(), inplace=false) + +# Save result directly to disk +ds = Dataset(result=time_mean) +compute_to_zarr(ds, "output.zarr") +``` +""" +function xmap end + import Base.mapslices function mapslices(f, d::YAXArray, addargs...; dims, kwargs...) !isa(dims, Tuple) && (dims = (dims,))