Skip to content

← JAX Numerical Computing step 21 of 25

Medium Primitives

Linear Interpolation

Jump to the task

Estimate between known points

Linear interpolation joins adjacent observations with straight line segments. To estimate a value between two coordinates, move the same fraction of the way between their observed values.

For endpoints (x0, y0) and (x1, y1) and a query q between them:

fraction = (q - x0) / (x1 - x0)
estimate = y0 + fraction * (y1 - y0)

jnp.interp finds the appropriate segment and applies this calculation. It is useful for lookup tables or resampling measurements onto a new grid. The resulting curve is continuous, but its slope can change abruptly at the known coordinates. It is not generally a smooth curve or an exact reconstruction of the process that produced the samples.

Worked example

import jax.numpy as jnp

x_known = jnp.array([0.0, 1.0, 2.0])
y_known = jnp.array([0.0, 2.0, 0.0])
x_query = jnp.array([-1.0, 0.5, 1.0, 1.5, 3.0])
result = jnp.interp(x_query, x_known, y_known)
# [0.0, 1.0, 2.0, 1.0, 0.0]

At 0.5, we are halfway from (0, 0) to (1, 2), so the result is 1. At 1, the query matches an observation and returns 2. Outside the known interval, the default behavior holds the nearest endpoint's value constant: the queries -1 and 3 both return 0 here.

Problem

Implement interp_at(x_query, x_known, y_known) using jnp.interp with its default boundary behavior.

  • Inputs are one-dimensional float32 arrays of finite values.
  • x_known is nonempty and strictly increasing. y_known has the same length and contains the corresponding observations.
  • x_query may be in any order. Return one value per query, preserving that order and shape. Empty queries give empty output.
  • Below the first known coordinate, return the first observed value. Above the last coordinate, return the last observed value.

With one known point, every query returns that point's observed value.

Things to watch

Keep each coordinate paired with its observation. If you sort unsorted data elsewhere, apply the same permutation to x_known and y_known. Sorting only the coordinates changes the data. Inputs here are already sorted; do not rely on JAX to validate or fix their ordering.

Endpoint values are held constant by default, but this is not the only behavior supported by the API. JAX also supports left, right, and period, including left="extrapolate" and right="extrapolate" for extending endpoint line segments. Leave those options at their defaults for this exercise. Interpolation between samples is different from extrapolation beyond them.