From 972e9916226ce8f75b4be8cc6dc9319fb8c1fa23 Mon Sep 17 00:00:00 2001
From: john verzani <jverzani@gmail.com>
Date: Tue, 7 Mar 2023 16:46:30 -0500
Subject: [PATCH] Issue 469 (#472)

* take 1

* fit using constrained least squares
---
 src/common.jl                     | 15 +++++-
 src/polynomials/standard-basis.jl | 89 +++++++++++++++++++++++++++++--
 test/StandardBasis.jl             | 10 ++++
 3 files changed, 107 insertions(+), 7 deletions(-)

diff --git a/src/common.jl b/src/common.jl
index 4692301c..7ba47e4c 100644
--- a/src/common.jl
+++ b/src/common.jl
@@ -107,6 +107,7 @@ the variance-covariance matrix.)
 
 For fitting with a large degree, the Vandermonde matrix is exponentially ill-conditioned. The [`ArnoldiFit`](@ref) type introduces an Arnoldi orthogonalization that fixes this problem.
 
+
 """
 function fit(P::Type{<:AbstractPolynomial},
              x::AbstractVector{T},
@@ -141,7 +142,7 @@ fit(x::AbstractVector,
 function _fit(P::Type{<:AbstractPolynomial},
              x::AbstractVector{T},
              y::AbstractVector{T},
-             deg::Integer = length(x) - 1;
+             deg = length(x) - 1;
              weights = nothing,
              var = :x,) where {T}
     x = mapdomain(P, x)
@@ -152,7 +153,17 @@ function _fit(P::Type{<:AbstractPolynomial},
         coeffs = vand \ y
     end
     R = float(T)
-    return P(R.(coeffs), var)
+    if isa(deg, Integer)
+        return P{R, Symbol(var)}(R.(coeffs))
+    else
+        cs = zeros(T, 1 + maximum(deg))
+        for (i,aᵢ) ∈ zip(deg, coeffs)
+            cs[1 + i] = aᵢ
+        end
+        return P{R, Symbol(var)}(cs)
+    end
+
+
 end
 
 
diff --git a/src/polynomials/standard-basis.jl b/src/polynomials/standard-basis.jl
index c0b26f03..e7bb5592 100644
--- a/src/polynomials/standard-basis.jl
+++ b/src/polynomials/standard-basis.jl
@@ -488,14 +488,34 @@ function  roots(p::P; kwargs...)  where  {T, P <: StandardBasisPolynomial{T}}
 end
 
 function vander(P::Type{<:StandardBasisPolynomial}, x::AbstractVector{T}, n::Integer) where {T <: Number}
-    A = Matrix{T}(undef, length(x), n + 1)
-    A[:, 1] .= one(T)
-    @inbounds for i in 1:n
-        A[:, i + 1] = A[:, i] .* x
+    vander(P, x, 0:n)
+    # A = Matrix{T}(undef, length(x), n + 1)
+    # A[:, 1] .= one(T)
+    # @inbounds for i in 1:n
+    #     A[:, i + 1] = A[:, i] .* x
+    # end
+    # return A
+end
+
+# skip some degrees
+function vander(P::Type{<:StandardBasisPolynomial}, x::AbstractVector{T}, degs) where {T <: Number}
+    A = Matrix{T}(undef, length(x),  length(degs))
+    Aᵢ = one.(x)
+
+    i′ = 1
+    for i ∈ 0:maximum(degs)
+        if i ∈ degs
+            A[:, i′] = Aᵢ
+            i′ += 1
+        end
+        for (i, xᵢ) ∈ enumerate(x)
+            Aᵢ[i] *= xᵢ
+        end
     end
-    return A
+    A
 end
 
+
 ## as noted at https://github.com/jishnub/PolyFit.jl, using method from SpecialMatrices is faster
 ## https://github.com/JuliaMatrices/SpecialMatrices.jl/blob/master/src/vandermonde.jl
 ## This is Algorithm 2 of https://www.maths.manchester.ac.uk/~higham/narep/narep108.pdf
@@ -531,6 +551,65 @@ function fit(P::Type{<:StandardBasisPolynomial},
     end
 end
 
+"""
+    fit(P::Type{<:StandardBasisPolynomial}, x, y, J, [cs::Dict{Int, T}]; weights, var)
+
+Using constrained least squares, fit a polynomial of the type
+`p = ∑_{i ∈ J} aᵢ xⁱ + ∑ cⱼxʲ` where `cⱼ` are fixed non-zero constants
+
+* `J`: a collection of degrees to find coefficients for
+* `cs`: If given, a `Dict` of key/values, `i => cᵢ`, which indicate the degree and value of the fixed non-zero constants.
+
+The degrees in `cs` and those in `J` should not intersect.
+
+Example
+```
+x = range(0, pi/2, 10)
+y = sin.(x)
+P = Polynomial
+p0 = fit(P, x, y, 5)
+p1 = fit(P, x, y, 1:2:5)
+p2 = fit(P, x, y, 3:2:5, Dict(1 => 1))
+[norm(p.(x) - y) for p ∈ (p0, p1, p2)] # 1.7e-5, 0.00016, 0.000248
+```
+
+"""
+function fit(P::Type{<:StandardBasisPolynomial},
+             x::AbstractVector{T},
+             y::AbstractVector{T},
+             J,
+             cs=nothing;
+             weights = nothing,
+             var = :x,) where {T}
+    _fit(P, x, y, J; weights=weights, var=var)
+end
+
+
+function fit(P::Type{<:StandardBasisPolynomial},
+             x::AbstractVector{T},
+             y::AbstractVector{T},
+             J,
+             cs::Dict{Int, S};
+             weights = nothing,
+             var = :x,) where {T,S}
+
+    for i ∈ J
+        haskey(cs, i) && throw(ArgumentError("cs can't overlap with deg"))
+    end
+
+    # we subtract off `∑cᵢ xⁱ`ⱼ from `y`;
+    # fit as those all degrees not in J are 0,
+    # then add back the constant coefficients
+
+    q = SparsePolynomial(cs)
+    y′ = y - q.(x)
+
+    p = fit(P, x, y′, J; weights=weights, var=var)
+
+    return p + q
+end
+
+
 function _polynomial_fit(P::Type{<:StandardBasisPolynomial}, x::AbstractVector{T}, y; var=:x) where {T}
     R = float(T)
     coeffs = Vector{R}(undef, length(x))
diff --git a/test/StandardBasis.jl b/test/StandardBasis.jl
index 4967183e..769212d1 100644
--- a/test/StandardBasis.jl
+++ b/test/StandardBasis.jl
@@ -642,8 +642,18 @@ end
         @test fit(P, 1:4, 1:4, var=:x) ≈ variable(P{Float64}, :x)
         @test fit(P, 1:4, 1:4, 1, var=:x) ≈ variable(P{Float64}, :x)
 
+        # issue #467, fit specific degrees only
+        p = fit(P, xs, ys, 1:2:9)
+        @test norm(p.(xs) - ys) ≤ 1e-4
+
+        # issue 467: with constants
+        p = fit(P, xs, ys, 3:2:9, Dict(1 => 1))
+        @test norm(p.(xs) - ys) ≤ 1e-3
+
     end
 
+
+
     f(x) = 1/(1 + 25x^2)
     N = 250; xs = [cos(j*pi/N) for j in N:-1:0];
     q = fit(ArnoldiFit, xs, f.(xs));