changeset 1273:7fab13c07412 refactor/grids

Allow multi-argument functions in eval_on
author Jonatan Werpers <jonatan@werpers.com>
date Sat, 25 Feb 2023 22:57:00 +0100
parents 3637daad71e8
children dcd8654ca33b
files src/Grids/grid.jl test/Grids/grid_test.jl
diffstat 2 files changed, 9 insertions(+), 5 deletions(-) [+]
line wrap: on
line diff
--- a/src/Grids/grid.jl	Sat Feb 25 22:49:43 2023 +0100
+++ b/src/Grids/grid.jl	Sat Feb 25 22:57:00 2023 +0100
@@ -54,8 +54,13 @@
 * Mention map(f,g) if you want a concrete array
 """
 eval_on(g::Grid, f) = eval_on(g, f, Base.IteratorSize(g)) # TBD: Borde f vara först som i alla map, sum, och dylikt
-eval_on(g::Grid, f, ::Base.HasShape) = LazyTensors.LazyFunctionArray((I...)->f(g[I...]), size(g))
-
+function eval_on(g::Grid, f, ::Base.HasShape)
+    if hasmethod(f, (Any,))
+        return LazyTensors.LazyFunctionArray((I...)->f(g[I...]), size(g))
+    else
+        return LazyTensors.LazyFunctionArray((I...)->f(g[I...]...), size(g))
+    end
+end
 
 """
     getcomponent(gfun, I::Vararg{Int})
--- a/test/Grids/grid_test.jl	Sat Feb 25 22:49:43 2023 +0100
+++ b/test/Grids/grid_test.jl	Sat Feb 25 22:57:00 2023 +0100
@@ -39,10 +39,9 @@
     @test eval_on(g, x̄ -> @SVector[x̄[2], x̄[1]]) isa LazyArray{SVector{2,Float64}}
     @test eval_on(g, x̄ -> @SVector[x̄[2], x̄[1]]) == map(x̄ -> @SVector[x̄[2], x̄[1]], g)
 
-
-
+    # Multi-argument functions
     f(x,y) = sin(x)*cos(y)
-    @test_broken eval_on(g, f) == map(p->f(p...), points(g))
+    @test eval_on(g, f) == map(x̄->f(x̄...), g)
 end
 
 @testset "getcomponent" begin