Differentiate complex function

[konstunn]

I have a complex function (method):

    def M(self, u, x0=None, th=None):
        if th is None:
            th = self.__th
        else:
            th = np.array(th)

        s = len(th)
        n = self.__n

        lst = list()
        lst.append(self.__F)
        lst.append(self.__C)
        lst.append(self.__G)
        lst.append(self.__H)
        lst.append(self.__w_cov)
        lst.append(self.__v_cov)
        lst.append(self.__x0_mean)
        lst.append(self.__x0_cov)

        # eval
        jlst = [autograd.jacobian(f)(th) for f in lst]

        jlst = [[np.squeeze(j, 2) for j in np.dsplit(jel, s)] for jel in jlst]

        dF, dC, dG, dH, dQ, dR, dX0, dP0 = jlst

        # eval
        F, C, G, H, Q, R, X0, P0 = [f(th) for f in lst]

        if x0 is not None:
            # on reshape fail exception will be raised
            X0 = np.array(x0).reshape([n, 1])

        C_A = np.vstack(dC)
        C_A = np.vstack([C, C_A])

        M = np.zeros([s, s])

        t = np.transpose
        inv = np.linalg.inv
        Sp = np.trace
        Pe = P0
        dPe = dP0
        Inn = np.eye(n)
        Onn = np.zeros([n, n])

        # TODO: cast every thing to np.matrix and use '*' multiplication syntax

        def F_A_f(F, dF, H, K_):
            _1st_col = [dF_i - K_ @ dH_i for dF_i, dH_i in zip(dF, dH)]
            _1st_col = np.vstack(_1st_col)

            bdiag = scipy.linalg.block_diag(*[F - K_ @ H] * s)
            rez = np.hstack([_1st_col, bdiag])

            _1st_row = np.hstack([np.zeros([n, n])] * s)
            _1st_row = np.hstack([F, _1st_row])

            rez = np.vstack([_1st_row, rez])
            return rez

        def X_Ap_f(F_A, X_Ap, u, k):
            if k == 0:
                F_ = np.vstack(dF)
                F_ = np.vstack([F, F_])

                # force dX0_i to be 2D array
                FdX0 = [F @ np.array(dX0_i, ndmin=2) for dX0_i in dX0]
                FdX0 = np.vstack(FdX0)
                OFdX0 = np.vstack([Onn, FdX0])

                # u[:,[k]] - get k-th column as column vector
                return F_ @ X0 + OFdX0 + C_A @ u[:, [0]]
            elif k > 0:
                return F_A @ X_Ap + C_A @ u[:, [k]]

        def Cf(i):
            i = i + 1
            O = [np.zeros([n, n])] * i
            O = np.hstack(O) if i else []
            C = np.hstack([O, np.eye(n)]) if i else np.eye(n)
            O = [np.zeros([n, n])] * (s-i)
            O = np.hstack(O) if s-i else []
            C = np.hstack([C, O]) if s-i else C
            return C

        u = np.array(u, ndmin=2)
        N = u.shape[1]

        if u.shape[0] != C.shape[1]:
            raise Exception('invalid shape of 'u'')

        for k in range(N):
            if k == 0:
                E_A = np.zeros([n*(s+1), n*(s+1)])
                X_Ap = X_Ap_f(None, None, u, k)
            elif k > 0:
                E_A = F_A @ E_A @ t(F_A) + K_A @ B @ t(K_A)
                X_Ap = X_Ap_f(F_A, X_Ap, u, k)

            # Pp, B, K, Pu, K_
            Pp = F @ Pe @ t(F) + G @ Q @ t(G)
            B = H @ Pp @ t(H) + R
            invB = inv(B)
            K = Pp @ t(H) @ invB
            Pu = (Inn - K @ H) @ Pp
            K_ = F @ K

            F_A = F_A_f(F, dF, H, K_)

            dPp = [dF_i @ Pe @ t(F) + F @ dPe_i @ t(F) + F @ Pe @ t(dF_i) +
                dG_i @ Q @ t(G) + G @ dQ_i @ t(G) + G @ Q @ t(dG_i)
                for dF_i, dPe_i, dG_i, dQ_i in zip(dF, dPe, dG, dQ)]

            dB = [dH_i @ Pp @ t(H) + H @ dPp_i @ t(H) + H @ Pp @ t(dH_i) + dR_i
                for dH_i, dPp_i, dR_i in zip(dH, dPp, dR)]

            dK = [(dPp_i @ t(H) + Pp @ t(dH_i) - Pp @ t(H) @ invB @ dB_i) @ invB
                for dPp_i, dH_i, dB_i in zip(dPp, dH, dB)]

            dPu = [(Inn - K @ H) @ dPp_i - (dK_i @ H + K @ dH_i) @ Pp
                for dPp_i, dK_i, dH_i in zip(dPp, dK, dH)]

            dK_ = [dF_i @ K + F @ dK_i for dF_i, dK_i in zip(dF, dK)]

            K_A = np.vstack(dK_)
            K_A = np.vstack([K_, K_A])

            # 8: AM
            AM = list()

            EXX = E_A + X_Ap @ t(X_Ap)

            C0 = Cf(0)

            for i, j in itertools.product(range(s), range(s)):
                S1 = Sp(C0 @ EXX @ t(C0) @ t(dH[j]) @ invB @ dH[i])
                S2 = Sp(C0 @ EXX @ t(Cf(j)) @ t(H) @ invB @ dH[i])
                S3 = Sp(Cf(i) @ EXX @ t(C0) @ t(dH[j]) @ invB @ H)
                S4 = Sp(Cf(i) @ EXX @ t(Cf(j)) @ t(H) @ invB @ H)
                S5 = 0.5 * Sp(dB[i] @ invB @ dB[j] @ invB)
                AM.append(S1 + S2 + S3 + S4 + S5)

            AM = np.array(AM).reshape([s, s])
            M = M + AM

            # TODO: dont forget to update P, dP etc.
            Pe = Pu
            dPe = dPu

        return M

I could paste a full source text of my class, including full class definition and usage example or push it all on github and give a link.

Well, I have the following error when I try to differentiate my function with autograd:

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
/usr/local/lib/python3.5/dist-packages/autograd/errors.py in wrapped(*args, **kwargs)
     47   def wrapped(*args, **kwargs):
---> 48     try: return fun(*args, **kwargs)
     49     except Exception as e: add_extra_error_message(e)

/usr/local/lib/python3.5/dist-packages/autograd/convenience_wrappers.py in jacfun(*args, **kwargs)
     40         ans_vspace = vspace(ans)
---> 41         jacobian_shape = ans_vspace.shape + vspace(args[argnum]).shape
     42         grads = map(vjp, ans_vspace.standard_basis())

TypeError: can only concatenate tuple (not "list") to tuple

During handling of the above exception, another exception occurred:

TypeError                                 Traceback (most recent call last)
<ipython-input-3-f7ecbc824472> in <module>()
----> 1 jacobian(m.M)(u)

/usr/local/lib/python3.5/dist-packages/autograd/errors.py in wrapped(*args, **kwargs)
     47   def wrapped(*args, **kwargs):
     48     try: return fun(*args, **kwargs)
---> 49     except Exception as e: add_extra_error_message(e)
     50   return wrapped
     51

/usr/local/lib/python3.5/dist-packages/autograd/errors.py in add_extra_error_message(e)
     69         else:
     70             raise_(AutogradHint, (extra_message, etype, value), traceback)
---> 71     raise_(etype, value, traceback)

/usr/local/lib/python3.5/dist-packages/future/utils/__init__.py in raise_(tp, value, tb)
    411             exc = tp
    412         if exc.__traceback__ is not tb:
--> 413             raise exc.with_traceback(tb)
    414         raise exc
    415

/usr/local/lib/python3.5/dist-packages/autograd/errors.py in wrapped(*args, **kwargs)
     46   @wraps(fun)
     47   def wrapped(*args, **kwargs):
---> 48     try: return fun(*args, **kwargs)
     49     except Exception as e: add_extra_error_message(e)
     50   return wrapped

/usr/local/lib/python3.5/dist-packages/autograd/convenience_wrappers.py in jacfun(*args, **kwargs)
     39         vjp, ans = make_vjp(fun, argnum)(*args, **kwargs)
     40         ans_vspace = vspace(ans)
---> 41         jacobian_shape = ans_vspace.shape + vspace(args[argnum]).shape
     42         grads = map(vjp, ans_vspace.standard_basis())
     43         return np.reshape(np.stack(grads), jacobian_shape)

TypeError: can only concatenate tuple (not "list") to tuple

What might be the cause?

[konstunn]

I think I have found the reason. I used to diff the function as follows:

u = [2.0, 2.0]
jacobian(m.fim)(u)

Then I get the error described in my comment above. But if I do it as follows:

u = np.array([2.0, 2.0])
jacobian(m.fim)(u)

and it works. Must be stupid mistake. Sorry also for that I wasn't completely specific with my example. But I suppose and hope it's already possible to get the problem from I have provided. Feel free to ask me to provide more complete example to reproduce the problem if it is neccessary.