chained broadcasts

.github/workflows/CI.yml

@@ -4,7 +4,8 @@ on:
   push:
     branches:
       - master
-    tags: '*'
+    tags:
+      - '*'
 jobs:
   test:
     name: Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }} - ${{ github.event_name }}
@@ -39,10 +40,12 @@ jobs:
             ${{ runner.os }}-
       - uses: julia-actions/julia-buildpkg@v1
       - uses: julia-actions/julia-runtest@v1
+        with:
+          coverage: true
         env:
           JULIA_NUM_THREADS: '4'
       - uses: julia-actions/julia-processcoverage@v1
       - uses: codecov/codecov-action@v5
         with:
           token: ${{ secrets.CODECOV_TOKEN}}
-          file: lcov.info
+          files: lcov.info

docs/Project.toml

@@ -33,8 +33,8 @@ YAXArrayBase = "90b8fcef-0c2d-428d-9c56-5f86629e9d14"
 YAXArrays = "c21b50f5-aa40-41ea-b809-c0f5e47bfa5c"
 Zarr = "0a941bbe-ad1d-11e8-39d9-ab76183a1d99"
 
+[sources]
+YAXArrays = {path = ".."}
+
 [compat]
 DocumenterVitepress = "0.2"
-
-[sources]
-YAXArrays = { path = ".." }

docs/src/UserGuide/read.md

@@ -42,7 +42,7 @@ using YAXArrays
 using NetCDF
 using Downloads: download
 
-path = download("https://www.unidata.ucar.edu/software/netcdf/examples/tos_O1_2001-2002.nc", "example.nc")
+path = download("https://archive.unidata.ucar.edu/software/netcdf/examples/tos_O1_2001-2002.nc", "example.nc")
 ds = open_dataset(path)
 ````
 

docs/src/UserGuide/select.md

@@ -8,7 +8,7 @@ using YAXArrays
 using NetCDF
 using Downloads: download
 
-path = download("https://www.unidata.ucar.edu/software/netcdf/examples/tos_O1_2001-2002.nc", "example.nc")
+path = download("https://archive.unidata.ucar.edu/software/netcdf/examples/tos_O1_2001-2002.nc", "example.nc")
 ds = open_dataset(path)
 ````
 

docs/src/UserGuide/write.md

@@ -7,7 +7,7 @@ using YAXArrays
 using NetCDF
 using Downloads: download
 
-path = download("https://www.unidata.ucar.edu/software/netcdf/examples/tos_O1_2001-2002.nc", "example.nc")
+path = download("https://archive.unidata.ucar.edu/software/netcdf/examples/tos_O1_2001-2002.nc", "example.nc")
 ds = open_dataset(path)
 ````
 

src/DAT/broadcast.jl

@@ -2,10 +2,23 @@ struct XStyle <: Broadcast.BroadcastStyle end
 Base.BroadcastStyle(::Broadcast.AbstractArrayStyle, ::XStyle) = XStyle()
 Base.BroadcastStyle(::XStyle, ::Broadcast.AbstractArrayStyle) = XStyle()
 Base.BroadcastStyle(::Type{<:YAXArray}) = XStyle()
-to_yax(x::Number) = YAXArray((),fill(x))
+to_yax(x::Number) = YAXArray((), fill(x))
 to_yax(x::DD.AbstractDimArray) = x
-
 function Base.broadcasted(::XStyle, f, args...)
-    args2 = map(to_yax,args)
-    xmap(XFunction(f,inplace=false),args2...)
+    return Broadcast.Broadcasted{XStyle}(f, args)
+end
+function Base.materialize(bc::Broadcast.Broadcasted{XStyle})
+    args2 = map(arg -> arg isa Broadcast.Broadcasted ? Base.materialize(arg) : arg, bc.args)
+    args2 = map(to_yax, args2)
+    # determine output type by calling `eltype` on a dummy function call
+    dummy_args = map(a -> first(a.data), args2)
+    outtype = typeof(bc.f(dummy_args...))
+    return xmap(XFunction(bc.f; inplace=false), args2..., output=XOutput(; outtype))
+end
+function Base.materialize!(bc::Broadcast.Broadcasted{XStyle})
+    args2 = map(arg -> arg isa Broadcast.Broadcasted ? Base.materialize(arg) : arg, bc.args)
+    args2 = map(to_yax, args2)
+    dummy_args = map(a -> first(a.data), args2)
+    outtype = typeof(bc.f(dummy_args...))
+    return xmap(XFunction(bc.f; inplace=true), args2..., output=XOutput(; outtype))
 end

test/DAT/broadcast.jl

@@ -0,0 +1,63 @@
+using Test
+using YAXArrays
+
+function sample_arrays()
+    a = YAXArray(ones(3,4))
+    b = YAXArray(zeros(3,4))
+    c = YAXArray([1.0  NaN 3.0; 4.0 5.0 NaN])
+    return a, b, c
+end
+
+a, b, c = sample_arrays()
+
+@testset "YAXArray lazy broadcast tests" begin
+
+    # Scalar broadcasting
+    x1 = a .+ 1
+    @test eltype(x1) == eltype(a)
+    @test size(x1) == size(a)
+    @test all(x1[:] .== 2.0)
+
+    x2 = 1 .+ a
+    @test all(x2[:] .== 2.0)
+
+    # Element-wise addition
+    xadd = a .+ b
+    @test all(xadd[:] .== 1.0)
+
+    # Element-wise multiplication
+    xmul = a .* b
+    @test all(xmul[:] .== 0.0)
+
+    # Unary functions
+    xneg = -a
+    @test all(xneg[:] .== -1.0)
+
+    xabs = abs.(xneg)
+    @test all(xabs[:] .== 1.0)
+
+    # Logical / predicates
+    xisnan = isnan.(c)
+    @test eltype(xisnan) == Bool
+    @test xisnan[1,2] == true
+    @test xisnan[2,3] == true
+    @test xisnan[1,1] == false
+
+    xnotnan = .!isnan.(c)
+    @test xnotnan[1,1] == true
+    @test xnotnan[1,2] == false
+
+    # Mixed operations
+    xmix = (a .+ b) .* 2 .- 1
+    @test all(xmix[:] .== 1.0)
+
+    # Chained broadcasts
+    xchain = .!isnan.(c .+ 1)
+    @test xchain[1,2] == false
+    @test xchain[1,1] == true
+
+    # Mixed operations with numbers
+    xscalar = a .* 3 .+ 1
+    @test all(xscalar[:] .== 4.0)
+    @test isa(a .+ b, YAXArray)
+end