{"id":19645000,"url":"https://github.com/quantecon/jaxinterp","last_synced_at":"2025-07-02T03:35:32.801Z","repository":{"id":174846143,"uuid":"652870110","full_name":"QuantEcon/jaxinterp","owner":"QuantEcon","description":"JAX compatible multilinear interpolation","archived":false,"fork":false,"pushed_at":"2023-06-13T23:40:36.000Z","size":4,"stargazers_count":8,"open_issues_count":3,"forks_count":2,"subscribers_count":6,"default_branch":"main","last_synced_at":"2025-06-29T09:33:14.772Z","etag":null,"topics":[],"latest_commit_sha":null,"homepage":null,"language":null,"has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":null,"status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/QuantEcon.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":null,"code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null},"funding":{"github":"numfocus","custom":"https://numfocus.org/donate-to-quantecon"}},"created_at":"2023-06-13T00:55:42.000Z","updated_at":"2024-12-07T19:22:42.000Z","dependencies_parsed_at":null,"dependency_job_id":"f475eb82-2e33-4a43-83cd-f94c49cc4441","html_url":"https://github.com/QuantEcon/jaxinterp","commit_stats":{"total_commits":5,"total_committers":2,"mean_commits":2.5,"dds":"0.19999999999999996","last_synced_commit":"933e504f0aceabbe13600825fac0f42fec60c559"},"previous_names":["quantecon/jaxinterp"],"tags_count":0,"template":false,"template_full_name":null,"purl":"pkg:github/QuantEcon/jaxinterp","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/QuantEcon%2Fjaxinterp","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/QuantEcon%2Fjaxinterp/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/QuantEcon%2Fjaxinterp/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/QuantEcon%2Fjaxinterp/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/QuantEcon","download_url":"https://codeload.github.com/QuantEcon/jaxinterp/tar.gz/refs/heads/main","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/QuantEcon%2Fjaxinterp/sbom","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":263071117,"owners_count":23409245,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":[],"created_at":"2024-11-11T14:31:28.860Z","updated_at":"2025-07-02T03:35:32.762Z","avatar_url":"https://github.com/QuantEcon.png","language":null,"funding_links":["https://github.com/sponsors/numfocus","https://numfocus.org/donate-to-quantecon"],"categories":[],"sub_categories":[],"readme":"# JAX compatible multilinear interpolation\n\nWe need fast linear interpolation in arbitrary dimensions, compatible with\n`jax.jit`.\n\nSuppose we have a function $f$ mapping $\\mathbb R^n$ to $\\mathbb R$ and we evalute $f$ on a\nfinite set of grid points $v_1, \\ldots, v_k$, where each $v_i$ is a point in\n$\\mathbb R^n$.  \n\nLet $f_i = f(v_i)$.\n\nWe now want to be able to evaluate $\\hat f(x)$, which is a linear interpolation\nof the grid points $v_1, \\ldots, v_k$ and corresponding function values \n$f_1, \\ldots, f_k$, at the point $x$.\n\nWe want to do this where $n$ can be any integer.\n\nWe also want to be able to do this in a vectorized, parallelized manner, so that\nwe can pass a large number of evaluation points $x_1, \\ldots, x_m$, and compute\nthe points $\\hat f(x_1), \\ldots, \\hat f(x_m)$ efficiently.\n\nThis is already possible in Numba, using the Econforge\n[interpolation library](https://www.econforge.org/interpolation.py/):\n\n\n```python\nimport numpy as np\nfrom interpolation.splines import UCGrid, CGrid, nodes\nfrom interpolation.splines import eval_linear\n\n# function to interpolate\nf = lambda x,y: np.sin(np.sqrt(x**2+y**2+0.00001))/np.sqrt(x**2+y**2+0.00001)\n\n# uniform cartesian grid with a small number of points\ngrid = UCGrid((-1.0, 1.0, 10), (-1.0, 1.0, 10))\n\n# get grid points\ngp = nodes(grid)   # 100x2 matrix\n\n# compute values on grid points\nvalues = f(gp[:,0], gp[:,1]).reshape((10,10))\n\n# interpolate at one point\npoint = np.array([0.1,0.45]) # 1d array\nval = eval_linear(grid, values, point)  # float\n\n# interpolate at many points:\npoints = np.random.random((10000,2))\neval_linear(grid, values, points) # 10000 vector\n```\n\nThe question is, what is the most efficient way to replicate this in JAX?\n\nUsing `jax.scipy.ndimage.map_coordinates`, @junnanZ has implemented evaluation at a single point for $n=4$ in [this file](https://github.com/jstac/sdfs_via_autodiff/blob/main/code/ssy/continuous_junnan/utils.py).\n\nWe could then extend this from one `point` to `points` via `vmap`.\n\nThis discussion suggests that the vmap approach will be efficient: https://github.com/google/jax/issues/6312\n\nOther alternatives include https://www.tensorflow.org/probability/api_docs/python/tfp/substrates/jax/math/batch_interp_regular_nd_grid\n\nAre these our best options?  If so, which is best?\n\nOnce we have decided, let's \n\n1. test to make sure that the output agrees with the output from the Python code\n   above, and\n1. provide a nice interface, like the one in the Python code above.\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fquantecon%2Fjaxinterp","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fquantecon%2Fjaxinterp","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fquantecon%2Fjaxinterp/lists"}