jorbit 0.1.1__py3-none-any.whl

jorbit/__init__.py

@@ -0,0 +1,23 @@
+import os
+from astropy.utils.data import download_files_in_parallel, is_url_in_cache
+
+from jorbit.data.constants import (
+    DEFAULT_PLANET_EPHEMERIS_URL,
+    DEFAULT_ASTEROID_EPHEMERIS_URL,
+)
+
+if (not is_url_in_cache(DEFAULT_PLANET_EPHEMERIS_URL)) or (
+    not is_url_in_cache(DEFAULT_ASTEROID_EPHEMERIS_URL)
+):
+    print("JPL DE440 ephemeris files not found in astropy cache, downloading now...")
+    print(
+        "Files are approx. 765 MB, may take several minutes but will not be repeated."
+    )
+    download_files_in_parallel(
+        [DEFAULT_PLANET_EPHEMERIS_URL, DEFAULT_ASTEROID_EPHEMERIS_URL],
+        cache=True,
+        show_progress=True,
+    )
+
+
+from jorbit.particle import Particle

jorbit/accelerations/__init__.py

@@ -0,0 +1,108 @@
+import jax
+
+jax.config.update("jax_enable_x64", True)
+import jax.numpy as jnp
+
+from jorbit.utils.states import SystemState
+from jorbit.accelerations.newtonian import newtonian_gravity
+from jorbit.accelerations.gr import ppn_gravity
+
+
+def create_newtonian_ephemeris_acceleration_func(ephem_processor):
+
+    def func(inputs: SystemState) -> jnp.ndarray:
+        perturber_xs, perturber_vs = ephem_processor.state(inputs.time)
+        perturber_log_gms = ephem_processor.log_gms
+
+        new_state = SystemState(
+            massive_positions=jnp.concatenate([perturber_xs, inputs.massive_positions]),
+            massive_velocities=jnp.concatenate(
+                [perturber_vs, inputs.massive_velocities]
+            ),
+            tracer_positions=inputs.tracer_positions,
+            tracer_velocities=inputs.tracer_velocities,
+            log_gms=jnp.concatenate([perturber_log_gms, inputs.log_gms]),
+            time=inputs.time,
+            acceleration_func_kwargs=inputs.acceleration_func_kwargs,
+        )
+
+        accs = newtonian_gravity(new_state)
+
+        num_perturbers = perturber_xs.shape[0]
+        return accs[num_perturbers:]
+
+    return jax.tree_util.Partial(func)
+
+
+def create_gr_ephemeris_acceleration_func(ephem_processor):
+
+    def func(inputs: SystemState) -> jnp.ndarray:
+        perturber_xs, perturber_vs = ephem_processor.state(inputs.time)
+        perturber_log_gms = ephem_processor.log_gms
+
+        new_state = SystemState(
+            massive_positions=jnp.concatenate([perturber_xs, inputs.massive_positions]),
+            massive_velocities=jnp.concatenate(
+                [perturber_vs, inputs.massive_velocities]
+            ),
+            tracer_positions=inputs.tracer_positions,
+            tracer_velocities=inputs.tracer_velocities,
+            log_gms=jnp.concatenate([perturber_log_gms, inputs.log_gms]),
+            time=inputs.time,
+            acceleration_func_kwargs=inputs.acceleration_func_kwargs,
+        )
+
+        accs = ppn_gravity(new_state)
+
+        num_perturbers = perturber_xs.shape[0]
+        return accs[num_perturbers:]
+
+    return jax.tree_util.Partial(func)
+
+
+def create_default_ephemeris_acceleration_func(ephem_processor):
+
+    def func(inputs: SystemState) -> jnp.ndarray:
+        num_gr_perturbers = 10  # the "planets"
+        num_newtonian_perturbers = 16  # the asteroids
+
+        perturber_xs, perturber_vs = ephem_processor.state(inputs.time)
+        perturber_log_gms = ephem_processor.log_gms
+
+        gr_state = SystemState(
+            massive_positions=jnp.concatenate(
+                [perturber_xs[:num_gr_perturbers], inputs.massive_positions]
+            ),
+            massive_velocities=jnp.concatenate(
+                [perturber_vs[:num_gr_perturbers], inputs.massive_velocities]
+            ),
+            tracer_positions=inputs.tracer_positions,
+            tracer_velocities=inputs.tracer_velocities,
+            log_gms=jnp.concatenate(
+                [perturber_log_gms[:num_gr_perturbers], inputs.log_gms]
+            ),
+            time=inputs.time,
+            acceleration_func_kwargs=inputs.acceleration_func_kwargs,
+        )
+        gr_acc = ppn_gravity(gr_state)[num_gr_perturbers:]
+
+        newtonian_state = SystemState(
+            massive_positions=jnp.concatenate(
+                [perturber_xs[num_gr_perturbers:], inputs.massive_positions]
+            ),
+            massive_velocities=jnp.concatenate(
+                [perturber_vs[num_gr_perturbers:], inputs.massive_velocities]
+            ),
+            tracer_positions=inputs.tracer_positions,
+            tracer_velocities=inputs.tracer_velocities,
+            log_gms=jnp.concatenate(
+                [perturber_log_gms[num_gr_perturbers:], inputs.log_gms]
+            ),
+            time=inputs.time,
+            acceleration_func_kwargs=inputs.acceleration_func_kwargs,
+        )
+        newtonian_acc = newtonian_gravity(newtonian_state)[num_newtonian_perturbers:]
+
+        return gr_acc + newtonian_acc
+
+    return jax.tree_util.Partial(func)

jorbit/accelerations/gr.py

@@ -0,0 +1,152 @@
+# These are pythonized/jaxified versions of acceleration models within REBOUNDx,
+# Tamayo et al. (2020) (DOI: 10.1093/mnras/stz2870). The gr_full function is the
+# equivalent of rebx_calculate_gr_full in REBOUNDx, which is itself based on
+# Newhall et al. (1984) (bibcode: 1983A&A...125..150N)
+# The original code is available at https://github.com/dtamayo/reboundx/blob/502abf3066d9bae174cb20538294c916e73391cd/src/gr_full.c
+# Accessed Fall 2024.
+
+# Many thanks to the REBOUNDx developers for their work, and for making it open source!
+
+import jax
+
+jax.config.update("jax_enable_x64", True)
+import jax.numpy as jnp
+
+from functools import partial
+
+from jorbit.utils.states import SystemState
+from jorbit.accelerations.newtonian import newtonian_gravity
+from jorbit.data.constants import SPEED_OF_LIGHT
+
+
+# equivalent of rebx_calculate_gr_full in reboundx
+@partial(jax.jit, static_argnames=["max_iterations"])
+def ppn_gravity(
+    inputs: SystemState,
+    max_iterations: int = 10,
+) -> jnp.ndarray:
+
+    c2 = inputs.acceleration_func_kwargs.get("c2", SPEED_OF_LIGHT**2)
+
+    # surrendering on the efficient handling of tracers vs. massive particles-
+    # lots of unnecessary computation here if T > 0, but ah well for now
+    M = inputs.massive_positions.shape[0]
+    T = inputs.tracer_positions.shape[0]
+    N = M + T
+    positions = jnp.concatenate(
+        [inputs.massive_positions, inputs.tracer_positions], axis=0
+    )
+    velocities = jnp.concatenate(
+        [inputs.massive_velocities, inputs.tracer_velocities], axis=0
+    )
+    gms = jnp.concatenate([jnp.exp(inputs.log_gms), jnp.zeros(T)])
+
+    dx = positions[:, None, :] - positions[None, :, :]
+    r2 = jnp.sum(dx * dx, axis=-1)
+    r = jnp.sqrt(r2)
+    r3 = r2 * jnp.sqrt(r2)
+
+    # Mask for i!=j calculations
+    mask = ~jnp.eye(N, dtype=bool)
+    prefac = jnp.where(mask, 1.0 / r3, 0.0)
+
+    # Newtonian acceleration
+    a_newt = -jnp.sum(
+        prefac[:, :, None] * dx * gms[None, :, None],
+        axis=1,
+    )  # (N,3)
+
+    dv = velocities[:, None, :] - velocities[None, :, :]  # (N,N,3)
+
+    x_com = jnp.sum(positions * gms[:, None], axis=0) / jnp.sum(gms)
+    v_com = jnp.sum(velocities * gms[:, None], axis=0) / jnp.sum(gms)
+
+    positions = positions - x_com
+    velocities = velocities - v_com
+
+    # Compute velocity-dependent terms
+    v2 = jnp.sum(velocities * velocities, axis=-1)  # (N,)
+    vdot = jnp.sum(velocities[:, None, :] * velocities[None, :, :], axis=-1)  # (N,N)
+
+    a1 = jnp.sum((4.0 / c2) * gms / r, axis=1, where=mask)
+    a1 = jnp.broadcast_to(a1, (N, N)).T
+
+    a2 = jnp.sum((1.0 / c2) * gms / r, axis=1, where=mask)
+    a2 = jnp.broadcast_to(a2, (N, N))
+
+    a3 = jnp.broadcast_to(-v2 / c2, (N, N)).T
+    a4 = -2.0 * jnp.broadcast_to(v2, (N, N)) / c2
+    a5 = (4.0 / c2) * vdot
+
+    a6_0 = jnp.sum(dx * velocities[None, :, :], axis=-1)
+    a6 = (3.0 / (2 * c2)) * (a6_0**2) / r2
+
+    a7 = jnp.sum(dx * a_newt[None, :, :], axis=-1) / (2 * c2)
+
+    factor1 = a1 + a2 + a3 + a4 + a5 + a6 + a7
+    part1 = (
+        jnp.broadcast_to(gms, (N, N))[:, :, None]
+        * dx
+        * factor1[:, :, None]
+        / r3[:, :, None]
+    )
+
+    factor2 = jnp.sum(
+        dx * (4 * velocities[:, None, :] - 3 * velocities[None, :, :]), axis=-1
+    )
+    part2 = (
+        jnp.broadcast_to(gms, (N, N))[:, :, None]
+        * (
+            factor2[:, :, None] * dv / r3[:, :, None]
+            + 7.0 / 2.0 * a_newt[None, :, :] / r[:, :, None]
+        )
+        / c2
+    )
+
+    a_const = jnp.sum(part1 + part2, axis=1, where=mask[:, :, None])
+
+    def iteration_step(a_curr):
+        rdota = jnp.sum(dx * a_curr[None, :, :], axis=-1)  # (N, N)
+        non_const = jnp.sum(
+            (gms[None, :, None] / (2.0 * c2))
+            * (
+                (dx * rdota[:, :, None] / r3[:, :, None])
+                + (7.0 * a_curr[None, :, :] / r[:, :, None])
+            ),
+            axis=1,
+            where=mask[:, :, None],
+        )
+
+        return non_const
+
+    def do_nothing(carry):
+        return carry
+
+    def do_iteration(carry):
+        a_prev, a_curr, _ = carry
+        non_const = iteration_step(a_curr)
+        a_next = a_const + non_const
+
+        ratio = jnp.max(jnp.abs((a_next - a_curr) / a_next), initial=0.0)
+
+        return (a_curr, a_next, ratio)
+
+    def body_fn(carry, _):
+        a_prev, a_curr, ratio = carry
+        should_continue = ratio > jnp.finfo(jnp.float64).eps
+        new_carry = jax.lax.cond(should_continue, do_iteration, do_nothing, carry)
+        return new_carry, None
+
+    # Initialize with constant terms
+    init_a = jnp.zeros_like(a_const)
+    init_carry = (init_a, a_const, 1.0)
+
+    # Run fixed number of iterations using scan
+    final_carry, _ = jax.lax.scan(body_fn, init_carry, None, length=max_iterations)
+
+    _, a_final, _ = final_carry
+
+    # Combine Newtonian and GR terms
+    a = a_newt + a_final
+
+    return a

jorbit/accelerations/newtonian.py

@@ -0,0 +1,51 @@
+import jax
+
+jax.config.update("jax_enable_x64", True)
+import jax.numpy as jnp
+
+from jorbit.utils.states import SystemState
+
+
+@jax.jit
+def newtonian_gravity(inputs: SystemState) -> jnp.ndarray:
+    M = inputs.massive_positions.shape[0]  # number of massive particles
+    T = inputs.tracer_positions.shape[0]  # number of tracer particles
+
+    # 1. Compute accelerations on massive particles due to other massive particles
+    dx_massive = (
+        inputs.massive_positions[:, None, :] - inputs.massive_positions[None, :, :]
+    )  # (M,M,3)
+    r2_massive = jnp.sum(dx_massive * dx_massive, axis=-1)  # (M,M)
+    r3_massive = r2_massive * jnp.sqrt(r2_massive)  # (M,M)
+
+    # Mask for i!=j calculations among massive particles
+    mask_massive = ~jnp.eye(M, dtype=bool)  # (M,M)
+    prefac_massive = jnp.where(mask_massive, 1.0 / r3_massive, 0.0)
+
+    # Accelerations on massive particles from massive particles
+    a_massive = -jnp.sum(
+        prefac_massive[:, :, None]
+        * dx_massive
+        * jnp.exp(inputs.log_gms[None, :, None]),
+        axis=1,
+    )  # (M,3)
+
+    # 2. Compute accelerations on tracer particles due to massive particles
+    # (This will work even when T=0 due to JAX's shape polymorphism)
+    dx_tracers = (
+        inputs.tracer_positions[:, None, :] - inputs.massive_positions[None, :, :]
+    )  # (T,M,3)
+    r2_tracers = jnp.sum(dx_tracers * dx_tracers, axis=-1)  # (T,M)
+    r3_tracers = r2_tracers * jnp.sqrt(r2_tracers)  # (T,M)
+
+    # Accelerations on tracer particles from massive particles
+    a_tracers = -jnp.sum(
+        (1.0 / r3_tracers)[:, :, None]
+        * dx_tracers
+        * jnp.exp(inputs.log_gms[None, :, None]),
+        axis=1,
+    )  # (T,3)
+
+    # Combine accelerations for all particles
+    # This works even when T=0 thanks to JAX's handling of zero-sized arrays
+    return jnp.concatenate([a_massive, a_tracers], axis=0)

jorbit/astrometry/sky_projection.py

@@ -0,0 +1,113 @@
+import jax
+
+jax.config.update("jax_enable_x64", True)
+import jax.numpy as jnp
+
+from jorbit.utils.states import SystemState
+from jorbit.integrators import initialize_ias15_integrator_state, ias15_evolve
+from jorbit.data.constants import INV_SPEED_OF_LIGHT
+
+
+@jax.jit
+def sky_sep(ra1, dec1, ra2, dec2):
+    # all inputs are floats, ICRS positions in radians
+    # output is in arcsec
+
+    # following the astropy source on .separation, using Vincenty
+    delta = ra2 - ra1
+    numerator = jnp.sqrt(
+        (jnp.cos(dec2) * jnp.sin(delta)) ** 2
+        + (
+            jnp.cos(dec1) * jnp.sin(dec2)
+            - jnp.sin(dec1) * jnp.cos(dec2) * jnp.cos(delta)
+        )
+        ** 2
+    )
+
+    denominator = jnp.sin(dec1) * jnp.sin(dec2) + jnp.cos(dec1) * jnp.cos(
+        dec2
+    ) * jnp.cos(delta)
+
+    return jnp.arctan2(numerator, denominator) * 206264.80624709636
+
+
+@jax.jit
+def tangent_plane_projection(ra_ref, dec_ref, ra, dec):
+    # Convert to unit vectors
+    cos_dec = jnp.cos(dec)
+    sin_dec = jnp.sin(dec)
+    cos_ra = jnp.cos(ra)
+    sin_ra = jnp.sin(ra)
+
+    # Initial cartesian coordinates
+    x = cos_dec * cos_ra
+    y = cos_dec * sin_ra
+    z = sin_dec
+
+    # Rotation matrices (combined into single operation)
+    cos_ra_ref = jnp.cos(ra_ref)
+    sin_ra_ref = jnp.sin(ra_ref)
+    cos_dec_ref = jnp.cos(dec_ref)
+    sin_dec_ref = jnp.sin(dec_ref)
+
+    # Apply rotations (optimized matrix multiplication)
+    x_rot = (x * cos_ra_ref + y * sin_ra_ref) * cos_dec_ref + z * sin_dec_ref
+    y_rot = -x * sin_ra_ref + y * cos_ra_ref
+    z_rot = -(x * cos_ra_ref + y * sin_ra_ref) * sin_dec_ref + z * cos_dec_ref
+
+    # Project to plane
+    xi = y_rot / x_rot
+    eta = z_rot / x_rot
+
+    return jnp.array([xi, eta])
+
+
+@jax.jit
+def on_sky(
+    x,
+    v,
+    time,
+    observer_position,
+    acc_func,
+    acc_func_kwargs={},
+):
+    # has to be one particle at one time to get the light travel time right
+    state = SystemState(
+        massive_positions=jnp.empty((0, 3)),
+        massive_velocities=jnp.empty((0, 3)),
+        tracer_positions=jnp.array([x]),
+        tracer_velocities=jnp.array([v]),
+        log_gms=jnp.empty(0),
+        time=time,
+        acceleration_func_kwargs=acc_func_kwargs,
+    )
+    a0 = acc_func(state)
+    initial_integrator_state = initialize_ias15_integrator_state(a0)
+
+    def scan_func(carry, scan_over):
+        xz = carry
+        earth_distance = jnp.linalg.norm(xz - observer_position)
+        light_travel_time = earth_distance * INV_SPEED_OF_LIGHT
+
+        positions, velocities, final_system_state, final_integrator_state = (
+            ias15_evolve(
+                state,
+                acc_func,
+                jnp.array([state.time - light_travel_time]),
+                initial_integrator_state,
+            )
+        )
+
+        return final_system_state.tracer_positions[0], None
+
+    xz, _ = jax.lax.scan(
+        scan_func,
+        state.tracer_positions[0],
+        None,
+        length=3,
+    )
+
+    X = xz - observer_position
+    calc_ra = jnp.mod(jnp.arctan2(X[1], X[0]) + 2 * jnp.pi, 2 * jnp.pi)
+    calc_dec = jnp.pi / 2 - jnp.arccos(X[-1] / jnp.linalg.norm(X))
+    return calc_ra, calc_dec

jorbit/astrometry/transformations.py

@@ -0,0 +1,153 @@
+import jax
+
+jax.config.update("jax_enable_x64", True)
+import jax.numpy as jnp
+
+from jorbit.data.constants import (
+    ICRS_TO_HORIZONS_ECLIPTIC_ROT_MAT,
+    HORIZONS_ECLIPTIC_TO_ICRS_ROT_MAT,
+    TOTAL_SOLAR_SYSTEM_GM,
+)
+
+
+@jax.jit
+def icrs_to_horizons_ecliptic(xs):
+    rotated_xs = jnp.dot(xs, ICRS_TO_HORIZONS_ECLIPTIC_ROT_MAT.T)
+    return rotated_xs
+
+
+@jax.jit
+def horizons_ecliptic_to_icrs(xs):
+    rotated_xs = jnp.dot(xs, HORIZONS_ECLIPTIC_TO_ICRS_ROT_MAT.T)
+    return rotated_xs
+
+
+@jax.jit
+def elements_to_cartesian(a, ecc, nu, inc, Omega, omega):
+    # # Each of the elements are (n_particles, )
+    # # The angles are in *degrees*. Always assuming orbital element angles are in degrees
+
+    nu *= jnp.pi / 180
+    inc *= jnp.pi / 180
+    Omega *= jnp.pi / 180
+    omega *= jnp.pi / 180
+
+    t = (a * (1 - ecc**2))[:, None]
+    r_w = (
+        t
+        / (1 + ecc[:, None] * jnp.cos(nu[:, None]))
+        * jnp.column_stack((jnp.cos(nu), jnp.sin(nu), nu * 0.0))
+    )
+    v_w = (
+        jnp.sqrt(TOTAL_SOLAR_SYSTEM_GM)
+        / jnp.sqrt(t)
+        * jnp.column_stack((-jnp.sin(nu), ecc + jnp.cos(nu), nu * 0))
+    )
+
+    zeros = jnp.zeros_like(omega, dtype=jnp.float64)
+    ones = jnp.ones_like(omega, dtype=jnp.float64)
+    Rot1 = jnp.array(
+        [
+            [jnp.cos(-omega), -jnp.sin(-omega), zeros],
+            [jnp.sin(-omega), jnp.cos(-omega), zeros],
+            [zeros, zeros, ones],
+        ]
+    )
+
+    Rot2 = jnp.array(
+        [
+            [ones, zeros, zeros],
+            [zeros, jnp.cos(-inc), -jnp.sin(-inc)],
+            [zeros, jnp.sin(-inc), jnp.cos(-inc)],
+        ]
+    )
+
+    Rot3 = jnp.array(
+        [
+            [jnp.cos(-Omega), -jnp.sin(-Omega), zeros],
+            [jnp.sin(-Omega), jnp.cos(-Omega), zeros],
+            [zeros, zeros, ones],
+        ]
+    )
+
+    rot = jax.vmap(
+        lambda r1, r2, r3: jnp.matmul(jnp.matmul(r1, r2), r3), in_axes=(2, 2, 2)
+    )(Rot1, Rot2, Rot3)
+
+    x = jax.vmap(lambda x, y: jnp.matmul(x, y))(r_w, rot)
+    v = jax.vmap(lambda x, y: jnp.matmul(x, y))(v_w, rot)
+
+    return x, v
+
+
+@jax.jit
+def cartesian_to_elements(x, v):
+    r_mag = jnp.linalg.norm(x, axis=1)
+    v_mag = jnp.linalg.norm(v, axis=1)
+
+    # Specific angular momentum
+    h = jnp.cross(x, v)
+    h_mag = jnp.linalg.norm(h, axis=1)
+
+    # Eccentricity vector
+    e_vec = jnp.cross(v, h) / TOTAL_SOLAR_SYSTEM_GM - x / r_mag[:, jnp.newaxis]
+    ecc = jnp.linalg.norm(e_vec, axis=1)
+
+    # Specific orbital energy
+    specific_energy = v_mag**2 / 2 - TOTAL_SOLAR_SYSTEM_GM / r_mag
+
+    a = -TOTAL_SOLAR_SYSTEM_GM / (2 * specific_energy)
+
+    inc = jnp.arccos(h[:, 2] / h_mag) * 180 / jnp.pi
+
+    n = jnp.cross(jnp.array([0, 0, 1]), h)
+    n_mag = jnp.linalg.norm(n, axis=1)
+
+    Omega = jnp.where(
+        n[:, 1] >= 0,
+        jnp.arccos(n[:, 0] / n_mag) * 180 / jnp.pi,
+        360.0 - jnp.arccos(n[:, 0] / n_mag) * 180 / jnp.pi,
+    )
+    Omega = jnp.where(n_mag == 0, 0, Omega)
+
+    omega = jnp.where(
+        n_mag > 0,
+        jnp.where(
+            e_vec[:, 2] >= 0,
+            jnp.arccos(
+                jnp.clip(
+                    jnp.sum(n * e_vec, axis=1)
+                    / (n_mag * jnp.linalg.norm(e_vec, axis=1)),
+                    -1,
+                    1,
+                )
+            )
+            * 180
+            / jnp.pi,
+            360
+            - jnp.arccos(
+                jnp.clip(
+                    jnp.sum(n * e_vec, axis=1)
+                    / (n_mag * jnp.linalg.norm(e_vec, axis=1)),
+                    -1,
+                    1,
+                )
+            )
+            * 180
+            / jnp.pi,
+        ),
+        0,
+    )
+
+    nu = jnp.where(
+        jnp.sum(x * v, axis=1) >= 0,
+        jnp.arccos(jnp.clip(jnp.sum(e_vec * x, axis=1) / (ecc * r_mag), -1, 1))
+        * 180
+        / jnp.pi,
+        360
+        - jnp.arccos(jnp.clip(jnp.sum(e_vec * x, axis=1) / (ecc * r_mag), -1, 1))
+        * 180
+        / jnp.pi,
+    )
+
+    return a, ecc, nu, inc, Omega, omega