python-numerics

When to Load This Skill

Load this skill in addition to python-skill when the task touches:

• numpy / np.ndarray / tensor shapes / dtype discipline
• jaxtyping shape-annotated types
• pyarrow tables, schemas, Parquet, IPC, compute kernels
• polars dataframes (lazy or eager)
• pandas (legacy code, interop with external libraries that demand it)
• Numeric hot paths that need numpy / numba / orjson / compiled primitives
• Shape-dependent validation at system boundaries

python-skill still applies in full — this is a specialization, not a replacement. Every principle there (Result[T], pattern matching, no dynamic magic, strict typing) is load-bearing here too.

Version Policy

• Python 3.14+
• numpy >= 2.0 (uses np.bool_, PEP-585 generics, improved stubs)
• pyarrow >= 20
• polars >= 1.0
• jaxtyping >= 0.3
• beartype >= 0.20 (the runtime checker behind @jaxtyped)
• pandas-stubs >= 3.0 (if you're stuck on pandas)
• pandera >= 0.24 (if you're validating pandas at boundaries)

Philosophy

Strict types at the edges; numbers in the middle. Numeric code is dtype- and shape-sensitive, and both belong in the type signature. When a function takes a (batch, channels, h, w) tensor of float32, that contract should be written in the type and enforced at the boundary, not buried in a docstring or discovered mid-stacktrace.

• Shape + dtype contracts are non-negotiable. Use jaxtyping for shape-annotated numpy.ndarray / torch.Tensor / jax.Array, enforced at runtime boundaries with @jaxtyped(typechecker=beartype). A bare np.ndarray in a function signature is a contract hole. polars / pyarrow data is shape-checked by its schema, not jaxtyping.
• Schemas are the contract. pyarrow pa.Schema and polars pl.Schema are the source of truth for tabular data. Every load, every save, every boundary checks against the schema. Same discipline as Java record + migration history.
• Result[T] at every boundary. Shape mismatches, schema mismatches, malformed input, dtype casts — all return Result[T], not raised exceptions. Pattern-match the Err with a specific code.
• Vectorize; never loop for numerics. Python-level for loops over numpy / pyarrow / polars data destroy the point of those libraries. Every hot path uses compute kernels (pc.multiply, np.dot, pl.col(...).mean()).
• Big-O comments on non-trivial numeric functions — compute AND memory, derived through callees. Note when an op makes a copy vs. returns a view; memory complexity hides there.
  • Header comment, not body line. Any function that does real work gets a header comment above it: one line saying what it is for (matters most for public APIs) plus its big-O for compute and memory. The big-O lives in the header, not on a line inside the body. Skip this for tiny, obvious helpers (a one-line np.asarray cast).
  • Comment the code, not the work plan. Comments are for the person reading and reviewing: the public API, anything not obvious, and the reason behind the code as a plain sentence (no WHY: label). Keep out phases, plan steps, ticket numbers, and words like for now, temporary, or later — they go stale and someone has to find and delete them. Unfinished work belongs in plan.md.
  • Editing counts too. When you change a function, keep the comments already there, fix any your change made wrong, and add the missing header, big-O, and reason. Do this only for code you add or change, not for nearby code you did not touch.
• polars over pandas as the default. Polars has first-class static types, lazy evaluation, better performance, and a saner API. Pandas is a legacy concern — support it, but don't reach for it first.
• No np.random.seed() global state. Use np.random.default_rng(seed); pass the Generator through your API explicitly.

Shape + Dtype Contracts via jaxtyping

jaxtyping writes shape and dtype into the annotation. Be precise about what each layer checks:

• pyright sees Float[np.ndarray, "batch n"] as plain np.ndarray. Static checkers cannot do shape or dtype inference on jaxtyping annotations (jaxtyping FAQ). You still get array-vs-list-vs-Tensor errors at edit time — nothing shape-level.
• Runtime, with @jaxtyped(typechecker=beartype): full enforcement — dtype, rank, fixed sizes, and dim-name consistency across all arguments and the return value of each call.
• Always: machine-readable shape documentation that cannot drift silently once a runtime check covers it.

from beartype import beartype
from jaxtyping import Float, jaxtyped
import numpy as np

# BAD — shape is a runtime mystery; mismatches surface deep in the math
def normalize(x: np.ndarray) -> np.ndarray: ...

# GOOD — shape + dtype in the type, enforced at the boundary
@jaxtyped(typechecker=beartype)
def normalize(
    x: Float[np.ndarray, "batch channels h w"],
) -> Float[np.ndarray, "batch channels h w"]:
    mean: Float[np.ndarray, "1 channels 1 1"] = x.mean(axis=(0, 2, 3), keepdims=True)
    std: Float[np.ndarray, "1 channels 1 1"] = x.std(axis=(0, 2, 3), keepdims=True)
    return (x - mean) / (std + 1e-8)

Naming axes is load-bearing. Float[np.ndarray, "batch channels h w"] and Float[np.ndarray, "b c h w"] are both legal but the first documents intent. Within one @jaxtyped call, the same name must bind to the same size across all arguments and the return; across separate calls there is no relation — cross-function consistency is a convention you keep by reusing the same names codebase-wide.

Runtime checks at boundaries only — public entry points, post-load validation, tests. The checks are not free; keep them off hot inner loops.

See references/jaxtyping.md for the full cheatsheet: dtype table, variadic and symbolic shape syntax, @jaxtyped placement, beartype claw, hypothesis property testing, gotchas, and why polars / pyarrow data stays schema-checked instead.

numpy — Typed Discipline

dtype and shape in every signature

from jaxtyping import Float
import numpy as np

# np.asarray over np.array at boundaries — avoids copy if input already matches
def to_feature_vector(
    raw: list[float] | Float[np.ndarray, "n"],
) -> Float[np.ndarray, "n"]:
    return np.asarray(raw, dtype=np.float32)

# Validate shape at boundaries with Result[T]
def parse_features(
    payload: list[float],
    feature_dim: int,
) -> Result[Float[np.ndarray, "feat_dim"]]:
    if len(payload) != feature_dim:
        return err("SHAPE", f"expected {feature_dim} features, got {len(payload)}")
    arr: Float[np.ndarray, "feat_dim"] = np.asarray(payload, dtype=np.float32)
    return ok(arr)

Explicit axis, always

# BAD — axis implied, breaks silently on unexpected shapes
total = x.sum()

# GOOD — explicit axis
total: Float[np.ndarray, "n"] = x.sum(axis=-1)

Generator for randomness, never np.random.seed()

# BAD — global mutable state, not thread-safe, not typed
np.random.seed(42)
x = np.random.rand(100)

# GOOD — explicit generator, no global state, typed
rng: np.random.Generator = np.random.default_rng(seed=42)
x: Float[np.ndarray, "n"] = rng.random(100, dtype=np.float32)

Views vs. copies — intentional and documented

# A view — shares memory with the original
view: Float[np.ndarray, "n"] = x[::2]

# A copy — independent
copy: Float[np.ndarray, "n"] = x[::2].copy()

Document when you return a view so callers know mutations are shared. When in doubt, copy — views cause action-at-a-distance bugs that are painful to debug.

No np.errstate masking

# BAD — silently hides numerical errors, produces inf / nan
with np.errstate(divide='ignore', invalid='ignore'):
    result = x / y

# GOOD — validate inputs upstream, let errors surface as Err
if bool((y == 0).any()):
    return err("DIVIDE_BY_ZERO", "denominator contains zero")
result: Float[np.ndarray, "n"] = x / y

pyarrow — Schema-First Columnar Data

pyarrow is the memory format underlying polars, duckdb, and modern pandas. Schemas are the contract; every Table you construct is validated against its schema, every read from Parquet uses it.

Always declare schemas explicitly

import pyarrow as pa

# Schema as a module-level constant, like SQL DDL
ORDERS_SCHEMA: pa.Schema = pa.schema([
    pa.field("id", pa.int64(), nullable=False),
    pa.field("customer_id", pa.int64(), nullable=False),
    pa.field("amount", pa.decimal128(10, 2), nullable=False),
    pa.field("status", pa.dictionary(pa.int32(), pa.string()), nullable=False),
    pa.field("created_at", pa.timestamp("us", tz="UTC"), nullable=False),
    pa.field("notes", pa.string(), nullable=True),
])

Why nullable is first-class: Unlike numpy (where null = NaN = a lie), Arrow has a proper null bitmap per column. nullable=False is part of the type; pyarrow and polars both refuse to set null values in non-nullable columns.

Dictionary encoding for categoricals

# Low-cardinality string column → dictionary-encoded for memory + speed
pa.field("status", pa.dictionary(pa.int32(), pa.string()), nullable=False)

Memory savings for real-world categoricals are 10-100×. Same discipline as Java's enum for flat sets.

Parquet for on-disk, Arrow IPC for cross-process

import pyarrow.compute as pc
import pyarrow.parquet as pq

# Write with compression
pq.write_table(table, "orders.parquet", compression="zstd")

# Read with column pruning + row filtering (pushed down to Parquet reader).
# Expression filters are the typed, modern form — the legacy tuple-DNF
# list trips pyright.
filtered: pa.Table = pq.read_table(
    "orders.parquet",
    columns=["id", "amount"],
    filters=pc.field("status") == "CONFIRMED",
)

Parquet column pruning and predicate pushdown happen in C before the Python layer sees any data. If you're reading 2 columns out of 50, you're paying only for those 2.

pa.compute for vectorized ops — never Python loops

import pyarrow.compute as pc

# BAD — Python loop, thousands of Python-object roundtrips
totals: list[float] = [row["amount"].as_py() * 1.1 for row in table.to_pylist()]

# GOOD — vectorized compute kernel, zero Python overhead
with_tax: pa.Table = table.set_column(
    table.schema.get_field_index("amount"),
    "amount",
    pc.multiply(table["amount"], 1.1),
)

Zero-copy bridges to polars / duckdb / numpy

import polars as pl
import duckdb

# Arrow → polars (zero copy). from_arrow returns DataFrame | Series —
# assert to narrow for pyright.
df = pl.from_arrow(table)
assert isinstance(df, pl.DataFrame)

# Arrow → duckdb (zero copy, register as view)
con = duckdb.connect()
con.register("orders", table)
result: pa.Table = con.execute(
    "select status, sum(amount) from orders group by 1"
).arrow()

# Arrow → numpy (zero copy for contiguous numeric columns; decimal /
# string columns come back as object dtype — cast first for numerics)
ids: Int[np.ndarray, "n"] = table["id"].to_numpy()

Streaming for large data

import pyarrow.ipc as ipc

# memory_map gives zero-copy reads (and satisfies the stubs, which want
# a file object for open_stream, not a path)
with pa.memory_map(stream_path) as source, ipc.open_stream(source) as reader:
    for batch in reader:
        batch: pa.RecordBatch
        process(batch)   # one batch at a time — never load the whole stream

Validate at boundaries — schemas + Result[T]

from pathlib import Path

def load_orders(path: Path) -> Result[pa.Table]:
    try:
        table: pa.Table = pq.read_table(path)
    except Exception as e:
        return Err(Failure.from_exception(e, "READ_FAILED"))

    if not table.schema.equals(ORDERS_SCHEMA):
        return err(
            "SCHEMA_MISMATCH",
            f"expected {ORDERS_SCHEMA}, got {table.schema}",
        )
    return ok(table)

Polars — The Default Dataframe

Polars is the default dataframe library for new code. First-class static types, lazy query optimization, better performance than pandas, and a cleaner API.

Schema-first construction

import polars as pl

# Declare the schema explicitly
schema: pl.Schema = pl.Schema({
    "id": pl.Int64,
    "customer_id": pl.Int64,
    "amount": pl.Decimal(10, 2),
    "status": pl.Categorical(),   # pl.Schema needs instantiated dtypes
    "created_at": pl.Datetime(time_zone="UTC"),
})

df: pl.DataFrame = pl.DataFrame(rows, schema=schema)

Expressions, not string column references

# BAD — string column names, no pyright checking
df.filter(df["status"] == "CONFIRMED").select(["id", "amount"])

# GOOD — pl.col() expressions, composable, chainable
df.filter(pl.col("status") == "CONFIRMED").select(pl.col("id"), pl.col("amount"))

# BEST — lazy execution, query optimizer picks the plan
result: pl.DataFrame = (
    df.lazy()
    .filter(pl.col("status") == "CONFIRMED")
    .group_by("customer_id")
    .agg(pl.col("amount").sum().alias("total"))
    .sort("total", descending=True)
    .collect()
)

Lazy by default for large data

# Lazy — plan is built, no execution yet
q: pl.LazyFrame = pl.scan_parquet("orders/*.parquet")

# Query optimizer picks the best plan (predicate pushdown, projection pruning)
result: pl.DataFrame = (
    q.filter(pl.col("status") == "CONFIRMED")
     .group_by("customer_id")
     .agg(pl.col("amount").sum())
     .collect()
)

scan_parquet never reads the files until .collect(). The query optimizer pushes filters and column pruning down to the Parquet reader, so you touch only the rows and columns you actually need.

Boundary validation

Same Result[T] pattern as pyarrow's load_orders above: read, compare df.schema against the expected pl.Schema, err("SCHEMA_MISMATCH") on drift, ok(df) otherwise.

Pandas — Legacy Interop Only

Polars is the default; migrate pandas code when you can. When you must hold pandas (legacy code, an external library that demands pd.DataFrame): pandas-stubs for pyright, pandera models validating at boundaries, Arrow as the bridge to polars. pandas 3.0 is copy-on-write, permanently — chained assignment never lands (warns ChainedAssignmentError), single-step .loc writes are fine, defensive .copy() calls are dead weight.

See references/pandas-interop.md for the discipline: pandera DataFrameModel boundaries, pandas-3.0 CoW semantics, anti-patterns, and the polars-vs-pandas decision table.

Testing Numeric Code

TDD applies to numeric code unchanged: red, green, vacuous mutation check; the regression test lands before the fix. For property tests the vacuous check is concrete — mutate the function under test (drop the mean subtraction) and confirm hypothesis finds the counterexample.

Property-test array invariants with hypothesis.extra.numpy — constrain the element strategy (the default generates nan / inf) and pick tolerances that respect float conditioning. Full recipe in references/jaxtyping.md.

Shape and dtype assertions

def test_parse_features_shape() -> None:
    payload: list[float] = [1.0, 2.0, 3.0]
    match parse_features(payload, feature_dim=3):
        case Ok(arr):
            assert arr.shape == (3,)
            assert arr.dtype == np.float32
        case Err():
            raise AssertionError("expected Ok")

Schema roundtrip tests for pyarrow / polars

def test_orders_schema_roundtrip(tmp_path: Path) -> None:
    table: pa.Table = make_orders_table(rows=100)
    path: Path = tmp_path / "orders.parquet"
    pq.write_table(table, path)

    match load_orders(path):
        case Ok(loaded):
            assert loaded.schema.equals(ORDERS_SCHEMA)
            assert loaded.num_rows == 100
        case Err():
            raise AssertionError("expected Ok")

Tooling Notes

pyright stub packages. numpy and polars ship complete inline types. pyarrow ships py.typed but its bundled annotations are incomplete — add the community pyarrow-stubs dev dependency, which types the pyarrow.compute kernels (generated at runtime, otherwise invisible to pyright). Even then a few pyarrow.parquet / pyarrow.ipc signatures leak Unknown and trip strict reportUnknownMemberType; scope per-line ignores there rather than relaxing the project. For jaxtyping, no stubs are needed — pyright sees its annotations as the bare array class (enforcement is runtime, see above); under full strict the jaxtyped import itself is partially unknown (sentinel-typed default) and may need # pyright: ignore[reportUnknownVariableType].

ruff rules that fire on numeric code. Most of the discipline from python-skill's references/tooling.md applies here, plus:

• PD (pandas-vet rules) — if you're using pandas, enable them to catch common pitfalls (.ix, chained indexing, inplace=True)
• NPY (numpy-specific rules) — catches deprecated numpy APIs

Cross-References

• python-skill/SKILL.md — the core language discipline (types, Result, pattern matching, ADTs, protocols, tooling)
• python-skill/references/performance.md — general Python performance idioms, allocation model, orjson serialization
• python-skill/references/testing.md — general pytest / hypothesis patterns
• python-skill/references/tooling.md — base pyproject.toml with ruff / pyright / pytest config. Extend it with PD and NPY rule families for numeric projects.
• references/jaxtyping.md (this skill) — shape-typing cheatsheet
• references/pandas-interop.md (this skill) — legacy pandas discipline