squishyplanet 0.0.1__tar.gz

squishyplanet-0.0.1/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Ben Cassese
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

squishyplanet-0.0.1/PKG-INFO

@@ -0,0 +1,34 @@
+Metadata-Version: 2.1
+Name: squishyplanet
+Version: 0.0.1
+Summary: A package for dealing with non-spherical exoplanets
+Author-email: Ben Cassese <b.c.cassese@columbia.edu>
+License: MIT
+Project-URL: Homepage, https://github.com/ben-cassese/squishyplanet
+Project-URL: Documentation, https://squishyplanet.readthedocs.io/en/latest/
+Keywords: astronomy,exoplanets
+Classifier: Programming Language :: Python
+Classifier: Intended Audience :: Science/Research
+Classifier: Topic :: Scientific/Engineering :: Astronomy
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: requests
+Requires-Dist: jax
+Requires-Dist: jaxlib
+Requires-Dist: quadax
+Requires-Dist: tqdm
+Requires-Dist: black
+Requires-Dist: astropy>5
+Requires-Dist: matplotlib
+Requires-Dist: sympy
+Requires-Dist: ipykernel
+
+# squishyplanet
+A package for dealing with non-spherical exoplanets
+
+[![Documentation Status](https://readthedocs.org/projects/squishyplanet/badge/?version=latest)](https://squishyplanet.readthedocs.io/en/latest/?badge=latest)
+![Build Status](https://github.com/ben-cassese/squishyplanet/actions/workflows/tests.yml/badge.svg)
+[![Coverage Status](https://coveralls.io/repos/github/ben-cassese/squishyplanet/badge.svg?branch=main&service=github)](https://coveralls.io/github/ben-cassese/squishyplanet?branch=main)
+
+WIP!

squishyplanet-0.0.1/README.md

@@ -0,0 +1,8 @@
+# squishyplanet
+A package for dealing with non-spherical exoplanets
+
+[![Documentation Status](https://readthedocs.org/projects/squishyplanet/badge/?version=latest)](https://squishyplanet.readthedocs.io/en/latest/?badge=latest)
+![Build Status](https://github.com/ben-cassese/squishyplanet/actions/workflows/tests.yml/badge.svg)
+[![Coverage Status](https://coveralls.io/repos/github/ben-cassese/squishyplanet/badge.svg?branch=main&service=github)](https://coveralls.io/github/ben-cassese/squishyplanet?branch=main)
+
+WIP!

squishyplanet-0.0.1/pyproject.toml

@@ -0,0 +1,50 @@
+[build-system]
+requires = ["setuptools"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "squishyplanet"
+version = "0.0.1"
+requires-python = ">=3.9"
+dependencies = [
+    "requests",
+    "jax",
+    "jaxlib",
+    "quadax",
+    "tqdm",
+    "black",
+    "astropy>5",
+    "matplotlib",
+    "sympy",
+    "ipykernel",
+]
+authors = [
+  {name = "Ben Cassese", email = "b.c.cassese@columbia.edu"},
+]
+license = {text = "MIT"}
+keywords = ["astronomy", "exoplanets"]
+description = "A package for dealing with non-spherical exoplanets"
+readme = "README.md"
+classifiers = [
+  "Programming Language :: Python",
+  "Intended Audience :: Science/Research",
+  "Topic :: Scientific/Engineering :: Astronomy"
+]
+
+[project.urls]
+Homepage = "https://github.com/ben-cassese/squishyplanet"
+Documentation = "https://squishyplanet.readthedocs.io/en/latest/"
+
+[tool.setuptools]
+include-package-data = true
+
+[tool.black]
+line-length = 88
+
+[tool.coverage.run]
+omit = [
+    "tests/*",
+]
+exclude_lines = [
+    '.*@_defjvp*'
+]

squishyplanet-0.0.1/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

squishyplanet-0.0.1/setup.py

@@ -0,0 +1,5 @@
+# needed for editable installs when using pip<21.1
+
+from setuptools import setup
+
+setup()

squishyplanet-0.0.1/squishyplanet/__init__.py

@@ -0,0 +1,7 @@
+from .oblate_system import OblateSystem
+
+__all__ = ["OblateSystem"]
+__url__ = "https://github.com/ben-cassese/squishyplanet"
+__license__ = "MIT"
+__description__ = "Modeling transits of non-spherical exoplanets"
+__version__ = "0.1.0"

squishyplanet-0.0.1/squishyplanet/engine/development/lambertian_reflectance.py

@@ -0,0 +1,177 @@
+import jax
+
+jax.config.update("jax_enable_x64", True)
+
+import jax.numpy as jnp
+
+
+# for illustration only- not actually doing this 2D integral anywhere
+def profile(
+    x,
+    y,
+    p_xx,
+    p_xy,
+    p_xz,
+    p_x0,
+    p_yy,
+    p_yz,
+    p_y0,
+    p_zz,
+    p_z0,
+    p_00,
+    x_c,
+    y_c,
+    z_c,
+    **kwargs,
+):
+    return (
+        -(
+            x_c
+            * (
+                p_x0
+                + 2 * p_xx * x
+                + p_xy * y
+                - (
+                    p_xz
+                    * (
+                        p_z0
+                        + p_xz * x
+                        + p_yz * y
+                        - jnp.sqrt(
+                            (p_z0 + p_xz * x + p_yz * y) ** 2
+                            - 4
+                            * p_zz
+                            * (
+                                -1
+                                + p_00
+                                + p_x0 * x
+                                + p_xx * x**2
+                                + p_y0 * y
+                                + p_xy * x * y
+                                + p_yy * y**2
+                            )
+                        )
+                    )
+                )
+                / (2.0 * p_zz)
+            )
+        )
+        - (
+            p_y0
+            + p_xy * x
+            + 2 * p_yy * y
+            - (
+                p_yz
+                * (
+                    p_z0
+                    + p_xz * x
+                    + p_yz * y
+                    - jnp.sqrt(
+                        (p_z0 + p_xz * x + p_yz * y) ** 2
+                        - 4
+                        * p_zz
+                        * (
+                            -1
+                            + p_00
+                            + p_x0 * x
+                            + p_xx * x**2
+                            + p_y0 * y
+                            + p_xy * x * y
+                            + p_yy * y**2
+                        )
+                    )
+                )
+            )
+            / (2.0 * p_zz)
+        )
+        * y_c
+        - jnp.sqrt(
+            (p_z0 + p_xz * x + p_yz * y) ** 2
+            - 4
+            * p_zz
+            * (
+                -1
+                + p_00
+                + p_x0 * x
+                + p_xx * x**2
+                + p_y0 * y
+                + p_xy * x * y
+                + p_yy * y**2
+            )
+        )
+        * z_c
+    ) / (
+        jnp.sqrt(
+            (p_z0 + p_xz * x + p_yz * y) ** 2
+            - 4
+            * p_zz
+            * (
+                -1
+                + p_00
+                + p_x0 * x
+                + p_xx * x**2
+                + p_y0 * y
+                + p_xy * x * y
+                + p_yy * y**2
+            )
+            + (
+                p_x0
+                + 2 * p_xx * x
+                + p_xy * y
+                - (
+                    p_xz
+                    * (
+                        p_z0
+                        + p_xz * x
+                        + p_yz * y
+                        - jnp.sqrt(
+                            (p_z0 + p_xz * x + p_yz * y) ** 2
+                            - 4
+                            * p_zz
+                            * (
+                                -1
+                                + p_00
+                                + p_x0 * x
+                                + p_xx * x**2
+                                + p_y0 * y
+                                + p_xy * x * y
+                                + p_yy * y**2
+                            )
+                        )
+                    )
+                )
+                / (2.0 * p_zz)
+            )
+            ** 2
+            + (
+                p_y0
+                + p_xy * x
+                + 2 * p_yy * y
+                - (
+                    p_yz
+                    * (
+                        p_z0
+                        + p_xz * x
+                        + p_yz * y
+                        - jnp.sqrt(
+                            (p_z0 + p_xz * x + p_yz * y) ** 2
+                            - 4
+                            * p_zz
+                            * (
+                                -1
+                                + p_00
+                                + p_x0 * x
+                                + p_xx * x**2
+                                + p_y0 * y
+                                + p_xy * x * y
+                                + p_yy * y**2
+                            )
+                        )
+                    )
+                )
+                / (2.0 * p_zz)
+            )
+            ** 2
+        )
+        * jnp.sqrt(x_c**2 + y_c**2 + z_c**2)
+    )

squishyplanet-0.0.1/squishyplanet/engine/development/quadratic_limb_darkened_transit.py

@@ -0,0 +1,265 @@
+# this is a quadratic-only limb darkening model, uses a different method than the
+# Agol, Luger, Foreman-Mackey 2020 scheme. Should produce the same results as the
+# more generate polynomial_limb_darkened_transit.py model when all terms > u2 are 0
+
+import jax
+
+jax.config.update("jax_enable_x64", True)
+import jax.numpy as jnp
+
+from quadax import quadgk
+
+from squishyplanet.engine.planet_3d import planet_3d_coeffs
+from squishyplanet.engine.planet_2d import planet_2d_coeffs
+from squishyplanet.engine.parametric_ellipse import (
+    poly_to_parametric,
+    cartesian_intersection_to_parametric_angle,
+)
+from squishyplanet.engine.skypos import skypos
+from squishyplanet.engine.star_planet_intersection import single_intersection_points
+from squishyplanet.engine.kepler import kepler
+
+
+@jax.jit
+def planet_flux_integrand(alpha, u1, u2, c_x1, c_x2, c_x3, c_y1, c_y2, c_y3, **kwargs):
+    return (
+        (
+            c_x3 * (-(jnp.sin(alpha) * c_y1) + jnp.cos(alpha) * c_y2)
+            - c_x2 * (c_y1 + jnp.cos(alpha) * c_y3)
+            + c_x1 * (c_y2 + jnp.sin(alpha) * c_y3)
+        )
+        * (
+            3
+            * u2
+            * (
+                (jnp.cos(alpha) * c_x1 + jnp.sin(alpha) * c_x2 + c_x3) ** 2
+                + (jnp.cos(alpha) * c_y1 + jnp.sin(alpha) * c_y2 + c_y3) ** 2
+            )
+            ** 2
+            - 4
+            * (u1 + 2 * u2)
+            * (
+                -1
+                + jnp.sqrt(
+                    1
+                    - (jnp.cos(alpha) * c_x1 + jnp.sin(alpha) * c_x2 + c_x3) ** 2
+                    - (jnp.cos(alpha) * c_y1 + jnp.sin(alpha) * c_y2 + c_y3) ** 2
+                )
+            )
+            + 2
+            * (
+                (jnp.cos(alpha) * c_x1 + jnp.sin(alpha) * c_x2 + c_x3) ** 2
+                + (jnp.cos(alpha) * c_y1 + jnp.sin(alpha) * c_y2 + c_y3) ** 2
+            )
+            * (
+                3
+                - 3 * u1
+                - 6 * u2
+                + 2
+                * u1
+                * jnp.sqrt(
+                    1
+                    - (jnp.cos(alpha) * c_x1 + jnp.sin(alpha) * c_x2 + c_x3) ** 2
+                    - (jnp.cos(alpha) * c_y1 + jnp.sin(alpha) * c_y2 + c_y3) ** 2
+                )
+                + 4
+                * u2
+                * jnp.sqrt(
+                    1
+                    - (jnp.cos(alpha) * c_x1 + jnp.sin(alpha) * c_x2 + c_x3) ** 2
+                    - (jnp.cos(alpha) * c_y1 + jnp.sin(alpha) * c_y2 + c_y3) ** 2
+                )
+            )
+        )
+    ) / (
+        12.0
+        * (
+            (jnp.cos(alpha) * c_x1 + jnp.sin(alpha) * c_x2 + c_x3) ** 2
+            + (jnp.cos(alpha) * c_y1 + jnp.sin(alpha) * c_y2 + c_y3) ** 2
+        )
+    )
+
+
+# @jax.jit
+def star_flux_integral(alpha1, alpha2, u1, u2, c_x1, c_x2, c_x3, c_y1, c_y2, c_y3):
+    x1 = c_x1 * jnp.cos(alpha1) + c_x2 * jnp.sin(alpha1) + c_x3
+    y1 = c_y1 * jnp.cos(alpha1) + c_y2 * jnp.sin(alpha1) + c_y3
+    theta1 = jnp.arctan2(y1, x1)
+    theta1 = jnp.where(theta1 < 0, theta1 + 2 * jnp.pi, theta1)
+
+    x2 = c_x1 * jnp.cos(alpha2) + c_x2 * jnp.sin(alpha2) + c_x3
+    y2 = c_y1 * jnp.cos(alpha2) + c_y2 * jnp.sin(alpha2) + c_y3
+    theta2 = jnp.arctan2(y2, x2)
+    theta2 = jnp.where(theta2 < 0, theta2 + 2 * jnp.pi, theta2)
+
+    # print(theta1, theta2)
+    # jax.debug.print("theta1 {x}", x=theta1)
+    # jax.debug.print("theta2 {x}", x=theta2)
+
+    delta = jnp.abs(jnp.arctan2(jnp.sin(theta1 - theta2), jnp.cos(theta1 - theta2)))
+    # jax.debug.print("weird delta {x}", x=delta)
+
+    # jax.debug.print("star contribution {x}", x=delta / (2 * jnp.pi))
+    return delta / (2 * jnp.pi)
+
+
+def limb_darkening_profile(x, y, u1, u2, **kwargs):
+    return (
+        1
+        + u1 * (-1 + jnp.sqrt(1 - x**2 - y**2))
+        - u2 * (-1 + jnp.sqrt(1 - x**2 - y**2)) ** 2
+    )
+
+
+@jax.jit
+def lightcurve(state, times):
+    epsabs = epsrel = 1e-12
+
+    # kepler's eq for times -> f
+    time_deltas = times - state["t_peri"]
+    mean_anomalies = 2 * jnp.pi * time_deltas / state["period"]
+    f = kepler(mean_anomalies, state["e"])
+    state["f"] = f
+    # state["f"] = times
+
+    three = planet_3d_coeffs(**state)
+    two = planet_2d_coeffs(**three)
+    para = poly_to_parametric(**two)
+
+    fluxes = jnp.ones_like(times)
+
+    positions = skypos(**state)
+
+    possibly_in_transit = (
+        (jnp.abs(positions[0, :]) < 1.0 + state["r"])
+        * (jnp.abs(positions[1, :]) < 1.0 + state["r"])
+        * (positions[2, :] > 0)
+    )
+
+    # not thrilled with this implementation- you can feed all of para straight to
+    # quadgk without scanning or vmapping, which makes me think you should be able to
+    # feed it just all of the entries where it's transiting. but, I kept running into
+    # ConcretizationTypeError issues, so this is the workaround for now
+    def fully_transiting(X):
+        para, _, _ = X
+        func = jax.tree_util.Partial(
+            planet_flux_integrand, u1=state["u1"], u2=state["u2"], **para
+        )
+        blocked_flux = quadgk(func, [0, 2 * jnp.pi], epsabs=epsabs, epsrel=epsrel)[0]
+        blocked_flux = blocked_flux / (
+            -(1 / 6) * jnp.pi * (-6 + 2 * state["u1"] + state["u2"])
+        )
+        return blocked_flux
+
+    def partially_transiting(X):
+        # return jnp.nan
+        para, xs, ys = X
+        alphas = cartesian_intersection_to_parametric_angle(xs, ys, **para)
+        alphas = jnp.where(xs != 999, alphas, 2 * jnp.pi)
+        alphas = jnp.where(alphas < 0, alphas + 2 * jnp.pi, alphas)
+        alphas = jnp.where(alphas > 2 * jnp.pi, alphas - 2 * jnp.pi, alphas)
+        alphas = jnp.sort(alphas)
+        # jax.debug.print("xs {x}", x=xs)
+        # jax.debug.print("ys {x}", x=ys)
+        # jax.debug.print("alphas {x}", x=alphas)
+
+        test_ang = alphas[0] + (alphas[1] - alphas[0]) / 2
+        test_ang = jnp.where(test_ang > 2 * jnp.pi, test_ang - 2 * jnp.pi, test_ang)
+        # jax.debug.print("test_ang {x}", x=test_ang)
+
+        test_val = limb_darkening_profile(
+            x=para["c_x1"] * jnp.cos(test_ang)
+            + para["c_x2"] * jnp.sin(test_ang)
+            + para["c_x3"],
+            y=para["c_y1"] * jnp.cos(test_ang)
+            + para["c_y2"] * jnp.sin(test_ang)
+            + para["c_y3"],
+            u1=0.6,
+            u2=0.2,
+            **para,
+        )
+        # jax.debug.print("test_val {x}", x=test_val)
+
+        func = jax.tree_util.Partial(
+            planet_flux_integrand, u1=state["u1"], u2=state["u2"], **para
+        )
+        # full = quadgk(func, [0, 2 * jnp.pi], epsabs=epsabs, epsrel=epsrel)[
+        #     0
+        # ] / (-(1 / 6) * jnp.pi * (-6 + 2 * state["u1"] + state["u2"]))
+
+        def testval_is_not_nan(_):
+            planet_contribution = quadgk(
+                func, [alphas[0], alphas[1]], epsabs=epsabs, epsrel=epsrel
+            )[0] / (-(1 / 6) * jnp.pi * (-6 + 2 * state["u1"] + state["u2"]))
+            return planet_contribution
+
+        def testval_is_nan(_):
+            # planet_contribution =
+
+            leg1 = quadgk(func, [alphas[1], 2 * jnp.pi], epsabs=epsabs, epsrel=epsrel)[
+                0
+            ]
+            leg2 = quadgk(func, [0.0, alphas[0]], epsabs=epsabs, epsrel=epsrel)[0]
+            planet_contribution = (leg1 + leg2) / (
+                -(1 / 6) * jnp.pi * (-6 + 2 * state["u1"] + state["u2"])
+            )
+            return planet_contribution
+
+        planet_contribution = jax.lax.cond(
+            jnp.isnan(test_val), testval_is_nan, testval_is_not_nan, ()
+        )
+        # jax.debug.print("planet_contribution {x}", x=planet_contribution)
+        # #jax.debug.print("full {x}", x=full)
+        # #jax.debug.print("side1 {x}", x=side1)
+        # #jax.debug.print("side2 {x}", x=side2)
+
+        # planet_contribution = jnp.where(~jnp.isnan(test_val), side2, side1)
+        star_contribution = star_flux_integral(
+            alpha1=alphas[0],
+            alpha2=alphas[1],
+            u1=state["u1"],
+            u2=state["u2"],
+            **para,
+        )
+
+        total_blocked = planet_contribution + star_contribution
+
+        return total_blocked
+
+    def transiting(X):
+        indv_para, indv_two = X
+        (
+            xs,
+            ys,
+        ) = single_intersection_points(**indv_two)
+        on_limb = jnp.sum(xs) != 999 * 4
+
+        return jax.lax.cond(
+            on_limb,
+            partially_transiting,
+            fully_transiting,
+            (indv_para, xs, ys),
+        )
+
+    def not_transiting(X):
+        return 0.0
+
+    def scan_func(carry, scan_over):
+        indv_para, indv_two, mask = scan_over
+        return None, jax.lax.cond(
+            mask, transiting, not_transiting, (indv_para, indv_two)
+        )
+
+    # _, transit_fluxes = jax.lax.scan(
+    #     scan_func, None, (para, two, possibly_in_transit), None
+    # )
+    # fluxes -= transit_fluxes
+    two["rho_xx"] = jnp.ones_like(f) * two["rho_xx"]
+    two["rho_xy"] = jnp.ones_like(f) * two["rho_xy"]
+    two["rho_yy"] = jnp.ones_like(f) * two["rho_yy"]
+
+    transit_fluxes = jax.lax.scan(
+        scan_func, None, (para, two, possibly_in_transit), None
+    )[1]
+
+    return fluxes - transit_fluxes