Skip to content

← JAX Numerical Computing step 19 of 25

Medium Primitives

Polynomial Fitting via polyfit

Jump to the task

Fit a curve to observations

jnp.polyfit(x, y, deg) chooses polynomial coefficients that minimize the sum of squared differences between predictions and observations. Degree 1 fits a line, degree 2 fits a quadratic, and degree 0 fits a constant. With noisy data, the fitted curve generally does not pass through every observation.

The returned coefficients use descending powers. For degree 2, [a, b, c] represents a*x**2 + b*x + c. This is also the order expected by jnp.polyval, which evaluates the fitted curve.

Worked example

import jax.numpy as jnp

x = jnp.array([0.0, 1.0, 2.0])
y = jnp.array([1.0, 2.0, 2.0])
coeffs = jnp.polyfit(x, y, 1)
predictions = jnp.polyval(coeffs, x)
# coeffs: approximately [0.5, 1.166667]
# predictions: approximately [1.166667, 1.666667, 2.166667]

The fitted line is approximately 0.5*x + 1.166667. It misses each observation slightly, but its total squared error is as small as possible among lines. Fitting is different from evaluating an already-known polynomial: polyfit finds coefficients; polyval uses them.

Interpolation and numerical limits

In exact arithmetic, n distinct input coordinates determine one polynomial of degree at most n - 1 that passes through all observations. Repeated coordinates do not provide that guarantee, and floating-point fitting does not promise zero residuals or exact coefficient values.

High powers and poorly spaced coordinates can make a fit sensitive to small numerical errors. Centering and scaling the coordinates can help, but it changes the coordinate system of the returned coefficients. If you fit using z = (x - center) / scale, evaluate with transformed z values too, or convert the coefficients back. Do not silently scale x in this exercise, which asks for coefficients in the original variable.

Problem

Implement polyfit_coeffs(x, y, deg) using jnp.polyfit.

  • x and y are one-dimensional float32 arrays of the same nonzero length.
  • deg is a nonnegative whole-number scalar, supplied as a float by this exercise's input format. Convert it to a Python integer.
  • Return shape (deg + 1,), with coefficients in descending power order.

Compare floating-point coefficients with a tolerance. A mathematically zero coefficient can appear as a small nonzero value. NumPy's numpy.polynomial.Polynomial class uses ascending powers; it is a separate API, not a jnp.Polynomial class.

Degree and JIT

The degree determines the output shape and must be known when JAX traces a compiled function. Converting a concrete scalar with int(deg) works for the eager call here. It does not make a dynamic traced degree static. For JIT, pass a Python degree as a static argument, for example with jax.jit(polyfit_coeffs, static_argnames=("deg",)). Changing that static degree can require another compilation.