Merge branch 'carma' of https://github.com/ywx649999311/tinygp into ywx649999311-carma

Author: dfm

src/tinygp/kernels/quasisep.py

@@ -24,10 +24,11 @@
     "Matern52",
     "Cosine",
     "CARMA",
+    "carma",
 ]
 
 from abc import ABCMeta, abstractmethod
-from typing import Any, Optional, Union
+from typing import Any, Optional, Tuple, Union
 
 import jax
 import jax.numpy as jnp
@@ -39,6 +40,8 @@
 from tinygp.solvers.quasisep.core import DiagQSM, StrictLowerTriQSM, SymmQSM
 from tinygp.solvers.quasisep.general import GeneralQSM
 
+eta = 1e-20  # avoid nan
+
 
 class Quasisep(Kernel, metaclass=ABCMeta):
     """The base class for all quasiseparable kernels
@@ -643,21 +646,27 @@ class CARMA(Quasisep):
 
     .. code-block:: python
 
-        kernel = CARMA.init(alpha=..., beta=..., sigma=...)
+        kernel = CARMA.init(alpha=..., beta=...)
+
+    .. note::
+        To fit a CARMA model with p > 2, the :func:`from_fpoly` method needs to
+        be used to construct a valid model.
     """
     alpha: JAXArray
     beta: JAXArray
     sigma: JAXArray
-    roots: JAXArray
-    proj: JAXArray
-    proj_inv: JAXArray
-    stn: JAXArray
+    arroots: JAXArray
+    acf: JAXArray
+    real_mask: JAXArray
+    complex_mask: JAXArray
+    complex_select: JAXArray
+    obsmodel: JAXArray
 
     @classmethod
     def init(
         cls, alpha: JAXArray, beta: JAXArray, sigma: Optional[JAXArray] = None
     ) -> "CARMA":
-        r"""Construct a CARMA kernel
+        r"""Construct a CARMA kernel using the alpha, beta parameters
 
         Args:
             alpha: The parameter :math:`\alpha` in the definition above. This
@@ -674,78 +683,229 @@ def init(
         p = alpha.shape[0]
         assert beta.shape[0] <= p
 
-        # We find the roots of the autoregressive polynomial as a means to find
-        # the eigendecomposition of the design matrix.
-        alpha_ext = jnp.append(alpha, 1.0)
-        roots = jnp.roots(alpha_ext[::-1])
-        proj = roots[:, None] ** jnp.arange(p)[None, :]
-        proj_inv = jnp.linalg.inv(proj)
-
-        # Compute the stationary covariance - there is almost certainly a more
-        # elegant way, but this works! I worked this out kind of by trial and
-        # error using sympy. There is a lot of known structure in the P_inf
-        # matrix that can be exploited to "simplify" this calculation.
-        # Specifically, there are only `p` degrees of freedom, and P_inf has the
-        # following structure:
-        #
-        #   P_inf = [
-        #     [ p0   0   -p1   0    p2 ]
-        #     [ 0    p1   0   -p2   0  ]
-        #     [-p1   0    p2   0   -p3 ]
-        #     [ 0   -p2   0    p3   0  ]
-        #     [ p2   0   -p3   0    p4 ]
-        #   ]
-        #
-        # Using this structure, we get can solve the usual:
-        #
-        #  A @ P + P @ A.T + L @ L.T = 0
-        #
-        # for `P`, and we get something like the following. Kelly et al. (2104)
-        # also have an expression for this (their V_{ij}), but I prefer to use
-        # this since it is probably roughly just as fast to compute, and it is
-        # strictly real-valued.
-        f = 2 * ((np.arange(2 * p) // 2) % 2) - 1
-        x = f * jnp.append(alpha_ext, jnp.zeros(p - 1))
-        params = jnp.stack([np.roll(x, k)[::2] for k in range(p)], axis=0)
-        params = jnp.linalg.solve(
-            params, 0.5 * sigma**2 * jnp.eye(p, 1, k=-p + 1)
-        )[:, 0]
-        stn_ = []
-        for j in range(p):
-            stn_.append([jnp.zeros(()) for _ in range(p)])
-            for n, k in enumerate(range(j - 2, -1, -2)):
-                stn_[-1][k] = (2 * (n % 2) - 1) * params[j - n - 1]
-            for n, k in enumerate(range(j, p, 2)):
-                stn_[-1][k] = (1 - 2 * (n % 2)) * params[n + j]
-        stn = jnp.array(list(map(jnp.stack, stn_)))
+        # find acf
+        arroots = CARMA.roots(jnp.append(alpha, 1.0))
+        acf = CARMA.carma_acf(arroots, alpha, beta * sigma)
+        # masks for selecting entries in matrixes
+        real_mask = jnp.where(arroots.imag == 0.0, jnp.ones(p), jnp.zeros(p))
+        complex_mask = -real_mask + 1
+        complex_idx = jnp.cumsum(-real_mask + 1) * complex_mask
+        complex_select = complex_mask * complex_idx % 2
+
+        # compute obsmodel
+        om_real = jnp.sqrt(jnp.abs(acf.real))
+        a, b, c, d = (
+            2 * acf.real * complex_mask,
+            2 * acf.imag * complex_mask,
+            -arroots.real * complex_mask,
+            -arroots.imag * complex_mask,
+        )
+        c2 = jnp.square(c)
+        d2 = jnp.square(d)
+        s2 = c2 + d2
+        h2_2 = d2 * (a * c - b * d) / (2 * c * s2 + eta * real_mask)
+        h2 = jnp.sqrt(h2_2)
+        h1 = (c * h2 - jnp.sqrt(a * d2 - s2 * h2_2)) / (d + eta * real_mask)
+        om_complex = jnp.array([h1, h2])
+        obsmodel = (om_real * real_mask) + jnp.ravel(om_complex)[
+            ::2
+        ] * complex_mask
 
         return cls(
-            sigma=sigma,
             alpha=alpha,
             beta=beta,
-            roots=roots,
-            proj=proj,
-            proj_inv=proj_inv,
-            stn=stn,
+            sigma=sigma,
+            arroots=arroots,
+            acf=acf,
+            real_mask=real_mask,
+            complex_mask=complex_mask,
+            complex_select=complex_select,
+            obsmodel=obsmodel,
+        )
+
+    @classmethod
+    def from_fpoly(
+        cls, alpha_fpoly: JAXArray, beta_fpoly: JAXArray, beta_mult: JAXArray
+    ) -> "CARMA":
+        """Construct a CARMA kernel using the roots of the characteristic polynomials
+
+        The roots can be re-parameterized as the coefficients of a product
+        of quadratic equations each with the second-order term set to 1. The
+        input for this constructor are said coefficients. See Equation 30 in
+        the paper linked above for a reference.
+
+        Args:
+            alpha_fpoly: The coefficients of the auto-regressive quadratic
+                equations corresponding to the alpha parameters.
+            beta_fpoly: The coefficients of the moving-average quadratic
+                equations corresponding to the beta parameters.
+            beta_mult: Equivalent to beta[-1] used in the init constructor.
+        """
+
+        alpha_fpoly = jnp.atleast_1d(alpha_fpoly)
+        beta_fpoly = jnp.atleast_1d(beta_fpoly)
+        beta_mult = jnp.atleast_1d(beta_mult)
+
+        alpha = CARMA.fpoly2poly(jnp.append(alpha_fpoly, jnp.array([1.0])))[
+            :-1
+        ]
+        beta = CARMA.fpoly2poly(jnp.append(beta_fpoly, beta_mult))
+
+        return CARMA.init(alpha, beta)
+
+    @staticmethod
+    @jax.jit
+    def roots(poly_coeffs: JAXArray) -> JAXArray:
+        roots = jnp.roots(poly_coeffs[::-1], strip_zeros=False)
+        return roots[jnp.argsort(roots.real)]
+
+    @staticmethod
+    @jax.jit
+    def fpoly2poly(fpoly_coeffs: JAXArray) -> JAXArray:
+        """Expand the factorized characteristic polynomial"""
+
+        size = fpoly_coeffs.shape[0] - 1
+        remain = size % 2
+        nPair = size // 2
+        mult_f = fpoly_coeffs[
+            -1:
+        ]  # The coeff of highest order term in the output
+
+        poly = jax.lax.cond(
+            remain == 1,
+            lambda x: jnp.array([1.0, x]),
+            lambda x: jnp.array([0.0, 1.0]),
+            fpoly_coeffs[-2],
         )
+        poly = poly[-remain + 1 :]
+
+        for p in jnp.arange(nPair):
+            poly = jnp.convolve(
+                poly,
+                jnp.append(
+                    jnp.array([fpoly_coeffs[p * 2], fpoly_coeffs[p * 2 + 1]]),
+                    jnp.ones((1,)),
+                )[::-1],
+            )
+
+        # the returned is low->high following Kelly+14
+        return poly[::-1] * mult_f
+
+    @staticmethod
+    def poly2fpoly(poly_coeffs: JAXArray) -> Tuple[JAXArray, JAXArray]:
+        """Factorize a polynomial into product of quadratic equations"""
+
+        fpoly = jnp.empty((0))
+        mult_f = poly_coeffs[-1]
+        roots = CARMA.roots(poly_coeffs / mult_f)
+        odd = bool(len(roots) & 0x1)
+
+        rootsComp = roots[roots.imag != 0]
+        rootsReal = roots[roots.imag == 0]
+        nCompPair = len(rootsComp) // 2
+        nRealPair = len(rootsReal) // 2
+
+        for i in range(nCompPair):
+            root1 = rootsComp[i]
+            root2 = rootsComp[i + 1]
+            fpoly = jnp.append(fpoly, (root1 * root2).real)
+            fpoly = jnp.append(fpoly, -(root1.real + root2.real))
+
+        for i in range(nRealPair):
+            root1 = rootsReal[i]
+            root2 = rootsReal[i + 1]
+            fpoly = jnp.append(fpoly, (root1 * root2).real)
+            fpoly = jnp.append(fpoly, -(root1.real + root2.real))
+
+        if odd:
+            fpoly = jnp.append(fpoly, -rootsReal[-1].real)
+
+        return fpoly, jnp.array(mult_f)
+
+    @staticmethod
+    def carma_acf(
+        arroots: JAXArray, arparam: JAXArray, maparam: JAXArray
+    ) -> JAXArray:
+        """Get ACVF coefficients given CARMA parameters
+
+        Args:
+            arroots (array(complex)): AR roots in a numpy array
+            arparam (array(float)): AR parameters in a numpy array
+            maparam (array(float)): MA parameters in a numpy array
+        Returns:
+            array(complex): ACVF coefficients, each element correspond to a root.
+        """
+        arparam = jnp.atleast_1d(arparam)
+        maparam = jnp.atleast_1d(maparam)
+        p = arparam.shape[0]
+        q = maparam.shape[0] - 1
+        sigma = maparam[0]
+
+        # MA param into Kelly's notation
+        maparam = maparam / sigma
+
+        # init acf product terms
+        num_left = jnp.zeros(p, dtype=jnp.complex128)
+        num_right = jnp.zeros(p, dtype=jnp.complex128)
+        denom = -2 * arroots.real + jnp.zeros_like(arroots) * 1j
+
+        for k in range(q + 1):
+            num_left += maparam[k] * jnp.power(arroots, k)
+            num_right += maparam[k] * jnp.power(jnp.negative(arroots), k)
+
+        root_idx = jnp.arange(p)
+        for j in range(1, p):
+            root_k = arroots[jnp.roll(root_idx, j)]
+            denom *= (root_k - arroots) * (jnp.conj(root_k) + arroots)
+
+        return sigma**2 * num_left * num_right / denom
 
     def design_matrix(self) -> JAXArray:
-        p = self.alpha.shape[0]
-        return jnp.concatenate((jnp.eye(p - 1, p, k=1), -self.alpha[None]))
+        dm_real = jnp.diag(self.arroots.real * self.real_mask)
+        dm_complex_diag = jnp.diag(self.arroots.real * self.complex_mask)
+        dm_complex_u = jnp.diag(
+            (self.arroots.imag * self.complex_select)[:-1], k=1
+        )
+
+        return dm_real + dm_complex_diag + -dm_complex_u.T + dm_complex_u
 
     def stationary_covariance(self) -> JAXArray:
-        return self.stn
+        p = self.acf.shape[0]
+        diag = jnp.diag(
+            jnp.where(self.acf.real > 0, jnp.ones(p), -jnp.ones(p))
+        )
+        diag_complex = jnp.diag(
+            (
+                2
+                * jnp.square(-self.arroots.real)
+                / jnp.square(-self.arroots.imag + eta)
+            )
+            * jnp.roll(self.complex_select, 1)
+            * self.complex_mask
+        )
+        c_over_d = self.arroots.real / (self.arroots.imag + eta)
+        sc_complex_u = jnp.diag((-c_over_d * self.complex_select)[:-1], k=1)
+
+        return diag + diag_complex + sc_complex_u + sc_complex_u.T
 
     def observation_model(self, X: JAXArray) -> JAXArray:
-        return jnp.append(
-            self.beta, jnp.zeros(self.alpha.shape[0] - self.beta.shape[0])
-        )
+        return self.obsmodel
 
     def transition_matrix(self, X1: JAXArray, X2: JAXArray) -> JAXArray:
         dt = X2 - X1
-        return (
-            self.proj_inv @ (jnp.exp(self.roots * dt)[:, None] * self.proj)
-        ).real
+        c = -self.arroots.real
+        d = -self.arroots.imag
+        decay = jnp.exp(-c * dt)
+        sin = jnp.sin(d * dt)
+
+        tm_real = jnp.diag(decay * self.real_mask)
+        tm_complex_diag = jnp.diag(decay * jnp.cos(d * dt) * self.complex_mask)
+        tm_complex_u = jnp.diag(
+            (decay * sin * self.complex_select)[:-1],
+            k=1,
+        )
+
+        return tm_real + tm_complex_diag + -tm_complex_u.T + tm_complex_u
 
 
 def _prod_helper(a1: JAXArray, a2: JAXArray) -> JAXArray: