Skip to content

Commit d65475d

Browse files
authored
Merge pull request #369 from fslaborg/repo-assist/improve-weighted-slr-20260406-4b1ca01b9d045c66
[Repo Assist] feat: implement OLS.Linear.Univariable.fitWithWeighting (weighted simple linear regression)
2 parents 9d2635c + 6c542d9 commit d65475d

2 files changed

Lines changed: 116 additions & 1 deletion

File tree

src/FSharp.Stats/Fitting/LinearRegression.fs

Lines changed: 50 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -288,6 +288,51 @@ module LinearRegression =
288288
let coefficientConstrained (xData : Vector<float>) (yData : Vector<float>) ((xC,yC): float*float) =
289289
(fitConstrained xData yData (xC,yC)).Coefficients
290290

291+
/// <summary>
292+
/// Calculates the intercept and slope for a weighted straight line fitting the data.
293+
/// Weighted least squares minimizes Σ w_i * (y_i - (a + b*x_i))².
294+
/// </summary>
295+
/// <param name="weighting">vector of non-negative weights, one per observation</param>
296+
/// <param name="xData">vector of x values</param>
297+
/// <param name="yData">vector of y values</param>
298+
/// <returns>Coefficients of [intercept; slope]</returns>
299+
/// <example>
300+
/// <code>
301+
/// let xData = vector [|1.;2.;3.;4.;5.;6.|]
302+
/// let yData = vector [|4.;7.;9.;10.;11.;15.|]
303+
/// // down-weight the last observation
304+
/// let weights = vector [|1.;1.;1.;1.;1.;0.1|]
305+
/// let coefficients =
306+
/// Univariable.fitWithWeighting weights xData yData
307+
/// </code>
308+
/// </example>
309+
let fitWithWeighting (weighting: Vector<float>) (xData: Vector<float>) (yData: Vector<float>) =
310+
if xData.Length <> yData.Length || xData.Length <> weighting.Length then
311+
raise (System.ArgumentException("Vectors x, y and weighting must have the same length!"))
312+
// Closed-form WLS normal equations for y = a + b*x:
313+
// [Σw Σwx ] [a] [Σwy ]
314+
// [Σwx Σwx²] [b] = [Σwxy]
315+
let mutable sw = 0.
316+
let mutable swx = 0.
317+
let mutable swy = 0.
318+
let mutable swxx = 0.
319+
let mutable swxy = 0.
320+
for i = 0 to xData.Length - 1 do
321+
let wi = weighting.[i]
322+
let xi = xData.[i]
323+
let yi = yData.[i]
324+
sw <- sw + wi
325+
swx <- swx + wi * xi
326+
swy <- swy + wi * yi
327+
swxx <- swxx + wi * xi * xi
328+
swxy <- swxy + wi * xi * yi
329+
let denom = sw * swxx - swx * swx
330+
if abs denom < System.Double.Epsilon then
331+
raise (System.ArgumentException("Degenerate weighting: all weight is concentrated at a single x value."))
332+
let slope = (sw * swxy - swx * swy) / denom
333+
let intercept = (swy - slope * swx) / sw
334+
Coefficients([|intercept; slope|])
335+
291336
/// <summary>
292337
/// Takes intercept and slope of simple linear regression to predict the corresponding y value.
293338
/// </summary>
@@ -954,7 +999,11 @@ type LinearRegression() =
954999
LinearRegression.OLS.Linear.RTO.fit xData yData
9551000
| Constraint.RegressionThroughXY coordinate ->
9561001
LinearRegression.OLS.Linear.Univariable.fitConstrained xData yData coordinate
957-
| _ -> failwithf "Weighted simple linear regression is not yet implemented! Use polynomial weighted regression with degree 1 instead."
1002+
| Some w ->
1003+
match _constraint with
1004+
| Constraint.Unconstrained ->
1005+
LinearRegression.OLS.Linear.Univariable.fitWithWeighting w xData yData
1006+
| _ -> failwithf "Constrained weighted simple linear regression is not yet implemented!"
9581007

9591008
| Method.Polynomial o ->
9601009
match _constraint with

tests/FSharp.Stats.Tests/Fitting.fs

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,72 @@ let leastSquaresCholeskyTests =
5555
)
5656
]
5757
open FSharp.Stats.Fitting.Spline
58+
59+
[<Tests>]
60+
let weightedSimpleLinearRegressionTests =
61+
// Known exact case: y = 2 + 3x, equal weights → coefficients must be exact
62+
let xData = vector [|1.; 2.; 3.; 4.; 5.|]
63+
let yExact = vector [|5.; 8.; 11.; 14.; 17.|] // 2 + 3x
64+
65+
testList "Weighted Simple Linear Regression" [
66+
67+
testCase "Equal unit weights reproduce unweighted fit" (fun () ->
68+
let weights = vector [|1.; 1.; 1.; 1.; 1.|]
69+
let wCoef = LinearRegression.OLS.Linear.Univariable.fitWithWeighting weights xData yExact
70+
let uCoef = LinearRegression.OLS.Linear.Univariable.fit xData yExact
71+
Expect.floatClose Accuracy.high wCoef.Constant uCoef.Constant "Intercept should match unweighted"
72+
Expect.floatClose Accuracy.high wCoef.Linear uCoef.Linear "Slope should match unweighted"
73+
)
74+
75+
testCase "Exact line – intercept 2, slope 3" (fun () ->
76+
let weights = vector [|1.; 1.; 1.; 1.; 1.|]
77+
let coef = LinearRegression.OLS.Linear.Univariable.fitWithWeighting weights xData yExact
78+
Expect.floatClose Accuracy.high coef.Constant 2. "Intercept should be 2"
79+
Expect.floatClose Accuracy.high coef.Linear 3. "Slope should be 3"
80+
)
81+
82+
testCase "Down-weighting an outlier pulls fit toward true line" (fun () ->
83+
// y = 2 + 3x except the last point is a large outlier
84+
let yOutlier = vector [|5.; 8.; 11.; 14.; 100.|]
85+
let weightsFlat = vector [|1.; 1.; 1.; 1.; 1.|]
86+
let weightsDownLast = vector [|1.; 1.; 1.; 1.; 0.001|]
87+
let coefFlat = LinearRegression.OLS.Linear.Univariable.fitWithWeighting weightsFlat xData yOutlier
88+
let coefDown = LinearRegression.OLS.Linear.Univariable.fitWithWeighting weightsDownLast xData yOutlier
89+
// The down-weighted fit should be closer to slope=3, intercept=2
90+
let errorFlat = abs (coefFlat.Linear - 3.) + abs (coefFlat.Constant - 2.)
91+
let errorDown = abs (coefDown.Linear - 3.) + abs (coefDown.Constant - 2.)
92+
Expect.isTrue (errorDown < errorFlat) "Down-weighting outlier should give fit closer to true line"
93+
)
94+
95+
testCase "Doubling all weights does not change coefficients" (fun () ->
96+
let yData = vector [|4.; 7.; 9.; 10.; 11.|]
97+
let w1 = vector [|1.; 1.; 1.; 1.; 1.|]
98+
let w2 = vector [|2.; 2.; 2.; 2.; 2.|]
99+
let coef1 = LinearRegression.OLS.Linear.Univariable.fitWithWeighting w1 xData yData
100+
let coef2 = LinearRegression.OLS.Linear.Univariable.fitWithWeighting w2 xData yData
101+
Expect.floatClose Accuracy.high coef1.Constant coef2.Constant "Intercept invariant to weight scaling"
102+
Expect.floatClose Accuracy.high coef1.Linear coef2.Linear "Slope invariant to weight scaling"
103+
)
104+
105+
testCase "Agrees with Polynomial.fitWithWeighting order 1" (fun () ->
106+
let yData = vector [|4.; 7.; 9.; 10.; 11.|]
107+
let weights = vector [|2.; 1.; 0.5; 1.; 1.|]
108+
let coefUni = LinearRegression.OLS.Linear.Univariable.fitWithWeighting weights xData yData
109+
let coefPoly = LinearRegression.OLS.Polynomial.fitWithWeighting 1 weights xData yData
110+
Expect.floatClose Accuracy.high coefUni.Constant coefPoly.Constant "Intercept agrees with poly order-1"
111+
Expect.floatClose Accuracy.high coefUni.Linear coefPoly.Linear "Slope agrees with poly order-1"
112+
)
113+
114+
testCase "LinearRegressor.fit dispatches to weighted path" (fun () ->
115+
let yData = vector [|4.; 7.; 9.; 10.; 11.|]
116+
let weights = vector [|2.; 1.; 0.5; 1.; 1.|]
117+
let coefDirect = LinearRegression.OLS.Linear.Univariable.fitWithWeighting weights xData yData
118+
let coefDispatch = LinearRegression.fit(xData, yData, FittingMethod = Method.SimpleLinear, Weighting = weights)
119+
Expect.floatClose Accuracy.high coefDirect.Constant coefDispatch.Constant "Intercept matches dispatch"
120+
Expect.floatClose Accuracy.high coefDirect.Linear coefDispatch.Linear "Slope matches dispatch"
121+
)
122+
]
123+
58124
[<Tests>]
59125
let splineTests =
60126
testList "Fitting.Spline" [

0 commit comments

Comments
 (0)