GeoJax

GeoJax is a fast, JAX-powered geometry and statistics toolkit for scientific computing, robotics, and ML pipelines. It provides efficient, differentiable routines for vector operations, alignment, projection, rotations, circular statistics, outlier detection, and more — all accelerated with JAX.

---

Features

• ✅ Vector algebra, projection, reflection, affine transforms
• ✅ Coordinate alignment via PCA and Gram-Schmidt
• ✅ Rotation matrices from vectors or axis-angles
• ✅ Circular statistics: mean, std, variance (radians & degrees)
• ✅ Mahalanobis distance and outlier detection
• ✅ Bounding shapes: AABB, bounding sphere, OBB
• ✅ Geometric predicates: unit, orthogonal, collinear
• ✅ Fully JIT-compiled with typed inputs via JAX

---

Installation

Install GeoJax and JAX (CPU backend by default):

pip install -e .
pip install "jax[cpu]"

Or with CUDA (if using an NVIDEA GPU):

pip install "jax[cuda11]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

Quick Start

import jax.numpy as jnp
from GeoJax import *

# Rotate a vector
v = jnp.array([1.0, 0.0, 0.0])
axis = jnp.array([0.0, 0.0, 1.0])
v_rotated = rotate_around_axis(v, axis, jnp.pi / 2)

# Detect outliers
X = jnp.array([[1, 2], [1.1, 2.1], [0.9, 1.9], [50, 80]])
outliers = detect_outliers_mahalanobis(X, alpha=0.99)

# Circular mean of radians
angles = jnp.array([0.0, jnp.pi / 2, jnp.pi])
mean_angle = circmean(angles)

Module Overview

Module	Description
core	Linear algebra and vector operations
alignment	PCA alignment, Gram-Schmidt, eigendecomposition
angles	Angle and signed angle computation
rotation	Axis-angle and vector-based rotation matrices
projection	Project to planes, spheres, vectors, or axes
circstats	Circular mean, std, and variance
analysis	Mahalanobis distance and robust outlier detection
bounds	Bounding boxes and spheres
checks	Predicate functions for geometry (unit, orthogonal, etc.)
basis	Canonical basis vectors
utils	Misc tools: normalize angles, origin flipping

Testing

Run all tests and check coverage:

python runtests.py

or manually:

pytest --cov=GeoJax --cov-report=term-missing

Documentation

Each function has it's own docstring - full documentation will be provided soon!

Licence

GeoJax is released under the MIT License