14 laplacian tests

.gitignore

@@ -15,19 +15,6 @@ deps/src/
 
 # Build artifacts for creating documentation generated by the Documenter package
 docs/build/
-docs/site/
+.DS_Store
 
-# File generated by Pkg, the package manager, based on a corresponding Project.toml
-# It records a fixed state of all packages used by the project. As such, it should not be
-# committed for packages, but should be committed for applications that require a static
-# environment.
 Manifest.toml
-
-# Python Env
-.venv
-env_setup.sh
-.python-version
-
-# Misc
-.DS_Store
-.vscode

src/GridTools.jl

@@ -19,6 +19,7 @@ export Dimension,
     VERTICAL,
     LOCAL,
     Field,
+    similar_field,
     Connectivity,
     FieldOffset,
     neighbor_sum,
@@ -207,6 +208,47 @@ function Field(
     return Field(Tuple(dim), data, Tuple(broadcast_dims), origin = origin)
 end
 
+"""
+    similar_field(f::Field, ::Type{T}=eltype(f.data), dims::Dims=size(f.data)) where {T}
+
+Create a new `Field` object similar to `f` but with a different element type `T` or different dimensions `dims`.
+
+# Arguments
+- `f::Field`: The original `Field` object to be replicated.
+- `T`: The element type of the new `Field` object (default: element type of `f.data`).
+- `dims::Dims`: The dimensions of the new `Field` object (default: size of `f.data`).
+
+# Returns
+- `Field`: A new `Field` object with the specified element type and dimensions.
+
+# Examples
+```julia-repl
+julia> A1Dim_ = Dimension(:A1_, LOCAL);
+
+julia> A1Dim = A1Dim_();
+
+julia> in_field = Field((A1Dim, A1Dim), 3*ones(Float64, 4, 4))
+4×4 Float64 Field with dimensions ("A1", "A1") with indices 1:4×1:4:
+ 3.0  3.0  3.0  3.0
+ 3.0  3.0  3.0  3.0
+ 3.0  3.0  3.0  3.0
+ 3.0  3.0  3.0  3.0
+
+julia> out_field = similar(in_field, Int);
+
+julia> size(out_field)
+(4, 4)
+
+julia> eltype(out_field)
+Int64
+
+```
+"""
+function similar_field(f::Field, ::Type{T}=eltype(f.data), dims::Dims=size(f.data)) where {T}
+    new_data = similar(f.data, T, dims)
+    return Field(f.dims, new_data, f.broadcast_dims, f.origin)
+end
+
 # TODO(tehrengruber): There is no need to have FieldOffset and FieldOffsetTS, remove FieldOffset
 struct FieldOffsetTS{
     Name,

src/gt2py/gt2py.jl

@@ -218,7 +218,8 @@ end
 function preprocess_definiton(expr::Expr, closure_vars::Dict)
     sat = single_assign_target_pass(expr)
     ucc = unchain_compairs_pass(sat)
-    ssa = single_static_assign_pass(ucc)
+    nop = canonicalize_arithmetic_ops(ucc)
+    ssa = single_static_assign_pass(nop)
     py_closure_vars = translate_closure_vars(closure_vars)
     annotations = get_annotation(ssa, py_closure_vars)
     return (ssa, py_closure_vars, annotations)

src/gt2py/jast_to_foast.jl

@@ -266,6 +266,7 @@ end
 
 function visit_(sym::Val{:call}, args::Array, outer_loc)
     if args[1] in bin_op
+        @assert length(args)==3 "Expected a binary operation. AST must be canonicalized using `canonicalize_arithmetic_ops` first."
         return foast.BinOp(
             op = visit(args[1]),
             left = visit(args[2], outer_loc),

src/gt2py/preprocessing.jl

@@ -28,6 +28,21 @@ function unchain_compairs_pass(expr::Expr)::Expr
     end
 end
 
+"""
+Transform all arithmetic operations with more than two arguments into binary operations.
+"""
+function canonicalize_arithmetic_ops(expr)
+    if expr isa Expr
+        expr = Expr(expr.head, map(canonicalize_arithmetic_ops, expr.args)...)  # visit all children
+
+        if expr.head == :call && expr.args[1] in bin_op && length(expr.args) > 3
+            op, a, b, tail... = expr.args
+            return canonicalize_arithmetic_ops(Expr(:call, op, Expr(:call, op, a, b), tail...))
+        end
+    end
+    return expr
+end
+
 function recursive_unchain(args::Array)::Expr
     if length(args) == 3
         return Expr(:call, args[2], args[1], args[3])