📜 Conventions

📜 Conventions#

1. Class Design Patterns#

Abstract-Final Pattern#

coordinax follows the abstract-final pattern, a design approach that cleanly separates interface from implementation and avoids deep inheritance hierarchies.

  • Abstract base classes define the interface using abc.ABC and @abc.abstractmethod. These are never instantiated.

  • Concrete (final) classes implement abstract interfaces and are marked with @final decorator. These should not be further subclassed.

  • One inheritance level only: Abstract classes never inherit from concrete classes, and concrete classes never inherit from other concrete classes.

Example:

from abc import ABC, abstractmethod
from typing import final


class AbstractChart(ABC):
    """Abstract base defining chart interface."""

    @abstractmethod
    def components(self) -> tuple[str, ...]: ...


@final
class Cartesian3D(AbstractChart):
    """Concrete chart implementation."""

    def components(self) -> tuple[str, ...]:
        return ("x", "y", "z")

This pattern improves code clarity and avoids the fragile base class problem. See Equinox documentation for more background.

2. Type System & JAX Integration#

PyTree Registration#

All coordinax objects are PyTrees—JAX’s abstraction for hierarchical data structures that can be traced through transformations like jit, vmap, and grad.

  • Objects register via @jax.tree_util.register_static (marker indicating the entire object is static—doesn’t change during JAX transformations).

  • PyTree registration is handled automatically by equinox.Module (used as base class via Equinox).

  • This enables seamless use with JAX: jax.vmap(my_function)(vector_array) works automatically.

Quax & ArrayValue#

To integrate custom types with JAX operations, coordinax uses Quax—a multiple-dispatch layer enabling custom array-like types in JAX.

  • quax.ArrayValue is the protocol: custom types inherit from it to work with JAX.

  • Distance, Angle, and Point are ArrayValue subclasses, so they work naturally with jnp.sqrt, jnp.sin, etc.

  • Operations like +, * are implemented via Quax dispatch on JAX primitives (see Multiple Dispatch section below).

Why this matters: Users can write JAX code treating Distance and Point like arrays—no special handling needed.

See Glossary: Quax, ArrayValue, PyTree.

Immutability & Functional Design#

All coordinax objects are immutable—they don’t change state; instead, operations return new objects.

  • Required for JAX compatibility and functional programming paradigm.

  • Use dataclassish.replace() to update attributes: new_vector = dataclassish.replace(vector, x=new_x).

  • Immutability enables safe use with JAX transformations (no hidden state mutations).

3. API Organization & Design Philosophy#

Scalar-First Design#

Functions in coordinax operate on scalar (0-dimensional) objects—individual points, single vectors. Batching is left to the user via jax.vmap.

  • Why?: Scalar operations JIT-compile more efficiently; users can vmap along any axis they choose.

  • Pattern: Define function(point, static_arg, ...) returning a single point. Users batch via:

    transform_many = jax.vmap(function, in_axes=(0, None, ...))
batched_result = transform_many(point_array, ...)

  • Performance: The scalar body JIT-compiles; vmap efficiently maps over batches.

This design maximizes flexibility and performance.

Functional vs Object-Oriented APIs#

coordinax provides both functional and object-oriented APIs:

  • Functional API (primary): Pure functions taking arguments. Returns new objects; never mutates. Example: pt_map(chart_from, chart_to, point).

  • Object-Oriented API (convenience): Methods on objects wrapping functional APIs. Example: point.transition_to(chart_to).

Both are equally powerful; OOP wraps functional. Choose based on readability.

Module Organization#

Source code (/src/coordinax/) uses this structure:

  • main: User-facing re-exports of primary functionality. Most users start here.

  • Alphabetic submodules: angles, charts, distances, frames, manifolds, representations, vectors. Organized by semantic concept.

  • _src/ subdirectories: Implementation details. Less stable; avoid importing directly.

  • Internal modules: internal folder for utilities not intended for public use.

Import patterns: Always import explicitly; use from_ constructors for flexibility.

See Glossary: Functional API, OOP API, Module Organization.

4. Coordinate Transformations & Conversions#

coordinax distinguishes between vector transformations (coordinate changes) and representation conversions (different forms of same data).

Vector Transformation: cconvert#

cconvert: Change vector components under a coordinate change; preserves role semantics.

  • Input: cconvert(current_vector, target_chart)

  • Output: New vector with components expressed in target_chart

  • Role-aware: Position and velocity transform differently (this is why cconvert exists, not just simple coordinate swaps)

  • Example:

    import coordinax.main as cx
import coordinax.charts as cxc
import unxt as u

# Create a position vector in Cartesian coordinates
pos_cartesian = cx.Point.from_(
    {"x": u.Q(1, "m"), "y": u.Q(0, "m"), "z": u.Q(0, "m")}, cxc.cart3d, cx.point
)
# Transform to spherical coordinates
pos_spherical = cx.cconvert(pos_cartesian, cxc.sph3d)
# Now pos_spherical has (r, theta, phi) components

See Glossary: Vector Transformation, cconvert; spec.md § Transformations.

Representation Conversion: cconvert#

cconvert: Convert object representation without changing underlying data; enables different forms for same object.

  • Input: cconvert(target_representation, current_object)

  • Output: Semantically equivalent object in new form

  • Example: Convert between Distance units (meters → kilometers) or Angle representations (radians → degrees)

  • Contrast with cconvert: Same object (e.g., same distance), different representation

See Glossary: Representation Conversion, cconvert; spec.md § Representations.

5. Multiple Dispatch Patterns#

coordinax uses plum-dispatch for flexible, type-aware function implementations.

Core Pattern: Type Routing#

import plum


@plum.dispatch
def add(x: int, y: int):
    return x + y


@plum.dispatch
def add(x: str, y: str):
    return f"{x}_{y}"

Plum selects implementation based on runtime types of all arguments (not just the first). This enables coordinax to seamlessly handle mixed types (e.g., Distance + Quantity).

Discovering All Implementations#

When working with a dispatched function, use the .methods attribute to see all registered implementations:

from coordinax.main import Distance

print(Distance.from_.methods)

This is essential for understanding what types are supported and avoiding duplicate registrations.

Generic Type Handling in Signatures#

Critical: Plum does NOT support parameterized generic types in function signatures. Always use the base class without type parameters:

# CORRECT
@plum.dispatch
def process(obj: cxc.AbstractChart, /):
    ...


# WRONG - causes plum dispatch warnings
from typing import Any


@plum.dispatch
def process(obj: cxc.AbstractChart[Any, Any, Any], /):
    ...

Add # type: ignore[type-arg] comment to suppress type-checker warnings about missing type parameters.

Promotion Pattern: Handling Mixed Types#

Common pattern for binary operations:

from jaxtyping import ArrayLike
import coordinax.distances as cxd


@plum.dispatch
def add(x: cxd.Distance, y: cxd.Distance):
    return cxd.Distance(x.value + y.value, x.unit)


@plum.dispatch
def add(x: ArrayLike, y: cxd.Distance):  # Promote array to Distance
    return cxd.Distance(x + y.value, y.unit)

Promotion dispatches handle mixed types by converting simpler types to richer ones, then redispatching.

See Glossary: Multiple Dispatch, Promotion; plum documentation for full reference.

6. Naming Conventions#

Class Name Prefixes#

  • Abstract...: Abstract base class defining an interface (e.g., AbstractChart, AbstractVector, AbstractDistance). See Abstract-Final Pattern.

  • PhysDisp: Shorthand for “position”, indicating a point/location vector (e.g., CartesianPhysDisp3D).

  • PhysVel: Shorthand for “velocity”, indicating velocity vector (e.g., CartesianPhysVel3D).

  • PhysAcc: Shorthand for “acceleration”, indicating acceleration vector (e.g., CartesianPhysAcc3D).

  • 0D, 1D, 2D, 3D, N-D: Chart dimension (e.g., AbstractChart3D for 3D manifold charts).

Method Naming Patterns#

  • from_(...): Constructor method accepting diverse input types. Flexible alternative to overloading __init__. Example: Distance.from_(10 * u.m) or Distance.from_((10, "m")).

Pre-Defined Chart Instances#

For convenience, modules provide lowercase instances:

  • cart3d: Instance of Cartesian3D

  • sph3d: Instance of Spherical3D

  • lonlatsph3d: Instance of LonLatSpherical3D

See Glossary: Chart Instance, Chart Class.

Functional vs Object-Oriented APIs#

As JAX is function-oriented, but Python is generally object-oriented, coordinax provides both functional and object-oriented APIs. The functional APIs are the primary APIs, but the object-oriented APIs are easy to use and call the functional APIs, so lose none of the power.

Multiple Dispatch#

coordinax uses multiple dispatch to hook into quax’s flexible and extensible system to enable custom array-ish objects, like Quantity, in JAX. Also, coordinax uses multiple dispatch to enable deep interoperability between coordinax and other libraries, like astropy (and anything user-defined).

For more information on multiple dispatch, see the plum documentation.