Support iterable of points

src/FirstFinDiffPlans.jl

@@ -12,9 +12,9 @@ struct FirstFinDiffPlan{T, D} end
 FirstFinDiffPlan{T, D}(θ, N) where {T, D} = FirstFinDiffPlan{T, D}()
 
 function compute!(
-    pinv::FirstPoincareInvariant{T, D, <:Any, <:FirstFinDiffPlan}, zs, t, p
+    pinv::FirstPoincareInvariant{T, D, <:Any, <:FirstFinDiffPlan}, t, p
 ) where {T, D}
-
+    zs = pinv.points
     N = getpointnum(pinv)
     θ = getform(pinv)
     I = zero(T)

src/FirstFourierPlans.jl

@@ -31,9 +31,10 @@ function FirstFourierPlan{T, D}(θ, N) where {T, D}
 end
 
 function compute!(
-    pinv::FirstPoincareInvariant{T, D, <:Any, <:FirstFourierPlan}, zs, t, p
+    pinv::FirstPoincareInvariant{T, D, <:Any, <:FirstFourierPlan}, t, p
 ) where {T, D}
 
+    zs = pinv.points
     N = getpointnum(pinv)
     plan = getplan(pinv)
 

src/PoincareInvariants.jl

@@ -17,6 +17,10 @@ export SecondChebyshevPlan, SecondFinDiffPlan
 
 export canonical_one_form, CanonicalSymplecticMatrix, canonical_two_form
 
+include("utils.jl")
+
+## AbstractPoincareInvariant ##
+
 """
     AbstractPoincareInvariant{T, D}
 
@@ -25,12 +29,84 @@ T for calculations.
 """
 abstract type AbstractPoincareInvariant{T, D} end
 
+# these types will be accepted as a point cloud representing an area or line
+# Any other type will get fed to a PoincareInvariantIter constructor
+# The rows of an AbstractMatrix are points and
+# the elements of the AbstractVector{AbstractVector} are points
+const PointCollection{T} = Union{AbstractMatrix{T}, AbstractVector{AbstractVector{T}}}
+
+checkpoints(T::Type) = @argcheck T <: PointCollection "Points must be a matrix or vector of vectors"
+
+writepoints!(pinv::AbstractPoincareInvariant, points::AbstractMatrix) = (pinv.points .= points)
+function writepoints!(pinv::AbstractPoincareInvariant, points::AbstractVector{AbstractVector})
+    for i in 1:length(points)
+        pinv.points[:, i] .= points[i]
+    end
+end
+
 """
     compute!(pinv::AbstractPoincareInvariant, args...)
 
 computes a Poincaré invariant.
+
+implementations should define a method `compute!(pinv, t, p)`, which acts on the internal
+points storage.
 """
-function compute! end
+function compute!(
+    pinv::AbstractPoincareInvariant,
+    points::PointCollection,
+    t::Real, p
+)
+    writepoints!(pinv, points)
+    compute!(pinv, t, p)
+end
+
+function compute!(pinv::AbstractPoincareInvariant, data, times::AbstractVector, p)
+    _compute_iter(pinv, data, times, p, Base.IteratorEltype(data))
+end
+
+function _compute_iter(pinv, data, times, p, ::Base.HasEltype)
+    checkpoints(eltype(data))
+    map(enumerate(data)) do (i, points)
+        writepoints!(pinv, points)
+        compute!(pinv, times[i], p)
+    end
+end
+
+function _compute_iter(pinv, data, times, p, ::Base.EltypeUnknown)
+    map(enumerate(data)) do (i, points)
+        checkpoints(typeof(points))
+        writepoints!(pinv, points)
+        compute!(pinv, times[i], p)
+    end
+end
+
+#=
+function compute!(picache::ThreadCache{<:AbstractPoincareInvariant{T}}, data, ts, p) where T
+    n = datalength(data)
+    lk = ReentrantLock()
+
+    out = Vector{T}(undef, n)
+    @threads for i in 1:n
+        # each thread gets one pinv
+        pinv = picache[Threads.threadid()]
+
+        # lock in case data has mutable state
+        lock(_ -> writepoints!(pinv, data, i), lk)
+
+        out[i] = compute!(pinv, ts[i], p)
+    end
+
+    out
+end
+
+function compute!(pinv::AbstractPoincareInvariant, data, t, p)
+    iter = PoincareInvariantIter(pinv, pointsiter, args...)
+    map(iter) do (points, t, p)
+        compute!(pinv, points, t, p)
+    end
+end
+=#
 
 """
     getplan(pinv::AbstractPoincareInvariant)
@@ -80,9 +156,8 @@ get invariant one- or two-form.
 function getform end
 
 
-## Utils ##
+## Canonical Forms ##
 
-include("utils.jl")
 include("CanonicalSymplecticForms.jl")
 
 using .CanonicalSymplecticForms: canonical_one_form, CanonicalSymplecticMatrix,
@@ -95,19 +170,22 @@ struct FirstPoincareInvariant{T, D, θT, P} <: AbstractPoincareInvariant{T, D}
     θ::θT
     N::Int
     plan::P
+    points::Matrix{T}
+end
+
+function FirstPoincareInvariant{T, D}(θ::θT, N::Integer, plan::P) where {T, D, θT, P}
+    n = getpointspec(N, P) |> getpointnum
+    points = Matrix{T}(undef, n, D)
+    FirstPoincareInvariant{T, D, θT, P}(θ, n, plan, points)
 end
 
 function FirstPoincareInvariant{T, D}(
     θ::θT, N::Integer, P::Type=DEFAULT_FIRST_PLAN
 ) where {T, D, θT}
-    ps = getpointspec(N, P)
-    plan = P{T, D}(θ, N)
-    FirstPoincareInvariant{T, D, θT, typeof(plan)}(θ, ps, plan)
+    plan = P{T, D}(θ, getpointspec(N, P))
+    FirstPoincareInvariant{T, D}(θ, N, plan)
 end
 
-FirstPoincareInvariant{T, D}(θ::θT, N::Integer, plan::P) where {T, D, θT, P} =
-    FirstPoincareInvariant{T, D, θT, P}(θ, N, plan)
-
 function FirstPoincareInvariant{T, D, typeof(canonical_one_form)}(
     N, P=DEFAULT_FIRST_PLAN
 ) where {T, D}
@@ -147,22 +225,24 @@ struct SecondPoincareInvariant{
     ω::ωT  # symplectic matrix or function returning one
     pointspec::PS  # specifies how many points
     plan::P  # plan for chebyshev transform, differentiation, etc...
+    points::Matrix{T}
 end
 
-function SecondPoincareInvariant{T, D}(
-    ω::ωT, N, P::Type=DEFAULT_SECOND_PLAN
-) where {T, D, ωT}
+function SecondPoincareInvariant{T, D}(ω::ωT, N, plan::P) where {T, D, ωT, P}
     ps = getpointspec(N, P)
-    plan = P{T, D}(ω, ps)
-    SecondPoincareInvariant{T, D, ωT, typeof(ps), typeof(plan)}(ω, ps, plan)
+    points = Matrix{T}(undef, getpointnum(ps), D)
+    SecondPoincareInvariant{T, D, ωT, typeof(ps), P}(ω, ps, plan, points)
 end
 
-function SecondPoincareInvariant{T, D}(ω::ωT, ps::PS, plan::P) where {T, D, ωT, PS, P}
-    SecondPoincareInvariant{T, D, ωT, PS, P}(ω, ps, plan)
+function SecondPoincareInvariant{T, D}(
+    ω, N, P::Type=DEFAULT_SECOND_PLAN
+) where {T, D}
+    plan = P{T, D}(ω, getpointspec(N, P))
+    SecondPoincareInvariant{T, D}(ω, N, plan)
 end
 
 function SecondPoincareInvariant{T, D, CanonicalSymplecticMatrix{T}}(
-    N, P=DEFAULT_SECOND_PLAN
+    N, P::Type=DEFAULT_SECOND_PLAN
 ) where {T, D}
     ω = CanonicalSymplecticMatrix{T}(D)
     SecondPoincareInvariant{T, D}(ω, N, P)
@@ -181,7 +261,7 @@ getpoints(pinv::SecondPoincareInvariant) = getpoints((x, y) -> (x, y), pinv)
 
 getpointspec(pinv::SecondPoincareInvariant) = pinv.pointspec
 
-## Implementations
+# Implementations
 
 include("SecondChebyshevPlans.jl")
 include("SecondFinDiffPlans.jl")

src/SecondChebyshevPlans.jl

@@ -95,7 +95,7 @@ end
 
 function getintegrand!(
     intcoeffs::AbstractMatrix, plan::CallIntPlan{T, D}, ω::ωT,
-    phasepoints, t, p, ∂xcoeffs, ∂ycoeffs
+    points, t, p, ∂xcoeffs, ∂ycoeffs
 ) where {T, D, ωT}
     invpaduatransform!(plan.∂xvals, plan.invpaduaplan, ∂xcoeffs)
     invpaduatransform!(plan.∂yvals, plan.invpaduaplan, ∂ycoeffs)
@@ -105,7 +105,7 @@ function getintegrand!(
         if ωT <: AbstractMatrix
             ωi = ω
         else
-            pnti = view(phasepoints, i, :)
+            pnti = view(points, i, :)
             ωi = ω(pnti, t, p)
         end
 
@@ -154,17 +154,18 @@ function SecondChebyshevPlan{T, D}(ω, N::Integer) where {T, D}
 end
 
 function compute!(
-    pinv::SecondPoincareInvariant{<:Any, <:Any, <:Any, <:Any, P}, phasepoints, t, p
+    pinv::SecondPoincareInvariant{<:Any, <:Any, <:Any, <:Any, P}, t, p
 ) where P <: SecondChebyshevPlan
     plan = pinv.plan
-    paduatransform!(plan.phasecoeffs, plan.paduaplan, phasepoints)
+    points = pinv.points
+    paduatransform!(plan.phasecoeffs, plan.paduaplan, points)
     differentiate!(plan.∂x, plan.∂y, plan.diffplan, plan.phasecoeffs)
-    getintegrand!(plan.intcoeffs, plan.intplan, pinv.ω, phasepoints, t, p, plan.∂x, plan.∂y)
+    getintegrand!(plan.intcoeffs, plan.intplan, pinv.ω, points, t, p, plan.∂x, plan.∂y)
     integrate(plan.intcoeffs, plan.intweights)
 end
 
-# function _compute!(plan::SecondChebyshevPlan, ω::AbstractMatrix, D, phasepoints)
-#     paduatransform!(plan.phasecoeffs, plan.paduaplan, phasepoints)
+# function _compute!(plan::SecondChebyshevPlan, ω::AbstractMatrix, D, points)
+#     paduatransform!(plan.phasecoeffs, plan.paduaplan, points)
 #     differentiate!(plan.∂x, plan.∂y, plan.diffplan, plan.phasecoeffs)
 #     getintegrand!(plan.intcoeffs, plan.intplan, ω, plan.∂x, plan.∂y)
 #     integrate(plan.intcoeffs, plan.intplan)

src/SecondFinDiffPlans.jl

@@ -14,17 +14,19 @@ function SecondFinDiffPlan{T, D}(ω, ps::NTuple{2, Int}) where {T, D}
 end
 
 function compute!(
-    pinv::SecondPoincareInvariant{T, D, ωT, <:Any, P}, vals, t, p
+    pinv::SecondPoincareInvariant{T, D, ωT, <:Any, P}, t, p
 ) where {T, D, ωT, P <: SecondFinDiffPlan}
     nx, ny = pinv.pointspec
-    @argcheck size(vals) == (nx * ny, D) "Expected points mtarix to have size $((nx * ny, D))"
+    points = pinv.points
+
+    @argcheck size(points) == (nx * ny, D) "Expected points mtarix to have size $((nx * ny, D))"
 
     I = zero(T)
 
     ∂xi = Vector{T}(undef, D)
     ∂yi = Vector{T}(undef, D)
 
-    colviews = [view(vals, :, d) for d in 1:D]
+    colviews = [view(points, :, d) for d in 1:D]
 
     i = 1
     for ix in 1:nx
@@ -37,7 +39,7 @@ function compute!(
             if ωT <: AbstractMatrix
                 ωi = pinv.ω
             else
-                pnti = view(vals, i, :)
+                pnti = view(points, i, :)
                 ωi = pinv.ω(pnti, t, p)
             end
 

test/runtests.jl

@@ -10,6 +10,6 @@ using SafeTestsets, Test
     @safetestset "SecondFinDiffPlans" begin include("test_SecondFinDiffPlans.jl") end
 end
 
-@safetestset "PoincareInvariants Function Tests" begin
+@safetestset "PoincareInvariants" begin
     include("test_PoincareInvariants.jl")
 end

test/test_PoincareInvariants.jl

@@ -16,6 +16,8 @@
         end
         @test pnts isa Matrix{Float64}
         @test size(pnts) == (getpointnum(pinv), D)
+
+        @test size(getpoints(pinv)::Vector{Float64}) == (getpointnum(pinv),)
     end
 
     @safetestset "SecondPoincareInvariant" begin
@@ -35,6 +37,8 @@
         end
         @test pnts isa Matrix{Float64}
         @test size(pnts) == (getpointnum(pinv), D)
+
+        @test size(getpoints(pinv)::Matrix{Float64}) == (getpointnum(pinv), 2)
     end
 end
 
@@ -218,6 +222,47 @@ end
     end
 end
 
+@safetestset "Iterable of Points" begin
+    using PoincareInvariants
+
+    D = 4
+    N = 398
+    Nsteps = 13
+    times = [2.0 + n * 0.15 for n in 0:Nsteps-1]
+
+    f1(z, t, p) = canonical_one_form(z, t, p) .* t
+    pi1 = FirstPI{Float64, D}(f1, N)
+
+    f2(z, t, p) = canonical_two_form(z, t, p) .* t
+    pi2 = SecondPI{Float64, D}(f2, N)
+
+    p1(θ) = (cospi(2θ), cospi(2θ), sinpi(2θ), sinpi(2θ))
+    p2(x, y) = (x, x, y, y)
+
+    testIs1 = [-2π * t for t in times]
+    testIs2 = [2 * t for t in times]
+
+    let data = [getpoints(p1, pi1) for t in times]
+        Is = compute!(pi1, data, times, nothing)
+        @test testIs1 ≈ Is atol=1e-13
+    end
+
+    let data = [getpoints(p2, pi2) for t in times]
+        Is = compute!(pi2, data, times, nothing)
+        @test testIs2 ≈ Is atol=1e-13
+    end
+
+    let data = (getpoints(p1, pi1) for t in times)
+        Is = compute!(pi1, data, times, nothing)
+        @test testIs1 ≈ Is atol=1e-13
+    end
+
+    let data = (getpoints(p2, pi2) for t in times)
+        Is = compute!(pi2, data, times, nothing)
+        @test testIs2 ≈ Is atol=1e-13
+    end
+end
+
 @safetestset "Linear Symplectic Maps" begin
     @safetestset "In 2D" begin
         using PoincareInvariants
@@ -398,3 +443,40 @@ end
     g3(x, y) = init2(x, y) |> f |> f |> f
     # neither of the two methods manage this
 end
+
+@safetestset "Pendulum" begin
+    using PoincareInvariants
+
+    # m, g, L = 1
+    function update(points::AbstractMatrix, dt)
+        out = copy(points)
+        for pnt in eachrow(out)
+            pnt[1] += dt * pnt[2]
+            pnt[2] -= dt * sin(pnt[1])
+        end
+        out
+    end
+
+    integrate(initial::AbstractMatrix, dt, n) = accumulate(1:n, init=initial) do points, _
+        update(points, dt)
+    end
+
+    init1(θ) = 2 .* sincospi(2θ)
+    I1 = 4π
+
+    init2(x, y) = (π * x - π/2, 4 * y - 2)
+    I2 = 4π
+
+    pi1 = CanonicalFirstPI{Float64, 2}(1_000)
+    pi2 = CanonicalSecondPI{Float64, 2}(1_00)
+
+    dt = 0.05
+    tmax = 2
+    times = 0:dt:tmax
+
+    data1 = integrate(getpoints(init1, pi1), dt, length(times))
+    data2 = integrate(getpoints(init2, pi2), dt, length(times))
+
+    @test maximum(abs, I1 .- compute!(pi1, data1, times, nothing)) < 1e-14
+    @test maximum(abs, I2 .- compute!(pi2, data2, times, nothing)) < 1e-3
+end