Derivative models

[dfm]

I've implemented a first pass at derivative models but there are some open questions so I'm going to remove it from #19 and just post it here for future work:

class _DerivativeHelperTerm(terms.Term):
    def __init__(self, term):
        self.term = term

    def get_coefficients(self):
        coeffs = self.term.get_coefficients()
        a, b, c, d = coeffs[2:]
        final_coeffs = [
            -coeffs[0] * coeffs[1],
            coeffs[1],
            b * d - a * c,
            -(a * d + b * c),
            c,
            d,
        ]
        return final_coeffs


class DerivativeLatentTerm(LatentTerm):
    def __init__(self, term, *, value_amplitude, gradient_amplitude):
        if term.__requires_general_addition__:
            raise TypeError(
                "You cannot perform operations on a term that requires general"
                " addition, it must be the outer term in the kernel"
            )

        val_amp = np.atleast_1d(value_amplitude)
        grad_amp = np.atleast_1d(gradient_amplitude)
        if val_amp.ndim != 1 or val_amp.shape != grad_amp.shape:
            raise ValueError("Dimension mismatch between amplitudes")
        amp = np.stack((val_amp, grad_amp), axis=-1)
        M = amp.shape[0]

        ar, cr, ac, bc, cc, dc = term.get_coefficients()

        Jr = len(cr)
        Jc = len(cc)
        J = Jr + 2 * Jc

        left = np.zeros((2, 4 * J))
        right = np.zeros((2, 4 * J))
        if Jr:
            left[0, :Jr] = 1.0
            left[0, 2 * Jr : 3 * Jr] = -1.0
            left[1, Jr : 2 * Jr] = 1.0
            left[1, 3 * Jr : 4 * Jr] = 1.0
            right[0, : 2 * Jr] = 1.0
            right[1, 2 * Jr : 4 * Jr] = 1.0

        if Jc:
            J0 = 4 * Jr
            for J0 in [4 * Jr, 4 * Jr + 4 * Jc]:
                left[0, J0 : J0 + Jc] = 1.0
                left[0, J0 + 2 * Jc : J0 + 3 * Jc] = -1.0
                left[1, J0 + Jc : J0 + 2 * Jc] = 1.0
                left[1, J0 + 3 * Jc : J0 + 4 * Jc] = 1.0
                right[0, J0 : J0 + 2 * Jc] = 1.0
                right[1, J0 + 2 * Jc : J0 + 4 * Jc] = 1.0

        self.left = np.dot(amp, left)[:, :, None]
        self.right = np.dot(amp, right)[:, :, None]

        # self.left = np.empty((M, 4 * J, 1))
        # self.left[:] = np.nan
        # self.left[:] = np.dot(amp, left)[:, :, None]

        # self.right = np.zeros((M, 4 * J, M))
        # right = np.dot(amp, right)
        # for m in range(M):
        #     self.right[m, :, m] = right[m]

        # self.left = np.reshape(self.left, (M, -1, 1))
        # self.right = np.reshape(self.right, (M, -1, 1))

        print(self.left)
        print(self.right)
        # print(self.right)
        # assert 0

        dterm = _DerivativeHelperTerm(term)
        d2term = terms.TermDiff(term)
        super().__init__(
            terms.TermSum(term, dterm, dterm, d2term),
            dimension=2,
            left_latent=self._left_latent,
            right_latent=self._right_latent,
        )

    def _left_latent(self, t, inds):
        return self.left[inds]

    def _right_latent(self, t, inds):
        return self.right[inds]