ManifoldFlows

Relaxation(alph::AbstractVector; T=Float32, m = matrix_that_maps_index_to_vector)

Creates a Relaxation struct, which maps discrete tokens to continuous points and back again.

angleaxisstack(R::AbstractArray{T, 3})

Convert a stack of 3-by-3-by-N rotation matrices to a row vector of angles and a 3-by-N matrix of axis...es?

batch_flowstate(statetuple::Tuple{Vararg{AbstractArray}}, flowtuple::Tuple{Vararg{Flow}})

Converts a tuple of abstract arrays and a tuple of Flows, into a tuple of Flow-appropriate FlowStates.

bcds2quats(bcd::AbstractArray{<: Real, 2})

Convert a 3xN array of partial quaternions to an array of (flat) unit quaternions.

flow(f::Flow, x0::FlowState, model; steps = 100, tracker = NullTracker())

Samples from the distribution implied by the model under the Flow f, starting from x0. f and x0 can also be tuples, with matches components. steps can be an integer, in which case a linear schedule is used, or a vector of times to specify the schedule. If a tracker is supplied, the sample paths are tracked.

interpolate(f::Flow, x0::A, x1::A, t::T) where A::FlowState{T}

Geodesic interpolation between two FlowStates. t` must be a scalar, or row vector. The flow states are interpolated regardless of the mask, and the output mask is the logical AND of the two input masks.

logrotstack(R::AbstractArray{T, 3})

Calculate the logarithmic map of each rotation in a stack of 3-by-3-by-N rotation matrices.

logrotstack(A) is calculating the same thing as stack([log(A[:,:,i]) for i in 1:size(A,3)])

loss(f::Flow, x̂1::A, x1::A, xt::A, t::T; masked = false, eps = T(0.01), pow = 2) where A::FlowState{T}

Compute a loss between the predicted end point x̂1 and the true end point x1, given the starting point xt and the time t. These should be considered as "default" losses, and you might need to adapt and adjust them for your problem.

perturb!(f::Flow, x::A, σ::T) where A::FlowState{T}

Perturb the flow by a random amount, respecting the manifold, but do not change states where mask is false.

quats2rots(q)

Convert a 4-by-N array of quaternions to a 3-by-3-by-N array of rotation matrices.

randrot(σ²)

Generate a random rotation matrix, with each element drawn from the exponential map of a normal distribution with variance σ².

relax(seq::AbstractVector, r::Relaxation)

Converts a sequence of discrete tokens to a matrix (where each column can be thought of as a multivariant "point").

rotidentitystack(T,N)

Generate a stack of 3-by-3-by-N identity matrices of type T.

sinkhorn(C, λ; iters=50, standardize = true)

Returns the "plan" from the Sinkhorn algorithm for optimal transport, given a cost matrix C and regularization parameter λ. If standardize is true, then the cost matrix is standardized by dividing by its standard deviation which can help with numerical stability.

slerp(a,b,t)

Interpolation between two 3-by-3-by-N stacks of rotation matrices,performing Spherical linear interpolation (Slerp) on each pair of rotation matrices.

unrelax(points::AbstractArray, r::Relaxation)

Converts a continuous matrix (where each column can be thought of as a multivariant "point") to a sequence of discrete tokens.