swept volume computation for GridVAMP

robot.jl

using LazySets: UnionSetArray, LineSegment, BallInf
using ReachabilityAnalysis
using Colors: @colorant_str
import LinearAlgebra
using CartesianIndexSets: CartesianIndexSet
import PlotUtils
using PlotUtils: RGBA
using RecipesBase

# https://github.com/JuliaReach/ReachabilityAnalysis.jl/issues/464
using LazySets: ConvexHull, concretize, Zonotope
concretize_skip(s::LazySets.ConvexHull) = ConvexHull(concretize(s.X), concretize(s.Y))
concretize_skip(s) = s

ReachabilityAnalysis._apply_setops(X::LazySet, ::Val{:custom_gng}) = concretize_skip(X)
ReachabilityAnalysis._alias(setops::Val{:custom_gng}) = setops


function LazySets.CartesianProduct(usa::UnionSetArray, other::Any)
    UnionSetArray(CartesianProduct.(usa.array, Ref(other)))
end

function sweep_twist(source_set::UnionSetArray, args...; kwargs...)
    UnionSetArray([
        sweep_twist(a, args...; kwargs...) for a in source_set.array])
end

function LazySets.UnionSetArray(a::Vector{<:UnionSetArray})
    UnionSetArray(reduce(vcat, map(usa->usa.array, a)))
end

# https://github.com/JuliaReach/LazySets.jl/issues/2679
function LazySets.cartesian_product(Z1::Zonotope, Z2::Zonotope)
    center = vcat(Z1.center, Z2.center)
    z1 = zeros(dim(Z1), size(Z2.generators, 2))
    z2 = zeros(dim(Z2), size(Z1.generators, 2))
    g1 = Z1.generators
    g2 = Z2.generators
    generators = [
        g1 z1;
        g2 z2
    ]
    return Zonotope(center, generators)
end

function _sweep_twist_prob(source_set, A, c)
    # center at t=0 is origin
    # and moves with velocity `c`
    # simultaneously, source_set is transformed incrementally by `A`

    # if `A` is a rotational field
    # then this does "decoupled" interpolation
    # e.g. lines in x, y, θ space
    # e.g. LERP for x, y and SLERP for θ
    A_ = fill(0.0, (4, 4))
    c_ = fill(0.0, (4, ))
    A_[1:2, 1:2] .= A
    A_[1:2, 3:4] .= -A
    c_[1:2] .= c
    c_[3:4] .= c
    # augment the center point
    source_set_ = CartesianProduct(source_set, convert(Zonotope, Singleton([0.0, 0.0])))
    source_set_ = concretize(source_set_)
    prob = @ivp(x' = A_ * x + c_, x(0) ∈ source_set_)
end

function sweep_twist(source_set, A, c; alg_δ=0.01)
    prob = _sweep_twist_prob(source_set, A, c)

    alg = GLGM06(δ=alg_δ, max_order=10,
        approx_model=ReachabilityAnalysis.Forward(sih=:concrete, exp=:base, setops=:custom_gng)
    )

    sol_ = solve(prob, tspan=(0.0, 1.0), alg=alg)
    #line = LazySets.project(sol_, [3, 4])

    sol = LazySets.project(sol_, [1, 2])
    reach_set_all = set(sol)
    # tspan_all = tspan(sol) # equivalent to `reduce(union, tspan.(sol.F))`
    # ReachSet(reach_set_all, tspan_all) # doesn't work
    return reach_set_all
end

function Base.:*(M, usa::UnionSetArray)
    UnionSetArray([M * a for a in usa.array])
end

function Base.:+(usa::UnionSetArray, c)
    UnionSetArray([a + c for a in usa.array])
end

LazySets.minkowski_sum(a, b::UnionSetArray) =  UnionSetArray(minkowski_sum.(Ref(a), b.array))

"""
return Set of integer coordinates S
p ∈ S iff BallInf(p, 0.5) ∩ S ≂̸ ∅

Outer Jordan digitization
supercover digitization
conservative rasterization (e.g. GPUs)

Digital Geometry by Reinhard Klette; Azriel Rosenfeld Published by Morgan Kaufmann, 2004
Section 2.3.5 Domain digitizations
"""
function supercover_grid(continuous_set)
    s = continuous_set
    D = LazySets.dim(s)
    sigma = BallInf(fill(0.0, D), 0.5)
    s_plus = minkowski_sum(sigma, s)

    bb = overapproximate(s_plus, Hyperrectangle)
    low = Int.(floor.(LazySets.low(bb) .- 0.0))
    hi = Int.(ceil.(LazySets.high(bb) .+ 0.0))
    lattice_range = CartesianIndex(low...):CartesianIndex(hi...)

    ci_to_vec(ci) = convert.(Float64, collect(Tuple(ci)))

    lattice_inside = filter(ci -> (ci_to_vec(ci) ∈ s_plus), lattice_range)
    grid = convert(CartesianIndexSet, Set(lattice_inside))
    return grid
end

propagator(A, c, set, t) = exp(A * t) * set + (c * t)

rot2d(θ) = [
    cos(θ) -sin(θ);
    sin(θ) cos(θ)
]

iso2d(x, y, θ) = [
    rot2d(θ)    [x; y]
    0 0         1
]

rot2d_field = [0 -1; 1 0]
# rot = exp(rot_field * angle_radians)

skeleton = UnionSetArray([
    LineSegment([0.0, -1.0], [0.0, 1.0]),
    LineSegment([0.0, 0.0], [4.0, 0.0])
])

offset = BallInf([0.0, 0.0], 0.25)
# continuum definition of robot geometry
body = UnionSetArray(minkowski_sum.(skeleton.array, Ref(offset)))

using Plots

if false
rotation = π/2
translation = [0.0, 0.0]

#source_set = BallInf([0.0, 0.0], 0.5)
source_set = body
swept = sweep_twist(source_set, rotation * rot2d_field, translation)
plt = plot(;aspect_ratio=:equal)

# plot continuous over-approximated swept volume
plot!(plt, swept, color=colorant"green")

# plot sampled swept volume
for t = 0.0:0.2:1.0
    plot!(plt, propagator(rotation * rot2d_field, translation, source_set, t), color=colorant"red", alpha=0.4)
end

scatter!(plt, map(x->tuple(x...), supercover_grid(swept)))
plt
end

using UnitfulAngles: turn
using UnitfulAngles.Unitful: rad
Δxy = 1
Nθ = 8
Δθ = inv(Nθ) * turn
wrapNθ(n) = mod(n, Nθ)

continuum(nc) = (xy=nc.nxy, θ=nc.nθ*Δθ)
continuumΔ(Δnc) = (Δxy=Δnc.Δnxy, Δθ=Δnc.Δnθ*Δθ)

function pose(set, configuration)
    c = configuration
    return rot2d(c.θ) * set + Vector{Float64}(c.xy)
end

function sweep(nc_from, Δnc)
    c_from = continuum(nc_from)
    Δc = continuumΔ(Δnc)

    source_set = pose(body, c_from)
    R = rad(Δc.Δθ) * rot2d_field
    t = Δc.Δxy

    swept = sweep_twist(source_set, R, t)
end

transitions = [
    (nc_from = (nxy=[0, 0], nθ=nθ), Δnc = (Δnxy=[Δnx, Δny], Δnθ=Δnθ))
    for nθ = 0:(Nθ-1), Δnx = [-1, 0, 1], Δny = [-1, 0, 1], Δnθ = [-2, -1, 0, 1, 2]
]

# plot these
nc_froms = [
    (nxy=[0, 0], nθ=0)
]

Δnc = [
    (Δnxy=[0, 0], Δnθ=0)
    (Δnxy=[1, 0], Δnθ=0)
    (Δnxy=[0, 1], Δnθ=0)
    (Δnxy=[0, 0], Δnθ=1)
]


quick_test = false
if quick_test
    transitions = [(;nc_from, Δnc) for nc_from in nc_froms for Δnc in Δncs]
end

println("compute swept")
swept_transitions = Dict(t => sweep(t.nc_from, t.Δnc)
    for t in transitions)

println("compute grid")
grid_swept_transitions = Dict(k => supercover_grid(v)
    for (k, v) in pairs(swept_transitions))

# avoid serialization relying on `CartesianIndexSets` or `OffsetArrays`
grid_swept_transitions_enumerated_set = sort([(transition = k, set = sort(map(Tuple, collect(v)))) for (k, v) in pairs(grid_swept_transitions)])

serialize_me = (;
    Nθ,
    grid_swept_transitions_enumerated_set,
    body
)

#=
import JSON3

open("swept1.json", "w") do f
    JSON3.pretty(f, JSON3.write(json_structure))
    println(f)
end
=#
mydir = @__DIR__
path_gridvamp = realpath(joinpath(mydir, "../../../"))
path_deps = joinpath(path_gridvamp, "deps")

using JLD2: save
save(joinpath(path_deps, "robot1.jld2"), Dict("value" => serialize_me))

# partial workaround for https://github.com/JuliaPlots/Plots.jl/issues/3448
# the workaround is partial because it will not apply the color cycling of `Plots.get_series_color` https://github.com/JuliaPlots/Plots.jl/blob/8cf268a3903f847ef3071019b3e77b99240ba9c4/src/args.jl#L1607-L1616
function _get_color_hack(plotattributes)
    full_color_default = RGB(0.0, 0.0, 1.0)
    full_alpha_default = 1.0
    empty = RGBA(0.0, 0.0, 0.0, 0.0)
    full = PlotUtils.plot_color(
        get(plotattributes, :seriescolor, full_color_default),
        get(plotattributes, :seriesalpha, full_alpha_default)
    )
    (;empty, full)
end

# partial workaround for https://github.com/JuliaPlots/Plots.jl/issues/2410
# this will work for `OffsetArray` and `Array`
# but might break for other types that e.g. wrap OffsetArray
_unwrap_offsetarray(a) = parent(a)

@recipe function f(s::CartesianIndexSet)
    axs = axes(s.x)
    f = first.(axs) .- 0.5
    l = last.(axs) .+ 0.5
    _range(start, step, stop) = start:step:stop
    extents = map(_range, f, [1.0, 1.0], l)
    (d1, d2) = extents
    colors = _get_color_hack(plotattributes)
    # recipe for plotting image sets yflip
    # set the default to `false` here.
    # it can be overridden, e.g. `plt = plot(set, yflip=true)`
    # This affects `plt.subplots[1][:yaxis][:flip]`
    yflip --> false
    d1, d2, ifelse.(transpose(_unwrap_offsetarray(s.x)), colors.full, colors.empty)
end

# TODO repair `nc_tos`

# plotting
gl = grid(length(nc_froms) + 1 , length(nc_tos) + 1)
p = plot(;layout = gl, aspect_ratio = 1.0, titlefontsize=6.0)

# setup subplots
for gi in CartesianIndices(axes(gl))
    gi = Tuple(gi)
    _subplot = gl[gi...]
    if gi == (1, 1)
        # bug, clears all subplots
        #plot!(_subplot, legend = false, grid = false, foreground_color_subplot=:white)
        continue
    end
    plot!(_subplot, aspect_ratio = 1.0, framestyle = :origin, link = :square, ticks=nothing, )
end

# plot column header
for (i, nc_from) = pairs(nc_froms)
    c_from = continuum(nc_from)
    plot!(gl[1 + i, 1], pose(body, c_from), title="q_from$i")
end

# plot row header
for (j, nc_to) = pairs(nc_tos)
    c_to = continuum(nc_to)
    plot!(gl[1, 1 + j], pose(body, c_to), title="q_to$j")
end

# plot swept volumes
for (i, nc_from) = pairs(nc_froms), (j, nc_to) = pairs(nc_tos)
    c_from = continuum(nc_from)
    c_to = continuum(nc_to)
    k = (;nc_from, nc_to)
    # this is an overapproximation, so plot it first / below
    plot!(gl[1 + i, 1 + j], grid_swept_transitions[k]; color=:green, alpha=0.9)

    plot!(gl[1 + i, 1 + j], swept_transitions[k])
    plot!(gl[1 + i, 1 + j]; title="f$i->t$j", titlefontsize=6.0)
end

display(p)

#=
_npose = (nxy=[0, 0], nθ=1)
_cpose = continuum(_npose)
object1 = pose(body, _cpose)
object2 = sweep(_npose, _npose)
p = plot(;aspect_ratio=1.0, gridlinewidth=2.0, gridalpha=1.0, xticks=-10:10, yticks=-10:10)
plot!(p, supercover_grid(object2), alpha=0.5)
plot!(p, object2, alpha=0.5)
plot!(p, object1, alpha=0.5)
display(p)
=#