Merge pull request #16 from EconForge/complete_ders

WIP: Add derivatives of complete polynomials

Author: sglyon

interpolation/complete_poly.py

@@ -73,6 +73,9 @@ def n_complete(n, d):
     return out
 
 
+#
+# Complete Polynomials Basis
+#
 def complete_polynomial(z, d):
     """
     Construct basis matrix for complete polynomial of degree `d`, given
@@ -111,7 +114,10 @@ def complete_polynomial(z, d):
     return out
 
 
-@jit(nopython=True, cache=True)
+# TODO: Currently turning off all cache arguments so that
+#       code works. This will be fixed in numba 0.32
+# @jit(nopython=True, cache=True)
+@jit(nopython=True)
 def _complete_poly_impl_vec(z, d, out):
     "out and z should be vectors"
     nvar = z.shape[0]
@@ -185,7 +191,8 @@ def _complete_poly_impl_vec(z, d, out):
         return
 
 
-@jit(nopython=True, cache=True)
+# @jit(nopython=True, cache=True)
+@jit(nopython=True)
 def _complete_poly_impl(z, d, out):
     nvar = z.shape[0]
     nobs = z.shape[1]
@@ -274,6 +281,257 @@ def _complete_poly_impl(z, d, out):
         return
 
 
+#
+# Complete Polynomials Derivative Basis
+#
+def complete_polynomial_der(z, d, der):
+    """
+    Construct basis matrix for complete polynomial of degree `d`, given
+    input data `z`.
+
+    Parameters
+    ----------
+    z : np.ndarray(size=(nvariables, nobservations))
+        The degree 1 realization of each variable. For example, if
+        variables are `q`, `r`, and `s`, then `z` should be
+        `z = np.row_stack([q, r, s])`
+
+    d : int
+        An integer specifying the degree of the complete polynomial
+
+    der : int
+        An integer specifying which variable to take derivative wrt --
+        a 0 means take derivative wrt first variable in z etc...
+
+    Returns
+    -------
+    out : np.ndarray(size=(ncomplete(nvariables, d), nobservations))
+        The basis matrix for the derivative of a complete polynomial
+        of degree d with respect to variable der
+
+    """
+    # check inputs
+    assert d >= 0, "d must be non-negative"
+    assert der >= 0, "derivative must be non-negative"
+    z = np.asarray(z)
+
+    # compute inds allocate space for output
+    nvar, nobs = z.shape
+    assert der < nvar, "derivative integer must be smaller than nobs in z"
+    out = np.zeros((n_complete(nvar, d), nobs))
+
+    if d > 5:
+        raise ValueError("Complete polynomial only implemeted up to degree 5")
+
+    # populate out with jitted function
+    _complete_poly_der_impl(z, d, der, out)
+
+    return out
+
+
+# @jit(nopython=True, cache=True)
+@jit(nopython=True)
+def _complete_poly_der_impl_vec(z, d, der, out):
+    "out and z should be vectors"
+    nvar = z.shape[0]
+
+    out[0] = 0.0
+
+    # fill first order stuff
+    if d >= 1:
+        # All linear terms except for one (the variable itself) are 0
+        for i in range(nvar):
+            out[i+1] = 0.0
+        out[der+1] = 1.0
+
+    if d == 1:
+        return
+
+    # now we need to fill in row nvar and beyond
+    ix = nvar
+    if d == 2:
+        for i1 in range(nvar):
+            # Get coefficients and values
+            (c1, t1) = (1, 1.0) if i1==der else (0, z[i1])
+            for i2 in range(i1, nvar):
+                (c2, t2) = (c1+1, 1.0) if i2==der else (c1, z[i2])
+
+                # Update index and out
+                ix += 1
+                out[ix] = c2 * t1*t2 * z[der]**(c2-1) if c2>0 else 0.0
+
+        return
+
+    if d == 3:
+        for i1 in range(nvar):
+            # Get coefficients and values
+            (c1, t1) = (1, 1.0) if i1==der else (0, z[i1])
+            for i2 in range(i1, nvar):
+                (c2, t2) = (c1+1, 1.0) if i2==der else (c1, z[i2])
+                ix += 1
+                out[ix] = c2 * t1*t2 * z[der]**(c2-1) if c2>0 else 0.0
+
+                for i3 in range(i2, nvar):
+                    (c3, t3) = (c2+1, 1.0) if i3==der else (c2, z[i3])
+                    ix += 1
+                    out[ix] = c3 * t1*t2*t3 * z[der]**(c3-1) if c3>0 else 0.0
+
+        return
+
+    if d == 4:
+        for i1 in range(nvar):
+            # Get coefficients and values
+            (c1, t1) = (1, 1.0) if i1==der else (0, z[i1])
+            for i2 in range(i1, nvar):
+                (c2, t2) = (c1+1, 1.0) if i2==der else (c1, z[i2])
+                ix += 1
+                out[ix] = c2 * t1*t2* z[der]**(c2-1) if c2>0 else 0.0
+
+                for i3 in range(i2, nvar):
+                    (c3, t3) = (c2+1, 1.0) if i3==der else (c2, z[i3])
+                    ix += 1
+                    out[ix] = c3*t1*t2*t3*z[der]**(c3-1) if c3>0 else 0.0
+
+                    for i4 in range(i3, nvar):
+                        (c4, t4) = (c3+1, 1.0) if i4==der else (c3, z[i4])
+                        ix += 1
+                        out[ix] = c4*t1*t2*t3*t4*z[der]**(c4-1) if c4>0 else 0.0
+
+        return
+
+    if d == 5:
+        for i1 in range(nvar):
+            # Get coefficients and values
+            (c1, t1) = (1, 1.0) if i1==der else (0, z[i1])
+            for i2 in range(i1, nvar):
+                (c2, t2) = (c1+1, 1.0) if i2==der else (c1, z[i2])
+                ix += 1
+                out[ix] = c2 * t1*t2* z[der]**(c2-1) if c2>0 else 0.0
+
+                for i3 in range(i2, nvar):
+                    (c3, t3) = (c2+1, 1.0) if i3==der else (c2, z[i3])
+                    ix += 1
+                    out[ix] = c3 * t1*t2*t3* z[der]**(c3-1) if c3>0 else 0.0
+
+                    for i4 in range(i3, nvar):
+                        (c4, t4) = (c3+1, 1.0) if i4==der else (c3, z[i4])
+                        ix += 1
+                        out[ix] = c4*t1*t2*t3*t4*z[der]**(c4-1) if c4>0 else 0.0
+
+                        for i5 in range(i4, nvar):
+                            (c5, t5) = (c4+1, 1.0) if i5==der else (c4, z[i5])
+                            ix += 1
+                            out[ix] = c5*t1*t2*t3*t4*t5*z[der]**(c5-1) if c5>0 else 0.0
+
+        return
+
+
+# @jit(nopython=True, cache=True)
+@jit(nopython=True)
+def _complete_poly_der_impl(z, d, der, out):
+    nvar = z.shape[0]
+    nobs = z.shape[1]
+
+    for k in range(nobs):
+        out[0, k] = 0.0
+
+    # fill first order stuff
+    if d >= 1:
+        # Make sure everything has zeros in it
+        for i in range(nvar):
+            for k in range(nobs):
+                out[i+1, k] = 0.0
+
+        # Then place ones where they belong in variable
+        for k in range(nobs):
+            out[der+1, k] = 1.0
+
+    if d == 1:
+        return
+
+    # now we need to fill in row nvar and beyond
+    ix = nvar
+    if d == 2:
+        for i1 in range(nvar):
+            for i2 in range(i1, nvar):
+                ix += 1
+                for k in range(nobs):
+                    c1, t1 = (1, 1.0) if i1==der else (0, z[i1, k])
+                    c2, t2 = (c1+1, 1.0) if i2==der else (c1, z[i2, k])
+                    out[ix, k] = c2*t1*t2*z[der, k]**(c2-1) if c2>0 else 0.0
+
+        return
+
+    if d == 3:
+        for i1 in range(nvar):
+            for i2 in range(i1, nvar):
+                ix += 1
+                for k in range(nobs):
+                    c1, t1 = (1, 1.0) if i1==der else (0, z[i1, k])
+                    c2, t2 = (c1+1, 1.0) if i2==der else (c1, z[i2, k])
+                    out[ix, k] = c2*t1*t2*z[der, k]**(c2-1) if c2>0 else 0.0
+
+                for i3 in range(i2, nvar):
+                    ix += 1
+                    for k in range(nobs):
+                        c3, t3 = (c2+1, 1.0) if i3==der else (c2, z[i3, k])
+                        out[ix, k] = c3*t1*t2*t3*z[der, k]**(c3-1) if c3>0 else 0.0
+
+        return
+
+    if d == 4:
+        for i1 in range(nvar):
+            for i2 in range(i1, nvar):
+                ix += 1
+                for k in range(nobs):
+                    c1, t1 = (1, 1.0) if i1==der else (0, z[i1, k])
+                    c2, t2 = (c1+1, 1.0) if i2==der else (c1, z[i2, k])
+                    out[ix, k] = c2*t1*t2*z[der, k]**(c2-1) if c2>0 else 0.0
+
+                for i3 in range(i2, nvar):
+                    ix += 1
+                    for k in range(nobs):
+                        c3, t3 = (c2+1, 1.0) if i3==der else (c2, z[i3, k])
+                        out[ix, k] = c3*t1*t2*t3*z[der, k]**(c3-1) if c3>0 else 0.0
+
+                    for i4 in range(i3, nvar):
+                        ix += 1
+                        for k in range(nobs):
+                            c4, t4 = (c3+1, 1.0) if i4==der else (c3, z[i4, k])
+                            out[ix, k] = c4*t1*t2*t3*t4*z[der, k]**(c4-1) if c4>0 else 0.0
+
+        return
+
+    if d == 5:
+        for i1 in range(nvar):
+            for i2 in range(i1, nvar):
+                ix += 1
+                for k in range(nobs):
+                    c1, t1 = (1, 1.0) if i1==der else (0, z[i1, k])
+                    c2, t2 = (c1+1, 1.0) if i2==der else (c1, z[i2, k])
+                    out[ix, k] = c2*t1*t2*z[der, k]**(c2-1) if c2>0 else 0.0
+
+                for i3 in range(i2, nvar):
+                    ix += 1
+                    for k in range(nobs):
+                        c3, t3 = (c2+1, 1.0) if i3==der else (c2, z[i3, k])
+                        out[ix, k] = c3*t1*t2*t3*z[der, k]**(c3-1) if c3>0 else 0.0
+
+                    for i4 in range(i3, nvar):
+                        ix += 1
+                        for k in range(nobs):
+                            c4, t4 = (c3+1, 1.0) if i4==der else (c3, z[i4, k])
+                            out[ix, k] = c4*t1*t2*t3*t4*z[der, k]**(c4-1) if c4>0 else 0.0
+
+                        for i5 in range(i4, nvar):
+                            ix += 1
+                            for k in range(nobs):
+                                c5, t5 = (c4+1, 1.0) if i5==der else (c4, z[i5, k])
+                                out[ix, k] = c5*t1*t2*t3*t4*t5*z[der, k]**(c5-1) if c5>0 else 0.0
+
+        return
+
+
 class CompletePolynomial:
 
     def __init__(self, n, d):
@@ -289,7 +547,12 @@ def fit_values(self, s, x, damp=0.0):
             new_coefs = np.ascontiguousarray(lstsq(Phi, x)[0])
             self.coefs = (1 - damp) * new_coefs + damp * self.coefs
 
+    def der(self, s, der):
+        dPhi = complete_polynomial_der(s.T, self.d, der).T
+        return np.dot(dPhi, self.coefs)
+
     def __call__(self, s):
 
         Phi = complete_polynomial(s.T, self.d).T
         return np.dot(Phi, self.coefs)
+

interpolation/tests/test_complete.py

@@ -1,7 +1,12 @@
 import numpy as np
-import numpy as np
 
 from interpolation.complete_poly import CompletePolynomial
+from interpolation.complete_poly import (n_complete, complete_polynomial,
+                                         complete_polynomial_der,
+                                         _complete_poly_impl,
+                                         _complete_poly_impl_vec,
+                                         _complete_poly_der_impl,
+                                         _complete_poly_der_impl_vec)
 
 def test_complete_scalar():
 
@@ -43,7 +48,29 @@ def f2(x, y): return x**3 - y
 
     cp.fit_values(points, vals, damp=0.5)
 
+def test_complete_derivative():
+
+    # TODO: Currently if z has a 0 value then it breaks because occasionally
+    #       tries to raise 0 to a negative power -- This can be fixed by
+    #       checking whether coefficient is 0 before trying to do anything...
+
+    # Test derivative vector
+    z = np.array([1, 2, 3])
+    sol_vec = np.array([0.0, 1.0, 0.0, 0.0, 2.0, 2.0, 3.0, 0.0, 0.0, 0.0])
+    out_vec = np.empty(n_complete(3, 2))
+    _complete_poly_der_impl_vec(z, 2, 0, out_vec)
+    assert(abs(out_vec - sol_vec).max() < 1e-10)
+
+    # Test derivative matrix
+    z = np.arange(1, 7).reshape(3, 2)
+    out_mat = complete_polynomial_der(z, 2, 1)
+    assert(abs(out_mat[0, :]).max() < 1e-10)
+    assert(abs(out_mat[2, :] - np.ones(2)).max() < 1e-10)
+    assert(abs(out_mat[-2, :] - np.array([5.0, 6.0])).max() < 1e-10)
+
 
 if __name__ == '__main__':
     test_complete_scalar()
     test_complete_vector()
+    test_complete_derivative()
+