even faster linalg

starry/_core/core.py

@@ -1485,6 +1485,42 @@ def get_kT0(self, rT):
         # Normalize to preserve the unit baseline
         return kT0 / tt.sum(kT0[0])
 
+    @autocompile
+    def get_kT(self, inc, theta, veq, u):
+        """
+        Get the kernels at an array of angular phases `theta`.
+
+        """
+        # Compute the convolution kernels
+        vsini = self.enforce_bounds(veq * tt.sin(inc), 0.0, self.vsini_max)
+        x = self.get_x(vsini)
+        rT = self.get_rT(x)
+        kT0 = self.get_kT0(rT)
+
+        # Compute the limb darkening operator
+        if self.udeg > 0:
+            F = self.F(
+                tt.as_tensor_variable(u), tt.as_tensor_variable([np.pi])
+            )
+            L = ts.dot(ts.dot(self.A1Inv, F), self.A1)
+            kT0 = tt.dot(tt.transpose(L), kT0)
+
+        # Compute the kernels at each epoch
+        kT = tt.zeros((self.nt, self.Ny, self.nk))
+        for m in range(self.nt):
+            kT = tt.set_subtensor(
+                kT[m],
+                tt.transpose(
+                    self.right_project(
+                        tt.transpose(kT0),
+                        inc,
+                        tt.as_tensor_variable(0.0),
+                        theta[m],
+                    )
+                ),
+            )
+        return kT
+
     @autocompile
     def get_D_data(self, kT0, inc, theta_scalar):
         """
@@ -1541,6 +1577,52 @@ def get_D(self, inc, theta, veq, u):
             ]
         )
 
+    @autocompile
+    def get_D_fixed_spectrum(self, inc, theta, veq, u, spectrum):
+        """
+        Return the Doppler matrix for a fixed spectrum.
+
+        """
+        # Get the convolution kernels
+        kT = self.get_kT(inc, theta, veq, u)
+
+        # The dot product is just a 2d convolution!
+        product = tt.nnet.conv2d(
+            tt.reshape(tt.reshape(spectrum, (-1,)), (self.nc, 1, 1, self.nwp)),
+            tt.reshape(kT, (self.nt * self.Ny, 1, 1, self.nk)),
+            border_mode="valid",
+            filter_flip=False,
+            input_shape=(self.nc, 1, 1, self.nwp),
+            filter_shape=(self.nt * self.Ny, 1, 1, self.nk),
+        )
+        product = tt.reshape(product, (self.nc, self.nt, self.Ny, self.nw))
+        product = tt.swapaxes(product, 1, 2)
+        product = tt.reshape(product, (self.Ny * self.nc, self.nt * self.nw))
+        product = tt.transpose(product)
+        return product
+
+    @autocompile
+    def dot_design_matrix_into(self, inc, theta, veq, u, matrix):
+        """
+        Dot the full Doppler design matrix into an arbitrary dense `matrix`.
+        This is equivalent to tt.dot(get_D(), matrix), but computes the
+        product with a single `conv2d` operation.
+
+        """
+        # Get the convolution kernels
+        kT = self.get_kT(inc, theta, veq, u)
+
+        # The dot product is just a 2d convolution!
+        product = tt.nnet.conv2d(
+            tt.reshape(tt.transpose(matrix), (-1, self.Ny, 1, self.nwp)),
+            tt.reshape(kT, (self.nt, self.Ny, 1, self.nk)),
+            border_mode="valid",
+            filter_flip=False,
+            input_shape=(None, self.Ny, 1, self.nwp),
+            filter_shape=(self.nt, self.Ny, 1, self.nk),
+        )
+        return tt.transpose(tt.reshape(product, (-1, self.nt * self.nw)))
+
     @autocompile
     def get_flux_from_design(self, inc, theta, veq, u, a):
         """
@@ -1557,47 +1639,34 @@ def get_flux_from_design(self, inc, theta, veq, u, a):
     def get_flux_from_conv(self, inc, theta, veq, u, a):
         """
         Compute the flux via a single 2d convolution.
-        This is the *fast* way of computing the model.
+        This is the *faster* way of computing the model.
 
         """
-        # Compute the convolution kernels
-        vsini = self.enforce_bounds(veq * tt.sin(inc), 0.0, self.vsini_max)
-        x = self.get_x(vsini)
-        rT = self.get_rT(x)
-        kT0 = self.get_kT0(rT)
-
-        # Compute the limb darkening operator
-        if self.udeg > 0:
-            F = self.F(
-                tt.as_tensor_variable(u), tt.as_tensor_variable([np.pi])
-            )
-            L = ts.dot(ts.dot(self.A1Inv, F), self.A1)
-            kT0 = tt.dot(tt.transpose(L), kT0)
-
-        # Compute the kernels at each epoch
-        kT0_r = kT0[:, ::-1]
-        kT = tt.zeros((self.nt, self.Ny, self.nk))
-        for m in range(self.nt):
-            kT = tt.set_subtensor(
-                kT[m],
-                tt.transpose(
-                    self.right_project(
-                        tt.transpose(kT0_r),
-                        inc,
-                        tt.as_tensor_variable(0.0),
-                        theta[m],
-                    )
-                ),
-            )
+        # Get the convolution kernels
+        kT = self.get_kT(inc, theta, veq, u)
 
         # The flux is just a 2d convolution!
         flux = tt.nnet.conv2d(
             tt.reshape(a, (1, self.Ny, 1, self.nwp)),
             tt.reshape(kT, (self.nt, self.Ny, 1, self.nk)),
             border_mode="valid",
+            filter_flip=False,
+            input_shape=(1, self.Ny, 1, self.nwp),
+            filter_shape=(self.nt, self.Ny, 1, self.nk),
         )
         return flux[0, :, 0, :]
 
+    @autocompile
+    def get_flux_from_convdot(self, inc, theta, veq, u, y, spectrum):
+        """
+        Compute the flux via a 2d convolution followed by a dot product.
+        This is usually the *fastest* way of computing the model.
+
+        """
+        D = self.get_D_fixed_spectrum(inc, theta, veq, u, spectrum)
+        flux = tt.dot(D, tt.reshape(tt.transpose(y), (-1,)))
+        return tt.reshape(flux, (self.nt, self.nw))
+
 
 class OpsSystem(object):
     """Class housing ops for modeling Keplerian systems."""

starry/doppler.py

@@ -350,6 +350,7 @@ def spectrum(self):
         The rest frame spectrum at wavelengths ``wav`` for each component.
 
         """
+        # NOTE the transpose here!
         return self._math.transpose(self._spectrum)
 
     @spectrum.setter
@@ -482,8 +483,7 @@ def spot(self, *, component=0, **kwargs):
         else:
             self[:, :, component] = self._map[:, :]
 
-    def design_matrix(self, theta=None):
-        """Return the full Doppler imaging design matrix."""
+    def _get_default_theta(self, theta):
         if theta is None:
             theta = self._math.cast(
                 np.linspace(0, 2 * np.pi, self._nt, endpoint=False)
@@ -495,23 +495,30 @@ def design_matrix(self, theta=None):
                 )
                 * self._angle_factor
             )
-        D = self.ops.get_D(self._inc, theta, self._veq, self._u)
-        return D
-
-    def flux(self, theta=None, mode="conv"):
-        """Return the model for the full spectral timeseries."""
-        if theta is None:
-            theta = self._math.cast(
-                np.linspace(0, 2 * np.pi, self._nt, endpoint=False)
+        return theta
+
+    def design_matrix(self, theta=None, fix_spectrum=False):
+        """Return the Doppler imaging design matrix."""
+        theta = self._get_default_theta(theta)
+        if fix_spectrum:
+            D = self.ops.get_D_fixed_spectrum(
+                self._inc, theta, self._veq, self._u, self._spectrum
             )
         else:
-            theta = (
-                self.ops.enforce_shape(
-                    self._math.cast(theta), np.array([self._nt])
-                )
-                * self._angle_factor
+            D = self.ops.get_D(self._inc, theta, self._veq, self._u)
+        return D
+
+    def flux(self, theta=None, mode="convdot"):
+        """
+        Return the model for the full spectral timeseries.
+
+        """
+        theta = self._get_default_theta(theta)
+        if mode == "convdot":
+            flux = self.ops.get_flux_from_convdot(
+                self._inc, theta, self._veq, self._u, self._y, self._spectrum
             )
-        if mode == "conv":
+        elif mode == "conv":
             flux = self.ops.get_flux_from_conv(
                 self._inc, theta, self._veq, self._u, self.spectral_map
             )
@@ -523,6 +530,14 @@ def flux(self, theta=None, mode="conv"):
             raise ValueError("Keyword `mode` must be one of `conv`, `design`.")
         return flux
 
+    def dot(self, matrix, theta=None):
+        """Dot the Doppler design matrix into a given matrix or vector."""
+        theta = self._get_default_theta(theta)
+        product = self.ops.dot_design_matrix_into(
+            self._inc, theta, self._veq, self._u, self._math.cast(matrix)
+        )
+        return product
+
     def show(self, theta=None, res=150, file=None, **kwargs):
         """
 

tests/greedy/test_doppler.py

@@ -1,15 +1,60 @@
 import starry
 import numpy as np
+from scipy.linalg import block_diag
+import pytest
 
 
-def test_conv2d():
+@pytest.fixture
+def map():
+    map = starry.DopplerMap(ydeg=10, nt=3, nc=2, nw=199, veq=50000)
+    map.load(["spot", "earth"], force_psd=True)
+    yield map
+
+
+@pytest.fixture
+def random():
+    yield np.random.default_rng(0)
+
+
+def test_flux(map):
     """
-    Test that our design matrix and conv2d implementations of the flux
+    Test that our various implementations of the flux
     yield identical results.
 
     """
-    map = starry.DopplerMap(ydeg=10, nt=3, nc=2, nw=199, veq=50000)
-    map.load(["spot", "earth"], force_psd=True)
-    flux1 = map.flux(mode="conv")
-    flux2 = map.flux(mode="design")
+    flux1 = map.flux(mode="convdot")
+    flux2 = map.flux(mode="conv")
+    flux3 = map.flux(mode="design")
     assert np.allclose(flux1, flux2)
+    assert np.allclose(flux1, flux3)
+
+
+def test_dot(map, random):
+    """
+    Test that our fast dot product method yields the same result as
+    instantiating the full design matrix and dotting it in.
+
+    """
+    D = map.design_matrix().todense()
+    matrix = random.normal(size=(map.nws * map.Ny, 5))
+    product1 = D @ matrix
+    product2 = map.dot(matrix)
+    assert np.allclose(product1, product2)
+
+
+def test_D_fixed_spectrum(map, random):
+    """
+    Test that our fast method for computing the design matrix
+    for fixed input spectrum yields the same result as instantiating
+    the full design matrix and dotting the spectral block matrix in.
+
+    """
+    DS = np.zeros((map.nt * map.nwf, 0))
+    D = map.design_matrix().todense()
+    for k in range(map.nc):
+        S = block_diag(
+            *[map.spectrum[:, k].reshape(-1, 1) for n in range(map.Ny)]
+        )
+        DS = np.hstack((DS, D @ S))
+    DS_fast = map.design_matrix(fix_spectrum=True)
+    assert np.allclose(DS, DS_fast)