sajax.geometry

geometry.py — JAX rotation matrices and coordinate transforms.

Replaces the astropy.coordinates.matrix_utilities.rotation_matrix dependency from the original SAGE code with pure JAX, making all geometry operations differentiable and JIT-compilable.

Geometry convention (identical to original SAGE):

• Observer is at z → +∞. The plane of sky is X-Y.

• The stellar rotation axis is the y-axis.

• inc_star = 90° → equator-on (observer sees the equator).

• inc_star = 0° → pole-on (observer looks at the north pole).

Functions

rotation_matrix_y(→ jax.numpy.ndarray)	3x3 active rotation matrix around the y-axis.
rotation_matrix_x(→ jax.numpy.ndarray)	3x3 active rotation matrix around the x-axis.
rotate_active_region(→ jax.numpy.ndarray)	Apply stellar rotation (y-axis) then stellar inclination (x-axis)

Module Contents

sajax.geometry.rotation_matrix_y(angle_rad: float) → jax.numpy.ndarray

3x3 active rotation matrix around the y-axis.

Parameters:

angle_rad (float) – Rotation angle in radians.

Return type:

jnp.ndarray, shape (3, 3)

sajax.geometry.rotation_matrix_x(angle_rad: float) → jax.numpy.ndarray

3x3 active rotation matrix around the x-axis.

Parameters:

angle_rad (float) – Rotation angle in radians.

Return type:

jnp.ndarray, shape (3, 3)

sajax.geometry.rotate_active_region(cart: jax.numpy.ndarray, phase_deg: float, inc_deg: float) → jax.numpy.ndarray

Apply stellar rotation (y-axis) then stellar inclination (x-axis) to a Cartesian coordinate vector of an active region.

This replaces the two-step stellar_rotation + stellar_inc functions from the original SAGE code.

Parameters:

• cart (jnp.ndarray, shape (3,)) – [x, y, z] pixel-coordinate position of the active region on the stellar sphere.

• phase_deg (float) – Rotational phase in degrees.

• inc_deg (float) – Stellar inclination in degrees (90 = equator-on, 0 = pole-on).

Returns:

Rotated [x, y, z] coordinates.

Return type:

jnp.ndarray, shape (3,)