docs: add docstrings to NonlinearSolveHomotopyContinuation.jl

Author: AayushSabharwal

lib/NonlinearSolveHomotopyContinuation/Project.toml

@@ -8,6 +8,7 @@ ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 CommonSolve = "38540f10-b2f7-11e9-35d8-d573e4eb0ff2"
 ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471"
 DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
+DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 HomotopyContinuation = "f213a82b-91d6-5c5d-acf7-10f1c761b327"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 NonlinearSolveBase = "be0214bd-f91f-a760-ac4e-3421ce2b2da0"
@@ -16,11 +17,12 @@ SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5"
 TaylorSeries = "6aa5eb33-94cf-58f4-a9d0-e4b2c4fc25ea"
 
 [compat]
-Aqua = "0.8"
 ADTypes = "1.11.0"
+Aqua = "0.8"
 CommonSolve = "0.2.4"
 ConcreteStructs = "0.2.3"
 DifferentiationInterface = "0.6.27"
+DocStringExtensions = "0.9.3"
 HomotopyContinuation = "2.12.0"
 LinearAlgebra = "1.11.0"
 NonlinearSolve = "4"

lib/NonlinearSolveHomotopyContinuation/src/NonlinearSolveHomotopyContinuation.jl

@@ -7,6 +7,7 @@ using SymbolicIndexingInterface
 using LinearAlgebra
 using ADTypes
 using TaylorSeries
+using DocStringExtensions
 import CommonSolve
 import HomotopyContinuation as HC
 import DifferentiationInterface as DI
@@ -15,6 +16,35 @@ using ConcreteStructs: @concrete
 
 export HomotopyContinuationJL, HomotopyNonlinearFunction
 
+"""
+    HomotopyContinuationJL{AllRoots}(; autodiff = true, kwargs...)
+    HomotopyContinuationJL(; kwargs...) = HomotopyContinuationJL{false}(; kwargs...)
+
+This algorithm is an interface to `HomotopyContinuation.jl`. It is only valid for
+fully determined polynomial systems. The `AllRoots` type parameter can be `true` or
+`false` and controls whether the solver will find all roots of the polynomial
+or a single root close to the initial guess provided to the `NonlinearProblem`.
+The polynomial function must allow complex numbers to be provided as the state.
+
+If `AllRoots` is `true`, the initial guess in the `NonlinearProblem` is ignored.
+The function must be traceable using HomotopyContinuation.jl's symbolic variables.
+Note that higher degree polynomials and systems with multiple unknowns can increase
+solve time significantly.
+
+If `AllRoots` is `false`, a single path is traced during the homotopy. The traced path
+depends on the initial guess provided to the `NonlinearProblem` being solved. This method
+does not require that the polynomial function is traceable via HomotopyContinuation.jl's
+symbolic variables.
+
+HomotopyContinuation.jl requires the jacobian of the system. In case a jacobian function
+is provided, it will be used. Otherwise, the `autodiff` keyword argument controls the
+autodiff method used to compute the jacobian. A value of `true` refers to
+`AutoForwardDiff` and `false` refers to `AutoFiniteDiff`. Alternate algorithms can be
+specified using ADTypes.jl.
+
+HomotopyContinuation.jl requires the taylor series of the polynomial system for the single
+root method. This is automatically computed using TaylorSeries.jl.
+"""
 @concrete struct HomotopyContinuationJL{AllRoots} <: NonlinearSolveBase.AbstractNonlinearSolveAlgorithm
     autodiff
     kwargs

lib/NonlinearSolveHomotopyContinuation/src/interface_types.jl

@@ -4,6 +4,16 @@ struct Inplace <: HomotopySystemVariant end
 struct OutOfPlace <: HomotopySystemVariant end
 struct Scalar <: HomotopySystemVariant end
 
+"""
+    $(TYPEDEF)
+
+A simple struct that wraps a polynomial function which takes complex input and returns
+complex output in a form that supports automatic differentiation. If the wrapped
+function if ``f: \\mathbb{C}^n \\rightarrow \\mathbb{C}^n`` then it is assumed
+the input arrays are real-valued and have length ``2n``. They are `reinterpret`ed
+into complex arrays and passed into the function. This struct has an in-place signature
+regardless of the signature of ``f``.
+"""
 @concrete struct ComplexJacobianWrapper{variant <: HomotopySystemVariant}
     f
 end
@@ -32,14 +42,49 @@ function (cjw::ComplexJacobianWrapper{Scalar})(u::AbstractVector{T}, x::Abstract
     return u
 end
 
+"""
+    $(TYPEDEF)
+
+A struct which implements the `HomotopyContinuation.AbstractSystem` interface for
+polynomial systems specified using `NonlinearProblem`.
+
+# Fields
+
+$(FIELDS)
+"""
 @concrete struct HomotopySystemWrapper{variant <: HomotopySystemVariant} <: HC.AbstractSystem
+    """
+    The wrapped polynomial function.
+    """
     f
+    """
+    The jacobian function, if provided to the `NonlinearProblem` being solved. Otherwise,
+    a `ComplexJacobianWrapper` wrapping `f` used for automatic differentiation.
+    """
     jac
+    """
+    The parameter object.
+    """
     p
+    """
+    The ADType for automatic differentiation.
+    """
     autodiff
+    """
+    The result from `DifferentiationInterface.prepare_jacobian`.
+    """
     prep
+    """
+    HomotopyContinuation.jl's symbolic variables for the system.
+    """
     vars
+    """
+    The `TaylorSeries.Taylor1` objects used to compute the taylor series of `f`.
+    """
     taylorvars
+    """
+    Preallocated intermediate buffers used for calculating the jacobian.
+    """
     jacobian_buffers
 end
 
@@ -168,9 +213,38 @@ function HC.ModelKit.taylor!(u::AbstractVector, ::Val{N}, sys::HomotopySystemWra
     return u
 end
 
+"""
+    $(TYPEDEF)
+
+A `HomotopyContinuation.AbstractHomotopy` which uses an inital guess ``x_0`` to construct
+the start system for the homotopy. The homotopy is
+
+```math
+\\begin{aligned}
+H(x, t) &= t * (F(x) - F(x_0)) + (1 - t) * F(x)
+        &= F(x) - t * F(x_0)
+\\end{aligned}
+```
+
+Where ``F: \\mathbb{R}^n \\rightarrow \\mathbb{R}^n`` is the polynomial system and
+``x_0 \\in \\mathbb{R}^n`` is the guess provided to the `NonlinearProblem`.
+
+# Fields
+
+$(TYPEDFIELDS)
+"""
 @concrete struct GuessHomotopy <: HC.AbstractHomotopy
+    """
+    The `HomotopyContinuation.AbstractSystem` to solve.
+    """
     sys <: HC.AbstractSystem
+    """
+    The residual of the initial guess, i.e. ``F(x_0)``.
+    """
     fu0
+    """
+    A preallocated `TaylorVector` for efficient `taylor!` implementation.
+    """
     taylorbuffer::HC.ModelKit.TaylorVector{5, ComplexF64}
 end
 
@@ -180,8 +254,6 @@ end
 
 Base.size(h::GuessHomotopy) = size(h.sys)
 
-# H(x, t) = (1 - t) * F(x) + t * (F(x) - F(x0))
-#         = F(x) - t * F(x0)
 function HC.ModelKit.evaluate!(u, h::GuessHomotopy, x, t, p = nothing)
     HC.ModelKit.evaluate!(u, h.sys, x, p)
     @inbounds for i in eachindex(u)

lib/NonlinearSolveHomotopyContinuation/src/solve.jl

@@ -1,3 +1,9 @@
+"""
+    $(TYPEDSIGNATURES)
+
+Create and return the appropriate `HomotopySystemWrapper` to use for solving the given
+`prob` with `alg`.
+"""
 function homotopy_continuation_preprocessing(prob::NonlinearProblem, alg::HomotopyContinuationJL)
     # cast to a `HomotopyNonlinearFunction`
     f = if prob.f isa HomotopyNonlinearFunction
@@ -36,7 +42,6 @@ function homotopy_continuation_preprocessing(prob::NonlinearProblem, alg::Homoto
         HC.variables(:x, axes(u0)...)
     end
 
-    # TODO: Is there an upper bound for the order?
     taylorvars = if isscalar
         Taylor1(zeros(ComplexF64, 5), 4)
     elseif iip