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.
@@ -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)).