unxts.linalg 2.0.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- unxts_linalg-2.0.0/.gitignore +166 -0
- unxts_linalg-2.0.0/PKG-INFO +50 -0
- unxts_linalg-2.0.0/README.md +20 -0
- unxts_linalg-2.0.0/docs/index.md +85 -0
- unxts_linalg-2.0.0/docs/linear-algebra.md +124 -0
- unxts_linalg-2.0.0/docs/quantity-matrix.md +132 -0
- unxts_linalg-2.0.0/docs/sharp-bits.md +60 -0
- unxts_linalg-2.0.0/docs/tutorial-metric.md +76 -0
- unxts_linalg-2.0.0/docs/units-matrix.md +72 -0
- unxts_linalg-2.0.0/pyproject.toml +90 -0
- unxts_linalg-2.0.0/src/unxts/linalg/__init__.py +42 -0
- unxts_linalg-2.0.0/src/unxts/linalg/_src/__init__.py +52 -0
- unxts_linalg-2.0.0/src/unxts/linalg/_src/_det.py +225 -0
- unxts_linalg-2.0.0/src/unxts/linalg/_src/_inv.py +212 -0
- unxts_linalg-2.0.0/src/unxts/linalg/_src/_products.py +148 -0
- unxts_linalg-2.0.0/src/unxts/linalg/_src/_quantity_matrix.py +545 -0
- unxts_linalg-2.0.0/src/unxts/linalg/_src/_register_primitives.py +1080 -0
- unxts_linalg-2.0.0/src/unxts/linalg/_src/_units_matrix.py +493 -0
- unxts_linalg-2.0.0/src/unxts/linalg/_src/_utils.py +37 -0
- unxts_linalg-2.0.0/src/unxts/linalg/_version.py +24 -0
- unxts_linalg-2.0.0/src/unxts/linalg/py.typed +0 -0
- unxts_linalg-2.0.0/tests/test_quantity_matrix.py +2522 -0
|
@@ -0,0 +1,166 @@
|
|
|
1
|
+
# -----------------------------------------------
|
|
2
|
+
# Project-Specific
|
|
3
|
+
|
|
4
|
+
# Version files
|
|
5
|
+
_version.py
|
|
6
|
+
|
|
7
|
+
|
|
8
|
+
# -----------------------------------------------
|
|
9
|
+
# General
|
|
10
|
+
|
|
11
|
+
# Byte-compiled / optimized / DLL files
|
|
12
|
+
__pycache__/
|
|
13
|
+
*.py[cod]
|
|
14
|
+
*$py.class
|
|
15
|
+
|
|
16
|
+
# C extensions
|
|
17
|
+
*.so
|
|
18
|
+
|
|
19
|
+
# Distribution / packaging
|
|
20
|
+
.Python
|
|
21
|
+
build/
|
|
22
|
+
develop-eggs/
|
|
23
|
+
dist/
|
|
24
|
+
downloads/
|
|
25
|
+
eggs/
|
|
26
|
+
.eggs/
|
|
27
|
+
lib/
|
|
28
|
+
lib64/
|
|
29
|
+
parts/
|
|
30
|
+
sdist/
|
|
31
|
+
var/
|
|
32
|
+
wheels/
|
|
33
|
+
share/python-wheels/
|
|
34
|
+
*.egg-info/
|
|
35
|
+
.installed.cfg
|
|
36
|
+
*.egg
|
|
37
|
+
MANIFEST
|
|
38
|
+
|
|
39
|
+
# PyInstaller
|
|
40
|
+
# Usually these files are written by a python script from a template
|
|
41
|
+
# before PyInstaller builds the exe, so as to inject date/other infos into it.
|
|
42
|
+
*.manifest
|
|
43
|
+
*.spec
|
|
44
|
+
|
|
45
|
+
# Installer logs
|
|
46
|
+
pip-log.txt
|
|
47
|
+
pip-delete-this-directory.txt
|
|
48
|
+
|
|
49
|
+
# Unit test / coverage reports
|
|
50
|
+
htmlcov/
|
|
51
|
+
.tox/
|
|
52
|
+
.nox/
|
|
53
|
+
.coverage
|
|
54
|
+
.coverage.*
|
|
55
|
+
.cache
|
|
56
|
+
nosetests.xml
|
|
57
|
+
coverage.xml
|
|
58
|
+
*.cover
|
|
59
|
+
*.py,cover
|
|
60
|
+
.hypothesis/
|
|
61
|
+
.pytest_cache/
|
|
62
|
+
cover/
|
|
63
|
+
|
|
64
|
+
# Translations
|
|
65
|
+
*.mo
|
|
66
|
+
*.pot
|
|
67
|
+
|
|
68
|
+
# Django stuff:
|
|
69
|
+
*.log
|
|
70
|
+
local_settings.py
|
|
71
|
+
db.sqlite3
|
|
72
|
+
db.sqlite3-journal
|
|
73
|
+
|
|
74
|
+
# Flask stuff:
|
|
75
|
+
instance/
|
|
76
|
+
.webassets-cache
|
|
77
|
+
|
|
78
|
+
# Scrapy stuff:
|
|
79
|
+
.scrapy
|
|
80
|
+
|
|
81
|
+
# Sphinx documentation
|
|
82
|
+
docs/_build/
|
|
83
|
+
docs/guides/perf.ipynb
|
|
84
|
+
|
|
85
|
+
# PyBuilder
|
|
86
|
+
.pybuilder/
|
|
87
|
+
target/
|
|
88
|
+
|
|
89
|
+
# Jupyter Notebook
|
|
90
|
+
.ipynb_checkpoints
|
|
91
|
+
|
|
92
|
+
# IPython
|
|
93
|
+
profile_default/
|
|
94
|
+
ipython_config.py
|
|
95
|
+
|
|
96
|
+
# pyenv
|
|
97
|
+
# For a library or package, you might want to ignore these files since the code is
|
|
98
|
+
# intended to run in multiple environments; otherwise, check them in:
|
|
99
|
+
# .python-version
|
|
100
|
+
|
|
101
|
+
# pipenv
|
|
102
|
+
# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
|
|
103
|
+
# However, in case of collaboration, if having platform-specific dependencies or dependencies
|
|
104
|
+
# having no cross-platform support, pipenv may install dependencies that don't work, or not
|
|
105
|
+
# install all needed dependencies.
|
|
106
|
+
#Pipfile.lock
|
|
107
|
+
|
|
108
|
+
# PEP 582; used by e.g. github.com/David-OConnor/pyflow
|
|
109
|
+
__pypackages__/
|
|
110
|
+
|
|
111
|
+
# Celery stuff
|
|
112
|
+
celerybeat-schedule
|
|
113
|
+
celerybeat.pid
|
|
114
|
+
|
|
115
|
+
# SageMath parsed files
|
|
116
|
+
*.sage.py
|
|
117
|
+
|
|
118
|
+
# Environments
|
|
119
|
+
.env
|
|
120
|
+
.venv
|
|
121
|
+
env/
|
|
122
|
+
venv/
|
|
123
|
+
ENV/
|
|
124
|
+
env.bak/
|
|
125
|
+
venv.bak/
|
|
126
|
+
|
|
127
|
+
# Spyder project settings
|
|
128
|
+
.spyderproject
|
|
129
|
+
.spyproject
|
|
130
|
+
|
|
131
|
+
# Rope project settings
|
|
132
|
+
.ropeproject
|
|
133
|
+
|
|
134
|
+
# mkdocs documentation
|
|
135
|
+
/site
|
|
136
|
+
|
|
137
|
+
# mypy
|
|
138
|
+
.mypy_cache/
|
|
139
|
+
.dmypy.json
|
|
140
|
+
dmypy.json
|
|
141
|
+
|
|
142
|
+
# Pyre type checker
|
|
143
|
+
.pyre/
|
|
144
|
+
|
|
145
|
+
# pytype static type analyzer
|
|
146
|
+
.pytype/
|
|
147
|
+
|
|
148
|
+
# Cython debug symbols
|
|
149
|
+
cython_debug/
|
|
150
|
+
|
|
151
|
+
# ruff
|
|
152
|
+
.ruff_cache/
|
|
153
|
+
|
|
154
|
+
# OS specific stuff
|
|
155
|
+
.DS_Store
|
|
156
|
+
.DS_Store?
|
|
157
|
+
._*
|
|
158
|
+
.Spotlight-V100
|
|
159
|
+
.Trashes
|
|
160
|
+
ehthumbs.db
|
|
161
|
+
Thumbs.db
|
|
162
|
+
|
|
163
|
+
# Common editor files
|
|
164
|
+
*~
|
|
165
|
+
*.swp
|
|
166
|
+
/.codspeed
|
|
@@ -0,0 +1,50 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: unxts.linalg
|
|
3
|
+
Version: 2.0.0
|
|
4
|
+
Summary: Heterogeneous-unit matrices and vectors for unxt (canonical: unxts.linalg)
|
|
5
|
+
Project-URL: Bug Tracker, https://github.com/GalacticDynamics/unxt/issues
|
|
6
|
+
Project-URL: Homepage, https://github.com/GalacticDynamics/unxt
|
|
7
|
+
Author-email: GalacticDynamics <nstarman@users.noreply.github.com>, Nathaniel Starkman <nstarman@users.noreply.github.com>
|
|
8
|
+
License: BSD-3-Clause
|
|
9
|
+
Classifier: Development Status :: 4 - Beta
|
|
10
|
+
Classifier: Intended Audience :: Developers
|
|
11
|
+
Classifier: Intended Audience :: Science/Research
|
|
12
|
+
Classifier: License :: OSI Approved :: BSD License
|
|
13
|
+
Classifier: Operating System :: OS Independent
|
|
14
|
+
Classifier: Programming Language :: Python
|
|
15
|
+
Classifier: Programming Language :: Python :: 3
|
|
16
|
+
Classifier: Programming Language :: Python :: 3 :: Only
|
|
17
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
18
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
19
|
+
Classifier: Programming Language :: Python :: 3.13
|
|
20
|
+
Classifier: Programming Language :: Python :: 3.14
|
|
21
|
+
Classifier: Topic :: Scientific/Engineering
|
|
22
|
+
Classifier: Typing :: Typed
|
|
23
|
+
Requires-Python: >=3.11
|
|
24
|
+
Requires-Dist: equinox>=0.11.8
|
|
25
|
+
Requires-Dist: jaxtyping>=0.2.34
|
|
26
|
+
Requires-Dist: plum-dispatch>=2.5.7
|
|
27
|
+
Requires-Dist: quax>=0.2.0
|
|
28
|
+
Requires-Dist: unxt
|
|
29
|
+
Description-Content-Type: text/markdown
|
|
30
|
+
|
|
31
|
+
# unxts.linalg
|
|
32
|
+
|
|
33
|
+
Heterogeneous-unit matrices and vectors for [unxt](https://github.com/GalacticDynamics/unxt).
|
|
34
|
+
|
|
35
|
+
This is the canonical package (`unxts.linalg`). It provides `QuantityMatrix` (alias `QM`): a quantity container whose elements may each carry a different unit, together with a static `UnitsMatrix` unit structure and unit-aware linear-algebra primitives (`det`, `inv`, the `matmul`/`matvec`/`vecmat`/`vecdot` products, and Quax-registered add/sub, mul/div, matmul, transpose, diag, and reduce-sum rules).
|
|
36
|
+
|
|
37
|
+
## Install
|
|
38
|
+
|
|
39
|
+
```bash
|
|
40
|
+
pip install unxts.linalg
|
|
41
|
+
```
|
|
42
|
+
|
|
43
|
+
## Usage
|
|
44
|
+
|
|
45
|
+
```python
|
|
46
|
+
import jax.numpy as jnp
|
|
47
|
+
import unxts.linalg as ul
|
|
48
|
+
|
|
49
|
+
qv = ul.QuantityMatrix(jnp.array([1.0, 2.0, 3.0]), unit=("m", "s", "kg"))
|
|
50
|
+
```
|
|
@@ -0,0 +1,20 @@
|
|
|
1
|
+
# unxts.linalg
|
|
2
|
+
|
|
3
|
+
Heterogeneous-unit matrices and vectors for [unxt](https://github.com/GalacticDynamics/unxt).
|
|
4
|
+
|
|
5
|
+
This is the canonical package (`unxts.linalg`). It provides `QuantityMatrix` (alias `QM`): a quantity container whose elements may each carry a different unit, together with a static `UnitsMatrix` unit structure and unit-aware linear-algebra primitives (`det`, `inv`, the `matmul`/`matvec`/`vecmat`/`vecdot` products, and Quax-registered add/sub, mul/div, matmul, transpose, diag, and reduce-sum rules).
|
|
6
|
+
|
|
7
|
+
## Install
|
|
8
|
+
|
|
9
|
+
```bash
|
|
10
|
+
pip install unxts.linalg
|
|
11
|
+
```
|
|
12
|
+
|
|
13
|
+
## Usage
|
|
14
|
+
|
|
15
|
+
```python
|
|
16
|
+
import jax.numpy as jnp
|
|
17
|
+
import unxts.linalg as ul
|
|
18
|
+
|
|
19
|
+
qv = ul.QuantityMatrix(jnp.array([1.0, 2.0, 3.0]), unit=("m", "s", "kg"))
|
|
20
|
+
```
|
|
@@ -0,0 +1,85 @@
|
|
|
1
|
+
# `unxts.linalg`
|
|
2
|
+
|
|
3
|
+
```{toctree}
|
|
4
|
+
:maxdepth: 1
|
|
5
|
+
:hidden:
|
|
6
|
+
|
|
7
|
+
quantity-matrix
|
|
8
|
+
units-matrix
|
|
9
|
+
linear-algebra
|
|
10
|
+
tutorial-metric
|
|
11
|
+
sharp-bits
|
|
12
|
+
```
|
|
13
|
+
|
|
14
|
+
`unxts.linalg` provides `QuantityMatrix` (alias `QM`): a quantity container whose elements may each carry a **different** unit. It is backed by a single JAX array plus a static `UnitsMatrix` describing the per-element units, and supports both 1-D (heterogeneous vector) and 2-D (heterogeneous matrix) structures.
|
|
15
|
+
|
|
16
|
+
It is useful for objects whose entries have mixed physical dimensions — Jacobians, metric tensors, and coordinate change-of-basis matrices — where a single scalar unit (as on `unxt.Quantity`) is not expressive enough.
|
|
17
|
+
|
|
18
|
+
## Install
|
|
19
|
+
|
|
20
|
+
::::{tab-set}
|
|
21
|
+
|
|
22
|
+
:::{tab-item} uv
|
|
23
|
+
|
|
24
|
+
```bash
|
|
25
|
+
uv add unxts.linalg
|
|
26
|
+
```
|
|
27
|
+
|
|
28
|
+
:::
|
|
29
|
+
|
|
30
|
+
:::{tab-item} pip
|
|
31
|
+
|
|
32
|
+
```bash
|
|
33
|
+
pip install unxts.linalg
|
|
34
|
+
```
|
|
35
|
+
|
|
36
|
+
:::
|
|
37
|
+
|
|
38
|
+
::::
|
|
39
|
+
|
|
40
|
+
Throughout these guides we import `unxt` as `u` and `unxts.linalg` as `ul` (so `QuantityMatrix` is `ul.QM`):
|
|
41
|
+
|
|
42
|
+
```{code-block} python
|
|
43
|
+
>>> import jax.numpy as jnp
|
|
44
|
+
>>> import unxt as u
|
|
45
|
+
>>> import unxts.linalg as ul
|
|
46
|
+
```
|
|
47
|
+
|
|
48
|
+
## At a glance
|
|
49
|
+
|
|
50
|
+
A 1-D `QuantityMatrix` is a vector whose entries each have their own unit:
|
|
51
|
+
|
|
52
|
+
```{code-block} python
|
|
53
|
+
>>> qv = ul.QuantityMatrix(jnp.array([1.0, 2.0, 3.0]), unit=("m", "s", "kg"))
|
|
54
|
+
>>> qv.unit.to_string()
|
|
55
|
+
'(m, s, kg)'
|
|
56
|
+
>>> 2 * qv
|
|
57
|
+
QuantityMatrix(Array([2., 4., 6.], dtype=float32), unit='(m, s, kg)')
|
|
58
|
+
```
|
|
59
|
+
|
|
60
|
+
Indexing a single element yields an ordinary `unxt.Quantity`:
|
|
61
|
+
|
|
62
|
+
```{code-block} python
|
|
63
|
+
>>> qv[0]
|
|
64
|
+
Quantity(Array(1., dtype=float32), unit='m')
|
|
65
|
+
```
|
|
66
|
+
|
|
67
|
+
## Guides
|
|
68
|
+
|
|
69
|
+
- [Quantity matrices](quantity-matrix.md) — constructing `QuantityMatrix`, indexing, unit conversion, and arithmetic.
|
|
70
|
+
- [Units matrices](units-matrix.md) — the immutable, hashable `UnitsMatrix` unit structure.
|
|
71
|
+
- [Linear algebra](linear-algebra.md) — matmul, transpose, `diag`, `det`, and `inv` with per-element unit tracking.
|
|
72
|
+
- [Tutorial: a heterogeneous metric](tutorial-metric.md) — a worked end-to-end example.
|
|
73
|
+
- [Sharp bits](sharp-bits.md) — the 1-D/2-D restriction and the uniform-unit requirements under `jax.jit`.
|
|
74
|
+
|
|
75
|
+
## Public API
|
|
76
|
+
|
|
77
|
+
`unxts.linalg` exposes:
|
|
78
|
+
|
|
79
|
+
- `QuantityMatrix` — the heterogeneous-unit matrix/vector (alias `QM`).
|
|
80
|
+
- `UnitsMatrix` — the immutable, hashable per-element unit structure.
|
|
81
|
+
- `matmul`, `matvec`, `vecmat`, `vecdot` — the four Array-API matrix/vector products, each resolving the batched matrix-vs-vector ambiguity that `@` alone cannot.
|
|
82
|
+
- `det`, `inv` — unit-tracking determinant and inverse (with their JAX primitives `det_p`, `inv_p`).
|
|
83
|
+
- `cdict_units` — extract per-key units from a component dictionary.
|
|
84
|
+
|
|
85
|
+
Importing `unxts.linalg` also registers, as import side effects, the Quax primitive rules (add/sub, mul/div, dot-general/matmul, transpose, gather/diag, reduce-sum) and the `plum` conversions/dispatch that let `QuantityMatrix` interoperate with the rest of `unxt`.
|
|
@@ -0,0 +1,124 @@
|
|
|
1
|
+
# Linear algebra with unit tracking
|
|
2
|
+
|
|
3
|
+
The whole point of `QuantityMatrix` is that linear-algebra operations carry the per-element units through the computation. `unxts.linalg` registers Quax rules for the underlying JAX primitives, so the [`quaxed`](https://github.com/GalacticDynamics/quaxed) drop-in `numpy` operators work directly on `QuantityMatrix` objects.
|
|
4
|
+
|
|
5
|
+
```{code-block} python
|
|
6
|
+
>>> import jax.numpy as jnp
|
|
7
|
+
>>> import quax
|
|
8
|
+
>>> import unxt as u
|
|
9
|
+
>>> import unxts.linalg as ul
|
|
10
|
+
```
|
|
11
|
+
|
|
12
|
+
## Matrix and vector products
|
|
13
|
+
|
|
14
|
+
`unxts.linalg` exposes the four NumPy / Array-API contraction functions, each of which contracts the shared axis and multiplies the corresponding per-element units:
|
|
15
|
+
|
|
16
|
+
| Function | Operands | Result |
|
|
17
|
+
| ----------------- | --------------- | ----------------- |
|
|
18
|
+
| `ul.matmul(A, B)` | matrix × matrix | matrix |
|
|
19
|
+
| `ul.matvec(A, v)` | matrix × vector | vector |
|
|
20
|
+
| `ul.vecmat(v, A)` | vector × matrix | vector |
|
|
21
|
+
| `ul.vecdot(a, b)` | vector · vector | scalar `Quantity` |
|
|
22
|
+
|
|
23
|
+
`matmul` contracts the shared axis and multiplies the corresponding units — here a dimensionless identity leaves a vector's units untouched:
|
|
24
|
+
|
|
25
|
+
```{code-block} python
|
|
26
|
+
>>> A = ul.QuantityMatrix(jnp.eye(3),
|
|
27
|
+
... unit=(("", "", ""), ("", "", ""), ("", "", "")))
|
|
28
|
+
>>> v = ul.QuantityMatrix(jnp.array([1.0, 2.0, 3.0]), unit=("m", "m", "m"))
|
|
29
|
+
>>> ul.matvec(A, v).value
|
|
30
|
+
Array([1., 2., 3.], dtype=float32)
|
|
31
|
+
```
|
|
32
|
+
|
|
33
|
+
Units on the contracted axis are combined and converted to a common reference, so mixed units are handled correctly:
|
|
34
|
+
|
|
35
|
+
```{code-block} python
|
|
36
|
+
>>> A2 = ul.QuantityMatrix(jnp.array([[1.0, 2.0], [3.0, 4.0]]),
|
|
37
|
+
... unit=(("m", "km"), ("m", "km")))
|
|
38
|
+
>>> v2 = ul.QuantityMatrix(jnp.array([1.0, 1.0]), unit=("s", "s"))
|
|
39
|
+
>>> ul.matvec(A2, v2).value
|
|
40
|
+
Array([2001., 4003.], dtype=float32)
|
|
41
|
+
```
|
|
42
|
+
|
|
43
|
+
A plain JAX array or an ordinary `unxt.Quantity` on one side is treated as dimensionless / uniform-unit respectively.
|
|
44
|
+
|
|
45
|
+
### Batches
|
|
46
|
+
|
|
47
|
+
A `QuantityMatrix` carries its logical 1-D/2-D units in the _trailing_ axes and treats _leading_ axes of the value array as batch dimensions. All four products — and the `@` operator — broadcast those batch axes. The `@` operator dispatches on the _logical_ rank (matrix vs vector), so a batch of matrices applied to a batch of vectors "just works":
|
|
48
|
+
|
|
49
|
+
```{code-block} python
|
|
50
|
+
>>> # A batch of 2 matrices applied to a batch of 2 vectors.
|
|
51
|
+
>>> Ab = ul.QuantityMatrix(
|
|
52
|
+
... jnp.array([[[1.0, 2.0], [3.0, 4.0]], [[5.0, 6.0], [7.0, 8.0]]]),
|
|
53
|
+
... unit=(("m", "m"), ("m", "m")),
|
|
54
|
+
... )
|
|
55
|
+
>>> vb = ul.QuantityMatrix(jnp.ones((2, 2)), unit=("s", "s"))
|
|
56
|
+
>>> (Ab @ vb).value
|
|
57
|
+
Array([[ 3., 7.],
|
|
58
|
+
[11., 15.]], dtype=float32)
|
|
59
|
+
```
|
|
60
|
+
|
|
61
|
+
There is one subtlety inherited from NumPy: a batched vector's value `(B, K)` is shape-indistinguishable from a matrix, so the raw **function** forms `jnp.matmul` / `quaxed.numpy.matmul` cannot express a batched matrix-vector product (and `ul.matmul` rejects it with a pointer to `matvec`). Use `@`, `ul.matvec`, or `ul.vecmat` — all of which name the ranks explicitly.
|
|
62
|
+
|
|
63
|
+
## Transpose and diagonal
|
|
64
|
+
|
|
65
|
+
`.T` swaps both the values and the unit structure of a 2-D matrix:
|
|
66
|
+
|
|
67
|
+
```{code-block} python
|
|
68
|
+
>>> a = ul.QuantityMatrix(jnp.array([[1.0, 2.0], [3.0, 4.0]]),
|
|
69
|
+
... unit=(("m", "s"), ("kg", "rad")))
|
|
70
|
+
>>> a.T.unit.to_string()
|
|
71
|
+
'((m, kg), (s, rad))'
|
|
72
|
+
```
|
|
73
|
+
|
|
74
|
+
`.diag()` extracts the diagonal as a 1-D `QuantityMatrix`, operating directly on the static unit structure so it works under `jax.jit`:
|
|
75
|
+
|
|
76
|
+
```{code-block} python
|
|
77
|
+
>>> M = ul.QuantityMatrix(jnp.diag(jnp.array([1.0, 2.0, 3.0])),
|
|
78
|
+
... unit=(("m", "s", "kg"),
|
|
79
|
+
... ("m", "s", "kg"),
|
|
80
|
+
... ("m", "s", "kg")))
|
|
81
|
+
>>> d = M.diag()
|
|
82
|
+
>>> d.unit.to_string()
|
|
83
|
+
'(m, s, kg)'
|
|
84
|
+
>>> d.value
|
|
85
|
+
Array([1., 2., 3.], dtype=float32)
|
|
86
|
+
```
|
|
87
|
+
|
|
88
|
+
## Determinant and inverse
|
|
89
|
+
|
|
90
|
+
`unxts.linalg` provides custom `det` and `inv` primitives with full JAX transform support (JIT, autodiff, vmap). On plain arrays they behave like `jnp.linalg.det` / `jnp.linalg.inv`; on a `QuantityMatrix` (wrapped with `quax.quaxify`) they additionally track units.
|
|
91
|
+
|
|
92
|
+
The determinant multiplies the main-diagonal units:
|
|
93
|
+
|
|
94
|
+
```{code-block} python
|
|
95
|
+
>>> from unxts.linalg import det, inv
|
|
96
|
+
>>> G = ul.QuantityMatrix(jnp.array([[2.0, 0.0], [0.0, 3.0]]),
|
|
97
|
+
... unit=(("m2", "m2"), ("m2", "m2")))
|
|
98
|
+
>>> quax.quaxify(det)(G)
|
|
99
|
+
Quantity(Array(6., dtype=float32), unit='m4')
|
|
100
|
+
```
|
|
101
|
+
|
|
102
|
+
The inverse carries the reciprocal units:
|
|
103
|
+
|
|
104
|
+
```{code-block} python
|
|
105
|
+
>>> B = ul.QuantityMatrix(jnp.array([[4.0, 0.0], [0.0, 1.0]]),
|
|
106
|
+
... unit=(("m2", "m2"), ("m2", "m2")))
|
|
107
|
+
>>> r = quax.quaxify(inv)(B)
|
|
108
|
+
>>> r.unit.to_string()
|
|
109
|
+
'((1 / m2, 1 / m2), (1 / m2, 1 / m2))'
|
|
110
|
+
>>> r.value
|
|
111
|
+
Array([[0.25, 0. ],
|
|
112
|
+
[0. , 1. ]], dtype=float32)
|
|
113
|
+
```
|
|
114
|
+
|
|
115
|
+
Because `det` and `inv` are real JAX primitives, they compose with `jax.jit`, `jax.grad`, and `jax.vmap` on plain arrays:
|
|
116
|
+
|
|
117
|
+
```{code-block} python
|
|
118
|
+
>>> import jax
|
|
119
|
+
>>> jax.jit(det)(jnp.array([[2.0, 0.0], [0.0, 3.0]]))
|
|
120
|
+
Array(6., dtype=float32)
|
|
121
|
+
>>> jax.grad(det)(jnp.array([[2.0, 0.0], [0.0, 3.0]]))
|
|
122
|
+
Array([[3., 0.],
|
|
123
|
+
[0., 2.]], dtype=float32)
|
|
124
|
+
```
|
|
@@ -0,0 +1,132 @@
|
|
|
1
|
+
# Quantity matrices
|
|
2
|
+
|
|
3
|
+
`QuantityMatrix` (alias `QM`) stores one numeric JAX array together with a static [`UnitsMatrix`](units-matrix.md) that gives the unit of **each** logical element. The shape of the unit structure decides whether the object behaves as a heterogeneous vector (1-D) or matrix (2-D).
|
|
4
|
+
|
|
5
|
+
```{code-block} python
|
|
6
|
+
>>> import jax.numpy as jnp
|
|
7
|
+
>>> import unxt as u
|
|
8
|
+
>>> import unxts.linalg as ul
|
|
9
|
+
```
|
|
10
|
+
|
|
11
|
+
## Construction
|
|
12
|
+
|
|
13
|
+
A 1-D `QuantityMatrix` takes a flat array and a tuple of units, one per element:
|
|
14
|
+
|
|
15
|
+
```{code-block} python
|
|
16
|
+
>>> qv = ul.QuantityMatrix(jnp.array([1.0, 2.0, 3.0]), unit=("m", "s", "kg"))
|
|
17
|
+
>>> qv.value
|
|
18
|
+
Array([1., 2., 3.], dtype=float32)
|
|
19
|
+
>>> qv.unit.to_string()
|
|
20
|
+
'(m, s, kg)'
|
|
21
|
+
>>> qv.ndim, qv.shape
|
|
22
|
+
(1, (3,))
|
|
23
|
+
```
|
|
24
|
+
|
|
25
|
+
A 2-D `QuantityMatrix` takes a 2-D array and a nested tuple of units:
|
|
26
|
+
|
|
27
|
+
```{code-block} python
|
|
28
|
+
>>> qm = ul.QuantityMatrix(jnp.ones((2, 2)), unit=(("m", "s"), ("kg", "rad")))
|
|
29
|
+
>>> qm.unit.to_string()
|
|
30
|
+
'((m, s), (kg, rad))'
|
|
31
|
+
>>> qm.ndim, qm.shape
|
|
32
|
+
(2, (2, 2))
|
|
33
|
+
```
|
|
34
|
+
|
|
35
|
+
Leading dimensions of the value array are treated as batch dimensions; only the trailing 1 or 2 axes are "logical" and must match the unit structure.
|
|
36
|
+
|
|
37
|
+
`QM` is a short alias for `QuantityMatrix` (mirroring `u.Q` for `Quantity`):
|
|
38
|
+
|
|
39
|
+
```{code-block} python
|
|
40
|
+
>>> ul.QM is ul.QuantityMatrix
|
|
41
|
+
True
|
|
42
|
+
```
|
|
43
|
+
|
|
44
|
+
### From a component dictionary
|
|
45
|
+
|
|
46
|
+
`QuantityMatrix.from_cdict` packs a dict of quantities into a 1-D matrix, preserving each entry's unit:
|
|
47
|
+
|
|
48
|
+
```{code-block} python
|
|
49
|
+
>>> v = {"x": u.Q(1.0, "m"), "y": u.Q(2.0, "s"), "z": u.Q(3.0, "kg")}
|
|
50
|
+
>>> ul.QuantityMatrix.from_cdict(v).unit.to_string()
|
|
51
|
+
'(m, s, kg)'
|
|
52
|
+
```
|
|
53
|
+
|
|
54
|
+
You may select and reorder a subset of keys:
|
|
55
|
+
|
|
56
|
+
```{code-block} python
|
|
57
|
+
>>> ul.QuantityMatrix.from_cdict(v, keys=("z", "x")).unit.to_string()
|
|
58
|
+
'(kg, m)'
|
|
59
|
+
```
|
|
60
|
+
|
|
61
|
+
## Indexing
|
|
62
|
+
|
|
63
|
+
Indexing a 1-D `QuantityMatrix` returns an ordinary `unxt.Quantity` — the element and its scalar unit:
|
|
64
|
+
|
|
65
|
+
```{code-block} python
|
|
66
|
+
>>> qv[0]
|
|
67
|
+
Quantity(Array(1., dtype=float32), unit='m')
|
|
68
|
+
>>> qv[2]
|
|
69
|
+
Quantity(Array(3., dtype=float32), unit='kg')
|
|
70
|
+
```
|
|
71
|
+
|
|
72
|
+
Indexing a row of a 2-D matrix returns a 1-D `QuantityMatrix`, while a full `[i, j]` index returns a scalar `Quantity`:
|
|
73
|
+
|
|
74
|
+
```{code-block} python
|
|
75
|
+
>>> qm2 = ul.QuantityMatrix(jnp.ones((2, 3)),
|
|
76
|
+
... unit=(("m", "s", "kg"), ("rad", "deg", "m")))
|
|
77
|
+
>>> qm2[0]
|
|
78
|
+
QuantityMatrix(Array([1., 1., 1.], dtype=float32), unit='(m, s, kg)')
|
|
79
|
+
>>> qm2[1, 2]
|
|
80
|
+
Quantity(Array(1., dtype=float32), unit='m')
|
|
81
|
+
```
|
|
82
|
+
|
|
83
|
+
## Arithmetic and unit conversion
|
|
84
|
+
|
|
85
|
+
Addition and subtraction convert each element of the right operand into the corresponding unit of the left operand before combining; the result adopts the left operand's units:
|
|
86
|
+
|
|
87
|
+
```{code-block} python
|
|
88
|
+
>>> import quaxed.numpy as qnp
|
|
89
|
+
>>> a = ul.QuantityMatrix(jnp.ones(3), unit=("m", "s", "kg"))
|
|
90
|
+
>>> b = ul.QuantityMatrix(jnp.ones(3), unit=("km", "ms", "g"))
|
|
91
|
+
>>> result = qnp.add(a, b)
|
|
92
|
+
>>> result.unit.to_string()
|
|
93
|
+
'(m, s, kg)'
|
|
94
|
+
>>> result.value
|
|
95
|
+
Array([1001. , 1.001, 1.001], dtype=float32)
|
|
96
|
+
```
|
|
97
|
+
|
|
98
|
+
Scalar multiplication scales the values and leaves the units unchanged:
|
|
99
|
+
|
|
100
|
+
```{code-block} python
|
|
101
|
+
>>> (2 * a).value
|
|
102
|
+
Array([2., 2., 2.], dtype=float32)
|
|
103
|
+
```
|
|
104
|
+
|
|
105
|
+
You can convert a whole `QuantityMatrix` to a compatible unit structure with `uconvert`:
|
|
106
|
+
|
|
107
|
+
```{code-block} python
|
|
108
|
+
>>> q = ul.QuantityMatrix(jnp.array([[1.0, 2.0], [3.0, 4.0]]),
|
|
109
|
+
... unit=(("m", "rad"), ("m", "rad")))
|
|
110
|
+
>>> target = u.unit((("km", "deg"), ("km", "deg")))
|
|
111
|
+
>>> q.uconvert(target).unit.to_string()
|
|
112
|
+
'((km, deg), (km, deg))'
|
|
113
|
+
```
|
|
114
|
+
|
|
115
|
+
When every element shares the same unit, a `QuantityMatrix` converts to a plain `unxt.Quantity`:
|
|
116
|
+
|
|
117
|
+
```{code-block} python
|
|
118
|
+
>>> import plum
|
|
119
|
+
>>> uniform = ul.QuantityMatrix(jnp.array([1.0, 2.0, 3.0]), unit=("m", "m", "m"))
|
|
120
|
+
>>> plum.convert(uniform, u.Q)
|
|
121
|
+
Quantity(Array([1., 2., 3.], dtype=float32), unit='m')
|
|
122
|
+
```
|
|
123
|
+
|
|
124
|
+
Mixed units make that conversion ambiguous, so it is rejected:
|
|
125
|
+
|
|
126
|
+
```{code-block} python
|
|
127
|
+
>>> try:
|
|
128
|
+
... plum.convert(qv, u.Q)
|
|
129
|
+
... except ValueError as e:
|
|
130
|
+
... print(e)
|
|
131
|
+
Cannot convert QuantityMatrix to Quantity unless all units are identical.
|
|
132
|
+
```
|
|
@@ -0,0 +1,60 @@
|
|
|
1
|
+
# Sharp bits
|
|
2
|
+
|
|
3
|
+
`QuantityMatrix` covers the common heterogeneous-unit vector/matrix cases, but a few restrictions follow from its design. Keep these in mind.
|
|
4
|
+
|
|
5
|
+
```{code-block} python
|
|
6
|
+
>>> import jax.numpy as jnp
|
|
7
|
+
>>> import quaxed.numpy as qnp
|
|
8
|
+
>>> import unxt as u
|
|
9
|
+
>>> import unxts.linalg as ul
|
|
10
|
+
```
|
|
11
|
+
|
|
12
|
+
## Only 1-D and 2-D structures
|
|
13
|
+
|
|
14
|
+
The unit structure is limited to 1-D (vector) and 2-D (matrix); there is no support for higher-rank _logical_ structures. (Leading **batch** dimensions on the value array are fine — only the trailing one or two axes are logical.)
|
|
15
|
+
|
|
16
|
+
```{code-block} python
|
|
17
|
+
>>> from unxts.linalg import UnitsMatrix
|
|
18
|
+
>>> try:
|
|
19
|
+
... UnitsMatrix(jnp.zeros((2, 2, 2)))
|
|
20
|
+
... except (TypeError, ValueError) as e:
|
|
21
|
+
... print("rejected")
|
|
22
|
+
rejected
|
|
23
|
+
```
|
|
24
|
+
|
|
25
|
+
## `det` and `inv` and their unit assumptions
|
|
26
|
+
|
|
27
|
+
`det` uses the product of the **main-diagonal** units, and `inv` requires a **uniform** unit that it can reciprocate — it raises `ValueError` on a heterogeneous-unit matrix, because a matrix inverse mixes entries and the per-element reciprocal would be wrong. These are exactly right for diagonal metrics and for matrices whose cofactor products share one physical dimension (the common case for coordinate metrics), but they are not general heterogeneous-unit determinants/inverses. Both require a 2-D matrix:
|
|
28
|
+
|
|
29
|
+
```{code-block} python
|
|
30
|
+
>>> import quax
|
|
31
|
+
>>> v = ul.QuantityMatrix(jnp.array([1.0, 2.0]), unit=("m", "s"))
|
|
32
|
+
>>> try:
|
|
33
|
+
... quax.quaxify(ul.det)(v)
|
|
34
|
+
... except ValueError as e:
|
|
35
|
+
... print("needs a 2-D matrix")
|
|
36
|
+
needs a 2-D matrix
|
|
37
|
+
```
|
|
38
|
+
|
|
39
|
+
## `diag` under `jax.jit` needs uniform units
|
|
40
|
+
|
|
41
|
+
The `.diag()` **method** operates on the static unit structure and works for heterogeneous units, even under `jit`:
|
|
42
|
+
|
|
43
|
+
```{code-block} python
|
|
44
|
+
>>> M = ul.QuantityMatrix(jnp.diag(jnp.array([1.0, 2.0, 3.0])),
|
|
45
|
+
... unit=(("m", "s", "kg"),
|
|
46
|
+
... ("m", "s", "kg"),
|
|
47
|
+
... ("m", "s", "kg")))
|
|
48
|
+
>>> M.diag().unit.to_string()
|
|
49
|
+
'(m, s, kg)'
|
|
50
|
+
```
|
|
51
|
+
|
|
52
|
+
By contrast `qnp.diag` lowers to a `gather`, whose indices are traced under `jit`; there the unit of each output element cannot be resolved individually, so **all units must be equal**. Prefer the `.diag()` method for heterogeneous-unit matrices.
|
|
53
|
+
|
|
54
|
+
## The `matmul` _function_ can't do a batched matrix-vector product
|
|
55
|
+
|
|
56
|
+
A batched vector's value `(B, K)` is shape-indistinguishable from a matrix. The **`@` operator is fine** — `QuantityMatrix.__matmul__` dispatches on the logical rank, so `A @ v` correctly does a (batched) matrix-vector product. But the raw **function** forms `jnp.matmul` / `quaxed.numpy.matmul` infer their contraction from the value shapes and so cannot; they silently succeed only when the sizes coincide and otherwise raise. `unxts.linalg.matmul` refuses a batched vector operand outright, pointing you at `matvec`. Prefer `@`, `ul.matvec`, or `ul.vecmat`. See [Linear algebra](linear-algebra.md#batches).
|
|
57
|
+
|
|
58
|
+
## It is a Quax type, not a materialisable array
|
|
59
|
+
|
|
60
|
+
`QuantityMatrix` is a `quax` array-ish value: it flows through `quax.quaxify`-ed functions but refuses to _materialise_ into a plain array (its elements have no single dtype-plus-unit), so use `.value` / `.unit` to inspect it, and `plum.convert(..., u.Q)` only when every unit is identical (see [Quantity matrices](quantity-matrix.md)).
|