### Install jaxoplanet from source

Source: https://github.com/exoplanet-dev/jaxoplanet/blob/main/docs/install.md

Clone the repository and install jaxoplanet in editable mode.

```bash
git clone https://github.com/exoplanet-dev/jaxoplanet.git
cd jaxoplanet
python -m pip install -e .
```

--------------------------------

### Install JAX (CPU)

Source: https://github.com/exoplanet-dev/jaxoplanet/blob/main/README.md

Installs the CPU version of JAX. Follow JAX documentation for other installation options.

```bash
python -m pip install "jax[cpu]"
```

--------------------------------

### Setup Jaxoplanet Environment

Source: https://github.com/exoplanet-dev/jaxoplanet/blob/main/docs/tutorials/rv.ipynb

Configure the number of CPUs and enable double-precision for JAX. This is a common setup step for numerical computations.

```python
import jax
import numpyro

numpyro.set_host_device_count(2)
jax.config.update("jax_enable_x64", True)
```

--------------------------------

### Install Miniforge for ARM64

Source: https://github.com/exoplanet-dev/jaxoplanet/blob/main/docs/install.md

Run the Miniforge installer script for ARM64 architecture on macOS.

```bash
bash Mambaforge-23.11.0-0-MacOSX-arm64.sh
```

--------------------------------

### Install jaxoplanet with pip

Source: https://github.com/exoplanet-dev/jaxoplanet/blob/main/docs/install.md

Use this command to install the latest stable version of jaxoplanet using pip.

```bash
python -m pip install jaxoplanet
```

--------------------------------

### Jaxoplanet and NumPyro Setup

Source: https://github.com/exoplanet-dev/jaxoplanet/blob/main/docs/tutorials/relative-astrometry.ipynb

Initializes JAX and NumPyro, setting the number of CPU cores and enabling 64-bit precision. This is a prerequisite for running the astrometric fitting.

```python
import jax
import numpyro

# For multi-core parallelism (useful when running multiple MCMC chains in parallel)
numpyro.set_host_device_count(2)

# For CPU (use "gpu" for GPU)
numpyro.set_platform("cpu")

jax.config.update("jax_enable_x64", True)
```

--------------------------------

### Run unit tests with Nox

Source: https://github.com/exoplanet-dev/jaxoplanet/blob/main/docs/install.md

After installing from source, run the unit tests using Nox.

```bash
python -m pip install nox
python -m nox -s test
```

--------------------------------

### Setup Jaxoplanet and NumPyro

Source: https://github.com/exoplanet-dev/jaxoplanet/blob/main/docs/tutorials/transit.ipynb

Configure the number of CPUs, set the platform to CPU, and enable double-precision numbers for JAX. This is essential for numerical stability and performance.

```python
import jax
import numpyro

# For multi-core parallelism (useful when running multiple MCMC chains in parallel)
numpyro.set_host_device_count(2)

# For CPU (use "gpu" for GPU)
numpyro.set_platform("cpu")

# For 64-bit precision since JAX defaults to 32-bit
jax.config.update("jax_enable_x64", True)
```

--------------------------------

### Install JAX with Conda or Pip

Source: https://github.com/exoplanet-dev/jaxoplanet/blob/main/docs/install.md

Install the JAX library using either Conda or pip within your activated environment.

```bash
conda install jax
```

```bash
python -m pip install "jax[cpu]"
```

--------------------------------

### Setup for Multwavelength Transit Fitting

Source: https://github.com/exoplanet-dev/jaxoplanet/blob/main/docs/tutorials/mw_transit.ipynb

Configures the number of CPUs, enables double-precision numbers, and imports necessary libraries for JAX, NumPyro, and plotting. Use this for setting up your environment for advanced transit modeling.

```python
import jaxoplanet
from jaxoplanet.light_curves import limb_dark_light_curve
from jaxoplanet.orbits import TransitOrbit
import numpy as np
import matplotlib.pyplot as plt
import numpyro
import numpyro.distributions as dist
import numpyro_ext
import numpyro_ext.distributions as distx, numpyro_ext.optim as optimx
import jax
import jax.numpy as jnp
import arviz as az
import corner
import itertools

numpyro.set_host_device_count(
    2
)  # For multi-core parallelism (useful when running multiple MCMC chains in parallel)
numpyro.set_platform("cpu")  # For CPU (use "gpu" for GPU)
jax.config.update(
    "jax_enable_x64", True
)  # For 64-bit precision since JAX defaults to 32-bit
```

--------------------------------

### JAX NumPy-like API Example

Source: https://github.com/exoplanet-dev/jaxoplanet/blob/main/docs/tutorials/introduction-to-jax.ipynb

Illustrates the similarity between JAX's NumPy API (jnp) and standard NumPy (np) for basic array operations.

```python
from typing import Union

import jax
import jax.numpy as jnp
import numpy as np

x = jnp.linspace(0, 5, 11)
x
```

--------------------------------

### Check Python Architecture

Source: https://github.com/exoplanet-dev/jaxoplanet/blob/main/docs/install.md

Verify your Python installation's architecture, essential for ARM Mac users.

```python
import platform
platform.machine()
```

--------------------------------

### Generate and Visualize Prior Samples

Source: https://github.com/exoplanet-dev/jaxoplanet/blob/main/docs/tutorials/relative-astrometry.ipynb

Generates prior samples from the model and visualizes them using a corner plot. This helps in identifying potential issues with the model setup before proceeding with optimization.

```python
import arviz as az
import corner

n_prior_samples = 2000

key = jax.random.key(0)

key, subkey = jax.random.split(key)
prior_samples = infer.Predictive(model, num_samples=n_prior_samples)(subkey)

converted_prior_samples = {
    f"{p}": np.expand_dims(prior_samples[p], axis=0) for p in prior_samples
}
var_names = [
    "Omega",
    "logP",
    "a",
    "incl",
    "ecc",
    "log_rho_s",
    "log_theta_s",
    "omega",
    "phase",
]
prior_idata = az.from_dict(converted_prior_samples)

corner.corner(prior_idata, var_names=var_names)
plt.show()
```

--------------------------------

### Visualize MCMC Samples with ArviZ

Source: https://github.com/exoplanet-dev/jaxoplanet/blob/main/docs/tutorials/ttv.ipynb

Converts NumPyro MCMC samples to an ArviZ inference data object and plots the trace of specified variables. Requires ArviZ and Matplotlib to be installed.

```python
import arviz as az

inf_data = az.from_numpyro(sampler)

az.plot_trace(
    inf_data,
    var_names=["transit_times_0", "transit_times_1", "duration", "r", "b", "u"],
    backend_kwargs={"constrained_layout": True},
)
plt.show()
```

--------------------------------

### Get Planet Position and Star Velocity

Source: https://github.com/exoplanet-dev/jaxoplanet/blob/main/docs/tutorials/quickstart.ipynb

Calculate the relative position of a planet and the velocity of the star over a given time range. Units are in Rsun/day.

```python
import jax.numpy as jnp
from matplotlib import pyplot as plt

# Get the position of the planet and velocity of the star as a function of time
t = jnp.linspace(0, 730, 5000)
x, y, z = earth.relative_position(t)
vx, vy, vz = earth.central_velocity(t)
```

--------------------------------

### Create and Activate Conda Environment

Source: https://github.com/exoplanet-dev/jaxoplanet/blob/main/docs/install.md

Set up a new Conda environment for jaxoplanet and activate it.

```bash
conda create --name jaxoplanet
conda activate jaxoplanet
conda install pip
```

--------------------------------

### Initialize and Run MCMC Sampler

Source: https://github.com/exoplanet-dev/jaxoplanet/blob/main/docs/tutorials/ttv.ipynb

Initializes and runs an MCMC sampler with NUTS kernel, specifying warmup and sampling steps, number of chains, and progress bar. Use this to perform Bayesian inference.

```python
sampler = numpyro.infer.MCMC(
    numpyro.infer.NUTS(
        model,
        dense_mass=True,
        regularize_mass_matrix=True,
        init_strategy=numpyro.infer.init_to_value(values=opt_params),
    ),
    num_warmup=500,
    num_samples=1000,
    num_chains=2,
    progress_bar=True,
)

sampler.run(jax.random.PRNGKey(1), time, yerr, y=y)
```

--------------------------------

### Optimize Exoplanet Model Parameters

Source: https://github.com/exoplanet-dev/jaxoplanet/blob/main/docs/tutorials/ttv.ipynb

Initializes and runs the optimization process for the exoplanet model. Requires pre-defined 'truth' parameters and the 'numpyro_ext.optim.optimize' function.

```python
init_params = {
    **truth,
    "logD": jnp.log(truth["duration"]),
    "logR": jnp.log(truth["r"]),
    "_b": truth["b"] / (1 + truth["r"]),
}

run_optim = numpyro_ext.optim.optimize(
    model,
    init_strategy=numpyro.infer.init_to_value(values=init_params),
)

opt_params = run_optim(jax.random.PRNGKey(5), time, yerr, y=y)
```

--------------------------------

### JIT Compilation with jax.jit()

Source: https://github.com/exoplanet-dev/jaxoplanet/blob/main/docs/tutorials/introduction-to-jax.ipynb

Demonstrates JIT compilation of a function using the `@jax.jit` decorator. Note that impurities like `print` are ignored during compilation, and the function is re-traced if input array shapes change.

```python
import jax
import jax.numpy as jnp

@jax.jit
def add_one(x: jax.Array) -> jax.Array:
    print("I'm being traced!")
    y = x + 1
    return y


print("First call")
print("Result:", add_one(jnp.array([0.0])))

print("\nSubsequent call with array of the same shape")
print("Result:", add_one(jnp.array([0.0])))

print("\nSubsequent call with array of different shape")
print("Result:", add_one(jnp.array([1.0, 2.0, 3.0])))

```

--------------------------------

### Define Exoplanet System and Parameters

Source: https://github.com/exoplanet-dev/jaxoplanet/blob/main/docs/tutorials/rv.ipynb

Define a star and a single exoplanet system using jaxoplanet's keplerian module. This sets up the physical parameters for simulations.

```python
import numpy as np
import matplotlib.pyplot as plt
from jaxoplanet.orbits import keplerian
from jaxoplanet import constants

# Convert 1.5 Mjup to Msun for jaxoplanet
truth = dict(mass=1.5 * constants.M_jup, period=3.0)
star = keplerian.Central(mass=1.0, radius=1.0)
system = keplerian.System(star).add_body(**truth)
```

--------------------------------

### Run MCMC Sampler

Source: https://github.com/exoplanet-dev/jaxoplanet/blob/main/docs/tutorials/rv.ipynb

Initializes and runs NumPyro's MCMC NUTS sampler to sample the posterior likelihood of model parameters. Requires `jax` and `numpyro.infer`.

```python
sampler = infer.MCMC(
    infer.NUTS(model, init_strategy=infer.init_to_value(values=init_values)),
    num_warmup=2000,
    num_samples=10000,
    progress_bar=True,
    num_chains=2,
)

sampler.run(jax.random.PRNGKey(6), time, y=rv_obs)
```

--------------------------------

### Compute and Plot Rotational Light Curve

Source: https://github.com/exoplanet-dev/jaxoplanet/blob/main/docs/tutorials/quickstart.ipynb

Compute the rotational light curve of a star using a SurfaceSystem and plot the resulting flux over time. Requires jax.numpy (jnp) for time linspace and plotting setup.

```python
from jaxoplanet.starry.light_curves import light_curve

plt.figure(figsize=(6.5, 2.5))
time = jnp.linspace(-20, 20, 1000)
flux = light_curve(system)(time).T[0]

plt.plot(time, flux)
plt.xlabel("time (days)")
plt.ylabel("relative flux")
plt.tight_layout()
```

--------------------------------

### Plot Single Spherical Harmonic and Coordinate System

Source: https://github.com/exoplanet-dev/jaxoplanet/blob/main/docs/conventions.ipynb

Displays a specific spherical harmonic (e.g., Y_1,1) and illustrates a 3D Cartesian coordinate system with an example point and vector. Requires `matplotlib.pyplot` and `jaxoplanet.starry.show_surface`.

```python
plt.figure(figsize=(6, 4))

y = np.array([0.0, 0.0, 0.0, 1.0])
show_surface(y)
plt.xlim(-2, 2)

xo, yo, zo = 0.9, 0.5, 10
r = 0.2
circle = plt.Circle((xo, yo), r, color="k", zorder=10)
plt.gca().add_artist(circle)
color1 = "k"
arrow_style = dict(
    head_width=0.06, head_length=0.06, zorder=10, width=0.01, lw=0, color=color1
)
L = 1.1
plt.plot(0, 0, ".", c=arrow_style["color"], ms=6)
plt.arrow(0, -L, 0, 2 * L, **arrow_style)
plt.arrow(-L, 0, 2 * L, 0, **arrow_style)
plt.text(0.1, -0.2, "$(0,0,0)$", color=color1, fontsize=14)
plt.text(L * 1.01, -0.2, "$x$", color=color1, fontsize=14)
plt.text(-0.2, L * 1.01, "$y$", color=color1, fontsize=14)
plt.text(-0.15, -0.15, "$z$", color=color1, fontsize=14)
plt.text(
    xo, yo + 1.5 * r, f"$({xo},{yo},{zo})$", color=color1, fontsize=14, ha="center"
)
plt.text(-0.9, 0.8, "$Y_{1,1}$", fontsize=18, ha="center", color=color1)
plt.tight_layout()
```

--------------------------------

### Print Package Versions

Source: https://github.com/exoplanet-dev/jaxoplanet/blob/main/docs/tutorials/about.ipynb

This snippet prints the versions of jaxoplanet and its related dependencies used in the tutorials. Ensure these versions match for compatibility.

```python
import jax
import arviz
import numpy
import corner
import numpyro
import jaxoplanet
import numpyro_ext

print(f"jaxoplanet.__version__ = {jaxoplanet.__version__}")
print(f"numpy.__version__ = {numpy.__version__}")
print(f"numpyro.__version__ = {numpyro.__version__}")
print(f"numpyro_ext.__version__ = {numpyro_ext.__version__}")
print(f"jax.__version__ = {jax.__version__}")
print(f"corner.__version__ = {corner.__version__}")
print(f"arviz.__version__ = {arviz.__version__}")
```

--------------------------------

### Optimizing linear model and plotting results

Source: https://github.com/exoplanet-dev/jaxoplanet/blob/main/docs/tutorials/introduction-to-jax.ipynb

Initializes model parameters, performs gradient descent optimization, and visualizes the fitted model against the original data. Prints the final optimized parameters.

```python
initial_params = {"slope": jnp.array([1.0]), "offset": jnp.array([0.0])}

params = gradient_descent(initial_params, data, lr=0.1, num_steps=50)

fig, ax = plt.subplots(figsize=(8, 4))
ax.scatter(data["x"], data["y"], marker=".", color="black", label="Data")
ax.plot(
    data["x"],
    linear_model(params, data["x"]),
    ls="dashed",
    color="orange",
    label="Optimised model",
)
ax.set(xlabel="x", ylabel="y")
ax.legend(loc="upper left", frameon=False)

print(f"slope: {params['slope'][0]:<.2f}, offset: {params['offset'][0]:<.2f}")
```

--------------------------------

### Jaxoplanet Kepler Solver Implementation (Python + JAX)

Source: https://github.com/exoplanet-dev/jaxoplanet/blob/main/docs/tutorials/core-from-scratch.ipynb

Combines the starter and refiner functions to solve Kepler's equation, including input range reduction and re-wrapping. Decorated with `@jax.jit` and `@jnp.vectorize` for JAX acceleration.

```python
@jax.jit
@jnp.vectorize
def kepler_solver_impl(mean_anom, ecc):
    mean_anom = mean_anom % (2 * jnp.pi)

    # We restrict to the range [0, pi)
    high = mean_anom > jnp.pi
    mean_anom = jnp.where(high, 2 * jnp.pi - mean_anom, mean_anom)

    # Solve
    ecc_anom = kepler_starter(mean_anom, ecc)
    ecc_anom = kepler_refiner(mean_anom, ecc, ecc_anom)

    # Re-wrap back into the full range
    ecc_anom = jnp.where(high, 2 * jnp.pi - ecc_anom, ecc_anom)

    return ecc_anom
```

--------------------------------

### Initialize and Run NUTS Sampler

Source: https://github.com/exoplanet-dev/jaxoplanet/blob/main/docs/tutorials/transit.ipynb

Initializes and runs the No-U-Turn Sampler (NUTS) for a given model and data. It configures the sampler with specific parameters like dense mass matrix, regularization, and an initialization strategy. The sampler then runs for a specified number of warmup and sampling steps across multiple chains.

```python
sampler = numpyro.infer.MCMC(
    numpyro.infer.NUTS(
        model,
        dense_mass=True,
        regularize_mass_matrix=True,
        init_strategy=numpyro.infer.init_to_value(values=opt_params),
    ),
    num_warmup=1000,
    num_samples=2000,
    num_chains=2,
    progress_bar=True,
)

sampler.run(jax.random.PRNGKey(1), time, yerr, y=y)
```

--------------------------------

### Define Orbital Parameters and Calculate Total Mass

Source: https://github.com/exoplanet-dev/jaxoplanet/blob/main/docs/tutorials/relative-astrometry.ipynb

Sets up initial orbital parameters for a binary system and calculates the total mass using Kepler's third law. Note that jaxoplanet requires the primary mass to be specified, and semi-major axis and period cannot both be provided.

```python
import astropy.units as u
from astropy import constants

# conversion constant from au to R_sun
au_to_R_sun = (constants.au / constants.R_sun).value

# Just to get started, let's take a look at the orbit using the parameter estimates from Pourbaix et al. 1998

# Orbital elements from Pourbaix et al. 1998
# For the relative astrometric fit, we only need the following parameters
a_ang = 0.324  # arcsec
parallax = 1.0  # arcsec (meaningless choice for now)
a = a_ang * au_to_R_sun / parallax
e = 0.798
i = 96.0 * u.deg  # [rad]
omega = 251.6 * u.deg - np.pi  # Pourbaix reports omega_2, but we want omega_1
Omega = 159.6 * u.deg
P = 28.8 * 365.25  # days
G_grav = constants.G.to(u.R_sun**3 / u.M_sun / u.day**2).value
# m_tot = (4 * np.pi * np.pi * (a * u.Rsun)**3 / (constants.G * (P * u.d)**2)).to(u.Msun).value
m_tot = 4 * np.pi * np.pi * a**3 / (G_grav * P**2)

T0 = Time(1989.92, format="decimalyear")
T0.format = "jd"
T0 = T0.value  # [Julian Date]
```

--------------------------------

### Kepler's Equation Starter Function (Python + JAX)

Source: https://github.com/exoplanet-dev/jaxoplanet/blob/main/docs/tutorials/core-from-scratch.ipynb

Provides an initial estimate for the eccentric anomaly using Markley's approximation. Requires JAX and JAX NumPy. Enable double precision with `jax.config.update('jax_enable_x64', True)` for literature comparisons.

```python
import jax
import jax.numpy as jnp

# Enabling double precision for comparison with standard methods from
# the literature, but everything here also works at lower precision
jax.config.update("jax_enable_x64", True)


def kepler_starter(mean_anom, ecc):
    ome = 1 - ecc
    M2 = jnp.square(mean_anom)
    alpha = 3 * jnp.pi / (jnp.pi - 6 / jnp.pi)
    alpha += 1.6 / (jnp.pi - 6 / jnp.pi) * (jnp.pi - mean_anom) / (1 + ecc)
    d = 3 * ome + alpha * ecc
    alphad = alpha * d
    r = (3 * alphad * (d - ome) + M2) * mean_anom
    q = 2 * alphad * ome - M2
    q2 = jnp.square(q)
    w = jnp.square(jnp.cbrt(jnp.abs(r) + jnp.sqrt(q2 * q + r * r)))
    return (2 * r * w / (jnp.square(w) + w * q + q2) + mean_anom) / d
```

--------------------------------

### Optimize for MAP Solution

Source: https://github.com/exoplanet-dev/jaxoplanet/blob/main/docs/tutorials/transit.ipynb

Find the Maximum A Posteriori (MAP) solution using numpyro_ext.optim.optimize. Supports initialization from prior medians or true values.

```python
init_param_method = "true_values"  # "prior_median" or "true_values"

if init_param_method == "prior_median":
    print("Starting from the prior medians")
    run_optim = numpyro_ext.optim.optimize(
        model, init_strategy=numpyro.infer.init_to_median()
    )
elif init_param_method == "true_values":
    print("Starting from the true values")
    init_params = {
        "t0": T0,
        "logP": jnp.log(PERIOD),
        "logD": jnp.log(DURATION),
        "logR": jnp.log(ROR),
        "_b": B / (1 + ROR),
        "u": U,
    }
    run_optim = numpyro_ext.optim.optimize(
        model,
        init_strategy=numpyro.infer.init_to_value(values=init_params),
    )

opt_params = run_optim(jax.random.PRNGKey(3), time, yerr, y=y)
```

--------------------------------

### Benchmarking JAX Functions with and without JIT

Source: https://github.com/exoplanet-dev/jaxoplanet/blob/main/docs/tutorials/introduction-to-jax.ipynb

Compares the execution speed of a custom function (`selu`) with and without JAX's JIT compilation. Use `.block_until_ready()` to ensure the timing includes the actual computation.

```python
import jax
import jax.numpy as jnp

def selu(x: jax.Array, alpha: float = 1.1, lmbda: float = 1.05) -> jax.Array:
    return lmbda * jnp.where(x > 0, x, alpha * jnp.exp(x) - alpha)


x = jnp.arange(0, 10, 10)

# Without jit compilation
# %timeit -n 2000 selu(x).block_until_ready()

# With jit compilation
selu_jit = jax.jit(selu)
# %timeit -n 2000 selu_jit(x).block_until_ready()

```

--------------------------------

### Set Initial Model Values and Plot

Source: https://github.com/exoplanet-dev/jaxoplanet/blob/main/docs/tutorials/rv.ipynb

Sets initial values for model parameters and plots the initial model against observed data. Requires `matplotlib.pyplot`.

```python
init_values = {"mass": 1.0, "period": 3.01, "error": rv_err}
init_model = rv_model(over_time, init_values)

plt.plot(over_time, init_model, "C0")
plot_data()
```

--------------------------------

### Instantiate a Keplerian System

Source: https://github.com/exoplanet-dev/jaxoplanet/blob/main/docs/tutorials/quickstart.ipynb

Instantiate a Keplerian system with a central object. Default parameters are used for the central object.

```python
from jaxoplanet.orbits.keplerian import System, Central

system = System(Central())  # a central object with some default parameters
```

--------------------------------

### Create a Surface System for Rotational Light Curves

Source: https://github.com/exoplanet-dev/jaxoplanet/blob/main/docs/tutorials/quickstart.ipynb

Attach a defined surface to a central body to create a SurfaceSystem for computing rotational light curves. Ensure the central body object is defined prior to this step.

```python
from jaxoplanet.starry.orbit import SurfaceSystem

system = SurfaceSystem(sun, surface)
```

--------------------------------

### Compute Second Derivative with Composed jax.grad

Source: https://github.com/exoplanet-dev/jaxoplanet/blob/main/docs/tutorials/introduction-to-jax.ipynb

Demonstrates how to compute higher-order derivatives by composing `jax.grad` with itself. This is useful for calculating second or higher derivatives of a function.

```python
def some_polynomial(x: float) -> float:
    return x**2 + 1


# Second derivative
print(f"f''(x) = {jax.grad(jax.grad(some_polynomial))(0.0)}")
```

--------------------------------

### Print MAP Parameters

Source: https://github.com/exoplanet-dev/jaxoplanet/blob/main/docs/tutorials/transit.ipynb

Iterate through the optimized parameters and print their values, excluding specific keys like 'light_curve', 'obs', and '_b'.

```python
for k, v in opt_params.items():
    if k in ["light_curve", "obs", "_b"]:
        continue
    print(f"{k}: {v}")
```

--------------------------------

### JAX Random Number Generation with PRNG Keys

Source: https://github.com/exoplanet-dev/jaxoplanet/blob/main/docs/tutorials/introduction-to-jax.ipynb

Demonstrates JAX's approach to random number generation using PRNG keys, including key splitting for generating subkeys.

```python
# To generate random numbers you need a PRNG key
key = jax.random.PRNGKey(seed=42)
print(f"Initial key = {key}")

# You can split this key into a `subkey` when needed
key, subkey = jax.random.split(key)
print(f"Key after split = {key}")
print(f"Subkey = {subkey}")

# Generate an array of random numbers
random_jnp = jax.random.uniform(key=subkey, shape=(5,))
```

--------------------------------

### Define Binary System with Jaxoplanet

Source: https://github.com/exoplanet-dev/jaxoplanet/blob/main/docs/tutorials/relative-astrometry.ipynb

Constructs a binary system using jaxoplanet's System, Central, and Body objects. The primary mass is derived from Kepler's third law, and orbital parameters are used to define the secondary body.

```python
from jaxoplanet.orbits.keplerian import System, Central, Body

primary = Central(mass=m_tot, radius=1.0)
system = System(primary).add_body(
    Body(
        time_peri=T0,
        period=P,
        inclination=i,
        eccentricity=e,
        omega_peri=omega,
        sin_asc_node=np.sin(Omega),
        cos_asc_node=np.cos(Omega),
        parallax=parallax,
    )
)
secondary = system.bodies[0]
```

--------------------------------

### Sample Posterior with MCMC

Source: https://github.com/exoplanet-dev/jaxoplanet/blob/main/docs/tutorials/relative-astrometry.ipynb

Defines and runs a Markov Chain Monte Carlo (MCMC) sampler to obtain posterior samples for the model parameters. Requires JAX and NumPyro.

```python
def sample_model(key, parallax=None):
    sampler = infer.MCMC(
        infer.NUTS(
            model,
            dense_mass=True,
            regularize_mass_matrix=True,
            init_strategy=infer.init_to_value(values=map_soln),
        ),
        num_warmup=1000,
        num_samples=1000,
        num_chains=2,
        progress_bar=True,
    )
    key, subkey = jax.random.split(key)
    sampler.run(subkey, parallax=parallax)
    return sampler


key, subkey = jax.random.split(key)
sampler = sample_model(subkey)
```

--------------------------------

### Sample Model with Parallax

Source: https://github.com/exoplanet-dev/jaxoplanet/blob/main/docs/tutorials/relative-astrometry.ipynb

Employ the `sample_model` function, providing parallax data, to generate posterior samples for the refined model. Ensure JAX random keys are split.

```python
key, subkey = jax.random.split(key)
plx_sampler = sample_model(subkey, parallax=parallax)
```

--------------------------------

### Summarize Model Parameters with Parallax

Source: https://github.com/exoplanet-dev/jaxoplanet/blob/main/docs/tutorials/relative-astrometry.ipynb

Generate a summary table of key model parameters, including those related to parallax, using ArviZ. This helps in assessing the quality and convergence of the model fit.

```python
az.summary(
    plx_idata,
    var_names=["P", "tperi", "a_ang", "omega", "Omega", "incl", "ecc", "m_tot", "plx"],
)
```

--------------------------------

### Create Stellar Surface with Spots

Source: https://github.com/exoplanet-dev/jaxoplanet/blob/main/docs/tutorials/rossiter.ipynb

Generates a stellar surface with features like spots using ylm_spot and defines its properties including limb darkening coefficients and oblateness.

```python
from jaxoplanet.starry.ylm import ylm_spot

y = ylm_spot(8)(0.8, 0.5, 0.0, 0.0)
surface = Surface(y=y, period=0.5, u=(0.2, 0.4), obl=0.2)
```

--------------------------------

### Summarize Posterior Inferences

Source: https://github.com/exoplanet-dev/jaxoplanet/blob/main/docs/tutorials/relative-astrometry.ipynb

Generates a summary table of the posterior distributions for specified parameters using ArviZ. Useful for quickly assessing parameter estimates and uncertainties.

```python
az.summary(
    idata,
    var_names=["P", "tperi", "a_ang", "omega", "Omega", "incl", "ecc"],
)
```

--------------------------------

### Print Sampler Summary

Source: https://github.com/exoplanet-dev/jaxoplanet/blob/main/docs/tutorials/transit.ipynb

Prints a summary of the MCMC sampler results directly, including convergence diagnostics and the number of divergences. This provides an alternative way to inspect the sampler's performance.

```python
# There's also a method to obtain similar results to `az.summary` but directly
# as a method with the MCMC sampler. It also gives us the number of divergences.
sampler.print_summary()
```

--------------------------------

### Generate and Plot Prior Samples

Source: https://github.com/exoplanet-dev/jaxoplanet/blob/main/docs/tutorials/mw_transit.ipynb

Generates prior samples from a model and visualizes them using a corner plot. Requires prior samples to be converted into an ArviZ InferenceData object.

```python
n_prior_samples = 2000
prior_samples = numpyro.infer.Predictive(model, num_samples=n_prior_samples)(
    jax.random.PRNGKey(0), t, yerr
)

# Let's make it into an arviz InferenceData object.
# To do so we'll first need to reshape the samples to be of shape (chains, draws, *shape)
converted_prior_samples = {
    f"{p}": np.expand_dims(prior_samples[p], axis=0) for p in prior_samples
}
prior_samples_inf_data = az.from_dict(converted_prior_samples)

# Plot the corner plot
fig = plt.figure(figsize=(20, 20))
_ = corner.corner(
    prior_samples_inf_data,
    fig=fig,
    var_names=["u", "t0", "duration", "b", "rors"],
    truths=[*U, T0, DURATION, B, *jnp.sqrt(DEPTHS)],
    show_titles=True,
    title_kwargs={\"fontsize\": 10},
    label_kwargs={\"fontsize\": 10},
)
```

--------------------------------

### Generate Trace Plots

Source: https://github.com/exoplanet-dev/jaxoplanet/blob/main/docs/tutorials/transit.ipynb

Generates trace plots for specified parameters using ArviZ. Trace plots visualize the sampled values over iterations for each chain, helping to assess convergence and explore the posterior distribution.

```python
_ = az.plot_trace(
    inf_data,
    var_names=["t0", "period", "duration", "r", "b", "u"],
    backend_kwargs={"constrained_layout": True},
)
```

--------------------------------

### Generate Corner Plot for Inferred Parameters

Source: https://github.com/exoplanet-dev/jaxoplanet/blob/main/docs/tutorials/rv.ipynb

Generates a corner plot to visualize the inferred model parameters ('mass', 'period') and their distributions. Requires `arviz`, `corner`, and `jaxoplanet.constants`.

```python
import arviz as az
import corner
import arviz as az

idata = az.from_numpyro(sampler)

idata = az.from_numpyro(sampler)

_ = corner.corner(
    idata,
    var_names=["mass", "period"],
    truths=[truth["mass"] / constants.M_jup, truth["period"]],
    show_titles=True,
    quantiles=[0.16, 0.5, 0.84],
    title_fmt=".4f",
)
```

--------------------------------

### Define Orbital Parameters and Wavelengths

Source: https://github.com/exoplanet-dev/jaxoplanet/blob/main/docs/tutorials/mw_transit.ipynb

Sets up essential orbital parameters like period, duration, and impact parameter, along with limb-darkening coefficients and the number of spectroscopic light curves. It also defines the wavelengths for each light curve and calculates wavelength-dependent transit depths.

```python
PERIOD = 10.0  # day
DURATION = 0.3  # day
T0 = 0.0  # day
B = 0.3  # impact parameter
U = np.array([0.4, 0.3])  # LDCs

num_lcs = 10
wavelengths = np.linspace(4.0, 4.9, num_lcs)

# Let's make the depth change for each spectroscopic light curve
DEPTHS = 0.01
DEPTHS += 1e-3 * np.exp(-(((wavelengths - 4.3) / 0.2) ** 2))

# Let's see what the theoretical transmission spectrum would look like
fig, ax = plt.subplots(dpi=150)
ax.plot(wavelengths, DEPTHS, marker=".", ms=10, ls="-")
ax.set_xlabel("wavelength [nm]", fontsize=10)
ax.set_ylabel("transit depth [unitless]", fontsize=10)
```

--------------------------------

### Performance Test: Direct Autodiff on Kepler Solver Implementation

Source: https://github.com/exoplanet-dev/jaxoplanet/blob/main/docs/tutorials/core-from-scratch.ipynb

Measures the performance of computing gradients by applying JAX's automatic differentiation directly to the original `kepler_solver_impl` function. This serves as a baseline to demonstrate the performance improvement gained from using a custom JVP rule.

```python
grad_func = jax.vmap(jax.grad(kepler_solver_impl, argnums=(0, 1)), in_axes=(0, None))
grad_func(mean_anom_bench, ecc)[1].block_until_ready()
%timeit grad_func(mean_anom_bench, ecc)[1].block_until_ready()
```

--------------------------------

### Optimize Model with Parallax

Source: https://github.com/exoplanet-dev/jaxoplanet/blob/main/docs/tutorials/relative-astrometry.ipynb

Use the `optimize_model` function with parallax values to find a more realistic solution. Requires JAX random keys and a list of parallax measurements.

```python
key, subkey = jax.random.split(key)
parallax = [24.05, 0.45]
plx_map_soln = optimize_model(subkey, parallax=parallax)
```

--------------------------------

### Verify Kepler Solver Implementation

Source: https://github.com/exoplanet-dev/jaxoplanet/blob/main/docs/tutorials/core-from-scratch.ipynb

Use this code to verify the accuracy of the implemented Kepler solver by comparing calculated eccentric anomalies with known values. Ensure matplotlib and jax.numpy are imported.

```python
import matplotlib.pyplot as plt

ecc = 0.5
true_ecc_anom = jnp.linspace(0, 2 * jnp.pi, 50_000)[:-1]
mean_anom = true_ecc_anom - ecc * jnp.sin(true_ecc_anom)

calc_ecc_anom = kepler_solver_impl(mean_anom, ecc)

plt.plot(true_ecc_anom, jnp.abs(calc_ecc_anom - true_ecc_anom), "k")
plt.axhline(0, color="k")
plt.xlabel("eccentric anomaly")
plt.ylabel(r"error from Kepler solver");
```

--------------------------------

### JAX Trigonometric Operations

Source: https://github.com/exoplanet-dev/jaxoplanet/blob/main/docs/tutorials/introduction-to-jax.ipynb

Demonstrates JAX's ability to perform trigonometric operations and calculate the mean of an array, similar to NumPy.

```python
unit_circle_arr = jnp.sin(x) ** 2 + jnp.cos(x) ** 2
print(jnp.round(jnp.mean(unit_circle_arr), 7))
```

--------------------------------

### Plot Phase Plots for TTV Analysis

Source: https://github.com/exoplanet-dev/jaxoplanet/blob/main/docs/tutorials/ttv.ipynb

Generates phase plots for individual planets to demonstrate issues with phasing planets exhibiting Transit Timing Variations (TTVs). Requires 'matplotlib.pyplot' and optimized parameters.

```python
fig, ax = plt.subplots(2, 1, figsize=(7.5, 4))

map_t0s = [opt_params[f"t0_{i}"] for i in range(2)]
map_periods = [opt_params[f"period_{i}"] for i in range(2)]

for i, (period, t0) in enumerate(zip(map_periods, map_t0s)):
    folded_time = (time - t0 + 0.5 * period) % period - 0.5 * period
    ax[i].errorbar(
        folded_time,
        y - map_light_curve[:, 1 - i],
        yerr=yerr,
        fmt=".k",
        label="data",
        zorder=-1000,
    )
    ax[i].set_xlim(-0.5 * period, 0.5 * period)
    ax[i].set_xlabel("time since transit [days]")
    ax[i].set_ylabel("relative flux")
    ax[i].set_title("Planet 1")
    ax[i].legend()
```

--------------------------------

### Plot MAP Model Fit

Source: https://github.com/exoplanet-dev/jaxoplanet/blob/main/docs/tutorials/transit.ipynb

Visualize the MAP model's light curve against the simulated data and the true light curve. This helps assess the quality of the MAP solution.

```python
plt.plot(time, y, ".", c="0.7", label="data")
plt.plot(time, y_true, "-k", label="truth")
plt.plot(time, opt_params["light_curve"], "--C0", label="MAP model")
plt.xlabel("time [days]")
plt.ylabel("relative flux")
plt.legend(fontsize=10, loc=4)
plt.xlim(time.min(), time.max())
plt.tight_layout()
```

--------------------------------

### Convert Sampler to ArviZ InferenceData

Source: https://github.com/exoplanet-dev/jaxoplanet/blob/main/docs/tutorials/relative-astrometry.ipynb

Convert the results from the `sample_model` function into an ArviZ `InferenceData` object for easier analysis and visualization. This step is crucial for using ArviZ's plotting and summary tools.

```python
plx_idata = az.from_numpyro(plx_sampler)
```

--------------------------------

### Define and Visualize a Non-uniform Star Surface

Source: https://github.com/exoplanet-dev/jaxoplanet/blob/main/docs/tutorials/quickstart.ipynb

Define a star's surface using spherical harmonics and visualize it. Requires importing necessary modules and setting a random seed for reproducibility.

```python
import numpy as np
from jaxoplanet.starry.ylm import Ylm
from jaxoplanet.starry.surface import Surface
from jaxoplanet.starry.visualization import show_surface

np.random.seed(42)
y = Ylm.from_dense([1.00, *np.random.normal(0.0, 2e-2, size=15)])

u_star = (0.1, 0.1)
surface = Surface(inc=1.0, obl=0.2, period=27.0, u=u_star, y=y)

plt.figure(figsize=(3, 3))
show_surface(surface)
```

--------------------------------

### Benchmark Kepler Solver Performance

Source: https://github.com/exoplanet-dev/jaxoplanet/blob/main/docs/tutorials/core-from-scratch.ipynb

Benchmark the performance of the Kepler solver using JAX's timing utilities. This code is designed to minimize Python overhead by using a large number of test values and blocking until computation is ready.

```python
# Using a large number of test values to minimize effects of Python overhead
mean_anom_bench = jax.device_put(
    jnp.linspace(0, 2 * jnp.pi, 500_000)[:-1]
).block_until_ready()
kepler_solver_impl(mean_anom_bench, ecc).block_until_ready()
%timeit kepler_solver_impl(mean_anom_bench, ecc).block_until_ready()
```

--------------------------------

### Generate ArviZ Inference Data and Summary

Source: https://github.com/exoplanet-dev/jaxoplanet/blob/main/docs/tutorials/transit.ipynb

Converts the sampler output to an ArviZ InferenceData object and generates a summary of the posterior samples. This summary includes diagnostics like the Gelman-Rubin R-hat statistic and effective sample size (ESS) for specified parameters.

```python
inf_data = az.from_numpyro(sampler)
samples = sampler.get_samples()
az.summary(inf_data, var_names=["t0", "period", "duration", "r", "b", "u"])
```

--------------------------------

### Generate Exoplanet Phase Plot

Source: https://github.com/exoplanet-dev/jaxoplanet/blob/main/docs/tutorials/transit.ipynb

This code snippet generates a phase plot for exoplanet transit data. It calculates the median inferred parameters from MCMC samples, models the light curve, and then plots the folded observed data along with the inferred model. Ensure you have 'numpy', 'matplotlib.pyplot', and 'jax.numpy' imported.

```python
inferred_params = {
    "t0": np.median(samples["t0"]),
    "period": np.median(samples["period"]),
    "duration": np.median(samples["duration"]),
    "r": np.median(samples["r"]),
    "b": np.median(samples["b"]),
    "u": np.median(samples["u"], axis=0),
}

y_model = light_curve_model(time, inferred_params)

fig, ax = plt.subplots()

# Plot the folded data
t_fold = (
    time - inferred_params["t0"] + 0.5 * inferred_params["period"])
% inferred_params["period"] - 0.5 * inferred_params["period"]
ax.errorbar(t_fold, y, yerr=yerr, fmt=".", color="0.6", label="data", zorder=-100)

# Plot the folded model
inds = np.argsort(t_fold)
ax.plot(t_fold[inds], y_model[inds], color="C0", label="inferred model")
ax.set_xlim(inferred_params["duration"] * jnp.array([-1, 1]) * 1.5)
ax.set_xlabel("time since transit [days]")
ax.set_ylabel("relative flux")
ax.legend(loc=4)
plt.tight_layout()
```

--------------------------------

### jaxoplanet BibTeX Entry

Source: https://github.com/exoplanet-dev/jaxoplanet/blob/main/README.md

BibTeX entry for citing the jaxoplanet GitHub repository. Use this if you find the code useful in your research.

```bibtex
@software{jaxoplanet,
  author       = {Soichiro Hattori and
                  Lionel Garcia and
                  Catriona Murray and
                  Jiayin Dong and
                  Shashank Dholakia and
                  David Degen and
                  Daniel Foreman-Mackey},
  title        = {{exoplanet-dev/jaxoplanet: Astronomical time series analysis with JAX}},
  month        = mar,
  year         = 2024,
  publisher    = {Zenodo},
  version      = {v0.0.2},
  doi          = {10.5281/zenodo.10736936},
  url          = {https://doi.org/10.5281/zenodo.10736936}
}
```

--------------------------------

### Convert Sampler Output to ArviZ InferenceData

Source: https://github.com/exoplanet-dev/jaxoplanet/blob/main/docs/tutorials/relative-astrometry.ipynb

Converts the results from the MCMC sampler into an ArviZ InferenceData object for easier analysis and visualization. Requires ArviZ.

```python
idata = az.from_numpyro(sampler)
```

--------------------------------

### Enable 64-bit Precision in JAX

Source: https://github.com/exoplanet-dev/jaxoplanet/blob/main/docs/commonissues.md

Enable double precision calculations in JAX to mitigate numerical precision issues. This is often necessary when encountering NaNs or infinities.

```python
import jax

jax.config.update("jax_enable_x64", True)
```

--------------------------------

### Compute Elementwise Derivatives with JAX

Source: https://github.com/exoplanet-dev/jaxoplanet/blob/main/docs/tutorials/core-from-scratch.ipynb

Creates a function to compute elementwise gradients of the `kepler_solver` with respect to specified arguments using `jax.vmap` and `jax.grad`. This is useful for applying gradient computations across arrays.

```python
def elementwise_grad(argnums):
    return jax.vmap(jax.grad(kepler_solver, argnums=argnums), in_axes=(0, None))
```

--------------------------------

### Sample Priors with Predictive

Source: https://github.com/exoplanet-dev/jaxoplanet/blob/main/docs/tutorials/transit.ipynb

Draw samples from the model's priors using numpyro.infer.Predictive and convert them into an ArviZ InferenceData object for analysis.

```python
import arviz as az

n_prior_samples = 3000
prior_samples = numpyro.infer.Predictive(model, num_samples=n_prior_samples)(
    jax.random.PRNGKey(0), time, yerr
)

# Let's make it into an arviz InferenceData object.
# To do so we'll first need to reshape the samples to be of shape (chains, draws, *shape)
converted_prior_samples = {
    f"{p}": np.expand_dims(prior_samples[p], axis=0) for p in prior_samples
}
prior_samples_inf_data = az.from_dict(converted_prior_samples)
```

--------------------------------

### Define a Two-Body System for RV Calculation

Source: https://github.com/exoplanet-dev/jaxoplanet/blob/main/docs/tutorials/rossiter.ipynb

Creates a jaxoplanet SurfaceSystem with a central body and a rotating stellar surface, adding a transiting exoplanet body. Note: jaxoplanet currently only supports two-body systems for dynamic calculations.

```python
from jaxoplanet.starry.orbit import Body, Central, SurfaceSystem

system = SurfaceSystem(Central(), surface).add_body(
    Body(period=20.0, radius=0.1, mass=0.1)
)
```