dimtensor 1.3.0

Unit-aware tensors for physics and scientific machine learning

dimtensor

Unit-aware tensors for physics and scientific machine learning.

dimtensor wraps NumPy, PyTorch, and JAX arrays with physical unit tracking, catching dimensional errors at operation time rather than after hours of computation.

Installation

pip install dimtensor

For PyTorch or JAX support:

pip install dimtensor[torch]  # PyTorch integration
pip install dimtensor[jax]    # JAX integration
pip install dimtensor[all]    # All optional dependencies

Quick Start

from dimtensor import DimArray, units

# Create dimension-aware arrays
velocity = DimArray([10, 20, 30], units.m / units.s)
time = DimArray([1, 2, 3], units.s)

# Operations preserve/check dimensions
distance = velocity * time
print(distance)  # [10 40 90] m

# Catch errors early
acceleration = DimArray([9.8], units.m / units.s**2)
velocity + acceleration  # DimensionError: cannot add m/s to m/s^2

Features

• Dimensional Safety: Operations between incompatible dimensions raise DimensionError immediately
• Unit Conversion: Convert between compatible units with .to()
• SI Units: Full support for SI base and derived units
• NumPy Integration: Works with NumPy arrays and supports common operations
• PyTorch Integration: DimTensor wraps PyTorch tensors with autograd support
• JAX Integration: DimArray registered as JAX pytree for jit/vmap/grad
• Physical Constants: CODATA constants with proper units and uncertainties
• Uncertainty Propagation: Track and propagate measurement uncertainties
• I/O Support: Save/load to JSON, HDF5, Parquet, NetCDF, pandas, xarray
• Visualization: Matplotlib and Plotly integration with auto unit labels
• Domain Units: Astronomy, chemistry, and engineering units

PyTorch Integration

import torch
from dimtensor.torch import DimTensor
from dimtensor import units

# Unit-aware tensors with autograd
v = DimTensor(torch.tensor([1.0, 2.0, 3.0], requires_grad=True), units.m / units.s)
t = DimTensor(torch.tensor([0.5, 1.0, 1.5]), units.s)
d = v * t  # distance in meters

# Gradients work
loss = d.sum()
loss.backward()
print(v.grad)  # Gradients flow through

# GPU support
v_cuda = v.cuda()  # Move to GPU, preserving units

JAX Integration

import jax
import jax.numpy as jnp
from dimtensor.jax import DimArray
from dimtensor import units

# Unit-aware arrays compatible with JAX transformations
@jax.jit
def kinetic_energy(mass, velocity):
    return 0.5 * mass * velocity**2

m = DimArray(jnp.array([1.0, 2.0]), units.kg)
v = DimArray(jnp.array([10.0, 20.0]), units.m / units.s)
E = kinetic_energy(m, v)  # JIT-compiled, units preserved
print(E)  # [50. 400.] J

Physical Constants

from dimtensor import constants

# CODATA physical constants with proper units
print(constants.c)           # Speed of light: 299792458.0 m/s
print(constants.h)           # Planck constant with uncertainty
print(constants.G)           # Gravitational constant

# Use in calculations
E = constants.c**2 * DimArray([1.0], units.kg)  # E = mc^2

Uncertainty Propagation

from dimtensor import DimArray, units

# Measurement with uncertainty
length = DimArray([10.0], units.m, uncertainty=[0.1])
width = DimArray([5.0], units.m, uncertainty=[0.05])

# Uncertainties propagate through operations
area = length * width
print(area)  # 50 +/- 0.71 m^2 (propagated in quadrature)

I/O Support

JSON

from dimtensor.io import save_json, load_json

arr = DimArray([1.0, 2.0, 3.0], units.m)
save_json(arr, "data.json")
loaded = load_json("data.json")  # Units preserved

HDF5

from dimtensor.io import save_hdf5, load_hdf5

arr = DimArray([1.0, 2.0, 3.0], units.m)
save_hdf5(arr, "data.h5", compression="gzip")
loaded = load_hdf5("data.h5")

Pandas

from dimtensor.io import to_dataframe, from_dataframe

data = {
    "position": DimArray([1, 2, 3], units.m),
    "velocity": DimArray([10, 20, 30], units.m / units.s)
}
df = to_dataframe(data)  # Unit metadata in attrs
arrays = from_dataframe(df)  # Reconstruct DimArrays

NetCDF

from dimtensor.io import save_netcdf, load_netcdf

arr = DimArray([1.0, 2.0, 3.0], units.m)
save_netcdf(arr, "data.nc")
loaded = load_netcdf("data.nc")

Parquet

from dimtensor.io import save_parquet, load_parquet

arr = DimArray([1.0, 2.0, 3.0], units.m)
save_parquet(arr, "data.parquet", compression="snappy")
loaded = load_parquet("data.parquet")

xarray

from dimtensor.io import to_xarray, from_xarray

arr = DimArray([1.0, 2.0, 3.0], units.m)
da = to_xarray(arr, name="distance", dims=("x",))
restored = from_xarray(da)  # Back to DimArray with units

Visualization

dimtensor integrates with matplotlib and plotly for automatic axis labeling.

Matplotlib

from dimtensor import DimArray, units
from dimtensor.visualization import setup_matplotlib, plot

# Option 1: Enable automatic unit labels globally
setup_matplotlib()
import matplotlib.pyplot as plt

time = DimArray([0, 1, 2, 3], units.s)
distance = DimArray([0, 10, 40, 90], units.m)
plt.plot(time, distance)  # Axes labeled automatically: [s], [m]

# Option 2: Use wrapper functions
plot(time, distance)  # Same result

# Convert units during plotting
plot(time, distance, y_unit=units.km)  # y-axis in kilometers

# Error bars from uncertainty
measurement = DimArray([10, 20, 30], units.m, uncertainty=[0.5, 0.5, 0.5])
from dimtensor.visualization import errorbar
errorbar(time[:3], measurement)  # Auto-extracts uncertainty

Plotly

from dimtensor import DimArray, units
from dimtensor.visualization import plotly

time = DimArray([0, 1, 2, 3], units.s)
distance = DimArray([0, 10, 40, 90], units.m)

# Line plot with automatic labels
fig = plotly.line(time, distance, title="Motion")
fig.show()

# Scatter with error bars (auto-extracts uncertainty)
measurement = DimArray([10, 20, 30], units.m, uncertainty=[0.5, 0.5, 0.5])
fig = plotly.scatter_with_errors(time[:3], measurement)

Domain-Specific Units

Astronomy

from dimtensor import DimArray
from dimtensor.domains.astronomy import parsec, AU, solar_mass, light_year

# Distance to Proxima Centauri
distance = DimArray([4.24], light_year)
distance_pc = distance.to(parsec)  # ~1.3 pc

# Stellar mass
mass = DimArray([1.0], solar_mass)  # 1 solar mass

Chemistry

from dimtensor import DimArray
from dimtensor.domains.chemistry import molar, dalton, ppm

# Solution concentration
concentration = DimArray([0.1], molar)  # 0.1 M

# Atomic mass
carbon_mass = DimArray([12.011], dalton)  # Carbon atomic weight

# Trace amounts
contamination = DimArray([50], ppm)  # 50 parts per million

Engineering

from dimtensor import DimArray
from dimtensor.domains.engineering import MPa, hp, BTU, kWh

# Material stress
yield_strength = DimArray([250], MPa)

# Engine power
power = DimArray([100], hp)  # 100 horsepower

# Energy consumption
electricity = DimArray([500], kWh)  # 500 kilowatt-hours

Examples

Kinematics

from dimtensor import DimArray, units

v = DimArray([10], units.m / units.s)  # velocity
t = DimArray([5], units.s)              # time
d = v * t                               # distance = 50 m

Force and Energy

m = DimArray([2], units.kg)              # mass
g = DimArray([9.8], units.m / units.s**2)  # gravity
h = DimArray([10], units.m)              # height

# Potential energy
PE = m * g * h  # 196 J

# Force
F = m * g  # 19.6 N

Unit Conversion

distance = DimArray([1], units.km)
in_meters = distance.to(units.m)  # 1000 m
in_miles = distance.to(units.mile)  # ~0.621 mi

Why dimtensor?

1. Catch errors early: Don't waste compute on dimensionally invalid calculations
2. Self-documenting code: Units make physics code clearer
3. ML-ready: PyTorch autograd and JAX transformations fully supported
4. Scientific computing: Physical constants and uncertainty propagation
5. Lightweight: Just metadata tracking, minimal overhead

License

MIT