Merge pull request #1878 from JuliaRobotics/25Q3/enh/keepsol

fast-forward keep solve changes (first use is extracting DERelative odes)

Author: dehann

IncrementalInference/ext/IncrInfrDiffEqFactorExt.jl

@@ -75,6 +75,7 @@ function DERelative(
   state1::AbstractVector{<:Real} = allocate(getPointIdentity(domain)), # zeros(getDimension(domain)),
   tspan::Tuple{<:Real, <:Real} = _calcTimespan(Xi),
   problemType = ODEProblem, # DiscreteProblem,
+  keepSolution::Bool = false
 )
   #
   datatuple = if 2 < length(Xi)
@@ -87,8 +88,21 @@ function DERelative(
   fproblem = problemType(f, state0, tspan, datatuple; dt)
   # backward time problem
   bproblem = problemType(f, state1, (tspan[2], tspan[1]), datatuple; dt = -dt)
+
+  _keepSolution = if keepSolution
+    Vector{typeof((;t=Vector{Float64}(),u=Vector{Vector{Float64}}()))}()
+  else
+    nothing
+  end 
+
   # build the IIF recognizable object
-  return DERelative(domain, fproblem, bproblem, datatuple) #, getSample)
+  return DERelative(
+    domain, 
+    fproblem, 
+    bproblem, 
+    datatuple, 
+    _keepSolution
+  )
 end
 
 function DERelative(
@@ -103,6 +117,7 @@ function DERelative(
   state0::AbstractVector{<:Real} = allocate(getPointIdentity(domain)), #zeros(getDimension(domain)),
   tspan::Tuple{<:Real, <:Real} = _calcTimespan(Xi),
   problemType = DiscreteProblem,
+  keepSolution::Bool = false
 )
   return DERelative(
     Xi,
@@ -114,6 +129,7 @@ function DERelative(
     state1,
     tspan,
     problemType,
+    keepSolution,
   )
 end
 #
@@ -301,12 +317,28 @@ function IncrementalInference.sampleFactor(cf::CalcFactor{<:DERelative}, N::Int
     cf._legacyParams[1], M_
   end
 
+  # solve ODE for N many particles transiting this factor
+  kS = cf.factor.keepSolution
+  if !isnothing( kS )
+    resize!(kS, N)
+    # @info "kS" length(kS)
+    for i in 1:N
+      kS[i] = (;t=Vector{Float64}(),u=Vector{Vector{Float64}}())
+    end
+  end
   # solve likely elements
   for i = 1:N
     # TODO, does this respect hyporecipe ???
     idxArr = (k -> cf._legacyParams[k][i]).(1:length(cf._legacyParams))
-    _solveFactorODE!(meas[i], prob, u0pts[i], _maketuplebeyond2args(idxArr...)...)
-    # _solveFactorODE!(meas, prob, u0pts, i, _maketuplebeyond2args(cf._legacyParams...)...)
+    oderes = _solveFactorODE!(meas[i], prob, u0pts[i], _maketuplebeyond2args(idxArr...)...)
+    if !isnothing( kS )
+      resize!(kS[i].t, length(oderes.t))
+      resize!(kS[i].u, length(oderes.u))
+      for j = 1:length(oderes.t)
+        kS[i].t[j] = oderes.t[j]
+        kS[i].u[j] = oderes.u[j]
+      end
+    end
   end
 
   # return meas, M

IncrementalInference/src/entities/ExtFactors.jl

@@ -19,11 +19,11 @@ DevNotes
 - FIXME Lots of consolidation and standardization to do, see RoME.jl #244 regarding Manifolds.jl.
 - TODO does not yet handle case where a factor spans across two timezones.
 """
-struct DERelative{T <: StateType, P, D} <: RelativeObservation
+struct DERelative{T <: StateType, P, D, O} <: RelativeObservation
   domain::Type{T}
   forwardProblem::P
   backwardProblem::P
   """ second element of this data tuple is additional variables that will be passed down as a parameter """
   data::D
-  # specialSampler::Function
+  keepSolution::O
 end

IncrementalInference/src/entities/FactorOperationalMemory.jl

@@ -26,7 +26,8 @@ Base.@kwdef struct CommonConvWrapper{
   AM <: AbstractManifold,
   HR <: HypoRecipeCompute,
   MT, 
-  G
+  G,
+  KCF <: Union{Nothing, <: Channel{<: CalcFactor}}
 } <: FactorCache
   # Basic factor topological info
   """ Values consistent across all threads during approx convolution """
@@ -67,6 +68,8 @@ Base.@kwdef struct CommonConvWrapper{
   res::Vector{Float64} = zeros(manifold_dimension(manifold))
   """ experimental feature to embed gradient calcs with ccw """
   _gradients::G = nothing
+    """ working memory to store residual from optimization routines """
+  keepCalcFactor::KCF
 end
 
 

IncrementalInference/src/services/ApproxConv.jl

@@ -10,6 +10,7 @@ function approxConvBelief(
   N::Int = length(measurement),
   nullSurplus::Real = 0,
   skipSolve::Bool = false,
+  keepCalcFactor::Union{Nothing, <:Channel} = nothing,
 )
   #
   v_trg = getVariable(dfg, target)
@@ -25,7 +26,8 @@ function approxConvBelief(
     solveKey, 
     N, 
     skipSolve, 
-    nullSurplus
+    nullSurplus,
+    keepCalcFactor
   )
 
   len = length(ipc)
@@ -85,6 +87,7 @@ function approxConvBelief(
   path::AbstractVector{Symbol} = Symbol[],
   skipSolve::Bool = false,
   nullSurplus::Real = 0,
+  keepCalcFactor::Union{Nothing, <:Channel} = nothing,
 )
   #
   # @assert isVariable(dfg, target) "approxConv(dfg, from, target,...) where `target`=$target must be a variable in `dfg`"
@@ -137,6 +140,7 @@ function approxConvBelief(
       N,
       skipSolve,
       nullSurplus,
+      keepCalcFactor,
     )
     if length(path) == 2
       return pts1Bel
@@ -155,7 +159,7 @@ function approxConvBelief(
       # this is a factor path[idx]
       fct = getFactor(dfg, path[idx])
       addFactor!(tfg, fct)
-      ptsBel = approxConvBelief(tfg, fct, path[idx + 1]; solveKey, N, skipSolve)
+      ptsBel = approxConvBelief(tfg, fct, path[idx + 1]; solveKey, N, skipSolve, keepCalcFactor)
       initVariable!(tfg, path[idx + 1], ptsBel)
       !setPPE ? nothing : setPPE!(tfg, path[idx + 1], solveKey, ppemethod)
     end
@@ -184,10 +188,11 @@ function calcProposalBelief(
   solveKey::Symbol = :default,
   nullSurplus::Real = 0,
   dbg::Bool = false,
+  keepCalcFactor::Union{Nothing, <:Channel} = nothing,
 )
   #
   # assuming it is properly initialized TODO
-  proposal = approxConvBelief(dfg, fct, target, measurement; solveKey, N, nullSurplus)
+  proposal = approxConvBelief(dfg, fct, target, measurement; solveKey, N, nullSurplus, keepCalcFactor)
 
   # return the proposal belief and inferdim, NOTE likely to be changed
   return proposal

IncrementalInference/src/services/CalcFactor.jl

@@ -376,6 +376,7 @@ function _createCCW(
   _blockRecursion::Bool = false,
   attemptGradients::Bool = true,
   userCache::CT = nothing,
+  keepCalcFactor::Bool = false,
 ) where {T <: AbstractFactor, CT}
   #
   if length(Xi) !== 0
@@ -394,7 +395,7 @@ function _createCCW(
   fullvariables = tuple(Xi...) # convert(Vector{VariableCompute}, Xi)
   # create a temporary CalcFactor object for extracting the first sample
 
-  _cf = CalcFactorNormSq(
+  _cf_nt = CalcFactorNormSq(
     usrfnc,
     1,
     _varValsAll,
@@ -408,12 +409,12 @@ function _createCCW(
   )
 
   # get a measurement sample
-  meas_single = sampleFactor(_cf, 1)[1]
+  meas_single = sampleFactor(_cf_nt, 1)[1]
   elT = typeof(meas_single)
   #TODO preallocate measurement?
   measurement = Vector{elT}()
 
-  #FIXME chicken and egg problem for getting measurement type, so creating twice.
+  # NOTE chicken and egg problem for getting measurement type, so creating twice.
   _cf = CalcFactorNormSq(
     usrfnc,
     1,
@@ -427,6 +428,11 @@ function _createCCW(
     nothing,
   )
 
+  keepCalcFactor_ = if keepCalcFactor
+    Channel{CalcFactor}(1024)
+  else
+    nothing
+  end
 
   # partialDims are sensitive to both which solvefor variable index and whether the factor is partial
   partial = hasfield(T, :partial) # FIXME, use isPartial function instead
@@ -441,6 +447,7 @@ function _createCCW(
 
   # as per struct CommonConvWrapper
   _gradients = if attemptGradients
+    # FIXME update to proper AD tools
     attemptGradientPrep(
       varTypes,
       usrfnc,
@@ -477,6 +484,7 @@ function _createCCW(
     ),
     measurement,
     _gradients,
+    keepCalcFactor = keepCalcFactor_
   )
 end
 
@@ -485,14 +493,15 @@ function updateMeasurement!(
   N::Int=1;
   measurement::AbstractVector = Vector{Tuple{}}(),
   needFreshMeasurements::Bool=true,
-  _allowThreads::Bool = true
+  _allowThreads::Bool = true,
+  keepCalcFactor::Union{Nothing, <:Channel} = nothing,
 )
   # FIXME do not divert Mixture for sampling
 
   # option to disable fresh samples or user provided
   if needFreshMeasurements
     # TODO this is only one thread, make this a for loop for multithreaded sampling
-    sampleFactor!(ccwl, N; _allowThreads)
+    sampleFactor!(ccwl, N; _allowThreads, keepCalcFactor)
   elseif 0 < length(measurement) 
     resize!(ccwl.measurement, length(measurement))
     ccwl.measurement[:] = measurement
@@ -520,6 +529,7 @@ function _beforeSolveCCW!(
   measurement = Vector{Tuple{}}(),
   needFreshMeasurements::Bool = true,
   solveKey::Symbol = :default,
+  keepCalcFactor::Union{Nothing, <:Channel} = nothing,
 ) where {F <: AbstractFactor} # F might be Mixture
   #
   if length(variables) !== 0
@@ -570,7 +580,7 @@ function _beforeSolveCCW!(
   _setCCWDecisionDimsConv!(ccwl, xDim)
 
   # FIXME do not divert Mixture for sampling
-  updateMeasurement!(ccwl, maxlen; needFreshMeasurements, measurement, _allowThreads=true)
+  updateMeasurement!(ccwl, maxlen; needFreshMeasurements, measurement, _allowThreads=true, keepCalcFactor)
 
   # used in ccw functor for AbstractRelativeMinimize
   resize!(ccwl.res, _getZDim(ccwl))
@@ -591,6 +601,7 @@ function _beforeSolveCCW!(
   measurement = Vector{Tuple{}}(),
   needFreshMeasurements::Bool = true,
   solveKey::Symbol = :default,
+  keepCalcFactor::Union{Nothing, <:Channel} = nothing,
 ) where {F <: AbstractFactor} # F might be Mixture
   # FIXME, NEEDS TO BE CLEANED UP AND WORK ON MANIFOLDS PROPER
 
@@ -606,7 +617,7 @@ function _beforeSolveCCW!(
 
   # FIXME do not divert Mixture for sampling
   # update ccwl.measurement values
-  updateMeasurement!(ccwl, maxlen; needFreshMeasurements, measurement, _allowThreads=true)
+  updateMeasurement!(ccwl, maxlen; needFreshMeasurements, measurement, _allowThreads=true, keepCalcFactor)
 
   return maxlen
 end

IncrementalInference/src/services/EvalFactor.jl

@@ -335,14 +335,15 @@ function evalPotentialSpecific(
   dbg::Bool = false,
   skipSolve::Bool = false,
   _slack = nothing,
+  keepCalcFactor::Union{Nothing, <:Channel} = nothing,
 ) where {T <: AbstractFactor}
   #
 
   # Prep computation variables
   # add user desired measurement values if 0 < length
   # 2023Q2, ccwl.varValsAll always points at the variable.VND.val memory locations
   #  remember when doing approxConv to make a deepcopy of the destination memory first.
-  maxlen = _beforeSolveCCW!(ccwl, variables, sfidx, N; solveKey, needFreshMeasurements, measurement)
+  maxlen = _beforeSolveCCW!(ccwl, variables, sfidx, N; solveKey, needFreshMeasurements, measurement, keepCalcFactor)
 
   # Check which variables have been initialized
   isinit = map(x -> isInitialized(x, solveKey), variables)
@@ -414,11 +415,12 @@ function evalPotentialSpecific(
   dbg::Bool = false,
   skipSolve::Bool = false,
   _slack = nothing,
+  keepCalcFactor::Union{Nothing, <:Channel} = nothing,
 ) where {T <: AbstractFactor}
   #
 
   # Prep computation variables
-  maxlen = _beforeSolveCCW!(ccwl, variables, sfidx, N; solveKey, needFreshMeasurements, measurement)
+  maxlen = _beforeSolveCCW!(ccwl, variables, sfidx, N; solveKey, needFreshMeasurements, measurement, keepCalcFactor)
 
   # # FIXME, NEEDS TO BE CLEANED UP AND WORK ON MANIFOLDS PROPER
   fnc = ccwl.usrfnc!
@@ -582,6 +584,7 @@ function evalFactor(
   dbg::Bool = false,
   skipSolve::Bool = false,
   _slack = nothing,
+  keepCalcFactor::Union{Nothing, <:Channel} = nothing,
 )
   #
   return evalPotentialSpecific(
@@ -598,6 +601,7 @@ function evalFactor(
     nullSurplus,
     skipSolve,
     _slack,
+    keepCalcFactor,
   )
   #
 end

IncrementalInference/src/services/FactorGraph.jl

@@ -723,6 +723,7 @@ function getDefaultFactorData(
   solveInProgress = 0,
   inflation::Real = getSolverParams(dfg).inflation,
   _blockRecursion::Bool = false,
+  keepCalcFactor::Bool = false,
 ) where {T <: AbstractFactor}
   #
 
@@ -741,6 +742,7 @@ function getDefaultFactorData(
     attemptGradients = getSolverParams(dfg).attemptGradients,
     _blockRecursion,
     userCache,
+    keepCalcFactor,
   )
 
   state = DFG.FactorState(
@@ -829,6 +831,7 @@ function DFG.addFactor!(
   inflation::Real = getSolverParams(dfg).inflation,
   namestring::Symbol = assembleFactorName(dfg, Xi),
   _blockRecursion::Bool = !getSolverParams(dfg).attemptGradients,
+  keepCalcFactor::Bool = false,
 )
   #
 
@@ -847,6 +850,7 @@ function DFG.addFactor!(
     # threadmodel,
     inflation,
     _blockRecursion,
+    keepCalcFactor,
   )
   #
   newFactor = FactorCompute(

IncrementalInference/src/services/NumericalCalculations.jl

@@ -157,6 +157,8 @@ function _solveLambdaNumeric(
     # TODO find good way for a solve to store diagnostics about number of failed converges etc.
     @warn "Optim did not converge (maxlog=10):" r maxlog=10
   end
+
+  # FIXME, how to use this exp when either Manifolds or LieGroups is used?
   return exp(M, ϵ, hat(M, ϵ, r.minimizer))
 end