Reorganization

Author: jonniedie

README.md

@@ -173,150 +173,3 @@ lotka_sol = solve(lotka_prob)
 Notice how cleanly the ```composed!``` function can pass variables from one function to another with no array index juggling in sight. This is especially useful for large models as it becomes harder to keep track top-level model array position when adding new or deleting old components from the model. We could go further and compose ```composed!``` with other components ad (practically) infinitum with no mental bookkeeping.
 
 The main benefit, however, is now our differential equations are unit testable. Both ```lorenz``` and ```lotka``` can be run as their own ```ODEProblem``` with ```f``` set to zero to see the unforced response.
-
-
-### Control of a sliding block
-In this example, we'll build a model of a block sliding on a surface and use ```ComponentArray```s to easily switch between coulomb and equivalent viscous damping models. The block is controlled by pushing and pulling a spring attached to it and we will use feedback through a PID controller to try to track a reference signal. For simplification, we are using the velocity of the block directly for the derivative term, rather than taking a filtered derivative of the error signal. We are also setting a deadzone on the friction force with exponential decay to zero velocity to get rid of simulation chatter during the static friction regime.
-
-```julia
-using ComponentArrays
-using DifferentialEquations
-using Interact: @manipulate
-using Parameters: @unpack
-using Plots
-
-## Setup
-const g = 9.80665
-
-maybe_apply(f::Function, x, p, t) = f(x, p, t)
-maybe_apply(f, x, p, t) = f
-
-# Allows functions of form f(x,p,t) to be applied and passed in as inputs
-function simulator(func; kwargs...)
-    simfun(dx, x, p, t) = func(dx, x, p, t; map(f->maybe_apply(f, x, p, t), (;kwargs...))...)
-    simfun(x, p, t) = func(x, p, t; map(f->maybe_apply(f, x, p, t), (;kwargs...))...)
-    return simfun
-end
-
-softsign(x) = tanh(1e3x)
-
-
-## Dynamics update functions
-# Sliding block with viscous friction
-function viscous_block!(D, vars, p, t; u=0.0)
-    @unpack m, c, k = p
-    @unpack v, x = vars
-
-    D.x = v
-    D.v = (-c*v + k*(u-x))/m
-    return x
-end
-
-# Sliding block with coulomb friction
-function coulomb_block!(D, vars, p, t; u=0.0)
-    @unpack m, μ, k = p
-    @unpack v, x = vars
-
-    D.x = v
-    a = -μ*g*softsign(v) + k*(u-x)/m
-    D.v = abs(a)<1e-3 && abs(v)<1e-3 ? -10v : a
-    return x
-end
-
-function PID_controller!(D, vars, p, t; err=0.0, v=0.0)
-    @unpack kp, ki, kd = p
-    @unpack x = vars
-
-    D.x = err
-    return ki*x + kp*err + kd*v
-end
-
-function feedback_sys!(D, components, p, t; ref=0.0)
-    @unpack ctrl, plant = components
-
-    u = p.ctrl.fun(D.ctrl, ctrl, p.ctrl.params, t; err=ref-plant.x, v=-plant.v)
-    return p.plant.fun(D.plant, plant, p.plant.params, t; u=u)
-end
-
-step_input(;time=1.0, mag=1.0) = (x,p,t) -> t>time ? mag : 0
-sine_input(;mag=1.0, period=10.0) = (x,p,t) -> mag*sin(t*2π/period)
-
-
-## Interactive GUI for switching out plant models and varying PID gains
-@manipulate for kp in 0:0.01:15,
-                ki in 0:0.01:15,
-                kd in 0:0.01:15,
-                damping in Dict(
-                    "Coulomb" => coulomb_block!,
-                    "Viscous" => viscous_block!,
-                ),
-                reference in Dict(
-                    "Sine" => sine_input,
-                    "Step" => step_input,
-                ),
-                magnitude in 0:0.01:10, # pop-pop!
-                period in 1:0.01:30,
-                plot_v in false
-
-    # Inputs
-    tspan = (0.0, 30.0)
-
-    ctrl_fun = PID_controller!
-    # plant_fun = coulomb_block!
-
-    ref = if reference==sine_input
-        reference(period=period, mag=magnitude)
-    else
-        reference(mag=magnitude)
-    end
-
-    m = 50.0
-    μ = 0.1
-    ω = 2π/period
-    c = 4*μ*m*g/(π*ω*magnitude) # Viscous equivalent damping
-    k = 50.0
-
-    plant_p = (m=m, μ=μ, c=c, k=k)
-    ctrl_p = (kp=kp, ki=ki, kd=kd)
-
-    plant_ic = (v=0, x=0)
-    ctrl_ic = (;x=0)
-
-
-
-    # Set up and solve
-    sys_p = (
-        ctrl = (
-            params = ctrl_p,
-            fun = ctrl_fun,
-        ),
-        plant = (
-            params = plant_p,
-            fun = damping,
-        ),
-    )
-    sys_ic = ComponentArray(ctrl=ctrl_ic, plant=plant_ic)
-    sys_fun = ODEFunction(simulator(feedback_sys!, ref=ref), syms=[:u, :v, :x])
-    sys_prob = ODEProblem(sys_fun, sys_ic, tspan, sys_p)
-
-    sol = solve(sys_prob, Tsit5())
-
-
-    # Plot
-    t = tspan[1]:0.1:tspan[2]
-    lims = magnitude*[-1, 1]
-    plotvars = plot_v ? [3, 2] : [3]
-    strip = plot(t, ref.(0, 0, t), ylim=1.2lims, label="r(t)")
-    plot!(strip, sol, vars=plotvars)
-    phase = plot(ref.(0, 0, t), map(x->x.plant.x, sol(t).u),
-        xlim=lims,
-        ylim=1.2lims,
-        legend=false,
-        xlabel="r(t)",
-        ylabel="x(t)",
-    )
-    plot(strip, phase, layout=(2, 1), size=(700, 800))
-
-end
-```
-<img src="assets/coulomb_control.png" alight="middle" />

src/ComponentArrays.jl

@@ -9,7 +9,7 @@ const FlatOrColonIdx = Union{FlatIdx, Colon}
 
 
 include("utils.jl")
-export fastindices
+export fastindices # Deprecated
 
 include("lazyarray.jl")
 
@@ -21,13 +21,19 @@ include("componentindex.jl")
 include("componentarray.jl")
 export ComponentArray, ComponentVector, ComponentMatrix, getaxes, getdata, valkeys
 
-include("set_get.jl")
+include("array_interface.jl")
+# Base methods: parent, size, elsize, axes, reinterpret, hcat, vcat, permutedims, IndexStyle, to_indices, to_index, getindex, setindex!, view, pointer, unsafe_convert, strides, stride
+# ArrayInterface methods: strides, size, lu_instance, parent_type
+
+include("namedtuple_interface.jl")
+# Base methods: hash, ==, keys, haskey, propertynames, getproperty, setproperty!
 
 include("similar_convert_copy.jl")
+# Base methods: similar, zero, copy, copyto!, deepcopy, convert (to Array and NamedTuple), promote
 
 include("broadcasting.jl")
-
-include("math.jl")
+# Base methods: BroadcastStyle, convert(to Broadcasted{Nothing}), similar, map, dataids
+# Broadcast methods: BroadcastStyle, broadcasted, broadcast_unalias
 
 include("show.jl")
 

src/array_interface.jl

@@ -0,0 +1,162 @@
+Base.parent(x::ComponentArray) = getfield(x, :data)
+
+Base.size(x::ComponentArray) = size(getdata(x))
+ArrayInterface.size(A::ComponentArray) = ArrayInterface.size(parent(A))
+
+Base.elsize(x::Type{<:ComponentArray{T,N,A,Axes}}) where {T,N,A,Axes} = Base.elsize(A)
+
+Base.axes(x::ComponentArray) = axes(getdata(x))
+
+Base.reinterpret(::Type{T}, x::ComponentArray, args...) where T = ComponentArray(reinterpret(T, getdata(x), args...), getaxes(x))
+
+Base.hcat(x::CV...) where {CV<:ComponentVector} = ComponentArray(reduce(hcat, getdata.(x)), getaxes(x[1])[1], FlatAxis())
+
+Base.vcat(x::ComponentVector, y::AbstractVector) = vcat(getdata(x), y)
+Base.vcat(x::AbstractVector, y::ComponentVector) = vcat(x, getdata(y))
+function Base.vcat(x::ComponentVector, y::ComponentVector)
+    if reduce((accum, key) -> accum || (key in keys(x)), keys(y); init=false)
+        return vcat(getdata(x), getdata(y))
+    else
+        data_x, data_y = getdata.((x, y))
+        ax_x, ax_y = only.(getaxes.((x, y)))
+        ax_y = reindex(ax_y, length(x))
+        idxmap_x, idxmap_y = indexmap.((ax_x, ax_y))
+        return ComponentArray(vcat(data_x, data_y), Axis((;idxmap_x..., idxmap_y...)))
+    end
+end
+Base.vcat(x::CV...) where {CV<:AdjOrTransComponentArray} = ComponentArray(reduce(vcat, map(y->getdata(y.parent)', x)), getaxes(x[1]))
+Base.vcat(x::ComponentVector...) = reduce(vcat, x)
+Base.vcat(x::ComponentVector, args...) = vcat(getdata(x), getdata.(args)...)
+Base.vcat(x::ComponentVector, args::Vararg{AbstractVector{T}, N}) where {T,N} = vcat(getdata(x), getdata.(args)...)
+
+function Base.permutedims(x::ComponentArray, dims)
+    axs = getaxes(x)
+    return ComponentArray(permutedims(getdata(x), dims), map(i->axs[i], dims)...)
+end
+
+## Indexing
+Base.IndexStyle(::Type{<:ComponentArray{T,N,<:A,<:Axes}}) where {T,N,A,Axes} = IndexStyle(A)
+
+# Since we aren't really using the standard approach to indexing, this will forward things to
+# the correct methods
+Base.to_indices(x::ComponentArray, i::Tuple) = i
+Base.to_indices(x::ComponentArray, i::Tuple{Vararg{Union{Integer, CartesianIndex}, N}}) where N = i
+Base.to_indices(x::ComponentArray, i::Tuple{Vararg{Int64}}) where N = i
+Base.to_index(x::ComponentArray, i) = i
+
+# Get AbstractAxis index
+@inline Base.getindex(::AbstractAxis, idx::FlatIdx) = ComponentIndex(idx)
+@inline Base.getindex(ax::AbstractAxis, ::Colon) = ComponentIndex(:, ax)
+@inline Base.getindex(::AbstractAxis{IdxMap}, s::Symbol) where IdxMap =
+    ComponentIndex(getproperty(IdxMap, s))
+
+# Get ComponentArray index
+Base.@propagate_inbounds Base.getindex(x::ComponentArray, idx::CartesianIndex) = getdata(x)[idx]
+Base.@propagate_inbounds Base.getindex(x::ComponentArray, idx::FlatIdx...) = getdata(x)[idx...]
+Base.@propagate_inbounds function Base.getindex(x::ComponentArray, idx::FlatOrColonIdx...)
+    axs = map((ax, i) -> getindex(ax, i).ax, getaxes(x), idx)
+    axs = remove_nulls(axs...)
+    return ComponentArray(getdata(x)[idx...], axs...)
+end
+Base.@propagate_inbounds Base.getindex(x::ComponentArray, ::Colon) = getdata(x)[:]
+@inline Base.getindex(x::ComponentArray, ::Colon...) = x
+Base.@propagate_inbounds Base.getindex(x::ComponentArray, idx...) = getindex(x, toval.(idx)...)
+@inline Base.getindex(x::ComponentArray, idx::Val...) = _getindex(x, idx...)
+
+# Set ComponentArray index
+@inline Base.setindex!(x::ComponentArray, v, idx::FlatIdx...) = setindex!(getdata(x), v, idx...)
+Base.@propagate_inbounds Base.setindex!(x::ComponentArray, v, ::Colon) = setindex!(getdata(x), v, :)
+Base.@propagate_inbounds Base.setindex!(x::ComponentArray, v, idx...) = setindex!(x, v, toval.(idx)...)
+@inline Base.setindex!(x::ComponentArray, v, idx::Val...) = _setindex!(x, v, idx...)
+
+# Explicitly view
+Base.@propagate_inbounds Base.view(x::ComponentArray, idx::ComponentArrays.FlatIdx...) = view(getdata(x), idx...)
+Base.@propagate_inbounds Base.view(x::ComponentArray, idx...) = _getindex(x, toval.(idx)...)
+
+# Generated get and set index methods to do all of the heavy lifting in the type domain
+@generated function _getindex(x::ComponentArray, idx...)
+    ci = getindex.(getaxes(x), getval.(idx))
+    inds = map(i -> i.idx, ci)
+    axs = map(i -> i.ax, ci)
+    axs = remove_nulls(axs...)
+    # the index must be valid after computing `ci`
+    :(Base.@_inline_meta; ComponentArray(Base.maybeview(getdata(x), $inds...), $axs...))
+end
+
+@generated function _setindex!(x::ComponentArray, v, idx...)
+    ci = getindex.(getaxes(x), getval.(idx))
+    inds = map(i -> i.idx, ci)
+    # the index must be valid after computing `ci`
+    return :(Base.@_inline_meta; setindex!(getdata(x), v, $inds...))
+end
+
+## Linear Algebra
+Base.pointer(x::ComponentArray{T,N,A,Axes}) where {T,N,A<:DenseArray,Axes} = pointer(getdata(x))
+
+Base.unsafe_convert(::Type{Ptr{T}}, x::ComponentArray{T,N,A,Axes}) where {T,N,A,Axes} = Base.unsafe_convert(Ptr{T}, getdata(x))
+
+Base.strides(x::ComponentArray) = strides(getdata(x))
+ArrayInterface.strides(A::ComponentArray) = ArrayInterface.strides(parent(A))
+for f in [:device, :stride_rank, :contiguous_axis, :contiguous_batch_size, :dense_dims] 
+    @eval ArrayInterface.$f(::Type{ComponentArray{T,N,A,Axes}}) where {T,N,A,Axes} = ArrayInterface.$f(A)
+end
+
+Base.stride(x::ComponentArray, k) = stride(getdata(x), k)
+Base.stride(x::ComponentArray, k::Int64) = stride(getdata(x), k)
+
+ArrayInterface.lu_instance(jac_prototype::ComponentArray) = ArrayInterface.lu_instance(getdata(jac_prototype))
+
+ArrayInterface.parent_type(::Type{ComponentArray{T,N,A,Axes}}) where {T,N,A,Axes} = A
+
+
+
+# While there are some cases where these were faster, it is going to be almost impossible to
+# to keep up with method ambiguity errors due to other array types overloading *, /, and \.
+# Leaving these here and commented out for now, but will delete them later.
+
+# # Avoid slower fallback
+# for f in [:(*), :(/), :(\)]
+#     @eval begin
+#         # The normal stuff
+#         Base.$f(x::ComponentArray, y::AbstractArray) = $f(getdata(x), y)
+#         Base.$f(x::AbstractArray, y::ComponentArray) = $f(x, getdata(y))
+#         Base.$f(x::ComponentArray, y::ComponentArray) = $f(getdata(x), getdata(y))
+
+#         # A bunch of special cases to avoid ambiguous method errors
+#         Base.$f(x::ComponentArray, y::AbstractMatrix) = $f(getdata(x), y)
+#         Base.$f(x::AbstractMatrix, y::ComponentArray) = $f(x, getdata(y))
+
+#         Base.$f(x::ComponentArray, y::AbstractVector) = $f(getdata(x), y)
+#         Base.$f(x::AbstractVector, y::ComponentArray) = $f(x, getdata(y))
+#     end
+# end
+
+# # Adjoint/transpose special cases
+# for f in [:(*), :(/)]
+#     @eval begin
+#         Base.$f(x::Adjoint, y::ComponentArray) = $f(getdata(x), getdata(y))
+#         Base.$f(x::Transpose, y::ComponentArray) = $f(getdata(x), getdata(y))
+
+#         Base.$f(x::Adjoint{T,<:AbstractVector{T}}, y::ComponentVector) where T = $f(x, getdata(y))
+#         Base.$f(x::Transpose{T,<:AbstractVector{T}}, y::ComponentVector) where T = $f(x, getdata(y))
+
+#         Base.$f(x::Adjoint{T,<:AbstractVector{T}}, y::ComponentMatrix{T,A,Axes}) where {T,A,Axes} = $f(x, getdata(y))
+#         Base.$f(x::Transpose{T,<:AbstractVector{T}}, y::ComponentMatrix{T,A,Axes}) where {T,A,Axes} = $f(x, getdata(y))
+
+#         Base.$f(x::Adjoint{T,<:AbstractMatrix{T}}, y::ComponentVector) where {T} = $f(x, getdata(y))
+#         Base.$f(x::Transpose{T,<:AbstractMatrix{T}}, y::ComponentVector) where {T} = $f(x, getdata(y))
+
+#         Base.$f(x::ComponentArray, y::Adjoint{T,<:AbstractVector{T}}) where T = $f(getdata(x), y)
+#         Base.$f(x::ComponentArray, y::Transpose{T,<:AbstractVector{T}}) where T = $f(getdata(x), y)
+
+#         Base.$f(x::ComponentArray, y::Adjoint{T,<:ComponentVector}) where T = $f(getdata(x), getdata(y))
+#         Base.$f(x::ComponentArray, y::Transpose{T,<:ComponentVector}) where T = $f(getdata(x), getdata(y))
+
+#         # There seems to be a new method in Julia > v.1.4 that specializes on this
+#         Base.$f(x::Adjoint{T,<:AbstractVector{T}}, y::ComponentVector{T,A,Axes}) where {T<:Number,A,Axes} = $f(x, getdata(y))
+#         Base.$f(x::Transpose{T,<:AbstractVector{T}}, y::ComponentVector{T,A,Axes}) where {T<:Number,A,Axes} = $f(x, getdata(y))
+
+#         Base.$f(x::Adjoint{T,<:AbstractVector{T}}, y::ComponentVector{T,A,Axes}) where {T<:Real,A,Axes} = $f(getdata(x), getdata(y))
+#         Base.$f(x::Transpose{T,<:AbstractVector{T}}, y::ComponentVector{T,A,Axes}) where {T<:Real,A,Axes} = $f(getdata(x), getdata(y))
+#     end
+# end

src/axis.jl

@@ -78,6 +78,7 @@ const NullOrFlatView{Inds,IdxMap} = ViewAxis{Inds,IdxMap,<:NullorFlatAxis}
 
 viewindex(::ViewAxis{Inds,IdxMap}) where {Inds,IdxMap} = Inds
 viewindex(::Type{<:ViewAxis{Inds,IdxMap}}) where {Inds,IdxMap} = Inds
+viewindex(i) = i
 
 
 
@@ -132,4 +133,10 @@ const NotPartitionedAxis = Union{Axis{IdxMap}, FlatAxis, NullAxis, ShapedAxis{Sh
 const NotShapedOrPartitionedAxis = Union{Axis{IdxMap}, FlatAxis, NullAxis} where {IdxMap}
 
 
-Base.merge(axs::Axis...) = Axis(merge(indexmap.(axs)...))
+Base.merge(axs::Axis...) = Axis(merge(indexmap.(axs)...))
+
+Base.lastindex(ax::AbstractAxis) = last(viewindex(last(indexmap(ax))))
+
+reindex(i, offset) = i .+ offset
+reindex(ax::Axis, offset) = Axis(map(x->reindex(x, offset), indexmap(ax)))
+reindex(ax::ViewAxis, offset) = ViewAxis(viewindex(ax) .+ offset, indexmap(ax))

src/broadcasting.jl

@@ -69,7 +69,7 @@ function Base.similar(bc::BC.Broadcasted{<:CAStyle{<:BC.Unknown, Axes, N}}, T::T
 end
 
 
-Base.Broadcast.broadcasted(f, x::ComponentArray) = ComponentArray(map(f, getdata(x)), getaxes(x))
+BC.broadcasted(f, x::ComponentArray) = ComponentArray(map(f, getdata(x)), getaxes(x))
 
 # Need a special case here because `map` doesn't follow same rules as normal broadcasting. To be safe and avoid ambiguities,
 # we'll just handle the case where everything is a ComponentArray. Else it falls back to a plain Array output.