Add better support for inplace vector fields (#4)

Author: dkarrasch

.github/workflows/CI.yml

@@ -0,0 +1,37 @@
+name: CI
+
+on:
+  push:
+    branches:
+      - master
+    tags: '*'
+  pull_request:
+
+jobs:
+  test:
+    name: Julia ${{ matrix.version }} - ${{ matrix.os }}
+    runs-on: ${{ matrix.os }}
+    continue-on-error: ${{ matrix.version == 'nightly' }}
+    strategy:
+      fail-fast: false
+      matrix:
+        version:
+          - '1.0'
+          - '1'
+          - 'nightly'
+        os:
+          - ubuntu-latest
+          - macOS-latest
+          - windows-latest
+    steps:
+      - uses: actions/checkout@v2
+      - uses: julia-actions/setup-julia@latest
+        with:
+          version: ${{ matrix.version }}
+      - uses: julia-actions/julia-buildpkg@latest
+      - uses: julia-actions/julia-runtest@latest
+      - uses: julia-actions/julia-processcoverage@v1
+      - uses: codecov/codecov-action@v1
+        if: ${{ matrix.version == '1' && matrix.os == 'ubuntu-latest' }}
+        with:
+          file: lcov.info

.travis.yml

@@ -1,7 +1,7 @@
 name = "StreamMacros"
 uuid = "1418a599-1564-4a75-98e9-b890cf8b552f"
-authors = ["Alvaro de Diego"]
-version = "0.1.2"
+authors = ["Alvaro de Diego", "Daniel Karrasch"]
+version = "0.2.0"
 
 [deps]
 DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
@@ -15,7 +15,8 @@ SymEngine = "0.6, 0.7, 0.8"
 julia = "1"
 
 [extras]
+InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test"]
+test = ["InteractiveUtils", "Test"]

Project.toml

@@ -1,30 +1,34 @@
 # StreamMacros.jl
 
-[![build status](https://travis-ci.org/CoherentStructures/StreamMacros.jl.svg?branch=master)](https://travis-ci.org/CoherentStructures/StreamMacros.jl)
-[![codecov.io](http://codecov.io/github/CoherentStructures/StreamMacros.jl/coverage.svg?branch=master)](http://codecov.io/github/CoherentStructures/StreamMacros.jl?branch=master)
-[![stable](https://img.shields.io/badge/docs-stable-blue.svg)](https://coherentstructures.github.io/StreamMacros.jl/stable)
-[![dev](https://img.shields.io/badge/docs-dev-blue.svg)](https://coherentstructures.github.io/StreamMacros.jl/dev)
+*Convenience macros for defining velocity fields via their stream function.*
 
-Provides some convenience macros for defining velocity fields via their stream function.
+[![build status](https://github.com/CoherentStructures/StreamMacros.jl/workflows/CI/badge.svg?branch=master)](https://github.com/CoherentStructures/StreamMacros.jl/actions?query=workflow%3ACI)
+[![code coverage](http://codecov.io/github/CoherentStructures/StreamMacros.jl/coverage.svg?branch=master)](http://codecov.io/github/CoherentStructures/StreamMacros.jl?branch=master)
+[![stable docs](https://img.shields.io/badge/docs-stable-blue.svg)](https://coherentstructures.github.io/StreamMacros.jl/stable)
+[![dev docs](https://img.shields.io/badge/docs-dev-blue.svg)](https://coherentstructures.github.io/StreamMacros.jl/dev)
+
+## Installation
 
-# Installation
 Install by typing
 
-    ]add https://github.com/CoherentStructures/StreamMacros.jl
+```julia
+]add https://github.com/CoherentStructures/StreamMacros.jl.git
+```
 
 in the Julia REPL.
 
-# Examples
+## Examples
+
 The basic usage pattern is
 
 ```julia
 output = @velo_from_stream <stream_name> begin
     <stream_name> = ...
-   	# additional definitions
+    # additional definitions
 end
 ```
 
-as in e.g.
+as in
 
 ```julia
 using StreamMacros

README.md

@@ -1,5 +1,6 @@
 [deps]
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
+StreamMacros = "1418a599-1564-4a75-98e9-b890cf8b552f"
 
 [compat]
-Documenter = "0.24"
+Documenter = "0.24, 0.25, 0.26"

docs/Project.toml

@@ -11,7 +11,9 @@ const ODE = DiffEqBase
 export
     @define_stream,
     @velo_from_stream,
+    @velo!_from_stream,
     @var_velo_from_stream,
+    @var_velo!_from_stream,
     @vorticity_from_stream
 
 include("macros.jl")

src/StreamMacros.jl

@@ -102,6 +102,31 @@ macro velo_from_stream(name::Symbol)
     end
 end
 
+"""
+    @velo!_from_stream(name::Symbol, [code::Expr])
+
+Same as [`@velo_from_stream`](@ref), but returns the in-place version of the velocity field.
+"""
+macro velo!_from_stream(H::Symbol, formulas::Expr)
+    V, _ = streamline_derivatives(H, formulas)
+    V = sym_subst.(V, [[:x, :y, :p]], [[:(u[1]), :(u[2]), :(p[1])]])
+    quote
+        ODE.ODEFunction{true}(
+            (du, u, p, t) -> begin
+                du[1] = $(V[1])
+                du[2] = $(V[2])
+                return du
+            end,
+        )
+    end
+end
+macro velo!_from_stream(name::Symbol)
+    haskey(stream_dict, name) || (@error "stream $name not defined")
+    quote
+        @velo!_from_stream $name $(stream_dict[name])
+    end
+end
+
 """
     @var_velo_from_stream(name::Symbol, [code::Expr])
 
@@ -167,6 +192,43 @@ macro var_velo_from_stream(name::Symbol)
     end
 end
 
+"""
+    @var_velo!_from_stream(name::Symbol, [code::Expr])
+
+Same as [`@var_velo_from_stream`](@ref), but returns the in-place version of the (extended)
+velocity field.
+"""
+macro var_velo!_from_stream(H::Symbol, formulas::Expr)
+    V, DV = streamline_derivatives(H, formulas)
+
+    # substitute :x and :y by access to the first column of a matrix u
+    V = sym_subst.(V, [[:x, :y, :p]], [[:(u[1, 1]), :(u[2, 1]), :(p[1])]])
+    DV = sym_subst.(DV, [[:x, :y, :p]], [[:(u[1, 1]), :(u[2, 1]), :(p[1])]])
+
+    quote
+        ODE.ODEFunction{true}(
+            (du, u, p, t) -> begin
+                # take as input a 2x3 matrix which is interpreted in the following way:
+                # [ x[1] X[1,1] X[1,2]
+                #   x[2] X[2,1] X[2,2] ]
+                du[1,1] = $(V[1])
+                du[2,1] = $(V[2])
+                du[1,2] = $(DV[1,1]) * u[1,2] + $(DV[1,2]) * u[2,2]
+                du[2,2] = $(DV[2,1]) * u[1,2] + $(DV[2,2]) * u[2,2]
+                du[1,3] = $(DV[1,1]) * u[1,3] + $(DV[1,2]) * u[2,3]
+                du[2,3] = $(DV[2,1]) * u[1,3] + $(DV[2,2]) * u[2,3]
+                return du
+            end,
+        )
+    end
+end
+macro var_velo!_from_stream(name::Symbol)
+    haskey(stream_dict, name) || (@error "stream $name not defined")
+    quote
+        @var_velo!_from_stream $(name) $(stream_dict[name])
+    end
+end
+
 """
     @vorticity_from_stream(name::Symbol, [code::Expr])
 
@@ -341,10 +403,6 @@ simple_simplifier(expr) = expr
 
 
 expr_grad(expr, coord_vars::Vector{Symbol}) = expr_diff.(expr, coord_vars)
-function hessian(expr, coord_vars)
-    ∇expr = expr_grad(expr, coord_vars)
-    ∇²expr = expr_grad(expr)
-end
 
 ####t####################################################################################
 #                 Functions for symbolic manipulation of expressions                    #

src/macros.jl

@@ -1,51 +1,44 @@
 using StreamMacros
+using InteractiveUtils
 using DiffEqBase
 using StaticArrays
 using Test
 
 @testset "StreamMacros.jl" begin
-    bickley = @velo_from_stream stream begin
-      stream = psi₀ + psi₁
-      psi₀   = - U₀ * L₀ * tanh(y / L₀)
-      psi₁   =   U₀ * L₀ * sech(y / L₀)^2 * re_sum_term
-
-      re_sum_term =  Σ₁ + Σ₂ + Σ₃
-
-      Σ₁  =  ε₁ * cos(k₁*(x - c₁*t))
-      Σ₂  =  ε₂ * cos(k₂*(x - c₂*t))
-      Σ₃  =  ε₃ * cos(k₃*(x - c₃*t))
-
-      k₁ = 2/r₀      ; k₂ = 4/r₀    ; k₃ = 6/r₀
-
-      ε₁ = 0.0075    ; ε₂ = 0.15    ; ε₃ = 0.3
-      c₂ = 0.205U₀   ; c₃ = 0.461U₀ ; c₁ = c₃ + (√5-1)*(c₂-c₃)
-
-      U₀ = 62.66e-6  ; L₀ = 1770e-3 ; r₀ = 6371e-3
+    du = zeros(2); u= rand(2); t = rand(); DU = zeros(2,3); U = rand(2,3)
+    @define_stream bstream begin
+        bstream = psi₀ + psi₁
+        psi₀   = - U₀ * L₀ * tanh(y / L₀)
+        psi₁   =   U₀ * L₀ * sech(y / L₀)^2 * re_sum_term
+  
+        re_sum_term =  Σ₁ + Σ₂ + Σ₃
+  
+        Σ₁  =  ε₁ * cos(k₁*(x - c₁*t))
+        Σ₂  =  ε₂ * cos(k₂*(x - c₂*t))
+        Σ₃  =  ε₃ * cos(k₃*(x - c₃*t))
+  
+        k₁ = 2/r₀      ; k₂ = 4/r₀    ; k₃ = 6/r₀
+  
+        ε₁ = 0.0075    ; ε₂ = 0.15    ; ε₃ = 0.3
+        c₂ = 0.205U₀   ; c₃ = 0.461U₀ ; c₁ = c₃ + (√5-1)*(c₂-c₃)
+  
+        U₀ = 62.66e-6  ; L₀ = 1770e-3 ; r₀ = 6371e-3
     end
-    @test bickley isa ODEFunction
+    bickley = @velo_from_stream bstream
+    bickley! = @velo!_from_stream bstream
+    @test bickley isa ODEFunction{false}
+    @test bickley! isa ODEFunction{true}
     @test bickley(SVector(0.0,0.0), nothing, 0.0) isa SVector
+    @test bickley(SVector(0.0,0.0), nothing, 0.0) == bickley!(ones(2), [0.0, 0.0], nothing, 0.0)
+    VERSION ≥ v"1.1" && @test iszero(@allocated bickley!(du, u, nothing, t))
 
-    bickley2 = @var_velo_from_stream stream begin
-      stream = psi₀ + psi₁
-      psi₀   = - U₀ * L₀ * tanh(y / L₀)
-      psi₁   =   U₀ * L₀ * sech(y / L₀)^2 * re_sum_term
-
-      re_sum_term =  Σ₁ + Σ₂ + Σ₃
-
-      Σ₁  =  ε₁ * cos(k₁*(x - c₁*t))
-      Σ₂  =  ε₂ * cos(k₂*(x - c₂*t))
-      Σ₃  =  ε₃ * cos(k₃*(x - c₃*t))
-
-      k₁ = 2/r₀      ; k₂ = 4/r₀    ; k₃ = 6/r₀
-
-      ε₁ = 0.0075    ; ε₂ = 0.15    ; ε₃ = 0.3
-      c₂ = 0.205U₀   ; c₃ = 0.461U₀ ; c₁ = c₃ + (√5-1)*(c₂-c₃)
-
-      U₀ = 62.66e-6  ; L₀ = 1770e-3 ; r₀ = 6371e-3
-    end
-    @test bickley2 isa ODEFunction
-    @test (bickley2(SMatrix{2,3}(0.0,0.0,0.0,0.0,0.0,0.0), nothing, 0.0)
-           isa SMatrix{2,3})
+    bickley2 = @var_velo_from_stream bstream
+    bickley2! = @var_velo!_from_stream bstream
+    @test bickley2 isa ODEFunction{false}
+    @test bickley2! isa ODEFunction{true}
+    @test (bickley2(SMatrix{2,3}(0.0,0.0,0.0,0.0,0.0,0.0), nothing, 0.0) isa SMatrix{2,3})
+    @test bickley2(SMatrix{2,3}(0.0,0.0,0.0,0.0,0.0,0.0), nothing, 0.0) == bickley2!(ones(2,3), zeros(2,3), nothing, 0.0)
+    VERSION ≥ v"1.1" && @test iszero(@allocated bickley2!(DU, U, nothing, t))
 
     @test StreamMacros.expr_diff(:(x^2 + sin(x) + y), :x) == :(2x + cos(x))
 
@@ -55,10 +48,13 @@ using Test
 
     t1 = @var_velo_from_stream test_stream
     t2 = @velo_from_stream test_stream
-    @test t1 isa ODEFunction
-    @test t2 isa ODEFunction
+    @test t1 isa ODEFunction{false}
+    @test t2 isa ODEFunction{false}
     ω = @vorticity_from_stream test_stream
     @test ω(rand(), rand(), rand()) == 4
+    # works only locally, not on github CI
+    # ct = code_typed(ω, (Float64, Float64, Float64))
+    # @test first(ct).first.code[1] == :(return 4)
 
     @define_stream Ψ_rot_dgyre begin
         st          = window(t, 0, 1)*t^2*(3-2*t) + heaviside(t-1)
@@ -69,6 +65,10 @@ using Test
         Ψ_rot_dgyre = (1-st) * Ψ_P + st * Ψ_F
     end
     rot_double_gyre = @velo_from_stream Ψ_rot_dgyre
+    rot_double_gyre! = @velo!_from_stream Ψ_rot_dgyre
+    @test rot_double_gyre(SVector{2}(0.0,0.0), nothing, 0.0) == rot_double_gyre!(ones(2), zeros(2), nothing, 0.0)
+    VERSION ≥ v"1.1" && @test iszero(@allocated rot_double_gyre!(du, u, nothing, t))
+
     dgyrepast = @velo_from_stream Ψ_P begin
         Ψ_P = sin(2π*x)*sin(π*y)
     end