From f7e1bde90e68c74f20de40674352acb0ecbbb8ec Mon Sep 17 00:00:00 2001
From: arthurmccray <amccray@stanford.edu>
Date: Thu, 30 Jul 2026 12:24:24 -0700
Subject: [PATCH 1/5] adding to_polars method for Vector and making polars an
 optional dependency not included in dev

---
 pyproject.toml                            |   3 +
 src/quantem/core/datastructures/vector.py |  84 +++++++++++++++++-
 tests/datastructures/test_vector.py       |  82 +++++++++++++++++
 uv.lock                                   | 102 ++++++++++++++--------
 4 files changed, 235 insertions(+), 36 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 4df9fc22..8673bb6e 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -60,6 +60,9 @@ dependencies = [
 widgets = [
     "quantem.widget",
 ]
+dataframe = [
+    "polars>=1.0",
+]
 
 [tool.hatch.build.targets.sdist]
 # hatchling always includes:
diff --git a/src/quantem/core/datastructures/vector.py b/src/quantem/core/datastructures/vector.py
index aa2627d2..5c4ee967 100644
--- a/src/quantem/core/datastructures/vector.py
+++ b/src/quantem/core/datastructures/vector.py
@@ -2,7 +2,7 @@
 
 import copy
 from pathlib import Path
-from typing import Any, Literal, Sequence
+from typing import TYPE_CHECKING, Any, Literal, Sequence
 
 import numpy as np
 from numpy.typing import NDArray
@@ -15,6 +15,9 @@
     validate_vector_units,
 )
 
+if TYPE_CHECKING:
+    import polars as pl
+
 
 class Vector(AutoSerialize):
     """Ragged cell data on a fixed grid.
@@ -701,6 +704,73 @@ def __abs__(self) -> "Vector":
     # I/O
     # ------------------------------------------------------------------ #
 
+    def to_polars(self, dim_names: Sequence[str] | None = None) -> "pl.DataFrame":
+        """Export the current selection to a polars DataFrame.
+
+        Every ragged row becomes one DataFrame row. The fixed-grid location of
+        that row is carried in leading integer columns -- one per fixed-grid
+        dimension -- followed by one column per selected field.
+
+        Grid coordinates are always reported in the *root* grid, not the
+        selection's local coordinates, so a view such as ``v[1, :]`` still
+        reports ``dim_0 == 1``. Both fixed-grid indexing and ``select_fields``
+        are honored, so ``v[1].select_fields("kx").to_polars()`` returns only
+        the ``kx`` rows belonging to cell 1.
+
+        Requires polars, which is an optional dependency:
+        ``pip install "quantem[dataframe]"``.
+
+        Parameters
+        ----------
+        dim_names : sequence of str, optional
+            Names for the fixed-grid index columns. Defaults to
+            ``("dim_0", ..., "dim_{ndim-1}")``. Must have one entry per
+            fixed-grid dimension.
+
+        Returns
+        -------
+        polars.DataFrame
+            Shape ``(total_rows, grid_ndim + num_fields)``.
+
+        Examples
+        --------
+        >>> v = Vector.from_shape((3, 2), fields=("kx", "ky"))
+        >>> v.to_polars().columns
+        ['dim_0', 'dim_1', 'kx', 'ky']
+        """
+        try:
+            import polars as pl
+        except ImportError as exc:  # pragma: no cover - depends on environment
+            raise ImportError(
+                "Vector.to_polars() requires polars, which is an optional dependency. "
+                'Install it with: pip install polars   (or: pip install "quantem[dataframe]")'
+            ) from exc
+
+        root_shape = self._state["shape"]
+        index_names = _resolve_dim_names(dim_names, len(root_shape))
+
+        collisions = [name for name in index_names if name in self.fields]
+        if collisions:
+            raise ValueError(
+                f"Fixed-grid column name(s) {collisions} collide with field name(s). "
+                "Pass dim_names=... to to_polars() to rename the index columns."
+            )
+
+        columns: dict[str, Any] = {}
+        if index_names:
+            counts = np.asarray(self.row_counts(), dtype=np.int64)
+            coords = np.unravel_index(self._selected_cell_indices(), root_shape)
+            for name, axis_coords in zip(index_names, coords):
+                columns[name] = pl.Series(
+                    name, np.repeat(axis_coords, counts).astype(np.int64, copy=False)
+                )
+
+        values = self.flatten()
+        for column, field in enumerate(self.fields):
+            columns[field] = pl.Series(field, values[:, column])
+
+        return pl.DataFrame(columns)
+
     def save(
         self,
         path: str | Path,
@@ -1028,6 +1098,18 @@ def _normalize_field_names(field_names: str | Sequence[str]) -> tuple[str, ...]:
     return normalized
 
 
+def _resolve_dim_names(dim_names: Sequence[str] | None, ndim: int) -> list[str]:
+    """Resolve fixed-grid index column names for DataFrame export."""
+    if dim_names is None:
+        return [f"dim_{i}" for i in range(ndim)]
+    resolved = [str(name) for name in dim_names]
+    if len(resolved) != ndim:
+        raise ValueError(f"Expected {ndim} dim_names, got {len(resolved)}")
+    if len(set(resolved)) != len(resolved):
+        raise ValueError("Duplicate dim_names are not allowed.")
+    return resolved
+
+
 def _normalize_units(units: str | Sequence[str] | None, count: int) -> list[str]:
     """Normalize field units to a list matching ``count``."""
     if units is None:
diff --git a/tests/datastructures/test_vector.py b/tests/datastructures/test_vector.py
index 954059d6..fbbad491 100644
--- a/tests/datastructures/test_vector.py
+++ b/tests/datastructures/test_vector.py
@@ -415,6 +415,88 @@ def test_from_data_supports_nested_fixed_grid(self):
         with pytest.raises(ValueError, match="same number of fields"):
             Vector.from_data(data=[np.array([[1.0, 2.0]]), np.array([[1.0, 2.0, 3.0]])])
 
+    def test_to_polars_line_vector(self):
+        pytest.importorskip("polars")
+        v = make_line_vector()
+
+        df = v.to_polars()
+        assert df.columns == ["dim_0", "intensity", "kx", "ky"]
+        assert df.height == v.total_rows == 6
+        assert df["dim_0"].to_list() == [0, 0, 1, 2, 2, 3]
+        np.testing.assert_array_equal(
+            df.select(v.fields).to_numpy(),
+            v.flatten(),
+        )
+
+    def test_to_polars_grid_and_dim_names(self):
+        pytest.importorskip("polars")
+        v = make_grid_vector()
+
+        df = v.to_polars()
+        assert df.columns == ["dim_0", "dim_1", "intensity", "kx", "ky"]
+        assert df.height == 6
+        assert list(zip(df["dim_0"], df["dim_1"])) == [
+            (0, 0),
+            (0, 1),
+            (1, 0),
+            (1, 1),
+            (2, 0),
+            (2, 1),
+        ]
+
+        renamed = v.to_polars(dim_names=("rx", "ry"))
+        assert renamed.columns == ["rx", "ry", "intensity", "kx", "ky"]
+
+    def test_to_polars_respects_current_selection(self):
+        pytest.importorskip("polars")
+        v = make_grid_vector()
+
+        # Field selection narrows the value columns, grid columns are kept.
+        field_view = v.select_fields("kx").to_polars()
+        assert field_view.columns == ["dim_0", "dim_1", "kx"]
+
+        # Fixed-grid selection reports *root* grid coordinates, not local ones.
+        cell_view = v[2].to_polars()
+        assert cell_view.height == 2
+        assert cell_view["dim_0"].to_list() == [2, 2]
+        assert cell_view["dim_1"].to_list() == [0, 1]
+
+        scalar_view = v[1, 1].select_fields("ky").to_polars()
+        assert scalar_view.columns == ["dim_0", "dim_1", "ky"]
+        assert scalar_view.to_dicts() == [{"dim_0": 1, "dim_1": 1, "ky": 211.0}]
+
+    def test_to_polars_handles_empty_cells_and_zero_dim_grid(self):
+        pytest.importorskip("polars")
+
+        sparse = Vector.from_shape(shape=(3,), fields=["a"])
+        sparse[0] = np.array([[1.0], [2.0]])
+        sparse[2] = np.array([[9.0]])
+        df = sparse.to_polars()
+        assert df["dim_0"].to_list() == [0, 0, 2]
+        assert df["a"].to_list() == [1.0, 2.0, 9.0]
+
+        empty = Vector.from_shape(shape=(2,), fields=["a", "b"]).to_polars()
+        assert empty.height == 0
+        assert empty.columns == ["dim_0", "a", "b"]
+
+        # A 0D fixed grid has no grid axes, so no index columns are emitted.
+        scalar_grid = Vector.from_shape(shape=(), fields=["a", "b"])
+        scalar_grid[...] = np.array([[1.0, 2.0], [3.0, 4.0]])
+        assert scalar_grid.to_polars().columns == ["a", "b"]
+
+    def test_to_polars_rejects_colliding_and_mismatched_names(self):
+        pytest.importorskip("polars")
+
+        with pytest.raises(ValueError, match="collide with field name"):
+            Vector.from_shape(shape=(2,), fields=["dim_0", "b"]).to_polars()
+
+        v = make_grid_vector()
+        with pytest.raises(ValueError, match="Expected 2 dim_names, got 1"):
+            v.to_polars(dim_names=("only_one",))
+
+        with pytest.raises(ValueError, match="Duplicate dim_names"):
+            v.to_polars(dim_names=("same", "same"))
+
     def test_save_and_load_round_trip(self, tmp_path):
         v = make_grid_vector()
         v.add_fields("extra", v.select_fields("intensity") + 1.0)
diff --git a/uv.lock b/uv.lock
index 36a6293e..2fe204ff 100644
--- a/uv.lock
+++ b/uv.lock
@@ -673,43 +673,43 @@ wheels = [
 
 [package.optional-dependencies]
 cublas = [
-    { name = "nvidia-cublas", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" },
-    { name = "nvidia-cuda-nvrtc", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" },
+    { name = "nvidia-cublas", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" },
+    { name = "nvidia-cuda-nvrtc", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" },
 ]
 cudart = [
-    { name = "nvidia-cuda-runtime", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" },
+    { name = "nvidia-cuda-runtime", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" },
 ]
 cufft = [
-    { name = "nvidia-cufft", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" },
-    { name = "nvidia-nvjitlink", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" },
+    { name = "nvidia-cufft", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" },
+    { name = "nvidia-nvjitlink", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" },
 ]
 cufile = [
-    { name = "nvidia-cufile", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" },
+    { name = "nvidia-cufile", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
 ]
 cupti = [
-    { name = "nvidia-cuda-cupti", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" },
+    { name = "nvidia-cuda-cupti", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" },
 ]
 curand = [
-    { name = "nvidia-curand", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" },
+    { name = "nvidia-curand", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" },
 ]
 cusolver = [
-    { name = "nvidia-cublas", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" },
-    { name = "nvidia-cusolver", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" },
-    { name = "nvidia-cusparse", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" },
-    { name = "nvidia-nvjitlink", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" },
+    { name = "nvidia-cublas", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" },
+    { name = "nvidia-cusolver", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" },
+    { name = "nvidia-cusparse", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" },
+    { name = "nvidia-nvjitlink", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" },
 ]
 cusparse = [
-    { name = "nvidia-cusparse", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" },
-    { name = "nvidia-nvjitlink", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" },
+    { name = "nvidia-cusparse", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" },
+    { name = "nvidia-nvjitlink", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" },
 ]
 nvjitlink = [
-    { name = "nvidia-nvjitlink", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" },
+    { name = "nvidia-nvjitlink", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" },
 ]
 nvrtc = [
-    { name = "nvidia-cuda-nvrtc", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" },
+    { name = "nvidia-cuda-nvrtc", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" },
 ]
 nvtx = [
-    { name = "nvidia-nvtx", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" },
+    { name = "nvidia-nvtx", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" },
 ]
 
 [[package]]
@@ -1224,7 +1224,7 @@ name = "importlib-metadata"
 version = "9.0.0"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "zipp" },
+    { name = "zipp", marker = "python_full_version < '3.12'" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/a9/01/15bb152d77b21318514a96f43af312635eb2500c96b55398d020c93d86ea/importlib_metadata-9.0.0.tar.gz", hash = "sha256:a4f57ab599e6a2e3016d7595cfd72eb4661a5106e787a95bcc90c7105b831efc", size = 56405, upload-time = "2026-03-20T06:42:56.999Z" }
 wheels = [
@@ -2540,6 +2540,34 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl", hash = "sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746", size = 20538, upload-time = "2025-05-15T12:30:06.134Z" },
 ]
 
+[[package]]
+name = "polars"
+version = "1.43.1"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "polars-runtime-32" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/32/79/720f4901230992f359653717e7cc3596731ed36e1a27be351d128ee0c3b7/polars-1.43.1.tar.gz", hash = "sha256:cb07ff3ad61c7b28043e6176e5fdb04a294346920b7f033deeb84116b4911883", size = 750058, upload-time = "2026-07-27T12:07:58.288Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/f3/7c/d74a56d5afae8e91924aad91cba621b286a05307e88569ecab666b3055b3/polars-1.43.1-py3-none-any.whl", hash = "sha256:f6ecd9184956f46442ddfcf6185423401475bbeffea59e238de8b9ecedacf16c", size = 846844, upload-time = "2026-07-27T12:06:33.913Z" },
+]
+
+[[package]]
+name = "polars-runtime-32"
+version = "1.43.1"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/60/44/390c9e9eef393991d907b7067264ddaa685b550a5bab94991765459a5e64/polars_runtime_32-1.43.1.tar.gz", hash = "sha256:2931c71fce2080ade2fc743207b3d70ea659f694e0273b6bacfe551ad6ce43e0", size = 3093618, upload-time = "2026-07-27T12:07:59.492Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/72/a5/72c075ff95b31807c3cf497757bdd87f81b5b8869e60a22ec532238bac99/polars_runtime_32-1.43.1-cp310-abi3-macosx_10_12_x86_64.whl", hash = "sha256:ddc7bca81e3616b74eba597bff09acbe3567e5b66798cc5305adc2e91a36e31e", size = 53084414, upload-time = "2026-07-27T12:06:36.283Z" },
+    { url = "https://files.pythonhosted.org/packages/b3/8f/a34347323459116becbc3b6223bc55b25a1be705d6064decac6b7aa0c769/polars_runtime_32-1.43.1-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:3bdaeb8b017be11d7d2f8a938ae98cb4135f90de451771c2bac4d21dca2f7cf4", size = 47514709, upload-time = "2026-07-27T12:06:39.63Z" },
+    { url = "https://files.pythonhosted.org/packages/83/7e/a0a22740388facc22f2faa81787b4150231a22dad42ad02b2c6efdba0d0a/polars_runtime_32-1.43.1-cp310-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:def1c19a339903ab39f1d134a69a9e8b02feb4f800f547809fbfed458ccb135e", size = 51351657, upload-time = "2026-07-27T12:06:44.701Z" },
+    { url = "https://files.pythonhosted.org/packages/34/f1/4a07318711eeb3a27c62c916751ca45f18df6b0891d18aba68ffd9c18a76/polars_runtime_32-1.43.1-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:4f323e5c4aa0f068c2c911bd01f1ecdc25c1f3a501879b29550a41ab95490965", size = 57289250, upload-time = "2026-07-27T12:06:48.173Z" },
+    { url = "https://files.pythonhosted.org/packages/26/74/c6b55dbb4db2574a9478bf1c2aa04624ea924fc44dc14b64050be462db7e/polars_runtime_32-1.43.1-cp310-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:b93d4bf59dcb0bac68ab2458890f1e8431d17369e59efc88625f0d92b22cad92", size = 51505657, upload-time = "2026-07-27T12:06:51.557Z" },
+    { url = "https://files.pythonhosted.org/packages/43/2e/12e987b0f311f20e41f904acc452824af3009579d60c878a5435ae1e9f0a/polars_runtime_32-1.43.1-cp310-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:977b4620d837f3ca0d131f858ac1bd72e2ae10a5e85a6cc27fddc71fbdb5005d", size = 55189721, upload-time = "2026-07-27T12:06:54.68Z" },
+    { url = "https://files.pythonhosted.org/packages/3a/f6/5bfba6f40a08b2b1bbe0507e9cc638ebada2dceaece45f9d87a905c57de5/polars_runtime_32-1.43.1-cp310-abi3-win_amd64.whl", hash = "sha256:fa557938e9113c12d59c56df8d7f7e1a411cc1954df0dafbb05900136c6329e7", size = 52566672, upload-time = "2026-07-27T12:06:57.725Z" },
+    { url = "https://files.pythonhosted.org/packages/f9/5e/41df901e2684e8857bdffd458c8157d237c701df9f3f922d101075f97724/polars_runtime_32-1.43.1-cp310-abi3-win_arm64.whl", hash = "sha256:cc186bad9f33b71f66ee6cc3ba2146d64315a120b0e988814d6fb1a37f0cc9e9", size = 46575356, upload-time = "2026-07-27T12:07:00.856Z" },
+]
+
 [[package]]
 name = "pre-commit"
 version = "4.6.1"
@@ -2940,6 +2968,9 @@ dependencies = [
 ]
 
 [package.optional-dependencies]
+dataframe = [
+    { name = "polars" },
+]
 widgets = [
     { name = "quantem-widget" },
 ]
@@ -2969,6 +3000,7 @@ requires-dist = [
     { name = "matplotlib" },
     { name = "numpy", specifier = ">2" },
     { name = "optuna", specifier = ">=4.5.0" },
+    { name = "polars", marker = "extra == 'dataframe'", specifier = ">=1.0" },
     { name = "quantem-widget", marker = "extra == 'widgets'", editable = "widget" },
     { name = "rosettasciio", specifier = ">=0.8.0" },
     { name = "scikit-image", specifier = ">=0.25.2" },
@@ -2981,7 +3013,7 @@ requires-dist = [
     { name = "tqdm" },
     { name = "zarr", specifier = ">3" },
 ]
-provides-extras = ["widgets"]
+provides-extras = ["widgets", "dataframe"]
 
 [package.metadata.requires-dev]
 dev = [
@@ -3363,7 +3395,7 @@ resolution-markers = [
     "python_full_version < '3.12'",
 ]
 dependencies = [
-    { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" } },
+    { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/7a/97/5a3609c4f8d58b039179648e62dd220f89864f56f7357f5d4f45c29eb2cc/scipy-1.17.1.tar.gz", hash = "sha256:95d8e012d8cb8816c226aef832200b1d45109ed4464303e997c5b13122b297c0", size = 30573822, upload-time = "2026-02-23T00:26:24.851Z" }
 wheels = [
@@ -3438,7 +3470,7 @@ resolution-markers = [
     "python_full_version >= '3.12' and python_full_version < '3.14'",
 ]
 dependencies = [
-    { name = "numpy", version = "2.5.1", source = { registry = "https://pypi.org/simple" } },
+    { name = "numpy", version = "2.5.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/a7/25/c2700dfaf6442b4effaa91af24ebce5dc9d31bb4a69706313aae70d72cd0/scipy-1.18.0.tar.gz", hash = "sha256:67b2ad2ad54c72ca6d04975a9b2df8c3638c34ddd5b28738e94fc2b57929d378", size = 30774447, upload-time = "2026-06-19T15:01:43.456Z" }
 wheels = [
@@ -3647,7 +3679,7 @@ resolution-markers = [
     "python_full_version < '3.12'",
 ]
 dependencies = [
-    { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" } },
+    { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/c5/cb/2f6d79c7576e22c116352a801f4c3c8ace5957e9aced862012430b62e14f/tifffile-2026.3.3.tar.gz", hash = "sha256:d9a1266bed6f2ee1dd0abde2018a38b4f8b2935cb843df381d70ac4eac5458b7", size = 388745, upload-time = "2026-03-03T19:14:38.134Z" }
 wheels = [
@@ -3663,7 +3695,7 @@ resolution-markers = [
     "python_full_version >= '3.12' and python_full_version < '3.14'",
 ]
 dependencies = [
-    { name = "numpy", version = "2.5.1", source = { registry = "https://pypi.org/simple" } },
+    { name = "numpy", version = "2.5.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/20/2f/e5fe51c8f782241d86fdf7251594b195f0d6c2fcf9d389079de212599246/tifffile-2026.7.14.tar.gz", hash = "sha256:ce2703e5ef22c868f1528d5f5b4ef75eefb019cf628a1c9ec0d17e0afeca8ef5", size = 437660, upload-time = "2026-07-14T23:41:31.737Z" }
 wheels = [
@@ -4017,12 +4049,12 @@ resolution-markers = [
     "python_full_version < '3.12'",
 ]
 dependencies = [
-    { name = "donfig" },
-    { name = "google-crc32c" },
-    { name = "numcodecs" },
-    { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" } },
-    { name = "packaging" },
-    { name = "typing-extensions" },
+    { name = "donfig", marker = "python_full_version < '3.12'" },
+    { name = "google-crc32c", marker = "python_full_version < '3.12'" },
+    { name = "numcodecs", marker = "python_full_version < '3.12'" },
+    { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" },
+    { name = "packaging", marker = "python_full_version < '3.12'" },
+    { name = "typing-extensions", marker = "python_full_version < '3.12'" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/31/5a/b8a0cf39a14c770c30bd1f2d120c54000c8cd9e84e8e79f38d9a7ce58071/zarr-3.1.6.tar.gz", hash = "sha256:d95e72cbea4b90e9a70679468b8266400331756232576ae2b43400ac5108d0eb", size = 386531, upload-time = "2026-03-23T17:25:18.748Z" }
 wheels = [
@@ -4038,12 +4070,12 @@ resolution-markers = [
     "python_full_version >= '3.12' and python_full_version < '3.14'",
 ]
 dependencies = [
-    { name = "donfig" },
-    { name = "google-crc32c" },
-    { name = "numcodecs" },
-    { name = "numpy", version = "2.5.1", source = { registry = "https://pypi.org/simple" } },
-    { name = "packaging" },
-    { name = "typing-extensions" },
+    { name = "donfig", marker = "python_full_version >= '3.12'" },
+    { name = "google-crc32c", marker = "python_full_version >= '3.12'" },
+    { name = "numcodecs", marker = "python_full_version >= '3.12'" },
+    { name = "numpy", version = "2.5.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" },
+    { name = "packaging", marker = "python_full_version >= '3.12'" },
+    { name = "typing-extensions", marker = "python_full_version >= '3.12'" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/93/8d/aeb164004f87543b06ef54f885d02c342c31ceb274e2bbec470a98927621/zarr-3.2.1.tar.gz", hash = "sha256:71565b738a0e7e8ed226f0516eba8c6bb53440ad7669a8c48ebb3534a161d035", size = 675161, upload-time = "2026-05-05T12:37:22.383Z" }
 wheels = [

From 435fbeac1bffa57388677b34be643ca0ca7f6fab Mon Sep 17 00:00:00 2001
From: arthurmccray <amccray@stanford.edu>
Date: Thu, 30 Jul 2026 13:52:12 -0700
Subject: [PATCH 2/5] converting vector to torch

---
 src/quantem/core/datastructures/vector.py | 806 ++++++++++++++--------
 src/quantem/core/io/serialize.py          |   8 +
 tests/datastructures/test_vector.py       | 360 ++++++----
 3 files changed, 777 insertions(+), 397 deletions(-)

diff --git a/src/quantem/core/datastructures/vector.py b/src/quantem/core/datastructures/vector.py
index 5c4ee967..37d598ca 100644
--- a/src/quantem/core/datastructures/vector.py
+++ b/src/quantem/core/datastructures/vector.py
@@ -1,12 +1,16 @@
 from __future__ import annotations
 
 import copy
+import math
+import numbers
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Literal, Sequence
 
 import numpy as np
+import torch
 from numpy.typing import NDArray
 
+from quantem.core import config
 from quantem.core.io.serialize import AutoSerialize
 from quantem.core.utils.validators import (
     validate_fields,
@@ -18,16 +22,18 @@
 if TYPE_CHECKING:
     import polars as pl
 
+DEFAULT_DTYPE = torch.float32
+
 
 class Vector(AutoSerialize):
-    """Ragged cell data on a fixed grid.
+    """Ragged cell data on a fixed grid, backed by torch.
 
     A ``Vector`` has two independent axes of structure:
     - fixed-grid dimensions given by ``shape``
     - ragged rows stored inside each fixed-grid cell
 
     Each ragged row has one value per named field, so each cell behaves like a
-    small 2D array with shape ``(n_rows, num_fields)``, where ``n_rows`` may
+    small 2D tensor with shape ``(n_rows, num_fields)``, where ``n_rows`` may
     vary from cell to cell.
 
     Parameters
@@ -43,57 +49,74 @@ class Vector(AutoSerialize):
         Descriptive name for the Vector.
     metadata : dict, optional
         Additional user metadata.
+    dtype : torch.dtype, optional
+        Row-buffer dtype. Defaults to ``torch.float32``.
+    device : str or torch.device, optional
+        Device for the row buffer. Defaults to ``"cpu"``.
 
     Notes
     -----
+    This class is torch-native: ``tensor``, ``flatten()`` and every arithmetic
+    result are ``torch.Tensor`` values living on ``device``. NumPy arrays,
+    Python sequences and scalars are accepted as *input* anywhere a payload is
+    taken and converted at the boundary; call ``numpy()`` to get NumPy back.
+
     The public API keeps fixed-grid indexing and field selection separate:
     - use ``[]`` for fixed-grid indexing
     - use ``select_fields(...)`` for field selection
 
     Fixed-grid indexing always returns a ``Vector``. A 0D selection exposes its
-    underlying cell array through ``.array``. Multi-cell selections can be
+    underlying cell tensor through ``.tensor``. Multi-cell selections can be
     concatenated with ``flatten()``.
 
     The internal representation is compact:
-    - ``_state["data"]`` stores all ragged rows in one numeric 2D array
+    - ``_state["data"]`` stores all ragged rows in one numeric 2D tensor
     - ``_state["cell_starts"]`` stores the start offset for each cell
     - ``_state["cell_lengths"]`` stores the row count for each cell
 
+    The offset bookkeeping is deliberately kept on the CPU even when the row
+    buffer lives on a GPU: it is read one scalar at a time, so keeping it on
+    the device would force a synchronization on every cell access.
+
     A ``Vector`` selection is a write-through view over shared storage. Views
     track only the selected fixed-grid shape, selected cell indices, and selected
-    field names.
+    field names. Because ``_state`` is shared, ``to(device)`` moves every view
+    of the same Vector.
 
     Examples
     --------
     Create a Vector and assign one cell:
 
-    >>> import numpy as np
+    >>> import torch
     >>> v = Vector.from_shape((2, 2), fields=("kx", "ky", "intensity"))
-    >>> v[0, 0] = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
-    >>> v[0, 0].array.shape
-    (2, 3)
+    >>> v[0, 0] = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
+    >>> v[0, 0].tensor.shape
+    torch.Size([2, 3])
 
     Select fields and apply in-place arithmetic:
 
     >>> kx = v.select_fields("kx")
     >>> kx += 16
     >>> kx.flatten().shape
-    (2, 1)
+    torch.Size([2, 1])
 
     Apply a rowwise transform with ``flatten()`` and ``set_flattened()``:
 
     >>> kx = v.select_fields("kx")
     >>> ky = v.select_fields("ky")
     >>> kx.set_flattened(
-    ...     np.where(
+    ...     torch.where(
     ...         ((kx.flatten() - 16) ** 2 + (ky.flatten() - 16) ** 2) < 12,
-    ...         10,
+    ...         10.0,
     ...         kx.flatten(),
     ...     )
     ... )
     """
 
-    __array_priority__ = 1000
+    # Opt out of NumPy's ufunc machinery entirely. This makes ``np.sin(vector)``
+    # raise a clear TypeError instead of building an object array, and makes
+    # ``ndarray + vector`` defer to ``Vector.__radd__``.
+    __array_ufunc__ = None
     _token = object()
 
     # ------------------------------------------------------------------ #
@@ -107,6 +130,8 @@ def __init__(
         units: Sequence[str] | None = None,
         name: str | None = None,
         metadata: dict[str, Any] | None = None,
+        dtype: torch.dtype | None = None,
+        device: str | int | torch.device | None = None,
         _token: object | None = None,
     ) -> None:
         if _token is not self._token:
@@ -119,6 +144,8 @@ def __init__(
             list(units) if units is not None else None,
             len(root_fields),
         )
+        root_dtype = DEFAULT_DTYPE if dtype is None else dtype
+        root_device = _resolve_device(device)
 
         self._state = {
             "shape": root_shape,
@@ -126,12 +153,12 @@ def __init__(
             "units": list(root_units),
             "name": name or f"{len(root_shape)}d ragged array",
             "metadata": dict(metadata or {}),
-            "data": np.empty((0, len(root_fields)), dtype=float),
-            "cell_starts": np.zeros(_cell_count(root_shape), dtype=np.int64),
-            "cell_lengths": np.zeros(_cell_count(root_shape), dtype=np.int64),
+            "data": torch.empty((0, len(root_fields)), dtype=root_dtype, device=root_device),
+            "cell_starts": torch.zeros(_cell_count(root_shape), dtype=torch.int64),
+            "cell_lengths": torch.zeros(_cell_count(root_shape), dtype=torch.int64),
         }
         self._selection_shape = root_shape
-        self._selection_indices: NDArray[np.int64] | None = None
+        self._selection_indices: torch.Tensor | None = None
         self._selected_fields: tuple[str, ...] | None = None
 
     @classmethod
@@ -139,16 +166,16 @@ def _from_view(
         cls,
         state: dict[str, Any],
         selection_shape: tuple[int, ...],
-        selection_indices: NDArray[np.int64] | None,
+        selection_indices: torch.Tensor | None,
         selected_fields: tuple[str, ...] | None,
     ) -> "Vector":
         """Build a view that shares backing storage with another Vector."""
         obj = cls.__new__(cls)
         obj._state = state
-        obj._selection_shape = selection_shape
         obj._selection_indices = (
-            None if selection_indices is None else selection_indices.astype(np.int64, copy=False)
+            None if selection_indices is None else selection_indices.to(torch.int64)
         )
+        obj._selection_shape = selection_shape
         obj._selected_fields = selected_fields
         return obj
 
@@ -161,6 +188,8 @@ def from_shape(
         units: Sequence[str] | None = None,
         name: str | None = None,
         metadata: dict[str, Any] | None = None,
+        dtype: torch.dtype | None = None,
+        device: str | int | torch.device | None = None,
     ) -> "Vector":
         """Create an empty Vector with the given fixed-grid shape and fields."""
         fields = _resolve_fields(fields, num_fields, None)
@@ -170,6 +199,8 @@ def from_shape(
             units=units,
             name=name,
             metadata=metadata,
+            dtype=dtype,
+            device=device,
             _token=cls._token,
         )
 
@@ -182,11 +213,16 @@ def from_data(
         units: Sequence[str] | None = None,
         name: str | None = None,
         metadata: dict[str, Any] | None = None,
+        dtype: torch.dtype | None = None,
+        device: str | int | torch.device | None = None,
     ) -> "Vector":
         """Create a Vector from nested fixed-grid data.
 
         The outer nesting defines the fixed-grid shape. Each leaf must coerce to a
-        2D cell array with consistent field count across all cells.
+        2D cell tensor with consistent field count across all cells. Leaves may be
+        tensors, NumPy arrays or nested sequences; they are cast to ``dtype``
+        (``torch.float32`` by default), so pass ``dtype=torch.float64`` to keep
+        double precision.
         """
         if not isinstance(data, (list, tuple)):
             raise TypeError(f"Data must be a list or tuple, got {type(data)}")
@@ -202,9 +238,11 @@ def from_data(
             units=units,
             name=name,
             metadata=metadata,
+            dtype=dtype,
+            device=device,
             _token=cls._token,
         )
-        vector._replace_cells(np.arange(len(cell_arrays), dtype=np.int64), cell_arrays)
+        vector._replace_cells(torch.arange(len(cell_arrays), dtype=torch.int64), cell_arrays)
         return vector
 
     # ------------------------------------------------------------------ #
@@ -255,64 +293,87 @@ def num_fields(self) -> int:
     @property
     def num_cells(self) -> int:
         """Return the number of fixed-grid cells in the current selection."""
-        return int(self._selected_cell_indices().size)
+        if self._selection_indices is None:
+            return _cell_count(self._state["shape"])
+        return int(self._selection_indices.numel())
 
     @property
     def total_rows(self) -> int:
         """Return the total ragged-row count in the current selection."""
-        return int(self._state["cell_lengths"][self._selected_cell_indices()].sum())
+        return int(self._selected_cell_lengths().sum())
 
     @property
-    def dtype(self) -> np.dtype[Any]:
-        """Return the NumPy dtype of the backing row buffer."""
+    def dtype(self) -> torch.dtype:
+        """Return the dtype of the backing row buffer."""
         return self._state["data"].dtype
 
+    @property
+    def device(self) -> str:
+        """Return the device string of the backing row buffer."""
+        return str(self._state["data"].device)
+
     # ------------------------------------------------------------------ #
     # Data access
     # ------------------------------------------------------------------ #
 
     @property
-    def array(self) -> NDArray[Any]:
-        """Return the selected cell as a NumPy array.
+    def tensor(self) -> torch.Tensor:
+        """Return the selected cell as a torch tensor.
+
+        Unlike ``Dataset.tensor``, which is the whole payload, this is *one
+        cell*: it is only valid for 0D selections and raises otherwise. Use
+        :meth:`flatten` to get every row of a multi-cell selection.
 
-        This is only valid for 0D selections. Single-field and contiguous
-        multi-field selections return writable views into the backing storage.
-        Non-contiguous multi-field selections return a copy because NumPy cannot
-        expose a writable column-subset view for that layout.
+        Contiguous field selections return writable views into the backing
+        storage. Reordered or non-contiguous selections return a copy, because
+        torch cannot expose a writable column-subset view for that layout.
         """
         if self.shape != ():
-            raise ValueError(".array is only valid when the selection contains exactly one cell.")
-        cell = self._cell_matrix(self._selected_cell_indices()[0])
-        cols = self._field_indices()
-        if cols.size == self._full_num_fields:
-            return cell
-        if cols.size == 1:
-            col = int(cols[0])
-            return cell[:, col : col + 1]
-        if _is_contiguous(cols):
-            return cell[:, int(cols[0]) : int(cols[-1]) + 1]
-        return cell[:, cols].copy()
-
-    def flatten(self) -> NDArray[Any]:
+            raise ValueError(".tensor is only valid when the selection contains exactly one cell.")
+        return self._selected_cell_matrix(int(self._selected_cell_indices()[0]))
+
+    def flatten(self) -> torch.Tensor:
         """Concatenate selected cells in row-major order.
 
-        Returns a 2D array with shape ``(total_rows, num_fields)`` even for
+        Returns a 2D tensor with shape ``(total_rows, num_fields)`` even for
         single-field selections.
         """
-        arrays = [
-            self._selected_cell_matrix(index)
-            for index in self._selected_cell_indices()
-            if self._cell_row_count(index) > 0
-        ]
-        if arrays:
-            return np.vstack(arrays)
-
-        dtype = self._state["data"].dtype if self._state["data"].ndim == 2 else float
-        return np.empty((0, self.num_fields), dtype=dtype)
+        data = self._state["data"]
+        gather = self._row_gather_index(self._selected_cell_indices())
+        if gather.numel() == 0:
+            return torch.empty((0, self.num_fields), dtype=data.dtype, device=data.device)
+        return _select_columns(data.index_select(0, gather.to(data.device)), self._field_indices())
+
+    def numpy(self) -> NDArray[Any]:
+        """Return the flattened selection as a NumPy array.
+
+        This is the NumPy counterpart of :meth:`flatten`, **not** of
+        :attr:`tensor` -- it covers the whole selection, not one cell. The
+        result is a detached CPU copy, and like ``Dataset.numpy()`` it is marked
+        read-only so accidental in-place writes raise instead of silently going
+        nowhere; use :meth:`set_flattened` to write values back.
+        """
+        array = self.flatten().cpu().numpy()
+        array.flags.writeable = False
+        return array
 
     def row_counts(self) -> list[int]:
         """Return per-cell row counts in the current selection order."""
-        return [self._cell_row_count(int(index)) for index in self._selected_cell_indices()]
+        return self._selected_cell_lengths().tolist()
+
+    def to(self, device: str | int | torch.device) -> "Vector":
+        """Move the backing row buffer to ``device`` and return ``self``.
+
+        ``device`` is normalized via :func:`quantem.core.config.validate_device`
+        so ``"cuda"``, ``0``, ``"cuda:0"`` and ``torch.device("cuda:0")`` all
+        resolve to the same canonical device.
+
+        Because all views share one ``_state``, this moves every view of the same
+        Vector, not just this one. The offset bookkeeping stays on the CPU by
+        design.
+        """
+        self._state["data"] = self._state["data"].to(_resolve_device(device))
+        return self
 
     # ------------------------------------------------------------------ #
     # Field management
@@ -449,24 +510,23 @@ def append_rows(self, idx: Any, rows: Any) -> None:
             raise ValueError("append_rows requires an index that selects exactly one cell.")
         target._require_full_field_view("append_rows")
 
-        new_rows = _coerce_cell_array(rows, target.num_fields)
+        new_rows = target._coerce_cell(rows, target.num_fields)
         if new_rows.shape[0] == 0:
             return
 
         cell_index = int(target._selected_cell_indices()[0])
-        existing = target._cell_matrix(cell_index)
-        combined = np.vstack((existing, new_rows)) if existing.shape[0] > 0 else new_rows.copy()
-        target._replace_cells(np.array([cell_index], dtype=np.int64), [combined])
+        combined = torch.cat((target._cell_matrix(cell_index), new_rows), dim=0)
+        target._replace_cells(torch.tensor([cell_index], dtype=torch.int64), [combined])
 
     def set_flattened(self, values: Any) -> None:
         """Write values back in flattened row-major order.
 
         This updates existing rows without changing per-cell row counts. It is
         the rowwise companion to ``flatten()`` and is especially useful for
-        NumPy-based transforms that operate on all selected rows at once.
+        tensor-based transforms that operate on all selected rows at once.
         """
         field_indices = self._field_indices()
-        targets = self._selected_cell_indices()
+        targets = self._selected_cell_indices().tolist()
         row_counts = self.row_counts()
         total_rows = sum(row_counts)
 
@@ -476,8 +536,9 @@ def set_flattened(self, values: Any) -> None:
             flat_values = values.flatten()
             if flat_values.shape[0] != total_rows:
                 raise ValueError(f"Expected {total_rows} rows, got {flat_values.shape[0]}")
+            flat_values = self._to_buffer(flat_values)
         else:
-            flat_values = _broadcast_field_values(values, total_rows, self.num_fields)
+            flat_values = self._broadcast_values(values, total_rows, self.num_fields)
 
         cursor = 0
         for target, rows in zip(targets, row_counts):
@@ -494,24 +555,20 @@ def compact(self) -> None:
         predictable at the cost of reallocating the backing buffer.
         """
         data = self._state["data"]
-        used_rows = int(self._state["cell_lengths"].sum())
-        if used_rows == 0:
-            self._state["data"] = np.empty((0, self._full_num_fields), dtype=data.dtype)
-            self._state["cell_starts"].fill(0)
-            return
-
-        compacted = np.empty((used_rows, self._full_num_fields), dtype=data.dtype)
-        starts = np.zeros_like(self._state["cell_starts"])
-        cursor = 0
-        for linear_index in range(_cell_count(self._state["shape"])):
-            length = self._cell_row_count(linear_index)
-            starts[linear_index] = cursor
-            if length > 0:
-                cell = self._cell_matrix(linear_index)
-                compacted[cursor : cursor + length] = cell
-                cursor += length
-        self._state["data"] = compacted
-        self._state["cell_starts"] = starts
+        lengths = self._state["cell_lengths"]
+        used_rows = int(lengths.sum())
+        if data.shape[0] == used_rows:
+            return  # already dense, nothing to reclaim
+
+        all_cells = torch.arange(_cell_count(self._state["shape"]), dtype=torch.int64)
+        gather = self._row_gather_index(all_cells)
+        if gather.numel() == 0:
+            self._state["data"] = torch.empty(
+                (0, self._full_num_fields), dtype=data.dtype, device=data.device
+            )
+        else:
+            self._state["data"] = data.index_select(0, gather.to(data.device))
+        self._state["cell_starts"] = torch.cumsum(lengths, 0) - lengths
 
     # ------------------------------------------------------------------ #
     # Python data model
@@ -529,6 +586,7 @@ def __repr__(self) -> str:
                 f"quantem.Vector, shape={self.shape}, name={self.name}",
                 f"  fields = {self.fields}",
                 f"  units: {self.units}",
+                f"  dtype: {self.dtype}, device: {self.device}",
             ]
         )
 
@@ -536,20 +594,7 @@ def __repr__(self) -> str:
 
     def copy(self) -> "Vector":
         """Return a deep copy of the current selection."""
-        copied = self.__class__(
-            shape=self.shape,
-            fields=self.fields,
-            units=self.units,
-            name=self.name,
-            metadata=copy.deepcopy(self.metadata),
-            _token=self.__class__._token,
-        )
-        target_cells = copied._selected_cell_indices()
-        source_arrays = [
-            self._selected_cell_matrix(index).copy() for index in self._selected_cell_indices()
-        ]
-        copied._replace_cells(target_cells, source_arrays)
-        return copied
+        return _vector_from_rows(self, self.flatten(), self.row_counts())
 
     def __getitem__(self, idx: Any) -> "Vector":
         """Return a fixed-grid selection as another Vector view."""
@@ -582,122 +627,129 @@ def __setitem__(self, idx: Any, value: Any) -> None:
     # Arithmetic operators
     # ------------------------------------------------------------------ #
 
-    def __array_ufunc__(self, ufunc: Any, method: str, *inputs: Any, **kwargs: Any) -> Any:
-        """Apply supported NumPy ufuncs elementwise.
-
-        Supported operations are limited to elementwise ``__call__`` ufuncs. The
-        result preserves the current selection shape and fields.
+    @classmethod
+    def __torch_function__(
+        cls,
+        func: Any,
+        types: Any,
+        args: tuple[Any, ...] = (),
+        kwargs: dict[str, Any] | None = None,
+    ) -> Any:
+        """Apply torch functions elementwise over the ragged rows.
+
+        Every ``Vector`` argument is replaced by its flattened rows, ``func`` is
+        applied, and any result shaped ``(total_rows, num_fields)`` is rebuilt
+        into a ``Vector`` preserving the selection shape and fields. Results of
+        any other shape -- reductions such as ``torch.sum`` and predicates such
+        as ``torch.allclose`` -- are returned as-is.
         """
-        if method != "__call__":
+        kwargs = {} if kwargs is None else kwargs
+        if kwargs.get("out") is not None:
             return NotImplemented
 
-        out = kwargs.get("out")
-        if out is not None:
-            return NotImplemented
-
-        vector_inputs = [value for value in inputs if isinstance(value, Vector)]
+        all_args = list(args) + list(kwargs.values())
+        vector_inputs = [value for value in all_args if isinstance(value, Vector)]
         if not vector_inputs:
             return NotImplemented
 
         template = vector_inputs[0]
         row_counts = template.row_counts()
-        total_rows = sum(row_counts)
 
         for other in vector_inputs[1:]:
             if other.shape != template.shape:
-                raise ValueError("Vector ufunc inputs must have matching fixed-grid shapes.")
+                raise ValueError("Vector inputs must have matching fixed-grid shapes.")
             if other.num_fields != template.num_fields:
-                raise ValueError("Vector ufunc inputs must have matching field counts.")
+                raise ValueError("Vector inputs must have matching field counts.")
             if other.row_counts() != row_counts:
-                raise ValueError("Vector ufunc inputs must have matching per-cell row counts.")
+                raise ValueError("Vector inputs must have matching per-cell row counts.")
+
+        flat_args = tuple(_flatten_torch_input(value) for value in args)
+        flat_kwargs = {key: _flatten_torch_input(value) for key, value in kwargs.items()}
+        result = func(*flat_args, **flat_kwargs)
 
-        flat_inputs = [
-            _normalize_ufunc_input(value, total_rows, template.num_fields) for value in inputs
-        ]
-        result = ufunc(*flat_inputs, **kwargs)
         if isinstance(result, tuple):
-            return tuple(_vector_from_flat_result(template, item, row_counts) for item in result)
-        return _vector_from_flat_result(template, result, row_counts)
+            return tuple(_maybe_wrap_result(template, item, row_counts) for item in result)
+        return _maybe_wrap_result(template, result, row_counts)
 
     def __add__(self, other: Any) -> "Vector":
-        return self._binary_op(other, np.add)
+        return self._binary_op(other, torch.add)
 
     def __sub__(self, other: Any) -> "Vector":
-        return self._binary_op(other, np.subtract)
+        return self._binary_op(other, torch.subtract)
 
     def __mul__(self, other: Any) -> "Vector":
-        return self._binary_op(other, np.multiply)
+        return self._binary_op(other, torch.multiply)
 
     def __truediv__(self, other: Any) -> "Vector":
-        return self._binary_op(other, np.divide)
+        return self._binary_op(other, torch.divide)
 
     def __floordiv__(self, other: Any) -> "Vector":
-        return self._binary_op(other, np.floor_divide)
+        return self._binary_op(other, torch.floor_divide)
 
     def __mod__(self, other: Any) -> "Vector":
-        return self._binary_op(other, np.mod)
+        return self._binary_op(other, torch.remainder)
 
     def __pow__(self, other: Any) -> "Vector":
-        return self._binary_op(other, np.power)
+        return self._binary_op(other, torch.pow)
 
     def __radd__(self, other: Any) -> "Vector":
-        return self._binary_op(other, np.add, reverse=True)
+        return self._binary_op(other, torch.add, reverse=True)
 
     def __rmul__(self, other: Any) -> "Vector":
-        return self._binary_op(other, np.multiply, reverse=True)
+        return self._binary_op(other, torch.multiply, reverse=True)
 
     def __rsub__(self, other: Any) -> "Vector":
-        return self._binary_op(other, np.subtract, reverse=True)
+        return self._binary_op(other, torch.subtract, reverse=True)
 
     def __rtruediv__(self, other: Any) -> "Vector":
-        return self._binary_op(other, np.divide, reverse=True)
+        return self._binary_op(other, torch.divide, reverse=True)
 
     def __rfloordiv__(self, other: Any) -> "Vector":
-        return self._binary_op(other, np.floor_divide, reverse=True)
+        return self._binary_op(other, torch.floor_divide, reverse=True)
 
     def __rmod__(self, other: Any) -> "Vector":
-        return self._binary_op(other, np.mod, reverse=True)
+        return self._binary_op(other, torch.remainder, reverse=True)
 
     def __rpow__(self, other: Any) -> "Vector":
-        return self._binary_op(other, np.power, reverse=True)
+        return self._binary_op(other, torch.pow, reverse=True)
 
     def __iadd__(self, other: Any) -> "Vector":
-        self._inplace_op(other, np.add)
+        self._inplace_op(other, torch.add)
         return self
 
     def __isub__(self, other: Any) -> "Vector":
-        self._inplace_op(other, np.subtract)
+        self._inplace_op(other, torch.subtract)
         return self
 
     def __imul__(self, other: Any) -> "Vector":
-        self._inplace_op(other, np.multiply)
+        self._inplace_op(other, torch.multiply)
         return self
 
     def __itruediv__(self, other: Any) -> "Vector":
-        self._inplace_op(other, np.divide)
+        self._inplace_op(other, torch.divide)
         return self
 
     def __ifloordiv__(self, other: Any) -> "Vector":
-        self._inplace_op(other, np.floor_divide)
+        self._inplace_op(other, torch.floor_divide)
         return self
 
     def __imod__(self, other: Any) -> "Vector":
-        self._inplace_op(other, np.mod)
+        self._inplace_op(other, torch.remainder)
         return self
 
     def __ipow__(self, other: Any) -> "Vector":
-        self._inplace_op(other, np.power)
+        self._inplace_op(other, torch.pow)
         return self
 
     def __neg__(self) -> "Vector":
-        return self._binary_op(-1, np.multiply)
+        return self._binary_op(-1, torch.multiply)
 
     def __pos__(self) -> "Vector":
         return self.copy()
 
     def __abs__(self) -> "Vector":
         result = self.copy()
-        result._inplace_unary(np.abs)
+        result._inplace_unary(torch.abs)
         return result
 
     # ------------------------------------------------------------------ #
@@ -759,13 +811,14 @@ def to_polars(self, dim_names: Sequence[str] | None = None) -> "pl.DataFrame":
         columns: dict[str, Any] = {}
         if index_names:
             counts = np.asarray(self.row_counts(), dtype=np.int64)
-            coords = np.unravel_index(self._selected_cell_indices(), root_shape)
+            cells = np.asarray(self._selected_cell_indices().tolist(), dtype=np.int64)
+            coords = np.unravel_index(cells, root_shape)
             for name, axis_coords in zip(index_names, coords):
                 columns[name] = pl.Series(
                     name, np.repeat(axis_coords, counts).astype(np.int64, copy=False)
                 )
 
-        values = self.flatten()
+        values = self.numpy()
         for column, field in enumerate(self.fields):
             columns[field] = pl.Series(field, values[:, column])
 
@@ -794,22 +847,77 @@ def save(
         skip : str, type, or list of (str or type)
             Attribute names/types to skip (by name or type) during serialization.
         compression_level : int or None
-            If set (0–9), applies Zstandard compression with Blosc backend at that level.
+            If set (0-9), applies Zstandard compression with Blosc backend at that level.
             Level 0 disables compression. Raises ValueError if > 9.
 
         Notes
         -----
         Skipped attribute names and types are also stored in the file metadata for correct
         round-trip skipping during load().
+
+        The row buffer and offset arrays are written as NumPy arrays rather than as
+        torch tensors, so they land in chunked, compressed Zarr arrays instead of
+        uncompressed ``torch.save`` blobs. :meth:`_post_load` converts them back to
+        CPU tensors on load, which also makes a GPU-saved Vector loadable without CUDA.
+
+        Note this only applies when the Vector is the *root* of the save.
+        ``AutoSerialize`` walks nested objects with ``_recursive_save`` rather than
+        calling their ``save()``, so a Vector held as an attribute of another
+        serializable object still round-trips correctly but writes its buffers as
+        uncompressed blobs. Fixing that properly means teaching
+        ``AutoSerialize._serialize_value`` to store plain non-grad tensors as Zarr
+        arrays, at which point this whole override collapses to ``compact()``.
         """
         self.compact()
-        super().save(
-            path,
-            mode=mode,
-            store=store,
-            skip=skip,
-            compression_level=compression_level,
-        )
+        buffer_keys = ("data", "cell_starts", "cell_lengths")
+        saved_state = {key: self._state[key] for key in buffer_keys}
+        saved_indices = self._selection_indices
+        try:
+            for key, tensor in saved_state.items():
+                self._state[key] = tensor.detach().cpu().numpy()
+            if saved_indices is not None:
+                self._selection_indices = saved_indices.detach().cpu().numpy()  # type: ignore[assignment]
+            super().save(
+                path,
+                mode=mode,
+                store=store,
+                skip=skip,
+                compression_level=compression_level,
+            )
+        finally:
+            for key, tensor in saved_state.items():
+                self._state[key] = tensor
+            self._selection_indices = saved_indices
+
+    def _post_load(self) -> None:
+        """Rehydrate NumPy-backed state into CPU tensors after deserialization.
+
+        Called by ``AutoSerialize._recursive_load``. This handles both files
+        written by :meth:`save` and older files written when Vector was
+        NumPy-backed; in both cases the stored dtype is preserved rather than
+        being coerced to the current default.
+        """
+        state = getattr(self, "_state", None)
+        if isinstance(state, dict):
+            if "data" in state:
+                state["data"] = _as_tensor(state["data"])
+            for key in ("cell_starts", "cell_lengths"):
+                if key in state:
+                    state[key] = _as_tensor(state[key], dtype=torch.int64)
+            if "shape" in state:
+                state["shape"] = tuple(int(dim) for dim in state["shape"])
+
+        selection_shape = getattr(self, "_selection_shape", None)
+        if selection_shape is not None:
+            self._selection_shape = tuple(int(dim) for dim in selection_shape)
+
+        indices = getattr(self, "_selection_indices", None)
+        if indices is not None:
+            self._selection_indices = _as_tensor(indices, dtype=torch.int64)
+
+        selected_fields = getattr(self, "_selected_fields", None)
+        if selected_fields is not None:
+            self._selected_fields = tuple(selected_fields)
 
     # ------------------------------------------------------------------ #
     # Private helpers — backing-store access
@@ -819,14 +927,18 @@ def save(
     def _full_num_fields(self) -> int:
         return len(self._state["fields"])
 
-    def _field_indices(self) -> NDArray[np.int64]:
-        """Map selected field names to column indices in the backing buffer."""
+    def _field_indices(self) -> list[int]:
+        """Map selected field names to column indices in the backing buffer.
+
+        Returned as a plain list so it can index a tensor on any device without
+        needing a matching index tensor there.
+        """
         if self._selected_fields is None:
-            return np.arange(self._full_num_fields, dtype=np.int64)
+            return list(range(self._full_num_fields))
 
         lookup = {field: i for i, field in enumerate(self._state["fields"])}
         try:
-            return np.array([lookup[field] for field in self._selected_fields], dtype=np.int64)
+            return [lookup[field] for field in self._selected_fields]
         except KeyError as exc:
             raise KeyError(f"Unknown field(s): {[str(exc.args[0])]}") from exc
 
@@ -835,36 +947,65 @@ def _require_full_field_view(self, operation: str) -> None:
         if self._selected_fields is not None:
             raise ValueError(f"{operation} is only allowed when all fields are selected.")
 
-    def _selected_cell_indices(self) -> NDArray[np.int64]:
+    def _selected_cell_indices(self) -> torch.Tensor:
         """Return linear cell indices for the current fixed-grid selection."""
         if self._selection_indices is None:
-            return np.arange(_cell_count(self._state["shape"]), dtype=np.int64)
+            return torch.arange(_cell_count(self._state["shape"]), dtype=torch.int64)
         return self._selection_indices
 
+    def _selected_cell_lengths(self) -> torch.Tensor:
+        """Return per-cell row counts for the current selection, in order."""
+        lengths = self._state["cell_lengths"]
+        if self._selection_indices is None:
+            return lengths
+        return lengths[self._selection_indices]
+
+    def _row_gather_index(self, cells: torch.Tensor) -> torch.Tensor:
+        """Buffer row indices for ``cells``, concatenated in row-major order.
+
+        This is the vectorized replacement for walking cells one at a time: the
+        result indexes ``_state["data"]`` directly, so gathering a whole
+        selection is a single ``index_select`` instead of one slice per cell.
+        """
+        lengths = self._state["cell_lengths"][cells]
+        total = int(lengths.sum())
+        if total == 0:
+            return torch.empty(0, dtype=torch.int64)
+        starts = self._state["cell_starts"][cells]
+        # Row r of output cell k comes from buffer row (start_k - out_start_k) + r.
+        offsets = starts - (torch.cumsum(lengths, 0) - lengths)
+        return torch.repeat_interleave(offsets, lengths) + torch.arange(total, dtype=torch.int64)
+
     def _cell_row_count(self, linear_index: int) -> int:
         """Return the row count for one cell in the backing buffer."""
         return int(self._state["cell_lengths"][linear_index])
 
-    def _cell_matrix(self, linear_index: int) -> NDArray[Any]:
+    def _cell_matrix(self, linear_index: int) -> torch.Tensor:
         """Return the full backing matrix for one cell."""
         start = int(self._state["cell_starts"][linear_index])
         length = int(self._state["cell_lengths"][linear_index])
         return self._state["data"][start : start + length]
 
-    def _selected_cell_matrix(self, linear_index: int) -> NDArray[Any]:
+    def _selected_cell_matrix(self, linear_index: int) -> torch.Tensor:
         """Return one cell with the current field selection applied."""
-        cell = self._cell_matrix(linear_index)
-        cols = self._field_indices()
-        if cols.size == self._full_num_fields:
-            return cell
-        if cols.size == 1:
-            col = int(cols[0])
-            return cell[:, col : col + 1]
-        if _is_contiguous(cols):
-            return cell[:, int(cols[0]) : int(cols[-1]) + 1]
-        return cell[:, cols].copy()
-
-    def _replace_cells(self, targets: NDArray[np.int64], arrays: Sequence[NDArray[Any]]) -> None:
+        return _select_columns(self._cell_matrix(linear_index), self._field_indices())
+
+    def _to_buffer(self, tensor: torch.Tensor) -> torch.Tensor:
+        """Cast a tensor to the backing buffer's dtype and device."""
+        data = self._state["data"]
+        return tensor.to(dtype=data.dtype, device=data.device)
+
+    def _coerce_cell(self, value: Any, num_fields: int) -> torch.Tensor:
+        """Normalize a single-cell payload onto this Vector's dtype/device."""
+        data = self._state["data"]
+        return _coerce_cell_array(value, num_fields, data.dtype, data.device)
+
+    def _broadcast_values(self, value: Any, total_rows: int, num_fields: int) -> torch.Tensor:
+        """Broadcast array-like input onto this Vector's dtype/device."""
+        data = self._state["data"]
+        return _broadcast_field_values(value, total_rows, num_fields, data.dtype, data.device)
+
+    def _replace_cells(self, targets: torch.Tensor, arrays: Sequence[Any]) -> None:
         """Replace complete cells in the compact row buffer.
 
         Whole-cell replacement is implemented by appending the new payload rows to
@@ -878,30 +1019,30 @@ def _replace_cells(self, targets: NDArray[np.int64], arrays: Sequence[NDArray[An
         if len(targets) == 0:
             return
 
-        normalized = [_coerce_cell_array(array, self._full_num_fields) for array in arrays]
+        normalized = [self._coerce_cell(array, self._full_num_fields) for array in arrays]
         payloads = [array for array in normalized if array.shape[0] > 0]
         if payloads:
-            appended = np.vstack(payloads)
-            self._state["data"] = np.concatenate((self._state["data"], appended), axis=0)
+            appended = torch.cat(payloads, dim=0)
+            self._state["data"] = torch.cat((self._state["data"], appended), dim=0)
 
-        cursor = self._state["data"].shape[0] - sum(array.shape[0] for array in normalized)
-        for target, array in zip(targets, normalized):
-            self._state["cell_starts"][target] = cursor
-            self._state["cell_lengths"][target] = array.shape[0]
-            cursor += array.shape[0]
+        lengths = torch.tensor([array.shape[0] for array in normalized], dtype=torch.int64)
+        cursor = self._state["data"].shape[0] - int(lengths.sum())
+        self._state["cell_starts"][targets] = cursor + torch.cumsum(lengths, 0) - lengths
+        self._state["cell_lengths"][targets] = lengths
 
         self._maybe_compact_storage()
 
     def _expand_storage(self, num_new_fields: int) -> None:
-        """Append new ``np.nan``-initialized columns for added fields."""
+        """Append new NaN-initialized columns for added fields."""
         data = self._state["data"]
-        dtype = np.result_type(data.dtype, float)
-        if data.shape[0] == 0:
-            self._state["data"] = np.empty((0, data.shape[1] + num_new_fields), dtype=dtype)
-            return
-
-        filler = np.full((data.shape[0], num_new_fields), np.nan, dtype=dtype)
-        self._state["data"] = np.concatenate((data.astype(dtype, copy=False), filler), axis=1)
+        # Promote to a float dtype first: torch.full(..., nan) rejects integer dtypes.
+        # This is a "smallest float that holds NaN" rule, independent of the
+        # new-Vector default in DEFAULT_DTYPE -- keep them separate.
+        dtype = torch.promote_types(data.dtype, torch.float32)
+        filler = torch.full(
+            (data.shape[0], num_new_fields), float("nan"), dtype=dtype, device=data.device
+        )
+        self._state["data"] = torch.cat((data.to(dtype), filler), dim=1)
 
     def _maybe_compact_storage(self) -> None:
         """Compact automatically once dead rows become materially larger than live rows."""
@@ -935,11 +1076,13 @@ def _assign_full_cells(self, value: Any) -> None:
                 raise ValueError(f"Expected {len(targets)} cells, got {len(source_cells)}")
             if value.num_fields != self.num_fields:
                 raise ValueError(f"Expected {self.num_fields} fields, got {value.num_fields}")
-            arrays = [value._selected_cell_matrix(index).copy() for index in source_cells]
+            arrays = [
+                value._selected_cell_matrix(index).clone() for index in source_cells.tolist()
+            ]
             self._replace_cells(targets, arrays)
             return
 
-        array = _coerce_cell_array(value, self.num_fields)
+        array = self._coerce_cell(value, self.num_fields)
         self._replace_cells(targets, [array] * len(targets))
 
     def _assign_selected_fields(self, value: Any) -> None:
@@ -950,35 +1093,39 @@ def _assign_selected_fields(self, value: Any) -> None:
         preserved, so each target cell keeps its existing row count and only the
         selected columns are overwritten.
         """
-        targets = self._selected_cell_indices()
+        targets = self._selected_cell_indices().tolist()
         field_indices = self._field_indices()
-        row_counts = [self._cell_row_count(index) for index in targets]
+        row_counts = self.row_counts()
         total_rows = sum(row_counts)
 
         if isinstance(value, Vector):
-            source_cells = value._selected_cell_indices()
+            source_cells = value._selected_cell_indices().tolist()
             if len(targets) != len(source_cells):
                 raise ValueError(f"Expected {len(targets)} cells, got {len(source_cells)}")
             if value.num_fields != self.num_fields:
                 raise ValueError(f"Expected {self.num_fields} fields, got {value.num_fields}")
-            source_counts = [value._cell_row_count(index) for index in source_cells]
+            source_counts = value.row_counts()
             if row_counts != source_counts:
                 raise ValueError("Per-cell row counts must match for field-selected assignment.")
-            snapshots = [value._selected_cell_matrix(index).copy() for index in source_cells]
+            snapshots = [
+                self._to_buffer(value._selected_cell_matrix(index)).clone()
+                for index in source_cells
+            ]
             for target, array in zip(targets, snapshots):
                 cell = self._cell_matrix(int(target))
                 if array.shape[0] > 0:
                     cell[:, field_indices] = array
             return
 
-        if np.isscalar(value):
+        if _is_scalar(value):
+            scalar = _scalar_value(value)
             for target in targets:
                 cell = self._cell_matrix(int(target))
                 if cell.shape[0] > 0:
-                    cell[:, field_indices] = value
+                    cell[:, field_indices] = scalar
             return
 
-        broadcast = _broadcast_field_values(value, total_rows, self.num_fields)
+        broadcast = self._broadcast_values(value, total_rows, self.num_fields)
         cursor = 0
         for target, rows in zip(targets, row_counts):
             chunk = broadcast[cursor : cursor + rows]
@@ -999,7 +1146,7 @@ def _binary_op(self, other: Any, op: Any, reverse: bool = False) -> "Vector":
 
     def _inplace_unary(self, op: Any) -> None:
         """Apply a unary elementwise operation in-place to the selected fields."""
-        targets = self._selected_cell_indices()
+        targets = self._selected_cell_indices().tolist()
         field_indices = self._field_indices()
         for target in targets:
             cell = self._cell_matrix(int(target))
@@ -1009,36 +1156,40 @@ def _inplace_unary(self, op: Any) -> None:
 
     def _inplace_op(self, other: Any, op: Any, reverse: bool = False) -> None:
         """Apply elementwise arithmetic in-place to the selected fields."""
-        targets = self._selected_cell_indices()
+        targets = self._selected_cell_indices().tolist()
         field_indices = self._field_indices()
-        row_counts = [self._cell_row_count(index) for index in targets]
+        row_counts = self.row_counts()
         total_rows = sum(row_counts)
 
         if isinstance(other, Vector):
-            source_cells = other._selected_cell_indices()
+            source_cells = other._selected_cell_indices().tolist()
             if len(targets) != len(source_cells):
                 raise ValueError(f"Expected {len(targets)} cells, got {len(source_cells)}")
             if other.num_fields != self.num_fields:
                 raise ValueError(f"Expected {self.num_fields} fields, got {other.num_fields}")
-            source_counts = [other._cell_row_count(index) for index in source_cells]
+            source_counts = other.row_counts()
             if row_counts != source_counts:
                 raise ValueError("Per-cell row counts must match for Vector arithmetic.")
-            snapshots = [other._selected_cell_matrix(index).copy() for index in source_cells]
+            snapshots = [
+                self._to_buffer(other._selected_cell_matrix(index)).clone()
+                for index in source_cells
+            ]
             for target, rhs in zip(targets, snapshots):
                 cell = self._cell_matrix(int(target))
                 lhs = cell[:, field_indices]
                 cell[:, field_indices] = op(rhs, lhs) if reverse else op(lhs, rhs)
             return
 
-        if np.isscalar(other):
+        if _is_scalar(other):
+            scalar = _scalar_value(other)
             for target in targets:
                 cell = self._cell_matrix(int(target))
                 lhs = cell[:, field_indices]
                 if lhs.shape[0] > 0:
-                    cell[:, field_indices] = op(other, lhs) if reverse else op(lhs, other)
+                    cell[:, field_indices] = op(scalar, lhs) if reverse else op(lhs, scalar)
             return
 
-        broadcast = _broadcast_field_values(other, total_rows, self.num_fields)
+        broadcast = self._broadcast_values(other, total_rows, self.num_fields)
         cursor = 0
         for target, rows in zip(targets, row_counts):
             chunk = broadcast[cursor : cursor + rows]
@@ -1049,6 +1200,81 @@ def _inplace_op(self, other: Any, op: Any, reverse: bool = False) -> None:
             cursor += rows
 
 
+def _resolve_device(device: str | int | torch.device | None) -> torch.device:
+    """Normalize a device specifier.
+
+    Note that ``None`` means CPU here, whereas ``config.validate_device(None)``
+    resolves to whatever accelerator is available. A data container that
+    silently lands on a GPU is surprising, so the default is explicit and
+    ``to()`` is how you move one.
+    """
+    if device is None:
+        return torch.device("cpu")
+    resolved, _ = config.validate_device(device)
+    return torch.device(resolved)
+
+
+def _as_tensor(
+    value: Any,
+    dtype: torch.dtype | None = None,
+    device: torch.device | str | None = None,
+) -> torch.Tensor:
+    """Coerce array-like input (tensor, ndarray, sequence, scalar) to a tensor.
+
+    Incoming tensors are detached: the row buffer is written in place, which is
+    not allowed on a tensor that requires grad.
+    """
+    if isinstance(value, torch.Tensor):
+        tensor = value.detach()
+    else:
+        if isinstance(value, np.ndarray) and any(stride < 0 for stride in value.strides):
+            # torch cannot wrap negatively-strided memory (e.g. arr[::-1]).
+            value = np.ascontiguousarray(value)
+        tensor = torch.as_tensor(value)
+    if dtype is not None and tensor.dtype != dtype:
+        tensor = tensor.to(dtype)
+    if device is not None and tensor.device != torch.device(device):
+        tensor = tensor.to(device)
+    return tensor
+
+
+def _is_scalar(value: Any) -> bool:
+    """Return True for values that broadcast as a single number."""
+    if isinstance(value, torch.Tensor):
+        return value.ndim == 0
+    return isinstance(value, (numbers.Number, np.generic))
+
+
+def _scalar_value(value: Any) -> Any:
+    """Unwrap a scalar-like value into something torch ops accept directly."""
+    # torch ops do not reliably accept numpy scalar types; python/tensor pass through.
+    return value.item() if isinstance(value, np.generic) else value
+
+
+def _flatten_torch_input(value: Any) -> Any:
+    """Replace a Vector argument with its flattened rows for torch dispatch."""
+    return value.flatten() if isinstance(value, Vector) else value
+
+
+def _maybe_wrap_result(template: "Vector", value: Any, row_counts: list[int]) -> Any:
+    """Rebuild a Vector from a rowwise torch result, or pass the result through.
+
+    The test is purely on shape: a tensor shaped like the flattened rows is
+    rebuilt, anything else is returned untouched. That covers the intent --
+    pointwise ops get wrapped, reductions and predicates do not -- but it is a
+    heuristic, not a knowledge of which ops are pointwise. A shape-preserving
+    non-pointwise op (``torch.t`` on a Vector whose row and field counts happen
+    to be equal) would be wrapped with its rows permuted. Swap this for an
+    explicit pointwise-op set if that ever bites.
+    """
+    if isinstance(value, torch.Tensor) and tuple(value.shape) == (
+        sum(row_counts),
+        template.num_fields,
+    ):
+        return _vector_from_rows(template, value, row_counts)
+    return value
+
+
 def _resolve_fields(
     fields: Sequence[str] | None,
     num_fields: int | None,
@@ -1083,7 +1309,7 @@ def _resolve_fields(
 
 def _cell_count(shape: tuple[int, ...]) -> int:
     """Return the number of fixed-grid cells in a shape."""
-    return int(np.prod(shape, dtype=np.int64)) if shape else 1
+    return math.prod(shape) if shape else 1
 
 
 def _normalize_field_names(field_names: str | Sequence[str]) -> tuple[str, ...]:
@@ -1135,20 +1361,25 @@ def _looks_like_field_selector(idx: Any) -> bool:
     return False
 
 
-def _coerce_cell_array(value: Any, num_fields: int) -> NDArray[Any]:
+def _coerce_cell_array(
+    value: Any,
+    num_fields: int,
+    dtype: torch.dtype,
+    device: torch.device,
+) -> torch.Tensor:
     """Normalize a single-cell payload to shape ``(n_rows, num_fields)``."""
     if isinstance(value, Vector):
         if value.shape != ():
             raise ValueError("Expected a 0D Vector for single-cell assignment.")
-        array = value.array.copy()
+        array = value.tensor.clone()
     else:
-        array = np.asarray(value)
+        array = _as_tensor(value)
 
     if array.ndim == 0:
         raise ValueError("Cell assignment requires a 2D array.")
     if array.ndim == 1:
-        if array.size == 0:
-            array = np.empty((0, num_fields), dtype=float)
+        if array.numel() == 0:
+            array = torch.empty((0, num_fields), dtype=dtype, device=device)
         elif num_fields == 1:
             array = array.reshape(-1, 1)
         else:
@@ -1157,12 +1388,12 @@ def _coerce_cell_array(value: Any, num_fields: int) -> NDArray[Any]:
         raise ValueError("Cell assignment requires a 2D array.")
     if array.shape[1] != num_fields:
         raise ValueError(f"Expected {num_fields} fields, got {array.shape[1]}")
-    return array
+    return array.to(dtype=dtype, device=device)
 
 
-def _flatten_fixed_grid(node: Any) -> tuple[tuple[int, ...], list[NDArray[Any]]]:
+def _flatten_fixed_grid(node: Any) -> tuple[tuple[int, ...], list[torch.Tensor]]:
     """Recursively flatten nested fixed-grid input into row-major cell order."""
-    if isinstance(node, np.ndarray):
+    if isinstance(node, (np.ndarray, torch.Tensor)):
         return (), [_coerce_inferred_cell_array(node)]
     if not isinstance(node, (list, tuple)):
         raise TypeError("Data must be a nested list/tuple of cell arrays or row sequences.")
@@ -1172,7 +1403,7 @@ def _flatten_fixed_grid(node: Any) -> tuple[tuple[int, ...], list[NDArray[Any]]]
         return (0,), []
 
     child_shape: tuple[int, ...] | None = None
-    cells: list[NDArray[Any]] = []
+    cells: list[torch.Tensor] = []
     for child in node:
         shape, child_cells = _flatten_fixed_grid(child)
         if child_shape is None:
@@ -1194,21 +1425,21 @@ def _looks_like_cell_rows(node: Sequence[Any]) -> bool:
 
 def _is_row_like(item: Any) -> bool:
     """Return True for a single row of scalar values."""
-    if isinstance(item, np.ndarray):
+    if isinstance(item, (np.ndarray, torch.Tensor)):
         return item.ndim == 1
     if not isinstance(item, (list, tuple)):
         return False
-    return all(np.isscalar(value) for value in item)
+    return all(_is_scalar(value) for value in item)
 
 
-def _coerce_inferred_cell_array(value: Any) -> NDArray[Any]:
-    """Infer a 2D cell array from row-like input during ``from_data``."""
-    array = np.asarray(value)
+def _coerce_inferred_cell_array(value: Any) -> torch.Tensor:
+    """Infer a 2D cell tensor from row-like input during ``from_data``."""
+    array = _as_tensor(value)
     if array.ndim == 0:
         raise ValueError("Cell data must be 1D or 2D.")
     if array.ndim == 1:
-        if array.size == 0:
-            return np.empty((0, 0), dtype=float)
+        if array.numel() == 0:
+            return torch.empty((0, 0), dtype=array.dtype)
         return array.reshape(1, -1)
     if array.ndim != 2:
         raise ValueError("Cell data must be 1D or 2D.")
@@ -1217,9 +1448,9 @@ def _coerce_inferred_cell_array(value: Any) -> NDArray[Any]:
 
 def _select_linear_indices(
     shape: tuple[int, ...],
-    current_indices: NDArray[np.int64],
+    current_indices: torch.Tensor,
     idx: Any,
-) -> tuple[tuple[int, ...], NDArray[np.int64]]:
+) -> tuple[tuple[int, ...], torch.Tensor]:
     """Apply fixed-grid indexing to a flattened cell-index view.
 
     ``current_indices`` stores the linear cell indices represented by the current
@@ -1230,13 +1461,13 @@ def _select_linear_indices(
     """
     if shape == ():
         if idx in ((), Ellipsis):
-            return (), np.array([int(current_indices[0])], dtype=np.int64)
+            return (), torch.tensor([int(current_indices[0])], dtype=torch.int64)
         raise IndexError("Too many indices for 0D Vector")
 
     index_tuple = _normalize_index_tuple(idx, len(shape))
     current_grid = current_indices.reshape(shape)
 
-    axis_positions: list[NDArray[np.int64]] = []
+    axis_positions: list[torch.Tensor] = []
     out_shape: list[int] = []
     scalar_axes: list[bool] = []
     for axis, axis_index in enumerate(index_tuple):
@@ -1249,14 +1480,14 @@ def _select_linear_indices(
     if all(scalar_axes):
         scalar_key = tuple(int(positions[0]) for positions in axis_positions)
         value = int(current_grid[scalar_key])
-        return (), np.array([value], dtype=np.int64)
+        return (), torch.tensor([value], dtype=torch.int64)
 
     mesh_inputs = [
         positions if not is_scalar else positions[:1]
         for positions, is_scalar in zip(axis_positions, scalar_axes)
     ]
-    grids = np.meshgrid(*mesh_inputs, indexing="ij")
-    selected = np.asarray(current_grid[tuple(grids)], dtype=np.int64).reshape(-1)
+    grids = torch.meshgrid(*mesh_inputs, indexing="ij")
+    selected = current_grid[tuple(grids)].reshape(-1).to(torch.int64)
     return tuple(out_shape), selected
 
 
@@ -1281,7 +1512,7 @@ def _normalize_index_tuple(idx: Any, ndim: int) -> tuple[Any, ...]:
     return idx
 
 
-def _positions_for_axis(axis_index: Any, size: int) -> tuple[NDArray[np.int64], bool]:
+def _positions_for_axis(axis_index: Any, size: int) -> tuple[torch.Tensor, bool]:
     """Resolve one axis index into concrete positions and scalar-vs-vector shape behavior."""
     if isinstance(axis_index, (bool, np.bool_)):
         raise TypeError("Boolean scalars are not valid Vector indices.")
@@ -1292,45 +1523,63 @@ def _positions_for_axis(axis_index: Any, size: int) -> tuple[NDArray[np.int64],
             index += size
         if index < 0 or index >= size:
             raise IndexError("Vector index out of range")
-        return np.array([index], dtype=np.int64), True
+        return torch.tensor([index], dtype=torch.int64), True
 
     if isinstance(axis_index, slice):
-        return np.arange(size, dtype=np.int64)[axis_index], False
+        return torch.arange(size, dtype=torch.int64)[axis_index], False
 
-    array = np.asarray(axis_index)
+    array = _as_index_tensor(axis_index)
     if array.ndim == 0:
-        if np.issubdtype(array.dtype, np.integer):
+        if _is_integer_dtype(array.dtype):
             return _positions_for_axis(int(array.item()), size)
         raise TypeError(f"Unsupported index type: {type(axis_index)!r}")
 
-    if array.dtype == bool or np.issubdtype(array.dtype, np.bool_):
+    if array.dtype == torch.bool:
         if array.ndim != 1:
             raise IndexError("Full-grid boolean masks are not supported.")
         if array.shape[0] != size:
             raise IndexError(
                 f"Boolean mask length {array.shape[0]} does not match axis length {size}"
             )
-        return np.flatnonzero(array).astype(np.int64, copy=False), False
+        return array.nonzero(as_tuple=False).reshape(-1).to(torch.int64), False
 
     if array.ndim != 1:
         raise IndexError("Fancy indexing arrays must be one-dimensional.")
-    if array.size == 0:
-        return np.array([], dtype=np.int64), False
-    if not np.issubdtype(array.dtype, np.integer):
+    if array.numel() == 0:
+        return torch.empty(0, dtype=torch.int64), False
+    if not _is_integer_dtype(array.dtype):
         raise TypeError("Fancy indices must be integers or booleans.")
 
-    positions = array.astype(np.int64, copy=True)
+    positions = array.to(torch.int64).clone()
     positions[positions < 0] += size
-    if np.any((positions < 0) | (positions >= size)):
+    if bool(((positions < 0) | (positions >= size)).any()):
         raise IndexError("Vector index out of range")
     return positions, False
 
 
-def _broadcast_field_values(value: Any, total_rows: int, num_fields: int) -> NDArray[Any]:
+def _as_index_tensor(value: Any) -> torch.Tensor:
+    """Coerce an index-like value into a CPU tensor without forcing a dtype."""
+    if isinstance(value, torch.Tensor):
+        return value.detach().cpu()
+    return _as_tensor(np.asarray(value))
+
+
+def _is_integer_dtype(dtype: torch.dtype) -> bool:
+    """Return True for signed/unsigned integer dtypes (excluding bool)."""
+    return not dtype.is_floating_point and not dtype.is_complex and dtype != torch.bool
+
+
+def _broadcast_field_values(
+    value: Any,
+    total_rows: int,
+    num_fields: int,
+    dtype: torch.dtype,
+    device: torch.device,
+) -> torch.Tensor:
     """Broadcast array-like input to flattened rowwise assignment shape."""
-    array = np.asarray(value)
+    array = _as_tensor(value, dtype=dtype, device=device)
     if array.ndim == 0:
-        return np.broadcast_to(array.reshape(1, 1), (total_rows, num_fields))
+        return array.reshape(1, 1).expand(total_rows, num_fields)
     if num_fields == 1 and array.ndim == 1:
         if total_rows == 0 and array.shape[0] == 0:
             return array.reshape(0, 1)
@@ -1338,55 +1587,60 @@ def _broadcast_field_values(value: Any, total_rows: int, num_fields: int) -> NDA
             raise ValueError(f"Expected {total_rows} values, got {array.shape[0]}")
         return array.reshape(total_rows, 1)
     try:
-        return np.broadcast_to(array, (total_rows, num_fields))
-    except ValueError as exc:
+        return torch.broadcast_to(array, (total_rows, num_fields))
+    except RuntimeError as exc:
         raise ValueError(
-            f"Cannot broadcast value with shape {array.shape} to ({total_rows}, {num_fields})"
+            f"Cannot broadcast value with shape {tuple(array.shape)} "
+            f"to ({total_rows}, {num_fields})"
         ) from exc
 
 
-def _normalize_ufunc_input(value: Any, total_rows: int, num_fields: int) -> Any:
-    """Normalize one ufunc input to flattened Vector-compatible form."""
-    if isinstance(value, Vector):
-        return value.flatten()
-    if np.isscalar(value):
-        return value
-    return _broadcast_field_values(value, total_rows, num_fields)
-
-
-def _vector_from_flat_result(
+def _vector_from_rows(
     template: Vector,
-    values: Any,
+    rows: torch.Tensor,
     row_counts: list[int],
 ) -> Vector:
-    """Build a Vector from flattened rowwise result data."""
-    total_rows = sum(row_counts)
-    flat_values = _broadcast_field_values(values, total_rows, template.num_fields)
+    """Build a Vector from an already row-major block of rows.
 
+    ``rows`` is exactly the buffer the result needs, so it is adopted directly
+    and the offsets are derived from ``row_counts`` -- no per-cell split and
+    re-concatenation.
+    """
+    source = _as_tensor(rows)
     result = Vector.from_shape(
         shape=template.shape,
         fields=template.fields,
         units=template.units,
         name=template.name,
+        dtype=source.dtype,
+        device=source.device,
     )
     result._state["metadata"] = copy.deepcopy(template.metadata)
 
-    if total_rows == 0:
-        result._state["data"] = np.empty((0, template.num_fields), dtype=flat_values.dtype)
-        return result
-
-    cursor = 0
-    cells: list[NDArray[Any]] = []
-    for rows in row_counts:
-        cells.append(flat_values[cursor : cursor + rows].copy())
-        cursor += rows
-
-    result._replace_cells(result._selected_cell_indices(), cells)
+    lengths = torch.tensor(row_counts, dtype=torch.int64)
+    result._state["data"] = source.contiguous()
+    result._state["cell_lengths"] = lengths
+    result._state["cell_starts"] = torch.cumsum(lengths, 0) - lengths
     return result
 
 
-def _is_contiguous(indices: NDArray[np.int64]) -> bool:
-    """Return True when integer column indices form one contiguous slice."""
-    if indices.size <= 1:
+def _is_contiguous(indices: Sequence[int]) -> bool:
+    """Return True when integer column indices form one ascending contiguous slice."""
+    if len(indices) <= 1:
         return True
-    return bool(np.all(indices[1:] - indices[:-1] == 1))
+    return all(after - before == 1 for before, after in zip(indices, indices[1:]))
+
+
+def _select_columns(rows: torch.Tensor, cols: list[int]) -> torch.Tensor:
+    """Apply a field selection to a 2D row block.
+
+    A contiguous ascending column run is returned as a writable view; any other
+    selection goes through advanced indexing, which copies. Reordered selections
+    must take the copying path so the columns come back in the requested order
+    rather than storage order.
+    """
+    if _is_contiguous(cols):
+        if not cols:
+            return rows[:, :0]
+        return rows[:, cols[0] : cols[-1] + 1]
+    return rows[:, cols]
diff --git a/src/quantem/core/io/serialize.py b/src/quantem/core/io/serialize.py
index 4dd03984..a71b33c3 100644
--- a/src/quantem/core/io/serialize.py
+++ b/src/quantem/core/io/serialize.py
@@ -740,6 +740,14 @@ def _recursive_load(
         if hasattr(obj, "__attrs_post_init__"):
             obj.__attrs_post_init__()
 
+        # Give the class a chance to normalize its restored state. Loading
+        # bypasses __init__, so classes that store data in a form different from
+        # what they serialize (e.g. Vector, which writes NumPy but works in
+        # torch) use this hook to rehydrate.
+        post_load = getattr(obj, "_post_load", None)
+        if callable(post_load):
+            post_load()
+
         # Fix PyTorch module set attributes after all loading is complete
         if isinstance(obj, torch.nn.Module):
             cls._fix_torch_module_sets(obj)
diff --git a/tests/datastructures/test_vector.py b/tests/datastructures/test_vector.py
index fbbad491..ddc9c435 100644
--- a/tests/datastructures/test_vector.py
+++ b/tests/datastructures/test_vector.py
@@ -2,10 +2,18 @@
 
 import numpy as np
 import pytest
+import torch
 
 from quantem.core.datastructures.vector import Vector
 from quantem.core.io.serialize import load
 
+requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
+
+
+def assert_rows(actual: torch.Tensor, expected: list[list[float]]) -> None:
+    """Compare a rowwise tensor result against literal expected values."""
+    torch.testing.assert_close(actual, torch.tensor(expected, dtype=actual.dtype))
+
 
 def make_line_vector() -> Vector:
     v = Vector.from_shape(
@@ -37,12 +45,13 @@ def test_initialization_and_len(self):
         assert len(v1) == 2
         assert v1.num_cells == 6
         assert v1.num_fields == 3
-        assert v1.dtype == np.dtype(float)
+        assert v1.dtype == torch.float32
+        assert v1.device == "cpu"
         assert v1.fields == ["a", "b", "c"]
         assert v1.units == ["none", "none", "none"]
         assert v1.name == "2d ragged array"
-        assert v1[0, 0].array.shape == (0, 3)
-        np.testing.assert_array_equal(v1[0, 0].flatten(), v1[0, 0].array)
+        assert v1[0, 0].tensor.shape == (0, 3)
+        torch.testing.assert_close(v1[0, 0].flatten(), v1[0, 0].tensor)
 
         v2 = Vector.from_shape(shape=(2, 3), num_fields=2)
         assert v2.fields == ["field_0", "field_1"]
@@ -62,26 +71,30 @@ def test_initialization_and_len(self):
         assert str(v1) == (
             "quantem.Vector, shape=(2, 3), name=2d ragged array\n"
             "  fields = ['a', 'b', 'c']\n"
-            "  units: ['none', 'none', 'none']"
+            "  units: ['none', 'none', 'none']\n"
+            "  dtype: torch.float32, device: cpu"
         )
 
-    def test_indexing_and_array_contract(self):
+    def test_indexing_and_tensor_contract(self):
         v = make_grid_vector()
 
         assert isinstance(v[:2, 1], Vector)
         assert v[:2, 1].shape == (2,)
         assert v[1].shape == (2,)
         assert v[1, 1].shape == ()
-        np.testing.assert_array_equal(v[-1, -1].array, np.array([[21.0, 121.0, 221.0]]))
+        assert_rows(v[-1, -1].tensor, [[21.0, 121.0, 221.0]])
 
         with pytest.raises(ValueError):
-            _ = v[:, 1].array
+            _ = v[:, 1].tensor
 
         result = v[[-1, 0], 1]
         assert result.shape == (2,)
         assert result.num_cells == 2
-        np.testing.assert_array_equal(result[0].array, np.array([[21.0, 121.0, 221.0]]))
-        np.testing.assert_array_equal(result[1].array, np.array([[1.0, 101.0, 201.0]]))
+        assert_rows(result[0].tensor, [[21.0, 121.0, 221.0]])
+        assert_rows(result[1].tensor, [[1.0, 101.0, 201.0]])
+
+        # Torch tensors work as fancy indices alongside lists and ndarrays.
+        assert v[torch.tensor([0, 2])].shape == (2, 2)
 
     def test_select_fields_and_chaining_equivalence(self):
         v = make_line_vector()
@@ -91,9 +104,9 @@ def test_select_fields_and_chaining_equivalence(self):
         assert selected.units == ["px"]
         assert selected.shape == v.shape
 
-        np.testing.assert_array_equal(
-            v.select_fields("kx")[2].array,
-            v[2].select_fields("kx").array,
+        torch.testing.assert_close(
+            v.select_fields("kx")[2].tensor,
+            v[2].select_fields("kx").tensor,
         )
 
         with pytest.raises(KeyError):
@@ -107,15 +120,36 @@ def test_select_fields_and_chaining_equivalence(self):
 
         multi = v.select_fields("intensity", "kx")
         assert multi.fields == ["intensity", "kx"]
-        assert multi.dtype == np.dtype(float)
+        assert multi.dtype == torch.float32
         assert multi.total_rows == 6
         assert multi.row_counts() == [2, 1, 2, 1]
 
-    def test_array_mutation_writes_through_for_single_field(self):
+    def test_select_fields_respects_requested_order(self):
+        v = Vector.from_shape(shape=(1,), fields=["a", "b", "c"])
+        v[0] = np.array([[1.0, 2.0, 3.0]])
+
+        # Data must come back in the order asked for, matching .fields --
+        # including the full-width reorder, which used to fall through a
+        # fast path that returned storage order.
+        reordered = v.select_fields("c", "b", "a")
+        assert reordered.fields == ["c", "b", "a"]
+        assert_rows(reordered[0].tensor, [[3.0, 2.0, 1.0]])
+        assert_rows(reordered.flatten(), [[3.0, 2.0, 1.0]])
+
+        subset = v.select_fields("c", "a")
+        assert subset.fields == ["c", "a"]
+        assert_rows(subset[0].tensor, [[3.0, 1.0]])
+
+        # Ascending contiguous selections keep the writable-view contract.
+        contiguous = v.select_fields("a", "b")
+        contiguous[0].tensor[0, 0] = -1.0
+        assert v[0].tensor[0, 0] == -1.0
+
+    def test_tensor_mutation_writes_through_for_single_field(self):
         v = make_line_vector()
-        cell = v.select_fields("kx")[1].array
+        cell = v.select_fields("kx")[1].tensor
         cell[0, 0] = 99.0
-        assert v[1].array[0, 1] == 99.0
+        assert v[1].tensor[0, 1] == 99.0
 
     def test_set_flattened_updates_rowwise(self):
         v = make_line_vector()
@@ -126,130 +160,185 @@ def test_set_flattened_updates_rowwise(self):
         flat_kx[mask[:, 0], 0] = -1.0
         kx.set_flattened(flat_kx)
 
-        np.testing.assert_array_equal(
-            kx.flatten(),
-            np.array([[10.0], [20.0], [-1.0], [-1.0], [-1.0], [-1.0]]),
-        )
+        assert_rows(kx.flatten(), [[10.0], [20.0], [-1.0], [-1.0], [-1.0], [-1.0]])
 
-    def test_field_arithmetic_with_scalar_and_ndarray(self):
+    def test_field_arithmetic_with_scalar_and_array(self):
         v = make_line_vector()
 
         kx = v.select_fields("kx")
         kx += 10
-        np.testing.assert_array_equal(
+        assert_rows(
             v.select_fields("kx").flatten(),
-            np.array([[20.0], [30.0], [40.0], [50.0], [60.0], [70.0]]),
+            [[20.0], [30.0], [40.0], [50.0], [60.0], [70.0]],
         )
 
+        # NumPy arrays are still accepted as input and converted at the boundary.
         v.select_fields("kx")[...] += np.arange(6)
-        np.testing.assert_array_equal(
+        assert_rows(
             v.select_fields("kx").flatten(),
-            np.array([[20.0], [31.0], [42.0], [53.0], [64.0], [75.0]]),
+            [[20.0], [31.0], [42.0], [53.0], [64.0], [75.0]],
+        )
+
+        # ... as are torch tensors.
+        v.select_fields("kx")[...] += torch.ones(6)
+        assert_rows(
+            v.select_fields("kx").flatten(),
+            [[21.0], [32.0], [43.0], [54.0], [65.0], [76.0]],
         )
 
         summed = v.select_fields("intensity") + v.select_fields("ky")
-        np.testing.assert_array_equal(
+        assert_rows(
             summed.flatten(),
-            np.array([[101.0], [202.0], [303.0], [404.0], [505.0], [606.0]]),
+            [[101.0], [202.0], [303.0], [404.0], [505.0], [606.0]],
         )
 
     def test_power_operations(self):
         v = make_line_vector()
 
         squared = v.select_fields("intensity") ** 2
-        np.testing.assert_array_equal(
-            squared.flatten(),
-            np.array([[1.0], [4.0], [9.0], [16.0], [25.0], [36.0]]),
-        )
+        assert_rows(squared.flatten(), [[1.0], [4.0], [9.0], [16.0], [25.0], [36.0]])
 
         intensity = v.select_fields("intensity")
         intensity **= 2
-        np.testing.assert_array_equal(
-            intensity.flatten(),
-            np.array([[1.0], [4.0], [9.0], [16.0], [25.0], [36.0]]),
-        )
+        assert_rows(intensity.flatten(), [[1.0], [4.0], [9.0], [16.0], [25.0], [36.0]])
 
         reverse = 2 ** v.select_fields("intensity")
-        np.testing.assert_array_equal(
+        assert_rows(
             reverse.flatten(),
-            np.array([[2.0], [16.0], [512.0], [65536.0], [33554432.0], [68719476736.0]]),
+            [[2.0], [16.0], [512.0], [65536.0], [33554432.0], [68719476736.0]],
         )
 
     def test_unary_mod_and_floor_division_operations(self):
         v = make_line_vector()
 
         negative = -v.select_fields("intensity")
-        np.testing.assert_array_equal(
-            negative.flatten(),
-            np.array([[-1.0], [-2.0], [-3.0], [-4.0], [-5.0], [-6.0]]),
-        )
+        assert_rows(negative.flatten(), [[-1.0], [-2.0], [-3.0], [-4.0], [-5.0], [-6.0]])
 
         absolute = abs(negative)
-        np.testing.assert_array_equal(
-            absolute.flatten(),
-            np.array([[1.0], [2.0], [3.0], [4.0], [5.0], [6.0]]),
-        )
+        assert_rows(absolute.flatten(), [[1.0], [2.0], [3.0], [4.0], [5.0], [6.0]])
 
         floored = v.select_fields("ky") // 150
-        np.testing.assert_array_equal(
-            floored.flatten(),
-            np.array([[0.0], [1.0], [2.0], [2.0], [3.0], [4.0]]),
-        )
+        assert_rows(floored.flatten(), [[0.0], [1.0], [2.0], [2.0], [3.0], [4.0]])
 
         modded = v.select_fields("ky") % 150
-        np.testing.assert_array_equal(
-            modded.flatten(),
-            np.array([[100.0], [50.0], [0.0], [100.0], [50.0], [0.0]]),
-        )
+        assert_rows(modded.flatten(), [[100.0], [50.0], [0.0], [100.0], [50.0], [0.0]])
 
         ky = v.select_fields("ky")
         ky //= 150
-        np.testing.assert_array_equal(
-            ky.flatten(),
-            np.array([[0.0], [1.0], [2.0], [2.0], [3.0], [4.0]]),
-        )
+        assert_rows(ky.flatten(), [[0.0], [1.0], [2.0], [2.0], [3.0], [4.0]])
 
         intensity = v.select_fields("intensity")
         intensity %= 2
-        np.testing.assert_array_equal(
-            intensity.flatten(),
-            np.array([[1.0], [0.0], [1.0], [0.0], [1.0], [0.0]]),
-        )
+        assert_rows(intensity.flatten(), [[1.0], [0.0], [1.0], [0.0], [1.0], [0.0]])
 
-    def test_numpy_ufunc_support(self):
+    def test_torch_function_support(self):
         v = make_line_vector()
+        kx = v.select_fields("kx")
 
-        sine = np.sin(v.select_fields("kx"))
-        np.testing.assert_allclose(
-            sine.flatten(),
-            np.sin(v.select_fields("kx").flatten()),
-        )
+        sine = torch.sin(kx)
+        assert isinstance(sine, Vector)
+        torch.testing.assert_close(sine.flatten(), torch.sin(kx.flatten()))
+
+        maximum = torch.maximum(kx, torch.tensor(35.0))
+        assert_rows(maximum.flatten(), [[35.0], [35.0], [35.0], [40.0], [50.0], [60.0]])
 
-        maximum = np.maximum(v.select_fields("intensity"), 3.0)  # type: ignore[arg-type]
+        # Multi-output functions return a tuple of Vectors.
+        mantissa, exponent = torch.frexp(v.select_fields("intensity"))
+        assert isinstance(mantissa, Vector)
+        assert_rows(mantissa.flatten(), [[0.5], [0.5], [0.75], [0.5], [0.625], [0.75]])
+
+        # Reductions do not have rowwise shape, so they pass through unwrapped.
+        total = torch.sum(kx)
+        assert isinstance(total, torch.Tensor)
+        assert not isinstance(total, Vector)
+        torch.testing.assert_close(total, torch.tensor(210.0))
+
+        # Mismatched ragged structure is rejected.
+        with pytest.raises(ValueError, match="matching per-cell row counts"):
+            torch.add(v[:2].select_fields("kx"), v[1:3].select_fields("kx"))
+
+    def test_numpy_interop_and_ufunc_optout(self):
+        v = make_line_vector()
+        kx = v.select_fields("kx")
+
+        # numpy() is the NumPy counterpart of flatten()
+        flat = kx.numpy()
+        assert isinstance(flat, np.ndarray)
         np.testing.assert_array_equal(
-            maximum.flatten(),
-            np.array([[3.0], [3.0], [3.0], [4.0], [5.0], [6.0]]),
+            flat, np.array([[10.0], [20.0], [30.0], [40.0], [50.0], [60.0]], dtype=np.float32)
         )
 
-        frac, whole = np.modf(v.select_fields("intensity") / 2.0)
-        np.testing.assert_allclose(
-            frac.flatten(),
-            np.array([[0.5], [0.0], [0.5], [0.0], [0.5], [0.0]]),
-        )
-        np.testing.assert_allclose(
-            whole.flatten(),
-            np.array([[0.0], [1.0], [1.0], [2.0], [2.0], [3.0]]),
+        # NumPy ufuncs are explicitly disabled; use the torch equivalent.
+        with pytest.raises(TypeError):
+            np.sin(kx)
+
+        # ndarray-on-the-left arithmetic still defers to Vector.__radd__
+        result = np.float64(1.0) + kx
+        assert isinstance(result, Vector)
+        assert_rows(result.flatten(), [[11.0], [21.0], [31.0], [41.0], [51.0], [61.0]])
+
+    def test_dtype_and_device_options(self):
+        default = Vector.from_shape(shape=(2,), fields=["a"])
+        assert default.dtype == torch.float32
+
+        doubled = Vector.from_shape(shape=(2,), fields=["a"], dtype=torch.float64)
+        doubled[0] = np.array([[1.0], [2.0]])
+        assert doubled.dtype == torch.float64
+        assert doubled[0].tensor.dtype == torch.float64
+
+        # The buffer dtype is authoritative: float64 input does not widen a
+        # float32 Vector.
+        narrow = Vector.from_shape(shape=(2,), fields=["a"])
+        narrow[0] = np.array([[1.0], [2.0]], dtype=np.float64)
+        assert narrow.dtype == torch.float32
+
+        from_data = Vector.from_data(
+            data=[np.array([[1.0, 2.0]])], fields=["a", "b"], dtype=torch.float64
         )
+        assert from_data.dtype == torch.float64
+
+    def test_device_property_and_to_cpu(self):
+        v = make_line_vector()
+        assert v.device == "cpu"
+        assert v.to("cpu") is v
+        assert v.device == "cpu"
+
+    @requires_cuda
+    def test_to_cuda_moves_shared_storage(self):
+        v = make_grid_vector()
+        view = v.select_fields("kx")
+
+        v.to("cuda")
+        assert v.device.startswith("cuda")
+        # Views share _state, so the move is visible through them too.
+        assert view.device.startswith("cuda")
+        assert v[1, 1].tensor.device.type == "cuda"
+        # Offset bookkeeping deliberately stays on the CPU.
+        assert v._state["cell_starts"].device.type == "cpu"
+
+        view += 1.0
+        assert view.flatten().device.type == "cuda"
+        assert isinstance(v.numpy(), np.ndarray)
+
+    @requires_cuda
+    def test_cuda_vector_saves_and_loads_to_cpu(self, tmp_path):
+        v = make_grid_vector().to("cuda")
+        path = tmp_path / "cuda_vector.zip"
+        v.save(path, mode="o")
+
+        # save() must leave the live object untouched on its original device.
+        assert v.device.startswith("cuda")
+
+        loaded = load(path)
+        assert loaded.device == "cpu"
+        torch.testing.assert_close(loaded[2, 1].tensor, v[2, 1].tensor.cpu())
 
     def test_field_assignment_from_vector_expression(self):
         v = make_line_vector()
         scale = 2.5
 
         v[:2].select_fields("intensity")[...] = v[2:4].select_fields("intensity") * scale
-        np.testing.assert_array_equal(
-            v[:2].select_fields("intensity").flatten(),
-            np.array([[10.0], [12.5], [15.0]]),
-        )
+        assert_rows(v[:2].select_fields("intensity").flatten(), [[10.0], [12.5], [15.0]])
 
     def test_field_assignment_requires_matching_per_cell_row_counts(self):
         v = make_line_vector()
@@ -260,28 +349,25 @@ def test_full_cell_assignment_allows_row_count_changes(self):
         v = make_line_vector()
 
         v[1] = v[0]
-        assert v[1].array.shape == (2, 3)
-        np.testing.assert_array_equal(v[1].array, v[0].array)
+        assert v[1].tensor.shape == (2, 3)
+        torch.testing.assert_close(v[1].tensor, v[0].tensor)
 
         v[0:2] = v[1:3]
-        assert v[0].array.shape == (2, 3)
-        assert v[1].array.shape == (2, 3)
+        assert v[0].tensor.shape == (2, 3)
+        assert v[1].tensor.shape == (2, 3)
 
         broadcast_cell = np.array([[9.0, 8.0, 7.0]])
         v[[0, 3]] = broadcast_cell
-        np.testing.assert_array_equal(v[0].array, broadcast_cell)
-        np.testing.assert_array_equal(v[3].array, broadcast_cell)
+        assert_rows(v[0].tensor, [[9.0, 8.0, 7.0]])
+        assert_rows(v[3].tensor, [[9.0, 8.0, 7.0]])
 
     def test_append_rows_and_compact(self):
         v = make_line_vector()
 
         v.append_rows(1, np.array([[7.0, 70.0, 700.0]]))
-        np.testing.assert_array_equal(
-            v[1].array,
-            np.array([[3.0, 30.0, 300.0], [7.0, 70.0, 700.0]]),
-        )
+        assert_rows(v[1].tensor, [[3.0, 30.0, 300.0], [7.0, 70.0, 700.0]])
 
-        v[1] = np.array([[8.0, 80.0, 800.0]])
+        v[1] = torch.tensor([[8.0, 80.0, 800.0]])
         assert v._state["data"].shape[0] > v.total_rows
 
         v.compact()
@@ -298,8 +384,11 @@ def test_boolean_indexing_is_axis_wise(self):
         selected = v[rows, cols]
 
         assert selected.shape == (2, 1)
-        np.testing.assert_array_equal(selected[0, 0].array, np.array([[1.0, 101.0, 201.0]]))
-        np.testing.assert_array_equal(selected[1, 0].array, np.array([[21.0, 121.0, 221.0]]))
+        assert_rows(selected[0, 0].tensor, [[1.0, 101.0, 201.0]])
+        assert_rows(selected[1, 0].tensor, [[21.0, 121.0, 221.0]])
+
+        # Torch bool masks work the same way.
+        assert v[torch.tensor([True, False, True]), cols].shape == (2, 1)
 
         with pytest.raises(IndexError):
             _ = v[np.array([[True, False], [False, True]])]
@@ -313,27 +402,27 @@ def test_empty_selection_is_valid_and_no_op_for_scalar_math(self):
         assert empty.flatten().shape == (0, 3)
 
         empty.select_fields("kx")[...] += 1
-        np.testing.assert_array_equal(v.flatten(), before)
+        torch.testing.assert_close(v.flatten(), before)
 
     def test_add_fields_defaults_expression_and_multiple_values(self):
         v = make_line_vector()
 
         v.add_fields(("h", "k"))
         assert v.fields == ["intensity", "kx", "ky", "h", "k"]
-        assert np.isnan(v[0].array[:, 3:5]).all()
+        assert torch.isnan(v[0].tensor[:, 3:5]).all()
 
         v.add_fields("field_out", v.select_fields("kx") + v.select_fields("ky"))
-        np.testing.assert_array_equal(
+        assert_rows(
             v.select_fields("field_out").flatten(),
-            np.array([[110.0], [220.0], [330.0], [440.0], [550.0], [660.0]]),
+            [[110.0], [220.0], [330.0], [440.0], [550.0], [660.0]],
         )
 
         v2 = make_line_vector()
         v2.add_fields(("h", "k"), (1.0, np.array([5.0, 6.0, 7.0, 8.0, 9.0, 10.0])))
-        np.testing.assert_array_equal(v2.select_fields("h").flatten(), np.ones((6, 1)))
-        np.testing.assert_array_equal(
+        assert_rows(v2.select_fields("h").flatten(), [[1.0]] * 6)
+        assert_rows(
             v2.select_fields("k").flatten(),
-            np.array([[5.0], [6.0], [7.0], [8.0], [9.0], [10.0]]),
+            [[5.0], [6.0], [7.0], [8.0], [9.0], [10.0]],
         )
 
         with pytest.raises(ValueError, match="all fields are selected"):
@@ -341,11 +430,11 @@ def test_add_fields_defaults_expression_and_multiple_values(self):
 
     def test_rename_fields(self):
         v = make_line_vector()
-        kx_data = v.select_fields("kx").flatten().copy()
+        kx_data = v.select_fields("kx").flatten().clone()
 
         v.rename_fields({"kx": "qx", "ky": "qy"})
         assert v.fields == ["intensity", "qx", "qy"]
-        np.testing.assert_array_equal(v.select_fields("qx").flatten(), kx_data)
+        torch.testing.assert_close(v.select_fields("qx").flatten(), kx_data)
 
         # Renaming through a field-selected view updates that view's selected names
         view = v.select_fields("qx")
@@ -366,17 +455,14 @@ def test_remove_fields_preserves_remaining_data(self):
         v.remove_fields(("kx", "extra"))
 
         assert v.fields == ["intensity", "ky"]
-        np.testing.assert_array_equal(
-            v[0].array,
-            np.array([[1.0, 100.0], [2.0, 200.0]]),
-        )
+        assert_rows(v[0].tensor, [[1.0, 100.0], [2.0, 200.0]])
 
     def test_copy_is_deep(self):
         v = make_line_vector()
         v_copy = v.select_fields(["intensity", "kx"]).copy()
 
-        v_copy[0].array[0, 0] = -1.0
-        assert v[0].array[0, 0] == 1.0
+        v_copy[0].tensor[0, 0] = -1.0
+        assert v[0].tensor[0, 0] == 1.0
         assert v_copy.fields == ["intensity", "kx"]
         assert v_copy.shape == (4,)
 
@@ -391,7 +477,7 @@ def test_from_data_supports_nested_fixed_grid(self):
         assert v.fields == ["a", "b"]
         assert v.units == ["u1", "u2"]
         assert v.name == "nested"
-        np.testing.assert_array_equal(v[0, 1].array, np.array([[3.0, 4.0], [5.0, 6.0]]))
+        assert_rows(v[0, 1].tensor, [[3.0, 4.0], [5.0, 6.0]])
 
         tuple_cells = [
             ([1.0, 2.0], [3.0, 4.0]),
@@ -399,16 +485,20 @@ def test_from_data_supports_nested_fixed_grid(self):
         ]
         tuple_vector = Vector.from_data(data=tuple_cells, fields=["a", "b"])
         assert tuple_vector.shape == (2,)
-        np.testing.assert_array_equal(tuple_vector[0].array, np.array([[1.0, 2.0], [3.0, 4.0]]))
-        np.testing.assert_array_equal(
-            tuple_vector[1].array,
-            np.array([[5.0, 6.0], [7.0, 8.0], [9.0, 10.0]]),
-        )
+        assert_rows(tuple_vector[0].tensor, [[1.0, 2.0], [3.0, 4.0]])
+        assert_rows(tuple_vector[1].tensor, [[5.0, 6.0], [7.0, 8.0], [9.0, 10.0]])
 
         tuple_data = (np.array([[1.0, 2.0]]), np.array([[3.0, 4.0]]))
         tuple_outer = Vector.from_data(data=tuple_data, fields=["a", "b"])
         assert tuple_outer.shape == (2,)
 
+        # Torch tensors are accepted as cell payloads too.
+        tensor_vector = Vector.from_data(
+            data=[torch.tensor([[1.0, 2.0]]), torch.tensor([[3.0, 4.0]])], fields=["a", "b"]
+        )
+        assert tensor_vector.shape == (2,)
+        assert_rows(tensor_vector[1].tensor, [[3.0, 4.0]])
+
         with pytest.raises(TypeError, match="Data must be a list or tuple"):
             Vector.from_data(data=np.array([1, 2, 3]))  # type: ignore[arg-type]
 
@@ -423,10 +513,7 @@ def test_to_polars_line_vector(self):
         assert df.columns == ["dim_0", "intensity", "kx", "ky"]
         assert df.height == v.total_rows == 6
         assert df["dim_0"].to_list() == [0, 0, 1, 2, 2, 3]
-        np.testing.assert_array_equal(
-            df.select(v.fields).to_numpy(),
-            v.flatten(),
-        )
+        np.testing.assert_array_equal(df.select(v.fields).to_numpy(), v.numpy())
 
     def test_to_polars_grid_and_dim_names(self):
         pytest.importorskip("polars")
@@ -507,12 +594,43 @@ def test_save_and_load_round_trip(self, tmp_path):
         with zipfile.ZipFile(path) as zf:
             names = [info.filename for info in zf.infolist()]
         assert len(names) < 30
+        # The row buffer is written as a compressed Zarr array, not a torch blob.
         assert "_state/data/zarr.json" in names
         assert all(not name.startswith("_selection_coords/") for name in names)
 
+        # save() must not leave the live object holding NumPy state.
+        assert isinstance(v._state["data"], torch.Tensor)
+
         loaded = load(path)
         assert isinstance(loaded, Vector)
         assert loaded.shape == v.shape
         assert loaded.fields == v.fields
         assert loaded.units == v.units
-        np.testing.assert_array_equal(loaded[2, 1].array, v[2, 1].array)
+        assert loaded.dtype == torch.float32
+        assert loaded.device == "cpu"
+        assert isinstance(loaded._state["cell_starts"], torch.Tensor)
+        torch.testing.assert_close(loaded[2, 1].tensor, v[2, 1].tensor)
+
+    def test_post_load_rehydrates_numpy_state(self):
+        """Vectors saved before the torch migration hold NumPy in _state."""
+        v = make_grid_vector()
+
+        # Simulate a NumPy-era file: every buffer restored as an ndarray, with
+        # float64 data and a list-valued shape.
+        v._state["data"] = v._state["data"].numpy().astype(np.float64)
+        v._state["cell_starts"] = v._state["cell_starts"].numpy()
+        v._state["cell_lengths"] = v._state["cell_lengths"].numpy()
+        v._state["shape"] = [3, 2]
+        v._selection_shape = [3, 2]
+        v._selection_indices = np.arange(6, dtype=np.int64)
+
+        v._post_load()
+
+        assert isinstance(v._state["data"], torch.Tensor)
+        assert isinstance(v._state["cell_starts"], torch.Tensor)
+        assert isinstance(v._selection_indices, torch.Tensor)
+        assert v._state["shape"] == (3, 2)
+        assert v._selection_shape == (3, 2)
+        # The stored precision is preserved rather than coerced to the default.
+        assert v.dtype == torch.float64
+        assert_rows(v[2, 1].tensor, [[21.0, 121.0, 221.0]])

From b00ea1dc7db2ea5d780c89d339b2697b402fcae9 Mon Sep 17 00:00:00 2001
From: arthurmccray <amccray@stanford.edu>
Date: Thu, 30 Jul 2026 16:55:57 -0700
Subject: [PATCH 3/5] cleaning up vector code

---
 src/quantem/core/datastructures/vector.py | 73 ++++++++++++++++++-----
 tests/datastructures/test_vector.py       | 61 +++++++++++++++++++
 2 files changed, 119 insertions(+), 15 deletions(-)

diff --git a/src/quantem/core/datastructures/vector.py b/src/quantem/core/datastructures/vector.py
index 37d598ca..502e156f 100644
--- a/src/quantem/core/datastructures/vector.py
+++ b/src/quantem/core/datastructures/vector.py
@@ -24,6 +24,21 @@
 
 DEFAULT_DTYPE = torch.float32
 
+# Only functions whose output preserves the meaning of each individual row may
+# be rebuilt as a Vector. Other torch functions still receive flattened tensor
+# inputs, but their results stay as ordinary tensors.
+_SAFE_ELEMENTWISE_TORCH_FUNCTIONS = frozenset(
+    getattr(torch, name)
+    for name in """
+        abs absolute acos acosh add asin asinh atan atan2 atanh
+        ceil clamp clip cos cosh divide erf erfc exp floor frexp
+        log log10 log1p log2 maximum minimum multiply neg pow
+        remainder round rsqrt sigmoid sin sinh sqrt square subtract
+        tan tanh trunc
+    """.split()
+    if hasattr(torch, name)
+)
+
 
 class Vector(AutoSerialize):
     """Ragged cell data on a fixed grid, backed by torch.
@@ -353,7 +368,7 @@ def numpy(self) -> NDArray[Any]:
         read-only so accidental in-place writes raise instead of silently going
         nowhere; use :meth:`set_flattened` to write values back.
         """
-        array = self.flatten().cpu().numpy()
+        array = self.flatten().detach().cpu().numpy()
         array.flags.writeable = False
         return array
 
@@ -525,6 +540,7 @@ def set_flattened(self, values: Any) -> None:
         the rowwise companion to ``flatten()`` and is especially useful for
         tensor-based transforms that operate on all selected rows at once.
         """
+        self._require_unique_cell_targets("set_flattened")
         field_indices = self._field_indices()
         targets = self._selected_cell_indices().tolist()
         row_counts = self.row_counts()
@@ -667,6 +683,8 @@ def __torch_function__(
         flat_kwargs = {key: _flatten_torch_input(value) for key, value in kwargs.items()}
         result = func(*flat_args, **flat_kwargs)
 
+        if func not in _SAFE_ELEMENTWISE_TORCH_FUNCTIONS:
+            return result
         if isinstance(result, tuple):
             return tuple(_maybe_wrap_result(template, item, row_counts) for item in result)
         return _maybe_wrap_result(template, result, row_counts)
@@ -1058,6 +1076,7 @@ def _maybe_compact_storage(self) -> None:
 
     def _assign(self, value: Any) -> None:
         """Dispatch assignment based on whether all fields or a subset are selected."""
+        self._require_unique_cell_targets("Assignment")
         if self._selected_fields is None:
             self._assign_full_cells(value)
         else:
@@ -1140,12 +1159,34 @@ def _assign_selected_fields(self, value: Any) -> None:
 
     def _binary_op(self, other: Any, op: Any, reverse: bool = False) -> "Vector":
         """Return a new Vector produced by elementwise arithmetic."""
-        result = self.copy()
-        result._inplace_op(other, op, reverse=reverse)
-        return result
+        row_counts = self.row_counts()
+        lhs = self.flatten()
+
+        if isinstance(other, Vector):
+            if other.num_cells != self.num_cells:
+                raise ValueError(f"Expected {self.num_cells} cells, got {other.num_cells}")
+            if other.num_fields != self.num_fields:
+                raise ValueError(f"Expected {self.num_fields} fields, got {other.num_fields}")
+            if other.row_counts() != row_counts:
+                raise ValueError("Per-cell row counts must match for Vector arithmetic.")
+            rhs: Any = other.flatten()
+        elif _is_scalar(other):
+            rhs = _scalar_value(other)
+        else:
+            rhs = _broadcast_field_values(
+                other,
+                sum(row_counts),
+                self.num_fields,
+                dtype=None,
+                device=lhs.device,
+            )
+
+        rows = op(rhs, lhs) if reverse else op(lhs, rhs)
+        return _vector_from_rows(self, rows, row_counts)
 
     def _inplace_unary(self, op: Any) -> None:
         """Apply a unary elementwise operation in-place to the selected fields."""
+        self._require_unique_cell_targets("In-place arithmetic")
         targets = self._selected_cell_indices().tolist()
         field_indices = self._field_indices()
         for target in targets:
@@ -1156,6 +1197,7 @@ def _inplace_unary(self, op: Any) -> None:
 
     def _inplace_op(self, other: Any, op: Any, reverse: bool = False) -> None:
         """Apply elementwise arithmetic in-place to the selected fields."""
+        self._require_unique_cell_targets("In-place arithmetic")
         targets = self._selected_cell_indices().tolist()
         field_indices = self._field_indices()
         row_counts = self.row_counts()
@@ -1199,6 +1241,16 @@ def _inplace_op(self, other: Any, op: Any, reverse: bool = False) -> None:
                 cell[:, field_indices] = op(chunk, lhs) if reverse else op(lhs, chunk)
             cursor += rows
 
+    def _require_unique_cell_targets(self, operation: str) -> None:
+        """Reject ambiguous write-through operations on repeated cell indices."""
+        indices = self._selection_indices
+        if indices is None:
+            return
+        if indices.numel() != torch.unique(indices).numel():
+            raise ValueError(
+                f"{operation} does not support repeated cell indices in a write selection."
+            )
+
 
 def _resolve_device(device: str | int | torch.device | None) -> torch.device:
     """Normalize a device specifier.
@@ -1257,16 +1309,7 @@ def _flatten_torch_input(value: Any) -> Any:
 
 
 def _maybe_wrap_result(template: "Vector", value: Any, row_counts: list[int]) -> Any:
-    """Rebuild a Vector from a rowwise torch result, or pass the result through.
-
-    The test is purely on shape: a tensor shaped like the flattened rows is
-    rebuilt, anything else is returned untouched. That covers the intent --
-    pointwise ops get wrapped, reductions and predicates do not -- but it is a
-    heuristic, not a knowledge of which ops are pointwise. A shape-preserving
-    non-pointwise op (``torch.t`` on a Vector whose row and field counts happen
-    to be equal) would be wrapped with its rows permuted. Swap this for an
-    explicit pointwise-op set if that ever bites.
-    """
+    """Rebuild a Vector from a shape-compatible result of an approved rowwise operation."""
     if isinstance(value, torch.Tensor) and tuple(value.shape) == (
         sum(row_counts),
         template.num_fields,
@@ -1573,7 +1616,7 @@ def _broadcast_field_values(
     value: Any,
     total_rows: int,
     num_fields: int,
-    dtype: torch.dtype,
+    dtype: torch.dtype | None,
     device: torch.device,
 ) -> torch.Tensor:
     """Broadcast array-like input to flattened rowwise assignment shape."""
diff --git a/tests/datastructures/test_vector.py b/tests/datastructures/test_vector.py
index ddc9c435..71ad4386 100644
--- a/tests/datastructures/test_vector.py
+++ b/tests/datastructures/test_vector.py
@@ -297,6 +297,46 @@ def test_dtype_and_device_options(self):
         )
         assert from_data.dtype == torch.float64
 
+    def test_out_of_place_arithmetic_uses_torch_dtype_promotion(self):
+        integer = Vector.from_shape(shape=(1,), fields=["a"], dtype=torch.int64)
+        integer[0] = [[3]]
+        divided = integer / 2
+        assert divided.dtype == torch.float32
+        assert_rows(divided.flatten(), [[1.5]])
+
+        real = Vector.from_shape(shape=(1,), fields=["a"])
+        real[0] = [[1.0]]
+
+        doubled = real + torch.tensor([[2.0]], dtype=torch.float64)
+        assert doubled.dtype == torch.float64
+        assert_rows(doubled.flatten(), [[3.0]])
+
+        complex_result = real + torch.tensor([[2.0j]], dtype=torch.complex64)
+        assert complex_result.dtype == torch.complex64
+        torch.testing.assert_close(
+            complex_result.flatten(), torch.tensor([[1.0 + 2.0j]], dtype=torch.complex64)
+        )
+
+    def test_unknown_shape_preserving_torch_function_returns_tensor(self):
+        v = make_line_vector()
+
+        flipped = torch.flip(v.select_fields("kx"), dims=(0,))
+
+        assert isinstance(flipped, torch.Tensor)
+        assert not isinstance(flipped, Vector)
+        torch.testing.assert_close(flipped, torch.flip(v.select_fields("kx").flatten(), dims=(0,)))
+
+    def test_numpy_detaches_tensor(self):
+        v = Vector.from_shape(shape=(1,), fields=["a"])
+        v._state["data"] = torch.tensor([[1.0]], requires_grad=True)
+        v._state["cell_starts"][0] = 0
+        v._state["cell_lengths"][0] = 1
+
+        array = v.numpy()
+
+        np.testing.assert_array_equal(array, np.array([[1.0]], dtype=np.float32))
+        assert not array.flags.writeable
+
     def test_device_property_and_to_cpu(self):
         v = make_line_vector()
         assert v.device == "cpu"
@@ -393,6 +433,27 @@ def test_boolean_indexing_is_axis_wise(self):
         with pytest.raises(IndexError):
             _ = v[np.array([[True, False], [False, True]])]
 
+    def test_repeated_fancy_indices_are_readable_but_not_writable(self):
+        v = make_line_vector()
+        repeated = v[[0, 0]].select_fields("intensity")
+
+        assert repeated.shape == (2,)
+        assert_rows(repeated.flatten(), [[1.0], [2.0], [1.0], [2.0]])
+        assert_rows((repeated + 1).flatten(), [[2.0], [3.0], [2.0], [3.0]])
+
+        before = v.flatten()
+        with pytest.raises(ValueError, match="repeated cell indices"):
+            repeated += 1
+        torch.testing.assert_close(v.flatten(), before)
+
+        with pytest.raises(ValueError, match="repeated cell indices"):
+            repeated.set_flattened(torch.zeros((4, 1)))
+        torch.testing.assert_close(v.flatten(), before)
+
+        with pytest.raises(ValueError, match="repeated cell indices"):
+            v[[0, 0]] = torch.tensor([[9.0, 9.0, 9.0]])
+        torch.testing.assert_close(v.flatten(), before)
+
     def test_empty_selection_is_valid_and_no_op_for_scalar_math(self):
         v = make_grid_vector()
         before = v.copy().flatten()

From cdf38b080b44984649a4429f524ab87692820327 Mon Sep 17 00:00:00 2001
From: arthurmccray <amccray@stanford.edu>
Date: Thu, 30 Jul 2026 17:00:32 -0700
Subject: [PATCH 4/5] fixing docstring

---
 src/quantem/core/datastructures/vector.py | 24 +++++++++++++++++------
 1 file changed, 18 insertions(+), 6 deletions(-)

diff --git a/src/quantem/core/datastructures/vector.py b/src/quantem/core/datastructures/vector.py
index 502e156f..0e4aaae5 100644
--- a/src/quantem/core/datastructures/vector.py
+++ b/src/quantem/core/datastructures/vector.py
@@ -651,13 +651,25 @@ def __torch_function__(
         args: tuple[Any, ...] = (),
         kwargs: dict[str, Any] | None = None,
     ) -> Any:
-        """Apply torch functions elementwise over the ragged rows.
+        """Apply torch functions over the ragged rows.
 
-        Every ``Vector`` argument is replaced by its flattened rows, ``func`` is
-        applied, and any result shaped ``(total_rows, num_fields)`` is rebuilt
-        into a ``Vector`` preserving the selection shape and fields. Results of
-        any other shape -- reductions such as ``torch.sum`` and predicates such
-        as ``torch.allclose`` -- are returned as-is.
+        Every ``Vector`` argument is replaced by its flattened rows and ``func``
+        is applied to those tensors. The result is rebuilt into a ``Vector``
+        -- preserving the selection shape and fields -- only when both hold:
+
+        - ``func`` is in :data:`_SAFE_ELEMENTWISE_TORCH_FUNCTIONS`, i.e. it maps
+          each row to a row and so keeps the ragged structure meaningful
+        - the result is a tensor shaped ``(total_rows, num_fields)``
+
+        Anything else is returned exactly as torch produced it. That covers
+        reductions (``torch.sum``), predicates (``torch.allclose``), and
+        shape-changing ops (``torch.t``) -- all of which still *work*, they just
+        hand back plain tensors rather than Vectors.
+
+        The allowlist is what makes this safe: shape alone is not a reliable
+        test, since ``torch.t`` on a Vector whose row and field counts happen to
+        be equal returns a same-shaped tensor that would otherwise be rewrapped
+        with its rows silently permuted.
         """
         kwargs = {} if kwargs is None else kwargs
         if kwargs.get("out") is not None:

From 03ea675a7ddef0ef0588c42fec71e13a8fdb8e04 Mon Sep 17 00:00:00 2001
From: Colin Ophus <cophus@gmail.com>
Date: Fri, 28 Aug 2026 15:55:00 -0700
Subject: [PATCH 5/5] Allow zero-length Vector axes and clarify torch sequence
 errors

Copying an empty fixed-grid selection such as v[[], :] raised
"Shape dimensions must be positive" because rebuilding the result
round-tripped its (0, n) shape through validate_shape. Zero-length
axes are reachable through indexing, so accept them and reject only
negative dimensions.

torch.cat/torch.stack take a sequence of tensors, so Vectors inside
that sequence were never seen as arguments and torch reported an
opaque "Multiple dispatch failed" error. Detect them and explain
that flatten() is the way to combine ragged rows.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
---
 src/quantem/core/datastructures/vector.py | 14 ++++++++++++++
 src/quantem/core/utils/validators.py      | 10 +++++++---
 tests/datastructures/test_vector.py       | 23 +++++++++++++++++++++++
 3 files changed, 44 insertions(+), 3 deletions(-)

diff --git a/src/quantem/core/datastructures/vector.py b/src/quantem/core/datastructures/vector.py
index 0e4aaae5..1aa21725 100644
--- a/src/quantem/core/datastructures/vector.py
+++ b/src/quantem/core/datastructures/vector.py
@@ -678,6 +678,20 @@ def __torch_function__(
         all_args = list(args) + list(kwargs.values())
         vector_inputs = [value for value in all_args if isinstance(value, Vector)]
         if not vector_inputs:
+            # Functions like torch.cat/torch.stack take a *sequence* of tensors,
+            # so the Vectors never appear as arguments in their own right. Say so,
+            # rather than letting torch report an opaque dispatch failure.
+            if any(
+                isinstance(item, Vector)
+                for value in all_args
+                if isinstance(value, (list, tuple))
+                for item in value
+            ):
+                raise TypeError(
+                    f"{func.__name__} takes a sequence of tensors, which cannot hold Vectors: "
+                    "ragged rows have no single shape to combine along. Pass flatten() "
+                    "results instead, e.g. torch.cat([a.flatten(), b.flatten()])."
+                )
             return NotImplemented
 
         template = vector_inputs[0]
diff --git a/src/quantem/core/utils/validators.py b/src/quantem/core/utils/validators.py
index 50361e25..b8c9b842 100644
--- a/src/quantem/core/utils/validators.py
+++ b/src/quantem/core/utils/validators.py
@@ -184,6 +184,10 @@ def validate_shape(shape: tuple[int, ...]) -> tuple[int, ...]:
     """
     Validate and convert shape to a tuple of integers.
 
+    Zero-length axes are allowed: an empty fixed-grid selection such as
+    ``vector[[], :]`` has shape ``(0, n)``, and copying or rebuilding it has to
+    round-trip that shape back through this validator.
+
     Parameters
     ----------
     shape : tuple[int, ...]
@@ -197,7 +201,7 @@ def validate_shape(shape: tuple[int, ...]) -> tuple[int, ...]:
     Raises
     ------
     ValueError
-        If shape contains non-positive integers
+        If shape contains negative integers
     TypeError
         If shape is not a tuple or contains non-integer values
     """
@@ -208,8 +212,8 @@ def validate_shape(shape: tuple[int, ...]) -> tuple[int, ...]:
     for dim in shape:
         if not isinstance(dim, int):
             raise TypeError(f"Shape dimensions must be integers, got {type(dim)}")
-        if dim <= 0:
-            raise ValueError(f"Shape dimensions must be positive, got {dim}")
+        if dim < 0:
+            raise ValueError(f"Shape dimensions must be non-negative, got {dim}")
         validated.append(dim)
 
     return tuple(validated)
diff --git a/tests/datastructures/test_vector.py b/tests/datastructures/test_vector.py
index 71ad4386..b8f4e007 100644
--- a/tests/datastructures/test_vector.py
+++ b/tests/datastructures/test_vector.py
@@ -518,6 +518,29 @@ def test_remove_fields_preserves_remaining_data(self):
         assert v.fields == ["intensity", "ky"]
         assert_rows(v[0].tensor, [[1.0, 100.0], [2.0, 200.0]])
 
+    def test_empty_selection_can_be_copied(self):
+        v = make_grid_vector()
+        empty = v[[], :]
+
+        copied = empty.copy()
+        assert copied.shape == (0, 2)
+        assert copied.fields == v.fields
+        assert copied.units == v.units
+        assert tuple(copied.flatten().shape) == (0, 3)
+        assert copied.total_rows == 0
+
+        with pytest.raises(ValueError, match="must be non-negative"):
+            Vector.from_shape(shape=(-1,), fields=["a"])
+
+    def test_torch_sequence_functions_report_a_clear_error(self):
+        v = make_line_vector()
+
+        with pytest.raises(TypeError, match="sequence of tensors"):
+            torch.cat([v, v])
+
+        # The suggested workaround does work
+        assert tuple(torch.cat([v.flatten(), v.flatten()]).shape) == (12, 3)
+
     def test_copy_is_deep(self):
         v = make_line_vector()
         v_copy = v.select_fields(["intensity", "kx"]).copy()