faster doppler by 50x

starry/_core/core.py

@@ -1510,6 +1510,9 @@ def get_D(self, inc, theta, veq, u):
         This is a horizontal stack of Toeplitz convolution matrices, one per
         spherical harmonic. These matrices are then stacked vertically for
         each rotational phase.
+
+        In general, instantiating this matrix (even in its sparse form) is not
+        a good idea: it's very slow, and can consume a ton of memory!
         """
         # Compute the convolution kernels
         vsini = self.enforce_bounds(veq * tt.sin(inc), 0.0, self.vsini_max)
@@ -1539,11 +1542,62 @@ def get_D(self, inc, theta, veq, u):
         )
 
     @autocompile
-    def get_flux(self, inc, theta, veq, u, a):
+    def get_flux_from_design(self, inc, theta, veq, u, a):
+        """
+        Compute the flux by dotting the design matrix into
+        the spectral map. This is the *slow* way of computing
+        the model.
+
+        """
         D = self.get_D(inc, theta, veq, u)
         flux = ts.dot(D, tt.reshape(a, (-1,)))
         return tt.reshape(flux, (self.nt, self.nw))
 
+    @autocompile
+    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.
+
+        """
+        # 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],
+                    )
+                ),
+            )
+
+        # 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",
+        )
+        return flux[0, :, 0, :]
+
 
 class OpsSystem(object):
     """Class housing ops for modeling Keplerian systems."""

starry/doppler.py

@@ -483,7 +483,7 @@ def spot(self, *, component=0, **kwargs):
             self[:, :, component] = self._map[:, :]
 
     def design_matrix(self, theta=None):
-        """Return the Doppler operator."""
+        """Return the full Doppler imaging design matrix."""
         if theta is None:
             theta = self._math.cast(
                 np.linspace(0, 2 * np.pi, self._nt, endpoint=False)
@@ -498,7 +498,7 @@ def design_matrix(self, theta=None):
         D = self.ops.get_D(self._inc, theta, self._veq, self._u)
         return D
 
-    def flux(self, theta=None):
+    def flux(self, theta=None, mode="conv"):
         """Return the model for the full spectral timeseries."""
         if theta is None:
             theta = self._math.cast(
@@ -511,9 +511,16 @@ def flux(self, theta=None):
                 )
                 * self._angle_factor
             )
-        flux = self.ops.get_flux(
-            self._inc, theta, self._veq, self._u, self.spectral_map
-        )
+        if mode == "conv":
+            flux = self.ops.get_flux_from_conv(
+                self._inc, theta, self._veq, self._u, self.spectral_map
+            )
+        elif mode == "design":
+            flux = self.ops.get_flux_from_design(
+                self._inc, theta, self._veq, self._u, self.spectral_map
+            )
+        else:
+            raise ValueError("Keyword `mode` must be one of `conv`, `design`.")
         return flux
 
     def show(self, theta=None, res=150, file=None, **kwargs):

tests/greedy/test_doppler.py

@@ -0,0 +1,15 @@
+import starry
+import numpy as np
+
+
+def test_conv2d():
+    """
+    Test that our design matrix and conv2d 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")
+    assert np.allclose(flux1, flux2)