From 50d3eac3ed64f8b34e27659a0e49e57b98b7426d Mon Sep 17 00:00:00 2001
From: Nathaniel Starkman <nstarman@users.noreply.github.com>
Date: Fri, 8 Dec 2023 20:09:27 -0500
Subject: [PATCH] typing hints (#21)

Signed-off-by: nstarman <nstarman@users.noreply.github.com>
---
 .../dynamics/mockstream/_df/base.py           |  6 ++---
 .../dynamics/mockstream/_df/fardal.py         |  4 +--
 .../mockstream/_mockstream_generator.py       |  8 +++---
 src/galdynamix/typing.py                      | 27 +++++++++++++++----
 4 files changed, 30 insertions(+), 15 deletions(-)

diff --git a/src/galdynamix/dynamics/mockstream/_df/base.py b/src/galdynamix/dynamics/mockstream/_df/base.py
index b0e34f56..7be7d427 100644
--- a/src/galdynamix/dynamics/mockstream/_df/base.py
+++ b/src/galdynamix/dynamics/mockstream/_df/base.py
@@ -13,11 +13,11 @@
 from galdynamix.dynamics._orbit import Orbit
 from galdynamix.dynamics.mockstream._core import MockStream
 from galdynamix.potential._potential.base import AbstractPotentialBase
-from galdynamix.typing import BatchVec3, FloatScalar, IntegerLike, Vec3, Vec6
+from galdynamix.typing import BatchVec3, FloatScalar, IntLike, Vec3, Vec6
 from galdynamix.utils import partial_jit
 
 Wif: TypeAlias = tuple[Vec3, Vec3, Vec3, Vec3]
-Carry: TypeAlias = tuple[IntegerLike, Vec3, Vec3, Vec3, Vec3]
+Carry: TypeAlias = tuple[IntLike, Vec3, Vec3, Vec3, Vec3]
 
 
 class AbstractStreamDF(eqx.Module):  # type: ignore[misc]
@@ -97,7 +97,7 @@ def _sample(
         prog_mass: FloatScalar,
         t: FloatScalar,
         *,
-        i: IntegerLike,
+        i: IntLike,
         seed_num: int,
     ) -> tuple[BatchVec3, BatchVec3, BatchVec3, BatchVec3]:
         """Generate stream particle initial conditions.
diff --git a/src/galdynamix/dynamics/mockstream/_df/fardal.py b/src/galdynamix/dynamics/mockstream/_df/fardal.py
index f2ddf7d4..a0b19066 100644
--- a/src/galdynamix/dynamics/mockstream/_df/fardal.py
+++ b/src/galdynamix/dynamics/mockstream/_df/fardal.py
@@ -8,7 +8,7 @@
 import jax.numpy as xp
 
 from galdynamix.potential._potential.base import AbstractPotentialBase
-from galdynamix.typing import FloatScalar, IntegerLike, Vec3, Vec6
+from galdynamix.typing import FloatScalar, IntLike, Vec3, Vec6
 from galdynamix.utils import partial_jit
 
 from .base import AbstractStreamDF
@@ -23,7 +23,7 @@ def _sample(
         prog_mass: FloatScalar,
         t: FloatScalar,
         *,
-        i: IntegerLike,
+        i: IntLike,
         seed_num: int,
     ) -> tuple[Vec3, Vec3, Vec3, Vec3]:
         """Generate stream particle initial conditions."""
diff --git a/src/galdynamix/dynamics/mockstream/_mockstream_generator.py b/src/galdynamix/dynamics/mockstream/_mockstream_generator.py
index e4c2cbf4..67a705a8 100644
--- a/src/galdynamix/dynamics/mockstream/_mockstream_generator.py
+++ b/src/galdynamix/dynamics/mockstream/_mockstream_generator.py
@@ -17,7 +17,7 @@
 from galdynamix.potential._potential.base import AbstractPotentialBase
 from galdynamix.typing import (
     FloatScalar,
-    IntegerScalar,
+    IntScalar,
     TimeVector,
     Vec6,
     VecN,
@@ -27,7 +27,7 @@
 
 from ._df import AbstractStreamDF
 
-Carry: TypeAlias = tuple[IntegerScalar, VecN, VecN]
+Carry: TypeAlias = tuple[IntScalar, VecN, VecN]
 
 
 def _converter_immutabledict_or_none(x: Any) -> ImmutableDict[Any] | None:
@@ -92,9 +92,7 @@ def _run_scan(
         qp0_lead = mock0_lead.qp
         qp0_trail = mock0_trail.qp
 
-        def scan_fn(
-            carry: Carry, idx: IntegerScalar
-        ) -> tuple[Carry, tuple[VecN, VecN]]:
+        def scan_fn(carry: Carry, idx: IntScalar) -> tuple[Carry, tuple[VecN, VecN]]:
             i, qp0_lead_i, qp0_trail_i = carry
             qp0_lead_trail = xp.vstack([qp0_lead_i, qp0_trail_i])
             t_i, t_f = ts[i], ts[-1]
diff --git a/src/galdynamix/typing.py b/src/galdynamix/typing.py
index cb10cd45..a0087a88 100644
--- a/src/galdynamix/typing.py
+++ b/src/galdynamix/typing.py
@@ -8,24 +8,37 @@
 # =============================================================================
 # Scalars
 
-IntegerScalar = Integer[Scalar, ""]
+IntScalar = Integer[Scalar, ""]
 """An integer scalar."""
 
-IntegerLike = IntegerScalar | int
+IntLike = IntScalar | int
+"""An integer or integer scalar."""
 
 FloatScalar = Float[Scalar, ""]
 """A float scalar."""
 
-FloatLike = FloatScalar | float
+FloatLike = FloatScalar | float | int
+"""A float(/int) or float scalar."""
+
+FloatOrIntScalar = FloatScalar | IntScalar
+"""A float or integer scalar."""
+
+FloatOrIntScalarLike = FloatLike | IntLike
+"""A float or integer or float(/int) scalar."""
+
 
 # =============================================================================
 # Vectors
 
 # -----------------------------------------------------------------------------
+# Shaped
 
 Vec3 = Float[Array, "3"]
 """A 3-vector, e.g. q=(x, y, z) or p=(vx, vy, vz)."""
 
+Matrix33 = Float[Array, "3 3"]
+"""A 3x3 matrix."""
+
 Vec6 = Float[Array, "6"]
 """A 6-vector e.g. qp=(x, y, z, vx, vy, vz)."""
 
@@ -43,6 +56,7 @@
 BatchableFloatLike = BatchFloatScalar | FloatLike
 
 # -----------------
+# Shaped
 
 BatchVec3 = Shaped[Vec3, "*batch"]
 """Zero or more batches of 3-vectors."""
@@ -53,11 +67,14 @@
 BatchVec7 = Shaped[Vec7, "*batch"]
 """Zero or more batches of 7-vectors."""
 
-VecN = Float[Array, "N"]
-
 ArrayAnyShape = Float[Array, "..."]
+"""An array with any shape."""
 
 # =============================================================================
 # Specific Vectors
 
+VecN = Float[Array, "N"]
+"""An (N,)-vector."""
+
 TimeVector = Float[Array, "time"]
+"""A vector of times."""