Unit-aware tensors for physics and scientific machine learning
Project description
dimtensor
Unit-aware tensors for physics and scientific machine learning
dimtensor catches dimensional errors at operation time, not after hours of computation.
from dimtensor import DimArray, units
# Operations check dimensions automatically
velocity = DimArray([10, 20, 30], units.m / units.s)
time = DimArray([1, 2, 3], units.s)
distance = velocity * time # [10 40 90] m
# Errors caught immediately
acceleration = DimArray([9.8], units.m / units.s**2)
velocity + acceleration # DimensionError: cannot add m/s to m/s^2
Why dimtensor?
| Problem | Solution |
|---|---|
| Silent unit errors waste compute time | Immediate DimensionError at operation time |
| PyTorch has no unit support | Native DimTensor with full autograd and GPU support |
| JAX incompatible with unit libraries | DimArray registered as pytree for jit/vmap/grad |
| Uncertainty handled separately | Built-in uncertainty propagation through all operations |
| Units lost during I/O | Save/load with units to JSON, HDF5, Parquet, NetCDF |
Features
- Dimensional Safety - Operations between incompatible dimensions raise
DimensionError - Unit Conversion - Convert between compatible units with
.to() - NumPy/PyTorch/JAX - Full integration with all three frameworks
- Physical Constants - CODATA 2022 constants with proper units and uncertainties
- Uncertainty Propagation - Track and propagate measurement uncertainties
- I/O Support - JSON, HDF5, Parquet, NetCDF, pandas, xarray
- Visualization - Matplotlib and Plotly with automatic unit labels
- Domain Units - Astronomy, chemistry, and engineering units
- Dimensional Inference - Infer dimensions from variable names and equations
- Dimensional Linting - Static analysis CLI for finding unit errors
- Optional Rust Backend - Accelerated operations when built from source
Installation
pip install dimtensor
For framework-specific support:
pip install dimtensor[torch] # PyTorch integration
pip install dimtensor[jax] # JAX integration
pip install dimtensor[all] # All optional dependencies
Optional: Rust Backend (v2.0+)
For improved performance, build the optional Rust backend:
# Install Rust (if not already installed)
curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh
pip install maturin
# Build and install the Rust extension
cd /path/to/dimtensor/rust
maturin build --release
pip install target/wheels/dimtensor_core-*.whl
The library automatically uses the Rust backend when available, with a pure Python fallback otherwise. Check availability:
from dimtensor._rust import HAS_RUST_BACKEND
print(f"Rust backend: {HAS_RUST_BACKEND}")
Quick Start
NumPy
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
PyTorch
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 flow through
d.sum().backward()
print(v.grad)
# GPU support
v_cuda = v.cuda()
JAX
import jax
import jax.numpy as jnp
from dimtensor.jax import DimArray
from dimtensor import units
@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: [50. 400.] J
Physical Constants
from dimtensor import constants, DimArray, units
print(constants.c) # Speed of light: 299792458.0 m/s
print(constants.h) # Planck constant with uncertainty
E = constants.c**2 * DimArray([1.0], units.kg) # E = mc^2
Uncertainty Propagation
from dimtensor import DimArray, units
length = DimArray([10.0], units.m, uncertainty=[0.1])
width = DimArray([5.0], units.m, uncertainty=[0.05])
area = length * width # 50 +/- 0.71 m^2 (propagated in quadrature)
I/O
from dimtensor import DimArray, units
from dimtensor.io import save_json, load_json, save_hdf5, load_hdf5
arr = DimArray([1.0, 2.0, 3.0], units.m)
# JSON
save_json(arr, "data.json")
loaded = load_json("data.json") # Units preserved
# HDF5
save_hdf5(arr, "data.h5", compression="gzip")
loaded = load_hdf5("data.h5")
Visualization
from dimtensor import DimArray, units
from dimtensor.visualization import plot
time = DimArray([0, 1, 2, 3], units.s)
distance = DimArray([0, 10, 40, 90], units.m)
plot(time, distance) # Axes labeled automatically: [s], [m]
Domain-Specific Units
from dimtensor import DimArray
from dimtensor.domains.astronomy import parsec, light_year, solar_mass
from dimtensor.domains.chemistry import molar, dalton
from dimtensor.domains.engineering import MPa, hp, kWh
# Astronomy
distance = DimArray([4.24], light_year).to(parsec) # ~1.3 pc
# Chemistry
concentration = DimArray([0.1], molar) # 0.1 M
# Engineering
stress = DimArray([250], MPa)
power = DimArray([100], hp)
Dataset Loaders (v3.3.0+)
from dimtensor.loaders import NistCodataLoader, NasaExoplanetLoader, PrismClimateLoader
# Load NIST CODATA 2022 fundamental constants
loader = NistCodataLoader()
constants_df = loader.load() # Returns DimArrays with units and uncertainties
print(constants_df['speed_of_light']) # 299792458.0 m/s (exact)
# Load NASA confirmed exoplanets
exo_loader = NasaExoplanetLoader()
exoplanets = exo_loader.load() # Orbital parameters with proper units
print(exoplanets['orbital_period']) # DimArray with unit: days
# Load PRISM climate data for a specific location
climate_loader = PrismClimateLoader(lat=40.7128, lon=-74.0060) # NYC
data = climate_loader.load(start_date='2020-01-01', end_date='2020-12-31')
print(data['temperature']) # DimArray with unit: degC
print(data['precipitation']) # DimArray with unit: mm
# All loaders support caching at ~/.dimtensor/cache/
Dimensional Inference (v2.0+)
from dimtensor.inference import infer_dimension, get_equations_by_domain
# Infer dimension from variable name
result = infer_dimension("velocity")
print(result.dimension) # L·T⁻¹ (length/time)
print(result.confidence) # 0.9
# Works with prefixes and suffixes
result = infer_dimension("initial_velocity_x")
print(result.dimension) # L·T⁻¹
# Query physics equations
mechanics = get_equations_by_domain("mechanics")
for eq in mechanics[:3]:
print(f"{eq.name}: {eq.formula}")
# Newton's Second Law: F = ma
# Kinetic Energy: KE = ½mv²
# Gravitational Force: F = Gm₁m₂/r²
Automatic Unit Inference (v3.3.0+)
from dimtensor.inference import infer_units
from dimtensor import units
# Infer unknown units from equations
equations = [
"F = m * a", # Newton's second law
"E = F * d" # Work equation
]
knowns = {
"m": units.kg,
"a": units.m / units.s**2,
"d": units.m
}
unknowns = ["F", "E"]
result = infer_units(equations, knowns, unknowns)
print(result["F"]) # kg·m·s⁻² (newton)
print(result["E"]) # kg·m²·s⁻² (joule)
# Detect dimensional inconsistencies
equations_bad = ["v = a + t"] # Invalid: cannot add acceleration to time
result = infer_units(equations_bad, {"a": units.m/units.s**2, "t": units.s}, ["v"])
# Raises DimensionError
Dimensional Linting (v2.1+)
# Lint a file for dimensional issues
dimtensor lint physics_simulation.py
# Example output:
# physics.py:15:4: W002 Potential dimension mismatch: LT⁻¹ + LT⁻²
# velocity + acceleration
# Suggestion: Cannot add/subtract LT⁻¹ and LT⁻². Check your units.
# JSON output for IDE integration
dimtensor lint --format=json src/
# Strict mode (report all inferences)
dimtensor lint --strict script.py
Useful Links
Call for Contributions
dimtensor is an open source project and welcomes contributions of all kinds. Here are ways to get involved:
- Report bugs - Open an issue
- Request features - Share ideas in discussions
- Contribute code - See our contributing guide
- Improve docs - Fix typos, add examples, clarify explanations
- Share use cases - Write tutorials or blog posts
Writing code isn't the only way to contribute. Good issues, documentation improvements, and community engagement are just as valuable.
License
MIT
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
Built Distribution
Filter files by name, interpreter, ABI, and platform.
If you're not sure about the file name format, learn more about wheel file names.
Copy a direct link to the current filters
File details
Details for the file dimtensor-3.3.0.tar.gz.
File metadata
- Download URL: dimtensor-3.3.0.tar.gz
- Upload date:
- Size: 116.5 kB
- Tags: Source
- Uploaded using Trusted Publishing? No
- Uploaded via: twine/6.2.0 CPython/3.11.11
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
0766ad3df223d6e6d4d8c0d9f5757e8b83b60082c71eac5cd7ffdbd1a80ea7c2
|
|
| MD5 |
fc208e912330d7e38f96810ec7bde965
|
|
| BLAKE2b-256 |
3d9705fc492611264dc9d133f9602bcece1744a70481b85bed1a80aaaea2ab3f
|
File details
Details for the file dimtensor-3.3.0-py3-none-any.whl.
File metadata
- Download URL: dimtensor-3.3.0-py3-none-any.whl
- Upload date:
- Size: 152.2 kB
- Tags: Python 3
- Uploaded using Trusted Publishing? No
- Uploaded via: twine/6.2.0 CPython/3.11.11
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
737753cff856c99270aafe02e384b5a025c80391acf6fa2b80f9c88641a265ca
|
|
| MD5 |
1c2a7c39c3ca9a159e332a7170fb1b58
|
|
| BLAKE2b-256 |
4082eadcee5b937147c0ebad70e105d12cd92d9b9667ebb48f3910c10980ff5b
|