feat: units

README.md

@@ -3,6 +3,7 @@
 [![image](https://img.shields.io/pypi/v/airtrafficsim.svg)](https://pypi.python.org/pypi/airtrafficsim)
 [![image](https://img.shields.io/pypi/l/airtrafficsim.svg)](https://pypi.python.org/pypi/airtrafficsim)
 [![image](https://img.shields.io/pypi/pyversions/airtrafficsim.svg)](https://pypi.python.org/pypi/airtrafficsim)
+[![image](https://img.shields.io/pypi/status/airtrafficsim)](https://pypi.python.org/pypi/airtrafficsim)
 
 <img src="docs/assets/img/Logo-full.png" width=50% />
 
@@ -42,6 +43,8 @@ Using the command above will install a version with very minimal footprint. Depe
 - `polars`: support for polars DataFrame (used in simulation and postprocessing third party data)
 - `networking`: support for downloading data from external third party sources
 - `era5`: support for parsing NetCDF for Google ARCO ERA5.
+- `jax`: support for automatic differentiation
+- `plot`: utils for nicer plotting
 
 For example:
 

examples/units.py

@@ -0,0 +1,3 @@
+import airtrafficsim.units as u
+
+print((u.KILOGRAM * u.METER * u.SECOND**-2).to_siunitx())

pyproject.toml

@@ -4,10 +4,10 @@ build-backend = "hatchling.build"
 
 [project]
 name = "AirTrafficSim"
-version = "0.2.0-alpha.1"
+version = "0.2.0-alpha.2"
 authors = [
     { name = "Frankie Hui", email = "kyhuiaf@connect.ust.hk" },
-    { name = "Abraham Cheung", email = "abraham@ylcheung.com" }
+    { name = "Abraham Cheung", email = "abraham@ylcheung.com" },
 ]
 description = "A lightweight collection of tools for air traffic management research."
 readme = "README.md"
@@ -19,7 +19,13 @@ classifiers = [
     "Intended Audience :: Information Technology",
     "Intended Audience :: Science/Research",
     "License :: OSI Approved :: MIT License",
+    "Operating System :: Microsoft :: Windows",
+    "Operating System :: MacOS",
+    "Operating System :: POSIX :: Linux",
+    "Topic :: Scientific/Engineering :: Physics",
     "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
     "Programming Language :: Python :: 3.12",
     "Programming Language :: Python :: 3.13",
     "Typing :: Typed",
@@ -42,6 +48,12 @@ networking = [
 era5 = [
     "xarray>=2024.11.0",
 ]
+jax = [
+    "jax>=0.5.0",
+]
+plot = [
+    "matplotlib>=3.10.0",
+]
 all = [
     "airtrafficsim[polars,networking,era5]",
 ]
@@ -57,7 +69,6 @@ dev = [
 test = [
     "pytest-cov>=6.0.0",
     "pytest>=8.3.4",
-    "jax>=0.4.38",
     "anyio>=4.7.0",
 ]
 docs = [

src/airtrafficsim/annotations.py

@@ -19,8 +19,10 @@
 At runtime, the types are effectively erased and static type checkers will
 not catch incompatible quantities.
 """
+
 # TODO: beartype integration
 # TODO: add optional LaTeX protocol
+from __future__ import annotations
 
 from dataclasses import dataclass
 from typing import Any, Generic, Literal, TypeAlias, TypeVar
@@ -34,40 +36,40 @@
 
 
 @dataclass(frozen=True, slots=True)
-class _Quantity(Generic[Units]):
-    unit: Units
+class Quantity(Generic[Units]):
+    unit: Units  # keeping it as a string for simplicity.
 
-    def __truediv__(self, other: "_Quantity[Any]") -> "Div":
+    def __truediv__(self, other: Quantity[Any]) -> Div:
         return Div(numerator=self, denominator=other)
 
 
 #
 # Base
 #
 
-Time: TypeAlias = _Quantity[Literal["s", "min", "hr"]]
-Length: TypeAlias = _Quantity[Literal["m", "ft", "nmi", "mi"]]
-Mass: TypeAlias = _Quantity[Literal["kg", "lbm"]]
-Temperature: TypeAlias = _Quantity[Literal["K", "°C", "°F", "°R"]]
-Angle: TypeAlias = _Quantity[Literal["rad", "deg"]]
+Time: TypeAlias = Quantity[Literal["s", "min", "hr"]]
+Length: TypeAlias = Quantity[Literal["m", "ft", "nmi", "mi"]]
+Mass: TypeAlias = Quantity[Literal["kg", "lbm"]]
+Temperature: TypeAlias = Quantity[Literal["K", "°C", "°F", "°R"]]
+Angle: TypeAlias = Quantity[Literal["rad", "deg"]]
 
 #
 # Derived
 #
 
-Force: TypeAlias = _Quantity[Literal["N", "lbf"]]
-Pressure: TypeAlias = _Quantity[Literal["Pa", "psi", "hPa", "inHg"]]
-Energy: TypeAlias = _Quantity[Literal["J"]]
-Power: TypeAlias = _Quantity[Literal["W"]]
-Velocity: TypeAlias = _Quantity[
+Force: TypeAlias = Quantity[Literal["N", "lbf"]]
+Pressure: TypeAlias = Quantity[Literal["Pa", "psi", "hPa", "inHg"]]
+Energy: TypeAlias = Quantity[Literal["J"]]
+Power: TypeAlias = Quantity[Literal["W"]]
+Velocity: TypeAlias = Quantity[
     Literal["m s⁻¹", "kt", "ft min⁻¹", "mi hr⁻¹", "km hr⁻¹"]
 ]
-Acceleration: TypeAlias = _Quantity[Literal["m s⁻²", "ft s⁻²"]]
-Density: TypeAlias = _Quantity[Literal["kg m⁻³", "slug ft⁻³"]]
-GasConstant: TypeAlias = _Quantity[Literal["J mol⁻¹ K⁻¹"]]
-MolarMass: TypeAlias = _Quantity[Literal["kg mol⁻¹"]]
-SpecificGasConstant: TypeAlias = _Quantity[Literal["J kg⁻¹ K⁻¹"]]
-ThrustSpecificFuelConsumption: TypeAlias = _Quantity[
+Acceleration: TypeAlias = Quantity[Literal["m s⁻²", "ft s⁻²"]]
+Density: TypeAlias = Quantity[Literal["kg m⁻³", "slug ft⁻³"]]
+GasConstant: TypeAlias = Quantity[Literal["J mol⁻¹ K⁻¹"]]
+MolarMass: TypeAlias = Quantity[Literal["kg mol⁻¹"]]
+SpecificGasConstant: TypeAlias = Quantity[Literal["J kg⁻¹ K⁻¹"]]
+ThrustSpecificFuelConsumption: TypeAlias = Quantity[
     Literal["kg s⁻¹ N⁻¹", "g s⁻¹ kN⁻¹", "lbm hr⁻¹ lbf⁻¹"]
 ]
 
@@ -82,11 +84,11 @@ def __truediv__(self, other: "_Quantity[Any]") -> "Div":
 
 
 class PressureAltitude(Length):
-    """Pressure altitude, as measured from altimeter"""
+    """Pressure altitude, as measured by altimeter"""
 
 
 class DensityAltitude(Length):
-    """Density altitude, as measured from altimeter"""
+    """Density altitude, as measured by altimeter"""
 
 
 class GeopotentialAltitude(Length):
@@ -159,7 +161,7 @@ class GS(Velocity):
 
 
 class WindSpeed(Velocity):
-    """Wind speed in the inertial reference frame"""
+    """Wind speed, inertial reference frame"""
 
 
 class SpeedOfSound(Velocity):
@@ -170,16 +172,16 @@ class GravitationalAcceleration(Acceleration):
     """Gravitational acceleration"""
 
 
-class TemperatureGradient(_Quantity[Literal["K m⁻¹"]]):
+class TemperatureGradient(Quantity[Literal["K m⁻¹"]]):
     """Lapse rate, below tropopause, ISA"""
 
 
-class MachNumber(_Quantity[None]):
+class MachNumber(Quantity[None]):
     """Mach number"""
 
 
-class AdiabaticIndex(_Quantity[None]):
-    """Ratio of specific heats, isentropic expansion factor"""
+class RatioOfSpecificHeats(Quantity[None]):
+    """Ratio of specific heats"""
 
 
 #
@@ -192,7 +194,7 @@ class AdiabaticIndex(_Quantity[None]):
 class Delta(Generic[Units]):
     """A difference between two quantities"""
 
-    quantity: _Quantity[Units]
+    quantity: Quantity[Units]
 
 
 @dataclass(frozen=True, slots=True)
@@ -203,5 +205,5 @@ class Div:
     Used in [unit conversion][airtrafficsim.unit_conversion].
     """
 
-    numerator: _Quantity[Any]
-    denominator: _Quantity[Any]
+    numerator: Quantity[Any]
+    denominator: Quantity[Any]

src/airtrafficsim/performance/airspeed.py

@@ -43,25 +43,25 @@
         CAS,
         EAS,
         TAS,
-        AdiabaticIndex,
         Density,
         ImpactPressure,
+        RatioOfSpecificHeats,
         StaticPressure,
     )
     from ..types import Array, ArrayOrScalarT
 
 
 def impact_pressure_from_cas(
     cas: Annotated[Array, CAS("m s⁻¹")],
-    gamma: Annotated[Array, AdiabaticIndex(None)] = GAMMA_DRY_AIR,
+    gamma: Annotated[Array, RatioOfSpecificHeats(None)] = GAMMA_DRY_AIR,
 ) -> Annotated[Array, ImpactPressure("Pa")]:
     """Impact pressure, compressible flow"""
     return impact_pressure(cas, RHO_0, P_0, gamma)
 
 
 def impact_pressure_from_cas_behind_normal_shock(
     cas: Annotated[Array, CAS("m s⁻¹")],
-    gamma: Annotated[Array, AdiabaticIndex(None)] = GAMMA_DRY_AIR,
+    gamma: Annotated[Array, RatioOfSpecificHeats(None)] = GAMMA_DRY_AIR,
 ) -> Annotated[Array, ImpactPressure("Pa")]:
     """Impact pressure, behind normal shock wave, supersonic flow"""
     return impact_pressure_behind_normal_shock(cas, A_0, P_0, gamma)
@@ -92,7 +92,7 @@ def tas_from_eas(
 def compressibility_factor(
     qc: Annotated[Array, ImpactPressure("Pa"), Gt(0)],
     p: Annotated[Array, StaticPressure("Pa")],
-    gamma: Annotated[Array, AdiabaticIndex(None)] = GAMMA_DRY_AIR,
+    gamma: Annotated[Array, RatioOfSpecificHeats(None)] = GAMMA_DRY_AIR,
 ) -> Array:
     """Assumption: subsonic speeds"""
     exponent = (gamma - 1) / gamma
@@ -103,7 +103,7 @@ def compressibility_factor(
 def eas_from_cas(
     cas: Annotated[Array, CAS("m s⁻¹"), Gt(0)],
     p: Annotated[Array, StaticPressure("Pa")],
-    gamma: Annotated[Array, AdiabaticIndex(None)] = GAMMA_DRY_AIR,
+    gamma: Annotated[Array, RatioOfSpecificHeats(None)] = GAMMA_DRY_AIR,
 ) -> Annotated[Array, EAS("m s⁻¹")]:
     """Assumption: subsonic speeds"""
     qc = impact_pressure_from_cas(cas)
@@ -126,7 +126,7 @@ def cas_from_tas(
     tas: Annotated[Array, TAS("m s⁻¹"), Gt(0)],
     rho: Annotated[Array, Density("kg m⁻³")],
     p: Annotated[Array, StaticPressure("Pa")],
-    gamma: Annotated[Array, AdiabaticIndex(None)] = GAMMA_DRY_AIR,
+    gamma: Annotated[Array, RatioOfSpecificHeats(None)] = GAMMA_DRY_AIR,
 ) -> Annotated[Array, CAS("m s⁻¹")]:
     """Assumption: subsonic speeds"""
     eas = eas_from_tas(tas, rho)

src/airtrafficsim/plot.py

@@ -0,0 +1,58 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Generic, overload
+
+from .annotations import Quantity, Units
+
+if TYPE_CHECKING:
+    from typing import Any
+
+    import polars as pl
+    from typing_extensions import Self
+
+
+def into_latex(unit: Units) -> str:
+    raise NotImplementedError
+
+
+@dataclass
+class Column(Generic[Units]):
+    quantity: Quantity[Units]
+    display_name: str | None = None
+    symbol: str | None = None
+    identifier: str | None = None
+    """
+    A unique identifier for retrieving the series in a dataframe (optional)
+    """
+
+    @classmethod
+    def from_quantity(cls, quantity: Quantity[Units]) -> Self:
+        if (doc := quantity.__doc__) is not None:
+            display_name = doc.split("\n")[0].split(",")[0]
+            print(display_name)
+        return cls(quantity)
+
+    @property
+    def label(self) -> str:
+        if self.display_name and self.symbol:
+            label = f"{self.display_name}, {self.symbol}"
+        elif self.display_name:
+            label = self.display_name
+        elif self.symbol:
+            label = f"{self.symbol}"
+        else:
+            raise ValueError("Either display_name or symbol must be provided.")
+        if self.quantity.unit is not None:  # hide for dimensionless
+            label += f" (${self.quantity.unit}$)"
+        return label
+
+    @overload
+    def __call__(self, df: pl.DataFrame) -> pl.Series: ...
+
+    @overload
+    def __call__(self, df: Any) -> Any: ...
+
+    def __call__(self, df: Any) -> Any:
+        """Returns series in the dataframe."""
+        return df[self.identifier]