feat: use limb darkening light curve in system starry light curve (#261)

* feat: use limb darkening light curve in system starry light curve if central.ydeg is 0

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix: use jax assert_allclose

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: lgrcia

src/jaxoplanet/starry/core/solution.py

@@ -121,9 +121,7 @@ def p_integral(order: int, l_max: int, b: Array, r: Array, kappa0: Array) -> Arr
 
     # low order variables
     low_order = np.min([order, 20])
-    zeros = jnp.zeros(order - low_order)
     roots, low_weights = roots_legendre(low_order)
-    low_weights = jnp.hstack((low_weights, zeros))
     phi = rng * (roots + 1)
     low_s2 = jnp.square(jnp.sin(phi))
     low_a1 = low_s2 - jnp.square(low_s2)
@@ -139,9 +137,12 @@ def p_integral(order: int, l_max: int, b: Array, r: Array, kappa0: Array) -> Arr
     high_a2 = jnp.where(r_cond, 0, delta + high_s2)
     high_a4 = 1 - 2 * high_s2
 
-    indices = []
-    weights = []
-    integrand = []
+    low_indices = []
+    low_integrand = []
+
+    high_indices = []
+    high_integrand = []
+
     n = 0
 
     for l in range(l_max + 1):  # noqa
@@ -160,8 +161,8 @@ def p_integral(order: int, l_max: int, b: Array, r: Array, kappa0: Array) -> Arr
                 result = (
                     2 * r * (r - b * c) * (1 - z2 * zero_safe_sqrt(z2)) / (3 * omz2)
                 )
-                integrand.append(jnp.where(cond, 0, 2 * result))
-                weights.append(high_weights)
+                high_integrand.append(jnp.where(cond, 0, 2 * result))
+                high_indices.append(n)
 
             elif mu % 2 == 0 and (mu // 2) % 2 == 0:
                 f = (
@@ -170,18 +171,18 @@ def p_integral(order: int, l_max: int, b: Array, r: Array, kappa0: Array) -> Arr
                     * low_a1 ** (0.25 * (mu + 4))
                     * low_a2 ** (0.5 * nu)
                 )
-                integrand.append(2 * jnp.hstack((f, zeros)))
-                weights.append(low_weights)
+                low_integrand.append(2 * f)
+                low_indices.append(n)
 
             elif mu == 1 and l % 2 == 0:
                 f = high_fa3 * high_a1 ** (l // 2 - 1) * high_a4
-                integrand.append(2 * f)
-                weights.append(high_weights)
+                high_integrand.append(2 * f)
+                high_indices.append(n)
 
             elif mu == 1:
                 f = high_fa3 * high_a1 ** ((l - 3) // 2) * high_a2 * high_a4
-                integrand.append(2 * f)
-                weights.append(high_weights)
+                high_integrand.append(2 * f)
+                high_indices.append(n)
 
             elif (mu - 1) % 2 == 0 and ((mu - 1) // 2) % 2 == 0:
                 f = (
@@ -190,23 +191,26 @@ def p_integral(order: int, l_max: int, b: Array, r: Array, kappa0: Array) -> Arr
                     * high_a1 ** ((mu - 1) // 4)
                     * high_a2 ** (0.5 * (nu - 1))
                 )
-                integrand.append(2 * f)
-                weights.append(high_weights)
+                high_integrand.append(2 * f)
+                high_indices.append(n)
 
             else:
                 n += 1
                 continue
 
-            indices.append(n)
             n += 1
 
-    indices = np.stack(indices)
-    weights = jnp.stack(weights)
+    low_indices = np.stack(low_indices)
+    high_indices = np.stack(high_indices)
+
+    low_P0 = rng * jnp.sum(jnp.stack(low_integrand) * low_weights, axis=1)
+    high_P0 = rng * jnp.sum(jnp.stack(high_integrand) * high_weights, axis=1)
 
-    P0 = rng * jnp.sum(jnp.stack(integrand) * weights, axis=1)
     P = jnp.zeros(l_max**2 + 2 * l_max + 1)
+    P = P.at[low_indices].set(low_P0)
+    P = P.at[high_indices].set(high_P0)
 
-    return P.at[indices].set(P0)
+    return P
 
 
 def rT(lmax: int) -> Array:

src/jaxoplanet/starry/light_curves.py

@@ -1,8 +1,11 @@
+from functools import partial
+
 import jax
 import jax.numpy as jnp
 import numpy as np
 import scipy
 
+from jaxoplanet.core.limb_dark import light_curve as _limb_dark_light_curve
 from jaxoplanet.starry.core.basis import A1, A2_inv, U
 from jaxoplanet.starry.core.polynomials import Pijk
 from jaxoplanet.starry.core.rotation import left_project
@@ -11,98 +14,6 @@
 from jaxoplanet.starry.system_observable import system_observable
 
 
-def _design_matrix_elements(
-    surface: Surface,
-    r: float | None = None,
-    x: float | None = None,
-    y: float | None = None,
-    z: float | None = None,
-    theta: float | None = None,
-    order: int = 20,
-    higher_precision: bool = False,
-    rv: bool = False,
-):
-    if higher_precision:
-        try:
-            from jaxoplanet.starry.multiprecision import (
-                basis as basis_mp,
-                utils as utils_mp,
-            )
-        except ImportError as e:
-            raise ImportError(
-                "The `mpmath` Python package is required for higher_precision=True."
-            ) from e
-
-    total_deg = surface.deg + (surface.vdeg if rv else 0)
-
-    rT_deg = rT(total_deg)
-
-    x = 0.0 if x is None else x
-    y = 0.0 if y is None else y
-    z = 0.0 if z is None else z
-
-    # no occulting body
-    if r is None:
-        b_rot = True
-        theta_z = 0.0
-        design_matrix_p = rT_deg
-
-    # occulting body
-    else:
-        b = jnp.sqrt(jnp.square(x) + jnp.square(y))
-        b_rot = jnp.logical_or(jnp.greater_equal(b, 1.0 + r), jnp.less_equal(z, 0.0))
-        b_occ = jnp.logical_not(b_rot)
-        theta_z = jnp.arctan2(x, y)
-
-        # trick to avoid nan `x=jnp.where...` grad caused by nan sT
-        r = jnp.where(b_rot, 1.0, r)
-        b = jnp.where(b_rot, 1.0, b)
-
-        sT = solution_vector(total_deg, order=order)(b, r)
-
-        if total_deg > 0:
-            if higher_precision:
-                A2 = np.atleast_2d(utils_mp.to_numpy(basis_mp.A2(total_deg)))
-            else:
-                A2 = scipy.sparse.linalg.inv(A2_inv(total_deg))
-                A2 = jax.experimental.sparse.BCOO.from_scipy_sparse(A2)
-        else:
-            A2 = jnp.array([[1]])
-
-        design_matrix_p = jnp.where(b_occ, sT @ A2, rT_deg)
-
-    if surface.ydeg == 0:
-        rotated_y = surface.y.todense()
-    else:
-        rotated_y = left_project(
-            surface.ydeg,
-            surface._inc,
-            surface._obl,
-            theta,
-            theta_z,
-            surface.y.todense(),
-        )
-
-    # limb darkening
-    if surface.udeg == 0:
-        p_u = Pijk.from_dense(jnp.array([1]))
-    else:
-        u = jnp.array([1, *surface.u])
-        p_u = Pijk.from_dense(u @ U(surface.udeg), degree=surface.udeg)
-
-    # surface map * limb darkening map
-    if higher_precision:
-        A1_val = np.atleast_2d(utils_mp.to_numpy(basis_mp.A1(surface.ydeg)))
-    else:
-        A1_val = jax.experimental.sparse.BCOO.from_scipy_sparse(A1(surface.ydeg))
-
-    p_y = Pijk.from_dense(A1_val @ rotated_y, degree=surface.ydeg)
-
-    norm = np.pi / (p_u.tosparse() @ rT(surface.udeg))
-
-    return design_matrix_p, p_y, p_u, norm, b_occ
-
-
 def surface_light_curve(
     surface: Surface,
     r: float | None = None,
@@ -165,13 +76,26 @@ def surface_light_curve(
         b = jnp.sqrt(jnp.square(x) + jnp.square(y))
         b_rot = jnp.logical_or(jnp.greater_equal(b, 1.0 + r), jnp.less_equal(z, 0.0))
         b_occ = jnp.logical_not(b_rot)
-        theta_z = jnp.arctan2(x, y)
 
         # trick to avoid nan `x=jnp.where...` grad caused by nan sT
         r = jnp.where(b_rot, 1.0, r)
         b = jnp.where(b_rot, 1.0, b)
 
-        sT = solution_vector(total_deg, order=order)(b, r)
+        if surface.ydeg == 0:
+            if surface.udeg == 0:
+                ld_u = jnp.array([])
+            else:
+                ld_u = jnp.concatenate(
+                    [jnp.atleast_1d(jnp.asarray(u_)) for u_ in surface.u], axis=0
+                )
+
+            lc_func = partial(_limb_dark_light_curve, ld_u, order=order)
+            lc = lc_func(b, r)
+            return 1.0 + jnp.where(b_occ, lc, 0)
+
+        else:
+            theta_z = jnp.arctan2(x, y)
+            sT = solution_vector(total_deg, order=order)(b, r)
 
         if total_deg > 0:
             if higher_precision:
@@ -210,10 +134,8 @@ def surface_light_curve(
         A1_val = jax.experimental.sparse.BCOO.from_scipy_sparse(A1(surface.ydeg))
 
     p_y = Pijk.from_dense(A1_val @ rotated_y, degree=surface.ydeg)
-
-    norm = np.pi / (p_u.tosparse() @ rT(surface.udeg))
-
     p_yu = p_y * p_u
+    norm = np.pi / (p_u.tosparse() @ rT(surface.udeg))
 
     return surface.amplitude * (p_yu.tosparse() @ design_matrix_p) * norm
 

tests/starry/light_curve_test.py

@@ -470,11 +470,7 @@ def test_EB_interchanged():
     flux_reversed = light_curve(system_2)(time + params["period"] / 2).sum(1)
 
     # for some reason the assert_allclose wasn't catching error here
-    np.testing.assert_allclose(
-        flux_ordered,
-        flux_reversed,
-        atol=1e-6 if flux_ordered.dtype.name == "float32" else 1e-12,
-    )
+    assert_allclose(flux_ordered, flux_reversed)
 
 
 @pytest.mark.parametrize("order", [20, 100, 500])