Merge pull request #23 from EconForge/albop/interp

Albop/interp: merging as is.

Author: albop

.travis.yml

@@ -1,8 +1,7 @@
 language: python
 
 python:
-    - "2.7"
-    - "3.5"
+    - "3.6"
 
 sudo: false
 
@@ -11,11 +10,7 @@ install:
   # conda line below will keep everything up-to-date.  We do this
   # conditionally because it saves us some downloading if the version is
   # the same.
-  - if [[ "$TRAVIS_PYTHON_VERSION" == "2.7" ]]; then
-      wget http://repo.continuum.io/miniconda/Miniconda-latest-Linux-x86_64.sh -O miniconda.sh;
-    else
-      wget http://repo.continuum.io/miniconda/Miniconda3-latest-Linux-x86_64.sh -O miniconda.sh;
-    fi
+  - wget http://repo.continuum.io/miniconda/Miniconda3-latest-Linux-x86_64.sh -O miniconda.sh;
   - bash miniconda.sh -b -p $HOME/miniconda
   - export PATH="$HOME/miniconda/bin:$PATH"
   - hash -r

README.md

@@ -49,6 +49,28 @@ In the near future:
 - JIT classes for all interpolation objects
 
 
+## jitted, non-uniform multilinear interpolation
+
+There is a simple `interp` function with a flexible API which does multinear on uniform or non uniform cartesian grids.
+
+```
+### 1d grid
+from interpolation import interp
+
+x = np.linspace(0,1,100)**2 # non-uniform points
+y = np.linspace(0,1,100)    # values
+
+# interpolate at one point:
+interp(x,y,0.5)
+
+# or at many points:
+u = np.linspace(0,1,1000)   # points
+interp(x,y,u)
+```
+
+
+
+
 
 
 ## smolyak interpolation

examples/example_mlinterp.py

@@ -0,0 +1,129 @@
+import numpy as np
+
+from numba import generated_jit
+import ast
+
+C = ((0.1,0.2),(0.1,0.2))
+l = (0.1,0.5)
+
+from interpolation.multilinear.fungen import extract_row, tensor_reduction
+
+tensor_reduction(C,l)
+
+ll = np.row_stack([1-np.array(l),l])
+ll
+np.einsum('ij,i,j', np.array(C), ll[0,:], ll[1,:])
+
+A = np.random.random((5,5))
+extract_row(A, 1, (2,2))
+
+from interpolation.multilinear.fungen import get_coeffs
+get_coeffs(A, (1,2))
+
+##########
+# interp #
+##########
+
+### 1d interpolation
+
+from interpolation import interp
+
+x = np.linspace(0,1,100)**2 # non-uniform points
+y = np.linspace(0,1,100)    # values
+
+# interpolate at one point:
+interp(x,y,0.5)
+
+# or at many points:
+u = np.linspace(0,1,1000)   # points
+interp(x,y,u)
+
+# one can iterate at low cost since the function is jitable:
+from numba import njit
+@njit
+def vec_eval(u):
+    N = u.shape[0]
+    out = np.zeros(N)
+    for n in range(N):
+        out[n] = interp(x,y,u)
+    return out
+
+print( abs(vec_eval(u) - interp(x,y,u)).max())
+
+
+### 2d interpolation (same for higher orders)
+
+
+from interpolation import interp
+
+x1 = np.linspace(0,1,100)**2 # non-uniform points
+x2 = np.linspace(0,1,100)**2 # non-uniform points
+y = np.array([[np.sqrt(u1**2 + u2**2) for u2 in x2] for u1 in x1])
+# (y[i,j] = sqrt(x1[i]**2+x2[j]**2)
+
+
+# interpolate at one point:
+
+interp(x1,x2,y,0.5,0.2)
+interp(x1,x2,y,(0.5,0.2))
+
+# or at many points: (each line corresponds to one observation)
+points = np.random.random((1000,2))
+interp(x1,x2,y,points)
+
+from numba import njit
+@njit
+def vec_eval(p):
+    N = p.shape[0]
+    out = np.zeros(N)
+    for n in range(N):
+        z1 = p[n,0]
+        z2 = p[n,1]
+        out[n] = interp(x1,x2,y,z1,z2)
+    return out
+
+print( abs(vec_eval(points) - interp(x1,x2,y,points)).max())
+
+
+# in the special case where the points at which one wants to interpolate
+# form a cartesian grid, one can use another call style:
+
+z1 = np.linspace(0,1,100)
+z2 = np.linspace(0,1,100)
+out = interp(x1,x2,y,z1,z2)
+# out[i,j] contains f(z1[i],z2[j])
+
+
+
+############
+# mlinterp #
+############
+
+# same as interp but with less flexible and more general API
+
+from interpolation import mlinterp
+
+x1 = np.linspace(0,1,100)**2 # non-uniform points for first dimensoin
+x2 = (0,1,100) # uniform points for second dimension
+grid = (x1,x2)
+y = np.array([[np.sqrt(u1**2 + u2**2) for u2 in x2] for u1 in x1])
+
+
+points = np.random.random((1000,2))
+
+# vectorized call:
+mlinterp(grid, y, points)
+
+# non-vectorized call (note third argument must be a tuple of floats of right size)
+mlinterp(grid, y, (0.4, 0.2))
+
+# arbitrary dimension
+
+d = 4
+K = 100
+N = 10000
+grid = (np.linspace(0,1,K),)*d
+y = np.random.random((K,)*d)
+z = np.random.random((N,d))
+
+mlinterp(grid,y,z)

interpolation/__init__.py

@@ -0,0 +1 @@
+from .multilinear.mlinterp import interp, mlinterp

interpolation/multilinear/__init__.py

@@ -0,0 +1,189 @@
+import numba
+import numpy as np
+from numba import float64, int64
+from numba import generated_jit, njit
+import ast
+
+from numba.extending import overload
+from numba.types.containers import Tuple, UniTuple
+
+
+####################
+# Dimension helper #
+####################
+
+t_coord = numba.typeof((2.3,2.4,1))           # type of an evenly spaced dimension
+t_array = numba.typeof(np.array([4.0, 3.9]))  # type of an unevenly spaced dimension
+
+# returns the index of a 1d point along a 1d dimension
+@generated_jit(nopython=True)
+def get_index(gc, x):
+    if gc == t_coord:
+        # regular coordinate
+        def fun(gc,x):
+            δ = (gc[1]-gc[0])/(gc[2]-1)
+            d = x-gc[0]
+            ii = d // δ
+            r = d-ii*δ
+            i = int(ii)
+            λ = r/δ
+            return (i, λ)
+        return fun
+    else:
+        # irregular coordinate
+        def fun(gc,x):
+            i = int(np.searchsorted(gc, x))-1
+            λ = (x-gc[i])/(gc[i+1]-gc[i])
+            return (i, λ)
+        return fun
+
+# returns number of dimension of a dimension
+@generated_jit(nopython=True)
+def get_size(gc):
+    if gc == t_coord:
+        # regular coordinate
+        def fun(gc):
+            return gc[2]
+        return fun
+    else:
+        # irregular coordinate
+        def fun(gc):
+            return len(gc)
+        return fun
+
+#####################
+# Generator helpers #
+#####################
+
+# the next functions replace the use of generators, with the difference that the
+# output is a tuple, which dimension is known by the jit guy.
+
+# example:
+# ```
+# def f(x): x**2
+# fmap(f, (1,2,3)) -> (1,3,9)
+# def g(x,y): x**2 + y
+# fmap(g, (1,2,3), 0.1) -> (1.1,3.1,9.1)  # (g(1,0.1), g(2,0.1), g(3,0.1))
+# def g(x,y): x**2 + y
+# fmap(g, (1,2,3), (0.1,0.2,0.3)) -> (1.1,3.0.12,9.3)
+# ```
+
+
+def fmap():
+    pass
+
+@overload(fmap)
+def _map(*args):
+
+    if len(args)==2 and isinstance(args[1], (Tuple, UniTuple)):
+        k = len(args[1])
+        s = "def __map(f, t): return ({}, )".format(str.join(', ',['f(t[{}])'.format(i) for i in range(k)]))
+    elif len(args)==3 and isinstance(args[1], (Tuple, UniTuple)):
+        k = len(args[1])
+        if isinstance(args[2], (Tuple, UniTuple)):
+            if len(args[2]) != k:
+                # we don't know what to do in this case
+                return None
+            s = "def __map(f, t1, t2): return ({}, )".format(str.join(', ',['f(t1[{}], t2[{}])'.format(i,i) for i in range(k)]))
+        else:
+            s = "def __map(f, t1, x): return ({}, )".format(str.join(', ',['f(t1[{}], x)'.format(i,i) for i in range(k)]))
+    else:
+        return None
+    d = {}
+    eval(compile(ast.parse(s),'<string>','exec'), d)
+    return d['__map']
+
+
+# not that `fmap` does nothing in python mode...
+# an alternative would be
+#
+# @njit
+# def _fmap():
+#     pass
+#
+# @overload(_fmap)
+# ...
+#
+# @njit
+# def fmap(*args):
+#     return _fmap(*args)
+#
+# but this seems to come with a performance cost.
+# It it is also possible to overload `map` but we would risk
+# a future conflict with the map api.
+
+
+#
+# @njit
+# def fmap(*args):
+#     return _fmap(*args)
+#
+# funzip(((1,2), (2,3), (4,3))) -> ((1,2,4),(2,3,3))
+
+@generated_jit(nopython=True)
+def funzip(t):
+    k = t.count
+    assert(len(set([e.count for e in t.types]))==1)
+    l = t.types[0].count
+    def print_tuple(t): return "({},)".format(str.join(", ", t))
+    tab =  [ [ 't[{}][{}]'.format(i,j) for i in range(k)] for j in range(l) ]
+    s = "def funzip(t): return {}".format(print_tuple( [print_tuple(e) for e in tab] ))
+    d = {}
+    eval(compile(ast.parse(s),'<string>','exec'), d)
+    return d['funzip']
+
+
+#####
+# array subscribing:
+# when X is a 2d array and I=(i,j) a 2d index, `get_coeffs(X,I)`
+# extracts `X[i:i+1,j:j+1]` but represents it as a tuple of tuple, so that
+# the number of its elements can be inferred by the compiler
+#####
+
+
+@generated_jit(nopython=True)
+def get_coeffs(X,I):
+    if X.ndim>len(I):
+        print("not implemented yet")
+    else:
+        from itertools import product
+        d = len(I)
+        mat = np.array( ["X[{}]".format(str.join(',', e)) for e in product(*[(f"I[{j}]", f"I[{j}]+1") for j in range(d)])] ).reshape((2,)*d)
+        mattotup = lambda mat: mat if isinstance(mat,str) else "({})".format(str.join(',',[mattotup(e) for e in mat]))
+        t = mattotup(mat)
+        s = "def get_coeffs(X,I): return {}".format(t)
+        dd = {}
+        eval(compile(ast.parse(s),'<string>','exec'), dd)
+        return dd['get_coeffs']
+        return fun
+
+# tensor_reduction(C,l)
+# (in 2d) computes the equivalent of np.einsum('ij,i,j->', C, [1-l[0],l[0]], [1-l[1],l[1]])`
+# but where l is a tuple and C a tuple of tuples.
+
+# this one is a temporary implementation (should be merged with old gen_splines* code)
+def gen_tensor_reduction(X, symbs, inds=[]):
+    if len(symbs) == 0:
+        return '{}[{}]'.format(X, str.join('][',[str(e) for e in inds]))
+    else:
+        h = symbs[0]
+        q = symbs[1:]
+        exprs = [  '{}*({})'.format(h if i==1 else '(1-{})'.format(h),gen_tensor_reduction(X, q,inds + [i])) for i in range(2)]
+        return str.join( ' + ', exprs )
+
+
+@generated_jit(nopython=True)
+def tensor_reduction(C,l):
+    ex = gen_tensor_reduction('C', ['l[{}]'.format(i) for i in range(len(l.types))])
+    dd = dict()
+    s = "def tensor_reduction(C,l): return {}".format(ex)
+    eval(compile(ast.parse(s),'<string>','exec'), dd)
+    return dd['tensor_reduction']
+
+@generated_jit(nopython=True)
+def extract_row(a, n, tup):
+    d = len(tup.types)
+    dd = {}
+    s = "def extract_row(a, n, tup): return ({},)".format(str.join(', ', [f"a[n,{i}]" for i in range(d)]))
+    eval(compile(ast.parse(s),'<string>','exec'), dd)
+    return dd['extract_row']