From 4732f159e7ec1db369ed5f95a7451bc12ccb0fa0 Mon Sep 17 00:00:00 2001 From: PBrdng Date: Tue, 11 Aug 2026 14:44:08 +0200 Subject: [PATCH 1/4] add a cache for moving points in the witness set --- src/critical_points.jl | 12 +++- src/gradient_cache.jl | 131 ++++++++++++++++++++++++++++++++++++- src/hypersurfaces.jl | 10 ++- src/pseudo_witness_sets.jl | 5 +- src/routing_functions.jl | 21 +++++- 5 files changed, 170 insertions(+), 9 deletions(-) diff --git a/src/critical_points.jl b/src/critical_points.jl index 1ca4ad9..e9c7cc5 100644 --- a/src/critical_points.jl +++ b/src/critical_points.jl @@ -24,6 +24,7 @@ function critical_points( start_grid_stepsize = 0.2, start_grid_center = nothing, monodromy_at_zero = false, + ntrackers::Int = Threads.nthreads(), options = MonodromyOptions( parameter_sampler = p -> 10 .* randn(ComplexF64, length(p)), max_loops_no_progress = 15 @@ -36,6 +37,7 @@ function critical_points( MS, H, S0, rhs0, k = _setup_monodromy_solver( ∇r, S0, rhs0; monodromy_at_zero = monodromy_at_zero, + ntrackers = ntrackers, options = options, ) @@ -69,6 +71,7 @@ function _setup_monodromy_solver( S0::Union{AbstractVector{<:AbstractVector{<:Number}},Nothing} = nothing, rhs0::Union{AbstractVector{<:Number},Nothing} = nothing; monodromy_at_zero = false, + ntrackers::Int = Threads.nthreads(), options = MonodromyOptions( parameter_sampler = p -> 10 .* randn(ComplexF64, length(p)), max_loops_no_progress = 15 @@ -80,8 +83,11 @@ function _setup_monodromy_solver( H = RoutingPointsHomotopy(∇r, p1, q1) ### Use monodromy to the system ∇r = rhs0 where we view the right-hand side are the parameters of the system - egtracker = EndgameTracker(H) - trackers = [egtracker] + ntrackers >= 1 || throw(ArgumentError("ntrackers must be positive")) + # Every routing evaluator owns mutable pseudo-witness trackers, moving fibres, + # derivative buffers, and LU workspaces. Give every outer tracker a deep copy + # so HomotopyContinuation can run paths concurrently without data races. + trackers = [EndgameTracker(i == 1 ? H : deepcopy(H)) for i = 1:ntrackers] x₀ = zeros(ComplexF64, size(H, k)) unique_points = UniquePoints(x₀, 1;) @@ -326,4 +332,4 @@ monodromy_result(R::RoutingPointsResult) = R.monodromy_result function Base.show(io::IO, R::RoutingPointsResult) npts = length(routing_points(R)) println(io, "Routing points result with $npts routing point(s)") -end \ No newline at end of file +end diff --git a/src/gradient_cache.jl b/src/gradient_cache.jl index 6d284f9..657f965 100644 --- a/src/gradient_cache.jl +++ b/src/gradient_cache.jl @@ -48,6 +48,21 @@ mutable struct GradientCache{T} M3::Matrix{T} gradient_temp::Vector{T} Hess_temp::Matrix{T} + # A path-local moving copy of the pseudo-witness fibre. Consecutive calls made by + # an outer path tracker are close, so continuing this fibre is much cheaper than + # restarting at the original witness slice for every evaluation. + fiber_point::Vector{T} + fiber_solutions::Vector{Vector{T}} + fiber_scratch::Vector{Vector{T}} + warm_fiber_tracking::Bool + fiber_valid::Bool + fiber_evaluations::Int + fiber_exact_hits::Int + fiber_warm_tracks::Int + fiber_cold_tracks::Int + fiber_fallbacks::Int + fiber_failures::Int + fiber_tracking_ns::UInt64 end function compute_systems(F, n, k, B) @@ -163,6 +178,10 @@ function GradientCache(PWS) gradient_temp = zeros(ComplexF64, k) Hess_temp = zeros(ComplexF64, k, k) + fiber_point = zeros(ComplexF64, k) + fiber_solutions = [copy(z) for z in PWS.tZ] + fiber_scratch = [similar(z) for z in PWS.tZ] + v0 = randn(ComplexF64, n+1) GradientCache{ComplexF64}(v0, @@ -213,12 +232,118 @@ function GradientCache(PWS) M2, M3, gradient_temp, - Hess_temp + Hess_temp, + fiber_point, + fiber_solutions, + fiber_scratch, + true, + false, + 0, 0, 0, 0, 0, 0, + UInt64(0) ) end +@inline function _same_fiber_point(a, b) + length(a) == length(b) || return false + @inbounds for i in eachindex(a, b) + a[i] == b[i] || return false + end + true +end + +function _track_fiber_from!(dest, PWS::PseudoWitnessSet, starts, p_start, p_target) + tracker = PWS.tracker + start_parameters!(tracker, p_start) + target_parameters!(tracker, p_target) + succeeded = true + for (i, start) in enumerate(starts) + code = HC.track!(tracker, start, 1) + copyto!(dest[i], tracker.tracker.state.x) + ok = HC.is_success(code) && all(isfinite, dest[i]) + PWS.track_report[i] = ok + succeeded &= ok + end + succeeded +end + +"""Track the pseudo-witness fibre to `p`, reusing the most recent fibre when possible. + +Updates are transactional: a failed warm track is discarded and retried from the +original witness slice. An incomplete fibre is never used for differentiation. +""" function track!(GC::GradientCache, PWS::PseudoWitnessSet, p) - track!(GC.line_hypersurface_intersections, PWS, p) + GC.fiber_evaluations += 1 + + if GC.warm_fiber_tracking && GC.fiber_valid && _same_fiber_point(GC.fiber_point, p) + GC.fiber_exact_hits += 1 + for i in eachindex(GC.line_hypersurface_intersections) + copyto!(GC.line_hypersurface_intersections[i], GC.fiber_solutions[i]) + PWS.track_report[i] = true + end + get_s_and_Uvals!(GC.Uvals, GC.S, GC, PWS) + return nothing + end + + t0 = time_ns() + success = false + if GC.warm_fiber_tracking && GC.fiber_valid + GC.fiber_warm_tracks += 1 + success = _track_fiber_from!( + GC.fiber_scratch, PWS, GC.fiber_solutions, GC.fiber_point, p, + ) + if !success + GC.fiber_fallbacks += 1 + end + end + + if !success + GC.fiber_cold_tracks += 1 + success = _track_fiber_from!(GC.fiber_scratch, PWS, PWS.tZ, PWS.L.point, p) + end + GC.fiber_tracking_ns += UInt64(time_ns() - t0) + + if !success + GC.fiber_failures += 1 + GC.fiber_valid = false + error("Failed to track the complete pseudo-witness fibre to the evaluation point.") + end + + copyto!(GC.fiber_point, p) + for i in eachindex(GC.fiber_solutions) + copyto!(GC.fiber_solutions[i], GC.fiber_scratch[i]) + copyto!(GC.line_hypersurface_intersections[i], GC.fiber_scratch[i]) + PWS.track_report[i] = true + end + GC.fiber_valid = GC.warm_fiber_tracking get_s_and_Uvals!(GC.Uvals, GC.S, GC, PWS) nothing -end \ No newline at end of file +end + +function set_warm_fiber_tracking!(GC::GradientCache, enabled::Bool) + GC.warm_fiber_tracking = enabled + GC.fiber_valid = false + GC +end + +function reset_fiber_cache!(GC::GradientCache) + GC.fiber_valid = false + GC.fiber_evaluations = 0 + GC.fiber_exact_hits = 0 + GC.fiber_warm_tracks = 0 + GC.fiber_cold_tracks = 0 + GC.fiber_fallbacks = 0 + GC.fiber_failures = 0 + GC.fiber_tracking_ns = UInt64(0) + GC +end + +fiber_tracking_stats(GC::GradientCache) = ( + enabled = GC.warm_fiber_tracking, + evaluations = GC.fiber_evaluations, + exact_hits = GC.fiber_exact_hits, + warm_tracks = GC.fiber_warm_tracks, + cold_tracks = GC.fiber_cold_tracks, + fallbacks = GC.fiber_fallbacks, + failures = GC.fiber_failures, + tracking_seconds = Float64(GC.fiber_tracking_ns) / 1e9, +) diff --git a/src/hypersurfaces.jl b/src/hypersurfaces.jl index 6304ea0..c23784e 100644 --- a/src/hypersurfaces.jl +++ b/src/hypersurfaces.jl @@ -5,7 +5,10 @@ hessian, degree, trace_test, sample_points, -decompose +decompose, +fiber_tracking_stats, +reset_fiber_cache!, +set_warm_fiber_tracking! @doc raw""" ProjectedHypersurface{TC} <: HC.AbstractSystem @@ -408,6 +411,11 @@ end hessian(h::ProjectedHypersurface{TC}, x, p = nothing) where {TC} = gradient_and_hessian(h, x, p)[2] +fiber_tracking_stats(h::ProjectedHypersurface) = fiber_tracking_stats(h.GC) +reset_fiber_cache!(h::ProjectedHypersurface) = reset_fiber_cache!(h.GC) +set_warm_fiber_tracking!(h::ProjectedHypersurface, enabled::Bool) = + set_warm_fiber_tracking!(h.GC, enabled) + # Helpers for the fused derivative systems in GradientCache. They unpack one flat evaluation diff --git a/src/pseudo_witness_sets.jl b/src/pseudo_witness_sets.jl index c9bd7b2..858abfd 100644 --- a/src/pseudo_witness_sets.jl +++ b/src/pseudo_witness_sets.jl @@ -207,6 +207,9 @@ The resulting points are stored in `u` and the success of each track is recorded """ function track!(u::Vector{Vector{ComplexF64}}, PWS::PseudoWitnessSet, p::AbstractVector) tracker = PWS.tracker + # Other clients may temporarily move the start parameters for warm continuation. + # A direct PWS track always starts at the defining witness slice. + start_parameters!(tracker, PWS.L.point) target_parameters!(tracker, p) # PWS.tZ contains the coordinates (t,Z) for the points where the line # (PWS.L.direction*t+PWS.L.point; Z) intersects V(F) @@ -328,4 +331,4 @@ function sample_points(PWS::PseudoWitnessSet, N::Int) # Return the desired number of sample points sample[1:N] -end \ No newline at end of file +end diff --git a/src/routing_functions.jl b/src/routing_functions.jl index bd3f41b..a5a1839 100644 --- a/src/routing_functions.jl +++ b/src/routing_functions.jl @@ -225,6 +225,25 @@ evaluate!(u, ∇r::RoutingGradient, x, p = nothing) = gradient!(u, ∇r.r, x, p) evaluate_and_jacobian(∇r::RoutingGradient, x, p = nothing) = gradient_and_hessian(∇r.r, x, p) evaluate_and_jacobian!(u, U, ∇r::RoutingGradient, x, p = nothing) = gradient_and_hessian!(u, U, ∇r.r, x, p) +function taylor!(u, ::Val{1}, F::RoutingGradient, x, p::TaylorVector) + # For a parameter homotopy the predictor asks for the coefficient of T in + # ∇log(r)(x) - p(T), with x held fixed. The routing gradient itself has no + # explicit parameters, hence this coefficient is simply -p₁. + _, p1 = vectors(p) + @inbounds for i in eachindex(u) + u[i] = -p1[i] + end + u +end + +function taylor!(u, ::Val{1}, F::RoutingGradient, x, p) + fill!(u, zero(eltype(u))) + u +end + function taylor!(u, ::Val, F::RoutingGradient, x, p) - fill!(u, zero(ComplexF64)) + # Higher coefficients require third and fourth derivatives of log(r). HC's + # predictor can fall back to its Hermite history; do not fabricate them. + fill!(u, zero(eltype(u))) + u end From 9afdbb3dc4e0e215c28130091c8b0eb3801dcb37 Mon Sep 17 00:00:00 2001 From: PBrdng Date: Tue, 11 Aug 2026 16:23:38 +0200 Subject: [PATCH 2/4] improve linear equations solving --- src/gradient_cache.jl | 149 +++++++++++--------------- src/hypersurfaces.jl | 241 ++++++++++++++---------------------------- 2 files changed, 140 insertions(+), 250 deletions(-) diff --git a/src/gradient_cache.jl b/src/gradient_cache.jl index 657f965..69d763a 100644 --- a/src/gradient_cache.jl +++ b/src/gradient_cache.jl @@ -4,10 +4,8 @@ mutable struct GradientCache{T} JsuF::HC.CompiledSystem JPF::HC.CompiledSystem JBF::HC.CompiledSystem - HF::HC.CompiledSystem - JxB::HC.CompiledSystem - JxP::HC.CompiledSystem - JPB::HC.CompiledSystem + adjoint_gradient_system::HC.CompiledSystem + contracted_hessian_system::HC.CompiledSystem S::Vector{T} X::Vector{T} Uvals::Matrix{T} @@ -15,37 +13,22 @@ mutable struct GradientCache{T} SB::Matrix{T} UP::Array{T,3} UB::Array{T,3} - A::Array{T,4} rhs1::Matrix{T} - rhs2::Vector{T} - rhs3::Vector{T} + adjoint_rhs::Vector{T} JsuF_vals::Vector{T} JPF_vals::Vector{T} JBF_vals::Vector{T} - HF_vals::Vector{T} - JxB_vals::Vector{T} - JxP_vals::Vector{T} - JPB_vals::Vector{T} + adjoint_gradient_input::Vector{T} + adjoint_gradient_vals::Vector{T} + contracted_hessian_input::Vector{T} + contracted_hessian_vals::Vector{T} JsuF_temp::Matrix{T} JPF_temp::Matrix{T} JBF_temp::Matrix{T} Jtu_temp::Matrix{T} # Temporary storage for evaluating JsuF JsuF_lu::Array{T,3} JsuF_ipiv::Matrix{LinearAlgebra.LAPACK.BlasInt} - JsuF_lu_success::Vector{Bool} - HF_temp::Array{T, 3} # Temporary storage for evaluating HF - JxB_temp::Array{T, 3} # Temporary storage for evaluating JxB - JxP_temp::Array{T, 3} # Temporary storage for evaluating JxP - JPB_temp::Array{T, 3} # Temporary storage for evaluating JPB - temp_Hi::Matrix{T} - temp_Jxpi::Matrix{T} - temp_Jxbi::Matrix{T} - temp_Jpbi::Matrix{T} ipiv::Vector{LinearAlgebra.LAPACK.BlasInt} # allocation for pivot for lu! in place linear solving - M::Matrix{T} - M1::Matrix{T} - M2::Matrix{T} - M3::Matrix{T} gradient_temp::Vector{T} Hess_temp::Matrix{T} # A path-local moving copy of the pseudo-witness fibre. Consecutive calls made by @@ -66,7 +49,8 @@ mutable struct GradientCache{T} end function compute_systems(F, n, k, B) - @unique_var uval[1:n-k] α[1:k] β[1:k] t + N = n - k + 1 + @unique_var uval[1:n-k] α[1:k] β[1:k] t λ[1:N] xp[1:N,1:k] xb[1:N,1:k] F_on_line = F([α + (1 / t) * β; uval]) v = vcat(t, uval) vars = vcat(t, uval, α) @@ -93,24 +77,48 @@ function compute_systems(F, n, k, B) JPF = CompiledSystem(System(reduce(vcat, JPF_exprs), variables = vars)) JBF = CompiledSystem(System(reduce(vcat, JBF_exprs), variables = vars)) - function J(x) - map(Iterators.product(∇v, x)) do (∇vi, xj) - evaluate(HC.ModelKit.differentiate(∇vi, xj), β => B) + # Compile the contractions that are actually needed. This avoids evaluating + # full G_xx, G_xp, G_xβ and G_pβ tensors at every fibre point. + adjoint_gradient_exprs = map(β) do βb + sum(1:N; init = 0) do i + λ[i] * HC.ModelKit.differentiate(F_on_line[i], βb) end end + adjoint_gradient_exprs = evaluate.(adjoint_gradient_exprs, Ref(β => B)) + adjoint_gradient_system = CompiledSystem(System( + adjoint_gradient_exprs, + variables = [vars; λ], + )) + + contracted_hessian_exprs = [begin + expression = 0 + for i = 1:N + residual_i = HC.ModelKit.differentiate( + HC.ModelKit.differentiate(F_on_line[i], α[a]), β[b], + ) + for r = 1:N + residual_i += HC.ModelKit.differentiate( + HC.ModelKit.differentiate(F_on_line[i], v[r]), β[b], + ) * xp[r, a] + residual_i += HC.ModelKit.differentiate( + HC.ModelKit.differentiate(F_on_line[i], v[r]), α[a], + ) * xb[r, b] + for c = 1:N + residual_i += HC.ModelKit.differentiate( + HC.ModelKit.differentiate(F_on_line[i], v[r]), v[c], + ) * xp[r, a] * xb[c, b] + end + end + expression += λ[i] * residual_i + end + evaluate(expression, β => B) + end for a = 1:k, b = 1:k] + contracted_hessian_system = CompiledSystem(System( + vec(contracted_hessian_exprs), + variables = [vars; λ; vec(xp); vec(xb)], + )) - HF_exprs = J(v) - JxB_exprs = J(β) - JxP_exprs = J(α) - JPB_exprs = map(Iterators.product(∇α, β)) do (∇αi, βj) - evaluate(HC.ModelKit.differentiate(∇αi, βj), β => B) - end - HF = CompiledSystem(System(reduce(vcat, vec(HF_exprs)), variables = vars)) - JxB = CompiledSystem(System(reduce(vcat, vec(JxB_exprs)), variables = vars)) - JxP = CompiledSystem(System(reduce(vcat, vec(JxP_exprs)), variables = vars)) - JPB = CompiledSystem(System(reduce(vcat, vec(JPB_exprs)), variables = vars)) - - return JsuF, JPF, JBF, HF, JxB, JxP, JPB + return JsuF, JPF, JBF, adjoint_gradient_system, contracted_hessian_system end @@ -135,22 +143,20 @@ function GradientCache(PWS) SB = zeros(ComplexF64, d, k) UP = zeros(ComplexF64, d, n - k, k) UB = zeros(ComplexF64, d, n - k, k) - A = zeros(ComplexF64, d, N, k, k) - - JsuF, JPF, JBF, HF, JxB, JxP, JPB = compute_systems(F, n, k, L.direction) + JsuF, JPF, JBF, adjoint_gradient_system, contracted_hessian_system = + compute_systems(F, n, k, L.direction) # rhs1 = zeros(ComplexF64, N, 2*k) - rhs2 = zeros(ComplexF64, N) - rhs3 = zeros(ComplexF64, k) + adjoint_rhs = zeros(ComplexF64, N) JsuF_vals = zeros(ComplexF64, N * N) JPF_vals = zeros(ComplexF64, N * k) JBF_vals = zeros(ComplexF64, N * k) - HF_vals = zeros(ComplexF64, N * N * N) - JxB_vals = zeros(ComplexF64, N * N * k) - JxP_vals = zeros(ComplexF64, N * N * k) - JPB_vals = zeros(ComplexF64, N * k * k) + adjoint_gradient_input = zeros(ComplexF64, N + k + N) + adjoint_gradient_vals = zeros(ComplexF64, k) + contracted_hessian_input = zeros(ComplexF64, N + k + N + 2 * N * k) + contracted_hessian_vals = zeros(ComplexF64, k * k) JsuF_temp = zeros(ComplexF64, N, 1+n-k) JPF_temp = zeros(ComplexF64, N, k) @@ -158,23 +164,9 @@ function GradientCache(PWS) Jtu_temp = zeros(ComplexF64, N, 1+n-k) # TODO: Maybe can reuse Jsu_temp.... JsuF_lu = zeros(ComplexF64, d, N, N) JsuF_ipiv = Matrix{LinearAlgebra.LAPACK.BlasInt}(undef, d, N) - JsuF_lu_success = zeros(Bool, d) - HF_temp = zeros(ComplexF64, N, N, N) - JxB_temp = zeros(ComplexF64, N, N, k) - JxP_temp = zeros(ComplexF64, N, N, k) - JPB_temp = zeros(ComplexF64, N, k, k) - temp_Hi = zeros(ComplexF64, N, N) - temp_Jxpi = zeros(ComplexF64, N, k) - temp_Jxbi = zeros(ComplexF64, N, k) - temp_Jpbi = zeros(ComplexF64, k, k) ipiv = Vector{LinearAlgebra.LAPACK.BlasInt}(undef, min(size(JsuF_temp,1), size(JsuF_temp,2))) - M = zeros(ComplexF64, k, k) - M1 = zeros(ComplexF64, k, n-k+1) - M2 = zeros(ComplexF64, n-k+1, k) - M3 = zeros(ComplexF64, k, n-k+1) - gradient_temp = zeros(ComplexF64, k) Hess_temp = zeros(ComplexF64, k, k) @@ -189,10 +181,8 @@ function GradientCache(PWS) JsuF, JPF, JBF, - HF, - JxB, - JxP, - JPB, + adjoint_gradient_system, + contracted_hessian_system, S, X, Uvals, @@ -200,37 +190,22 @@ function GradientCache(PWS) SB, UP, UB, - A, rhs1, - rhs2, - rhs3, + adjoint_rhs, JsuF_vals, JPF_vals, JBF_vals, - HF_vals, - JxB_vals, - JxP_vals, - JPB_vals, + adjoint_gradient_input, + adjoint_gradient_vals, + contracted_hessian_input, + contracted_hessian_vals, JsuF_temp, JPF_temp, JBF_temp, Jtu_temp, JsuF_lu, JsuF_ipiv, - JsuF_lu_success, - HF_temp, - JxB_temp, - JxP_temp, - JPB_temp, - temp_Hi, - temp_Jxpi, - temp_Jxbi, - temp_Jpbi, ipiv, - M, - M1, - M2, - M3, gradient_temp, Hess_temp, fiber_point, diff --git a/src/hypersurfaces.jl b/src/hypersurfaces.jl index c23784e..aac0596 100644 --- a/src/hypersurfaces.jl +++ b/src/hypersurfaces.jl @@ -136,15 +136,15 @@ function gradient!(u, h::ProjectedHypersurface{TC}, x, p = nothing) where {TC} # Use cached symbolic objects and arrays JsuF = GC.JsuF - JPF = GC.JPF - JBF = GC.JBF + adjoint_gradient_system = GC.adjoint_gradient_system v0 = GC.v0 S = GC.S Uvals = GC.Uvals - SB = GC.SB - rhs1, rhs2, rhs3 = GC.rhs1, GC.rhs2, GC.rhs3 - JsuF_vals, JPF_vals, JBF_vals = GC.JsuF_vals, GC.JPF_vals, GC.JBF_vals + adjoint_rhs = GC.adjoint_rhs + JsuF_vals = GC.JsuF_vals + adjoint_gradient_input = GC.adjoint_gradient_input + adjoint_gradient_vals = GC.adjoint_gradient_vals N, n = size(PWS.F) k = n_projection_variables(PWS) @@ -154,7 +154,9 @@ function gradient!(u, h::ProjectedHypersurface{TC}, x, p = nothing) where {TC} # `track!` restores or computes both the tracked intersections and the cached S/Uvals data. track!(GC, PWS, x) - #Obtain gradients of S and U with respect to p and β + # Adjoint implicit differentiation. If J = G_(s,u) and Jᵀλ = e₁, then + # -∂s/∂β = λᵀG_β, which is precisely one fibre contribution to + # ∇log|h|. This needs one solve rather than 2k forward sensitivity solves. for i = 1:length(S) if !PWS.track_report[i] # skip if i-th track failed @@ -167,30 +169,15 @@ function gradient!(u, h::ProjectedHypersurface{TC}, x, p = nothing) where {TC} JsuF_temp = GC.JsuF_temp _evaluate_fused_columns!(JsuF_temp, JsuF_vals, JsuF, v0, N, N) - JPF_temp = GC.JPF_temp - _evaluate_fused_columns!(JPF_temp, JPF_vals, JPF, v0, N, k) - - JBF_temp = GC.JBF_temp - _evaluate_fused_columns!(JBF_temp, JBF_vals, JBF, v0, k, N) - - _fill_rhs1!(rhs1, JPF_temp, JBF_temp) - - rhs1 .*= -1 - # In-place linear solving with pre-allocated pivot vector + fill!(adjoint_rhs, zero(ComplexF64)) + adjoint_rhs[1] = one(ComplexF64) _, ipiv, info = LinearAlgebra.LAPACK.getrf!(JsuF_temp, GC.ipiv) - if info == 0 # this indicates successful factorization - LinearAlgebra.LAPACK.getrs!('N', JsuF_temp, ipiv, rhs1) - else - fill!(rhs1, zero(ComplexF64)) - end + info == 0 || error("Singular implicit Jacobian while evaluating the projected hypersurface gradient.") + LinearAlgebra.LAPACK.getrs!('T', JsuF_temp, ipiv, adjoint_rhs) - # copy rhs1 row segment into SB row without creating slices - @inbounds @simd for jj = 1:k - SB[i, jj] = rhs1[1, k + jj] - end - @inbounds @simd for jj = 1:k - u[jj] -= SB[i, jj] - end + _fill_adjoint_gradient_input!(adjoint_gradient_input, v0, adjoint_rhs) + evaluate!(adjoint_gradient_vals, adjoint_gradient_system, adjoint_gradient_input) + u .+= adjoint_gradient_vals end @@ -218,20 +205,11 @@ function gradient_and_hessian!(u, U, h::ProjectedHypersurface{TC}, x, p = nothin JsuF = GC.JsuF JPF = GC.JPF JBF = GC.JBF - HF = GC.HF - JxB = GC.JxB - JxP = GC.JxP - JPB = GC.JPB + contracted_hessian_system = GC.contracted_hessian_system # Preallocated temporaries and cached LU data keep the Hessian path allocation-free. JsuF_lu = GC.JsuF_lu JsuF_ipiv = GC.JsuF_ipiv - JsuF_lu_success = GC.JsuF_lu_success - temp_Hi = GC.temp_Hi - temp_Jxpi = GC.temp_Jxpi - temp_Jxbi = GC.temp_Jxbi - temp_Jpbi = GC.temp_Jpbi - v0 = GC.v0 S = GC.S Uvals = GC.Uvals @@ -239,12 +217,10 @@ function gradient_and_hessian!(u, U, h::ProjectedHypersurface{TC}, x, p = nothin SB = GC.SB UP = GC.UP UB = GC.UB - A = GC.A - rhs1, rhs2, rhs3 = GC.rhs1, GC.rhs2, GC.rhs3 + rhs1, adjoint_rhs = GC.rhs1, GC.adjoint_rhs JsuF_vals, JPF_vals, JBF_vals = GC.JsuF_vals, GC.JPF_vals, GC.JBF_vals - HF_vals, JxB_vals, JxP_vals, JPB_vals = GC.HF_vals, GC.JxB_vals, GC.JxP_vals, GC.JPB_vals - - M, M1, M2, M3 = GC.M, GC.M1, GC.M2, GC.M3 + contracted_hessian_input = GC.contracted_hessian_input + contracted_hessian_vals = GC.contracted_hessian_vals k = n_projection_variables(PWS) N, n = size(PWS.F) @@ -279,18 +255,14 @@ function gradient_and_hessian!(u, U, h::ProjectedHypersurface{TC}, x, p = nothin rhs1 .*= -1 # In-place linear solving with pre-allocated pivot vector _, ipiv, info = LinearAlgebra.LAPACK.getrf!(JsuF_temp, GC.ipiv) - JsuF_lu_success[i] = (info == 0) + info == 0 || error("Singular implicit Jacobian while evaluating the projected hypersurface Hessian.") @inbounds for row = 1:N, col = 1:N JsuF_lu[i, row, col] = JsuF_temp[row, col] end @inbounds for jj = 1:N JsuF_ipiv[i, jj] = ipiv[jj] end - if info == 0 # this indicates successful factorization - LinearAlgebra.LAPACK.getrs!('N', JsuF_temp, ipiv, rhs1) - else - fill!(rhs1, zero(ComplexF64)) - end + LinearAlgebra.LAPACK.getrs!('N', JsuF_temp, ipiv, rhs1) _copy_rhs1_blocks!(SP, SB, UP, UB, rhs1, i) @inbounds @simd for jj = 1:k @@ -305,96 +277,34 @@ function gradient_and_hessian!(u, U, h::ProjectedHypersurface{TC}, x, p = nothin end end - # Compute the second-derivative contributions using the fused tensor systems. + # Evaluate the already-contracted second-order residual. The compiled system + # accepts (s,u,p,λ,x_p,x_β) and returns the k×k Hessian contribution directly. for j = 1:length(S) !PWS.track_report[j] && continue # skip if j-th track failed _fill_v0!(v0, S, Uvals, x, j) - HF_temp = GC.HF_temp - HF_nrows, HF_ncols = N, N - _evaluate_fused_tensor!(HF_temp, HF_vals, HF, v0, N, HF_nrows, HF_ncols) - - JxB_temp = GC.JxB_temp - JxB_nrows, JxB_ncols = N, k - _evaluate_fused_tensor!(JxB_temp, JxB_vals, JxB, v0, N, JxB_nrows, JxB_ncols) - - JxP_temp = GC.JxP_temp - JxP_nrows, JxP_ncols = N, k - _evaluate_fused_tensor!(JxP_temp, JxP_vals, JxP, v0, N, JxP_nrows, JxP_ncols) - - JPB_temp = GC.JPB_temp - JPB_nrows, JPB_ncols = k, k - _evaluate_fused_tensor!(JPB_temp, JPB_vals, JPB, v0, N, JPB_nrows, JPB_ncols) - - _fill_M1_M2!(M1, M2, SP, SB, UP, UB, j) - - for i = 1:N - - # copy slices into temporaries (avoids allocating SubArray objects) - @inbounds for r = 1:HF_nrows, c = 1:HF_ncols - temp_Hi[r, c] = HF_temp[i, r, c] - end - @inbounds for r = 1:JxP_nrows, c = 1:JxP_ncols - temp_Jxpi[r, c] = JxP_temp[i, r, c] - end - @inbounds for r = 1:JxB_nrows, c = 1:JxB_ncols - temp_Jxbi[r, c] = JxB_temp[i, r, c] - end - @inbounds for r = 1:JPB_nrows, c = 1:JPB_ncols - temp_Jpbi[r, c] = JPB_temp[i, r, c] - end - - # now step by step in-place matrix multiplications. - for a = 1:k, b = 1:k - A[j, i, a, b] = temp_Jpbi[b, a] # note the transpose here - end - mul!(M, transpose(temp_Jxpi), M2) - for a = 1:k, b = 1:k - A[j, i, a, b] += M[b, a] # note the transpose here - end - mul!(M, M1, temp_Jxbi) - for a = 1:k, b = 1:k - A[j, i, a, b] += M[b, a] # note the transpose here - end - mul!(M3, M1, temp_Hi) - mul!(M, M3, M2) - for a = 1:k, b = 1:k - A[j, i, a, b] += M[b, a] # note the transpose here - end - - end - end - - - # Reuse the LU factors of JsuF computed above when solving the final Hessian systems. - fill!(M, zero(ComplexF64)) # here M will get assigned the Hessian of log r - for j = 1:length(S) - - !PWS.track_report[j] && continue # skip if j-th track failed - !JsuF_lu_success[j] && continue - - Jtu = GC.Jtu_temp @inbounds for row = 1:N, col = 1:N Jtu[row, col] = JsuF_lu[j, row, col] end - @inbounds for jj = 1:N - GC.ipiv[jj] = JsuF_ipiv[j, jj] + @inbounds for ii = 1:N + GC.ipiv[ii] = JsuF_ipiv[j, ii] end - for a = 1:k, b = 1:k - for i = 1:N - rhs2[i] = A[j, i, a, b] - end - LinearAlgebra.LAPACK.getrs!('N', Jtu, GC.ipiv, rhs2) - M[a, b] += rhs2[1] + fill!(adjoint_rhs, zero(ComplexF64)) + adjoint_rhs[1] = one(ComplexF64) + LinearAlgebra.LAPACK.getrs!('T', Jtu, GC.ipiv, adjoint_rhs) + + _fill_contracted_hessian_input!( + contracted_hessian_input, v0, adjoint_rhs, SP, UP, SB, UB, j, + ) + evaluate!(contracted_hessian_vals, contracted_hessian_system, contracted_hessian_input) + @inbounds for b = 1:k, a = 1:k + U[a, b] += contracted_hessian_vals[(b - 1) * k + a] end end - for a = 1:k, b = 1:k - U[a, b] += M[a, b] - end nothing end @@ -431,23 +341,55 @@ set_warm_fiber_tracking!(h::ProjectedHypersurface, enabled::Bool) = v0 end -@inline function _unpack_fused_columns!(dest, vals, nrows, ncols) - for col = 1:ncols - offset = (col - 1) * nrows - @inbounds for row = 1:nrows - dest[row, col] = vals[offset + row] +@inline function _fill_adjoint_gradient_input!(dest, v0, λ) + offset = 0 + @inbounds for i in eachindex(v0) + dest[offset + i] = v0[i] + end + offset += length(v0) + @inbounds for i in eachindex(λ) + dest[offset + i] = λ[i] + end + dest +end + +@inline function _fill_contracted_hessian_input!(dest, v0, λ, SP, UP, SB, UB, idx) + offset = 0 + @inbounds for i in eachindex(v0) + dest[offset + i] = v0[i] + end + offset += length(v0) + @inbounds for i in eachindex(λ) + dest[offset + i] = λ[i] + end + offset += length(λ) + + # vec(x_p), column-major, with x=(s,u). + k = size(SP, 2) + N = size(UP, 2) + 1 + @inbounds for a = 1:k + dest[offset + (a - 1) * N + 1] = SP[idx, a] + for row = 2:N + dest[offset + (a - 1) * N + row] = UP[idx, row - 1, a] + end + end + offset += N * k + + # vec(x_β), column-major. + @inbounds for b = 1:k + dest[offset + (b - 1) * N + 1] = SB[idx, b] + for row = 2:N + dest[offset + (b - 1) * N + row] = UB[idx, row - 1, b] end end dest end -@inline function _unpack_fused_tensor!(dest, vals, nout, nrows, ncols) +@inline function _unpack_fused_columns!(dest, vals, nrows, ncols) for col = 1:ncols - for row = 1:nrows - offset = ((col - 1) * nrows + (row - 1)) * nout - @inbounds for out = 1:nout - dest[out, row, col] = vals[offset + out] - end + offset = (col - 1) * nrows + @inbounds for row = 1:nrows + dest[row, col] = vals[offset + row] end end dest @@ -458,11 +400,6 @@ end _unpack_fused_columns!(dest, vals, nrows, ncols) end -@inline function _evaluate_fused_tensor!(dest, vals, F, x, nout, nrows, ncols) - evaluate!(vals, F, x) - _unpack_fused_tensor!(dest, vals, nout, nrows, ncols) -end - @inline function _fill_rhs1!(rhs1, JPF_temp, JBF_temp) for col = 1:size(JPF_temp, 2) @inbounds for row = 1:size(rhs1, 1) @@ -491,25 +428,3 @@ end end nothing end - -@inline function _fill_M1_M2!(M1, M2, SP, SB, UP, UB, idx) - k = size(SP, 2) - N = size(M2, 1) - for a = 1:k - M1[a, 1] = SP[idx, a] - end - for a = 1:k - for b = 2:N - M1[a, b] = UP[idx, b - 1, a] - end - end - for b = 1:k - M2[1, b] = SB[idx, b] - end - for b = 1:k - for a = 2:N - M2[a, b] = UB[idx, a - 1, b] - end - end - nothing -end From fc3d9168900cb47a0e33d971b9210f4a08e6c9fd Mon Sep 17 00:00:00 2001 From: Oskar Henriksson Date: Wed, 12 Aug 2026 16:26:38 +0200 Subject: [PATCH 3/4] Make track! explanation internal Since it's an internal function, I don't thinkw e want it to appear in the docs. --- src/gradient_cache.jl | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/src/gradient_cache.jl b/src/gradient_cache.jl index 69d763a..1b1d722 100644 --- a/src/gradient_cache.jl +++ b/src/gradient_cache.jl @@ -241,11 +241,9 @@ function _track_fiber_from!(dest, PWS::PseudoWitnessSet, starts, p_start, p_targ succeeded end -"""Track the pseudo-witness fibre to `p`, reusing the most recent fibre when possible. - -Updates are transactional: a failed warm track is discarded and retried from the -original witness slice. An incomplete fibre is never used for differentiation. -""" +#Track the pseudo-witness fibre to `p`, reusing the most recent fibre when possible. +# Updates are transactional: a failed warm track is discarded and retried from the +# original witness slice. An incomplete fibre is never used for differentiation. function track!(GC::GradientCache, PWS::PseudoWitnessSet, p) GC.fiber_evaluations += 1 From b3df6b8f3142ff5eab6dd8903e85308c6933737f Mon Sep 17 00:00:00 2001 From: John Cobb Date: Wed, 12 Aug 2026 11:51:38 -0500 Subject: [PATCH 4/4] Aggregate parallel worker fiber stats after monodromy Remove the `ntrackers` parameter and always create one tracker per Julia thread. Add `_snapshot_worker_fiber_stats` and `_merge_worker_fiber_stats!` to collect fiber-tracking counters from deep-copied monodromy workers and fold their deltas back into the original hypersurface after each solve. Add helper `_fiber_tracking_counters` and `_add_fiber_tracking_delta!` to `GradientCache`. Improve docstring on `fiber_tracking_stats` to document the aggregation behavior. --- src/critical_points.jl | 44 ++++++++++++++++++++++++++++++++-------- src/gradient_cache.jl | 21 +++++++++++++++++++ src/hypersurfaces.jl | 6 ++++++ src/routing_functions.jl | 5 +++-- test/runtests.jl | 24 ++++++++++++++++++++++ 5 files changed, 90 insertions(+), 10 deletions(-) diff --git a/src/critical_points.jl b/src/critical_points.jl index e9c7cc5..74163ab 100644 --- a/src/critical_points.jl +++ b/src/critical_points.jl @@ -24,7 +24,6 @@ function critical_points( start_grid_stepsize = 0.2, start_grid_center = nothing, monodromy_at_zero = false, - ntrackers::Int = Threads.nthreads(), options = MonodromyOptions( parameter_sampler = p -> 10 .* randn(ComplexF64, length(p)), max_loops_no_progress = 15 @@ -37,7 +36,6 @@ function critical_points( MS, H, S0, rhs0, k = _setup_monodromy_solver( ∇r, S0, rhs0; monodromy_at_zero = monodromy_at_zero, - ntrackers = ntrackers, options = options, ) @@ -71,7 +69,6 @@ function _setup_monodromy_solver( S0::Union{AbstractVector{<:AbstractVector{<:Number}},Nothing} = nothing, rhs0::Union{AbstractVector{<:Number},Nothing} = nothing; monodromy_at_zero = false, - ntrackers::Int = Threads.nthreads(), options = MonodromyOptions( parameter_sampler = p -> 10 .* randn(ComplexF64, length(p)), max_loops_no_progress = 15 @@ -83,11 +80,10 @@ function _setup_monodromy_solver( H = RoutingPointsHomotopy(∇r, p1, q1) ### Use monodromy to the system ∇r = rhs0 where we view the right-hand side are the parameters of the system - ntrackers >= 1 || throw(ArgumentError("ntrackers must be positive")) # Every routing evaluator owns mutable pseudo-witness trackers, moving fibres, - # derivative buffers, and LU workspaces. Give every outer tracker a deep copy - # so HomotopyContinuation can run paths concurrently without data races. - trackers = [EndgameTracker(i == 1 ? H : deepcopy(H)) for i = 1:ntrackers] + # derivative buffers, and LU workspaces. HomotopyContinuation uses one tracker + # per Julia thread, so give each of those trackers an independent homotopy. + trackers = [EndgameTracker(i == 1 ? H : deepcopy(H)) for i = 1:Threads.nthreads()] x₀ = zeros(ComplexF64, size(H, k)) unique_points = UniquePoints(x₀, 1;) @@ -123,6 +119,32 @@ function _setup_monodromy_solver( return MS, H, S0, rhs0, k end +@inline _worker_hypersurfaces(tracker) = tracker.tracker.homotopy.∇r.r.H + +function _snapshot_worker_fiber_stats(MS::HomotopyContinuation.MonodromySolver) + [ + [_fiber_tracking_counters(h.GC) for h in _worker_hypersurfaces(tracker)] + for tracker in MS.trackers + ] +end + +function _merge_worker_fiber_stats!(targets, MS, before) + # The first worker owns the original routing function, so its counters are + # already visible. Add only the work performed by the deep-copied workers. + for worker_index = 2:length(MS.trackers) + workers = _worker_hypersurfaces(MS.trackers[worker_index]) + length(workers) == length(targets) || error("Worker routing function changed shape.") + for (target, worker, baseline) in zip(targets, workers, before[worker_index]) + _add_fiber_tracking_delta!( + target.GC, + baseline, + _fiber_tracking_counters(worker.GC), + ) + end + end + nothing +end + """ _expand_start_solutions(∇r, H, S0, rhs0, k; verbose, start_grid_width, start_stepsize, start_center, monodromy_at_zero) @@ -257,7 +279,13 @@ function _solve_and_trace( start_grid_width = 5, ) ### Monodromy - mon_result = monodromy_solve(MS, S0, rhs0, rand(UInt32)) + worker_stats_before = _snapshot_worker_fiber_stats(MS) + local mon_result + try + mon_result = monodromy_solve(MS, S0, rhs0, rand(UInt32)) + finally + _merge_worker_fiber_stats!(∇r.r.H, MS, worker_stats_before) + end ### Trace to ∇r=0 if !monodromy_at_zero diff --git a/src/gradient_cache.jl b/src/gradient_cache.jl index 1b1d722..76441ab 100644 --- a/src/gradient_cache.jl +++ b/src/gradient_cache.jl @@ -320,3 +320,24 @@ fiber_tracking_stats(GC::GradientCache) = ( failures = GC.fiber_failures, tracking_seconds = Float64(GC.fiber_tracking_ns) / 1e9, ) + +@inline _fiber_tracking_counters(GC::GradientCache) = ( + evaluations = GC.fiber_evaluations, + exact_hits = GC.fiber_exact_hits, + warm_tracks = GC.fiber_warm_tracks, + cold_tracks = GC.fiber_cold_tracks, + fallbacks = GC.fiber_fallbacks, + failures = GC.fiber_failures, + tracking_ns = GC.fiber_tracking_ns, +) + +function _add_fiber_tracking_delta!(GC::GradientCache, before, after) + GC.fiber_evaluations += after.evaluations - before.evaluations + GC.fiber_exact_hits += after.exact_hits - before.exact_hits + GC.fiber_warm_tracks += after.warm_tracks - before.warm_tracks + GC.fiber_cold_tracks += after.cold_tracks - before.cold_tracks + GC.fiber_fallbacks += after.fallbacks - before.fallbacks + GC.fiber_failures += after.failures - before.failures + GC.fiber_tracking_ns += after.tracking_ns - before.tracking_ns + GC +end diff --git a/src/hypersurfaces.jl b/src/hypersurfaces.jl index aac0596..8073e05 100644 --- a/src/hypersurfaces.jl +++ b/src/hypersurfaces.jl @@ -321,6 +321,12 @@ end hessian(h::ProjectedHypersurface{TC}, x, p = nothing) where {TC} = gradient_and_hessian(h, x, p)[2] +""" + fiber_tracking_stats(h::ProjectedHypersurface) + +Return cumulative fibre-tracking diagnostics. Work performed by parallel +monodromy workers is aggregated into the original hypersurface after each solve. +""" fiber_tracking_stats(h::ProjectedHypersurface) = fiber_tracking_stats(h.GC) reset_fiber_cache!(h::ProjectedHypersurface) = reset_fiber_cache!(h.GC) set_warm_fiber_tracking!(h::ProjectedHypersurface, enabled::Bool) = diff --git a/src/routing_functions.jl b/src/routing_functions.jl index a5a1839..869e177 100644 --- a/src/routing_functions.jl +++ b/src/routing_functions.jl @@ -242,8 +242,9 @@ function taylor!(u, ::Val{1}, F::RoutingGradient, x, p) end function taylor!(u, ::Val, F::RoutingGradient, x, p) - # Higher coefficients require third and fourth derivatives of log(r). HC's - # predictor can fall back to its Hermite history; do not fabricate them. + # Higher coefficients require third and fourth derivatives of log(r), which + # are not available yet. Preserve the legacy zero-coefficient behavior for + # compatibility; these values are not exact higher derivatives. fill!(u, zero(eltype(u))) u end diff --git a/test/runtests.jl b/test/runtests.jl index 8c6d535..bcf7e73 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -96,10 +96,34 @@ end ProjectedHypersurfaces.evaluate_and_jacobian!(u, U, H, x0, 0.0) @test norm(∇r_symbolic(x0)-q1 - u) < 1e-12 + # The parameter-homotopy Taylor coefficient is exactly the negative of the + # first parameter coefficient, independently of the current x series. + p0 = randn(ComplexF64, 2) + p1 = randn(ComplexF64, 2) + tp = TaylorVector{2}(ComplexF64, 2) + for i = 1:2 + tp[i] = (p0[i], p1[i]) + end + taylor_coefficient = zeros(ComplexF64, 2) + ProjectedHypersurfaces.taylor!(taylor_coefficient, Val(1), ∇r, x0, tp) + @test taylor_coefficient == -p1 # Test that the expansion of start solutions works ∇r = RoutingGradient(r) MS, H, S0, rhs0, k = ProjectedHypersurfaces._setup_monodromy_solver(∇r) + @test length(MS.trackers) == Threads.nthreads() + if Threads.nthreads() > 1 + worker_stats_before = ProjectedHypersurfaces._snapshot_worker_fiber_stats(MS) + target_before = ProjectedHypersurfaces._fiber_tracking_counters(h.GC) + copied_worker_h = ProjectedHypersurfaces._worker_hypersurfaces(MS.trackers[2])[1] + @test copied_worker_h.GC !== h.GC + copied_worker_h.GC.fiber_evaluations += 2 + copied_worker_h.GC.fiber_tracking_ns += UInt64(7) + ProjectedHypersurfaces._merge_worker_fiber_stats!(∇r.r.H, MS, worker_stats_before) + target_after = ProjectedHypersurfaces._fiber_tracking_counters(h.GC) + @test target_after.evaluations == target_before.evaluations + 2 + @test target_after.tracking_ns == target_before.tracking_ns + UInt64(7) + end S0, new_pts = ProjectedHypersurfaces._expand_start_solutions( ∇r, H, S0, rhs0, k; start_grid_width = 10,