FitLine.swift

// Copyright 2017 Todd Reed
//
// Permission is hereby granted, free of charge, to any person obtaining a copy of this software
// and associated documentation files (the “Software”), to deal in the Software without
// restriction, including without limitation the rights to use, copy, modify, merge, publish,
// distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the
// Software is furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all copies or
// substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING
// BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
// NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
// DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

// This code is part of the article at <https://www.toddreed.name/articles/line-fitting/>.

import Foundation

/// A line expressed in Hessian normal form: ρ = x cos θ + y sin θ.
///
/// `rho` is the length of the vector n from the origin to the line that is perpendicular to the
/// line. `theta` is the angle of between the x-axis and the vector n.
struct Line {
    let rho: Double
    let theta: Double
    
    init(rho: Double, theta: Double) {
        // Normalize rho and theta so that rho > 0 and theta in [0, 2π]
        var r = rho
        var t = theta

        if (r < 0) {
            r = -r
            t += Double.pi
        }
        if (t < 0 || t > 2*Double.pi) {
            t = normalizeAngle(t)
        }
        if (r == 0 && t > Double.pi) {
            t -= Double.pi
        }
        self.rho = r
        self.theta = t
    }
}

struct Point {
    let x: Double
    let y: Double
}

extension Point {
    static prefix func -(p: Point) -> Point {
        return Point(x: -p.x, y: -p.y)
    }

    static func +(a: Point, b: Point) -> Point {
        return Point(x: a.x+b.x, y: a.y+b.y)
    }

    static func -(a: Point, b: Point) -> Point {
        return Point(x: a.x-b.x, y: a.y-b.y)
    }

    static func *(s: Double, p: Point) -> Point {
        return Point(x: s*p.x, y: s*p.y)
    }

}

/// Returns the arithmetic mean of a sequence.
///
/// This implements the numerically stable algorithm detailed here:
/// https://diego.assencio.com/?index=c34d06f4f4de2375658ed41f70177d59. Thanks to Alan Nunns for
/// directing me to this algorithm, which avoids floating-point errors caused by the accumulation
/// of a large sum.
func mean<T>(_ values: T) -> T.Element where T: Sequence, T.Element: FloatingPoint {
    var mean: T.Element = 0
    for (i, v) in values.enumerated() {
        mean += (v-mean)/T.Element(i+1)
    }
    return mean
}

/// Normalizes an angle `theta` so that is lies in the 2π interval [center-π, center+π].
func normalizeAngle(_ theta: Double, center: Double = Double.pi) -> Double
{
    return theta-2*Double.pi*floor((theta+Double.pi-center)/(2*Double.pi))
}

/// Returns the best least-squares line fit to the provided data points using perpendicular
/// distances.
///
/// Returns nil if not enough points are provided or no unique solution exists.
func fitLine(points: [Point]) -> Line? {
    if points.count < 2 {
        return nil
    }

    let centroid = Point(x: mean(points.map { $0.x }), y: mean(points.map { $0.y }))
    let translatedPoints = points.map { p in p-centroid }

    // Compute a = ∑ (yᵢ²-xᵢ²). Instead of a direct summation, we use the factored form of
    // yᵢ²-xᵢ²: (yᵢ-xᵢ)×(yᵢ-xᵢ). This avoids “catastrophic cancellation”. Thanks to Alan
    // Nunns for identifying this improvement. Also see:
    //
    // Goldberg, David. 1991. “What Every Computer Scientist Should Know About
    // Floating-Point Arithmetic.” Computing Surveys (CSUR 23 (1): 5–48.
    // doi:10.1145/103162.103163.
    let a = translatedPoints
        .map({ p in (p.y-p.x) * (p.y+p.x) })
        .reduce(0.0, +)

    let b = translatedPoints
        .map({ p in p.x * p.y })
        .reduce(0.0, +)

    if (a == 0 && b == 0) {
        return nil
    } else {
        let beta = 2*b
        let gamma = hypot(a, beta)

        // Avoid subtracting two nearby floating-points to avoid “catastrophic
        // cancellation”. This could happen when a > 0 and a² ≫ b²; then a ≈ gamma.
        let theta = a > 0 ? atan2(-beta, a+gamma) : atan2(a-gamma, beta);
        let rho = centroid.x*cos(theta) + centroid.y*sin(theta)
        return Line(rho: rho, theta: theta)
    }
}