Pytensor-native interpolation functions

[jessegrabowski]

Description

Adds a new file, interpolation.py. Objective is to match a useful subset of scipy.interpolate. So far I have 1d interpolation with different strategies. Example:

import numpy as np
from pytensor.tensor.interpolate import interpolate1d
import matplotlib.pyplot as plt

x = np.linspace(-2, 6, 10)
y = np.sin(x)

f_linear = interpolate1d(x, y, method='linear')
f_nearest = interpolate1d(x, y, method='nearest')
f_first = interpolate1d(x, y, method='first')
f_last = interpolate1d(x, y, method='last')
f_mean = interpolate1d(x, y, method='mean')

x_hat = pt.dvector('x_hat')
ops = [f_linear, f_nearest, f_first, f_last, f_mean]
MODE = 'FAST_RUN'
funcs = [
    pytensor.function([x_hat], op(x_hat), mode=MODE) for op in ops
]

inputs = np.linspace(-2.5, 6.5, 1000)
labels = ['linear', 'nearest', 'first', 'last', 'mean']
fig, ax = plt.subplots()
ax.scatter(x, y, color='k')
for f, label in zip(funcs, labels):
    ax.plot(inputs, f(inputs), label=label)
ax.legend()
plt.show()

Everything is composed of pytensor primitives, so it ought to compile to JAX/NUMBA without additional changes. It seems there are so bugs though, so they don't work yet. Specifically, NUMBA raises the following error:

TypeError: pytensor.link.numba.dispatch.basic.numba_funcify() got multiple values for keyword argument 'parent_node'

And it also gives a warning about a missing Supervisor, which might or might not be related.

In addition to this more feature rich interpolate1d API, I also added pt.interp, which matches the numpy signature 1:1:

from pytensor.tensor.interpolate import interp

xp = [1., 2., 3.]
fp = [3., 2., 0.]
interp(2.5, xp, fp)
# 1.0
interp([0, 1, 1.5, 2.72, 3.14], xp, fp)
# array([3.  , 3.  , 2.5 , 0.56, 0.  ])

UNDEF = -99.0
interp(3.14, xp, fp, right=UNDEF)
# -99.0

I figure I'll also do a function that do polynomial interpolation (fit via least-squares). The holy grail here is BSplines, but how to do that is somewhat beyond me at the moment. I think just getting these simple methods into the library would be useful, then we can open issues to add more interesting methods.

Related Issue

• Closes Implement 1d interpolation #1140
• Related to Interpolated not supported in JAX/Numba backends pymc#6838

Checklist

• Checked that the pre-commit linting/style checks pass
• Included tests that prove the fix is effective or that the new feature works
• Added necessary documentation (docstrings and/or example notebooks)
• If you are a pro: each commit corresponds to a relevant logical change

Type of change

• New feature / enhancement
• Bug fix
• Documentation
• Maintenance
• Other (please specify):

---

📚 Documentation preview 📚: https://pytensor--1141.org.readthedocs.build/en/1141/

[ricardoV94]

Can you add labels to the PR?

[codecov]

Codecov Report

Attention: Patch coverage is 95.00000% with 4 lines in your changes missing coverage. Please review.

Project coverage is 82.14%. Comparing base (83c6b44) to head (d5daef1).
Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
pytensor/tensor/interpolate.py	94.59%	2 Missing and 2 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1141      +/-   ##
==========================================
+ Coverage   82.12%   82.14%   +0.02%     
==========================================
  Files         185      186       +1     
  Lines       48130    48210      +80     
  Branches     8669     8678       +9     
==========================================
+ Hits        39527    39603      +76     
- Misses       6438     6440       +2     
- Partials     2165     2167       +2     

Files with missing lines	Coverage Δ
pytensor/link/jax/dispatch/extra_ops.py	82.35% <100.00%> (+1.70%)	⬆️
pytensor/tensor/interpolate.py	94.59% <94.59%> (ø)