123 lines
3.4 KiB
TypeScript
123 lines
3.4 KiB
TypeScript
import * as math from 'npm:mathjs'
|
|
|
|
/**
|
|
* Performs non-linear least squares regression of data to an arbitrary function using math.js.
|
|
*
|
|
* @param {number[]} xData - Array of x-coordinates.
|
|
* @param {number[]} yData - Array of y-coordinates.
|
|
* @param {(x: number, params: number[]) => number} func - The function to fit to the data. Takes x and an array of parameters as input.
|
|
* @param {number[]} initialParams - Initial guesses for the parameters.
|
|
* @param {number} maxIterations - Maximum number of iterations.
|
|
* @param {number} tolerance - Tolerance for convergence.
|
|
* @returns {{ params: number[], residuals: number[], rSquared: number, iterations: number, converged: boolean }} - An object containing the fitted parameters and other information.
|
|
*/
|
|
export function curveFit<T extends number[]>(
|
|
xData: number[],
|
|
yData: number[],
|
|
func: (x: number, params: T) => number,
|
|
initialParams: T,
|
|
maxIterations: number = 1000,
|
|
tolerance: number = 1e-8,
|
|
): {
|
|
params: T
|
|
residuals: number[]
|
|
rSquared: number
|
|
iterations: number
|
|
converged: boolean
|
|
} {
|
|
const numPoints: number = xData.length
|
|
const numParams: number = initialParams.length
|
|
|
|
if (numPoints !== yData.length) {
|
|
throw new Error('xData and yData must have the same length.')
|
|
}
|
|
|
|
const params: number[] = [...initialParams]
|
|
let iterations: number = 0
|
|
let converged: boolean = false
|
|
|
|
while (iterations < maxIterations) {
|
|
const jacobian: number[][] = []
|
|
const residuals: number[] = []
|
|
|
|
for (let i: number = 0; i < numPoints; i++) {
|
|
const x: number = xData[i]
|
|
const y: number = yData[i]
|
|
const predicted: number = func(x, params as T)
|
|
const residual: number = y - predicted
|
|
residuals.push(residual)
|
|
|
|
const row: number[] = []
|
|
for (let j: number = 0; j < numParams; j++) {
|
|
const delta: number = 1e-6
|
|
const paramsPlusDelta = [...params] as T
|
|
paramsPlusDelta[j] += delta
|
|
const predictedPlusDelta: number = func(x, paramsPlusDelta)
|
|
const derivative: number = (predictedPlusDelta - predicted) / delta
|
|
row.push(derivative)
|
|
}
|
|
jacobian.push(row)
|
|
}
|
|
|
|
const jacobianTranspose: number[][] = math.transpose(jacobian)
|
|
const hessian: math.Matrix = math.matrix(
|
|
math.multiply(jacobianTranspose, jacobian),
|
|
)
|
|
const gradient: math.Matrix = math.matrix(
|
|
math.multiply(jacobianTranspose, residuals),
|
|
)
|
|
|
|
try {
|
|
const deltaParams: math.Matrix = math.usolve(
|
|
hessian,
|
|
gradient,
|
|
) as math.Matrix
|
|
|
|
let maxDelta: number = 0
|
|
for (let i: number = 0; i < numParams; i++) {
|
|
params[i] += deltaParams.get([i, 0])
|
|
maxDelta = Math.max(maxDelta, Math.abs(deltaParams.get([i, 0])))
|
|
}
|
|
|
|
if (maxDelta < tolerance) {
|
|
converged = true
|
|
break
|
|
}
|
|
} catch (_error) {
|
|
// console.error('Matrix solve failed:', error)
|
|
return {
|
|
params: params as T,
|
|
residuals: residuals,
|
|
rSquared: NaN,
|
|
iterations: iterations,
|
|
converged: false,
|
|
}
|
|
}
|
|
|
|
iterations++
|
|
}
|
|
|
|
const calculatedResiduals: number[] = xData.map((x: number, i: number) =>
|
|
yData[i] - func(x, params as T)
|
|
)
|
|
const ssRes: number = calculatedResiduals.reduce(
|
|
(sum: number, r: number) => sum + r * r,
|
|
0,
|
|
)
|
|
const yMean: number = yData.reduce((sum: number, y: number) => sum + y, 0) /
|
|
numPoints
|
|
const ssTot: number = yData.reduce(
|
|
(sum: number, y: number) => sum + (y - yMean) * (y - yMean),
|
|
0,
|
|
)
|
|
const rSquared: number = 1 - ssRes / ssTot
|
|
|
|
return {
|
|
params: params as T,
|
|
residuals: calculatedResiduals,
|
|
rSquared: rSquared,
|
|
iterations: iterations,
|
|
converged: converged,
|
|
}
|
|
}
|