Skip to content
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
99 changes: 99 additions & 0 deletions src/DAT/xmap.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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))
Expand Down Expand Up @@ -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,))
Expand Down
Loading