New poincare section using Roots and integrator stepping

old_vs_new_psos.jl

const PSOS_ERROR =
"the Poincaré surface of section did not have any points!"

function poincaresos2(ds::CDS{IIP, S, D}, plane, tfinal = 1000.0;
    direction = +1, Ttr::Real = 0.0, warning = true,
    diffeq...) where {IIP, S, D}

    integ = integrator(ds; diffeq...)

    planecrossing = PlaneCrossing{D}(plane, direction > 0 )

    psos = _poincaresos(integ, planecrossing, tfinal, Ttr)
    warning && length(psos) == 0 && warn(PSOS_ERROR)
    return psos
end

struct PlaneCrossing{D, P}
    plane::P
    dir::Bool
end
PlaneCrossing{D}(p::P, b) where {P, D} = PlaneCrossing{D, P}(p, b)
function (hp::PlaneCrossing{D, P})(u::AbstractVector) where {D, P<:Tuple}
    @inbounds x = u[hp.plane[1]] - hp.plane[2]
    hp.dir ? x : -x
end
function (hp::PlaneCrossing{D, P})(u::AbstractVector) where {D, P<:AbstractVector}
    x = zero(eltype(u))
    @inbounds for i in 1:D
        x += u[i]*hp.plane[i]
    end
    @inbounds x -= hp.plane[D+1]
    hp.dir ? x : -x
end

function _poincaresos(
    integ, planecrossing::PlaneCrossing{D}, tfinal, Ttr, atol = 1e-6, xrtol = 1e-6,
    ) where {D}

    Ttr != 0 && step!(integ, Ttr)

    psos = Dataset{D, eltype(integ.u)}()

    f = (t) -> planecrossing(integ(t))
    side = planecrossing(integ.u)

    while integ.t < tfinal + Ttr
        while side < 0
            integ.t > tfinal + Ttr && break
            step!(integ)
            side = planecrossing(integ.u)
        end
        while side > 0
            integ.t > tfinal + Ttr && break
            step!(integ)
            side = planecrossing(integ.u)
        end
        integ.t > tfinal + Ttr && break

        # I am now guaranteed to have `t` in negative and `tprev` in positive
        tcross = find_zero(f, (integ.tprev, integ.t), Bisection(),
                           xrtol = xrtol, atol = atol)

        ucross = integ(tcross)
        push!(psos.data, SVector{D}(ucross))

    end
    return psos
end

----------------------------------------------------------------------------------------------

# Benchmarks
ds = Systems.henonheiles([0., 0.1, 0.5, 0.])

# compile: 
output = poincaresos2(ds, (3, 0.0), 1000.0, Ttr = 200.0);
output = poincaresos(ds, (3, 0.0), 1000.0, Ttr = 200.0);

@btime  output = poincaresos2($ds, (3, 0.0), 1000.0, Ttr = 200.0);
  10.572 ms (91395 allocations: 2.42 MiB)

@btime  output = poincaresos($ds, (3, 0.0), 1000.0, Ttr = 200.0);
  21.430 ms (75122 allocations: 4.53 MiB)

# Benchmark compile time when given a new system
ds = Systems.lorenz()

julia> @time  output = poincaresos2(ds, (2, 0.0), 1000.0, Ttr = 200.0);
 30.234987 seconds (15.06 M allocations: 795.535 MiB, 2.56% gc time)

julia> @time  output = poincaresos(ds, (2, 0.0), 1000.0, Ttr = 200.0);
 32.305996 seconds (8.30 M allocations: 471.558 MiB, 2.11% gc time)