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/4] 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 4df9fc229..8673bb6ea 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 aa2627d27..5c4ee9676 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 954059d6b..fbbad4919 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 36a6293e2..2fe204ff1 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/4] 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 5c4ee9676..37d598ca3 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 4dd039840..a71b33c36 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 fbbad4919..ddc9c435f 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/4] 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 37d598ca3..502e156f3 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 ddc9c435f..71ad4386e 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/4] 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 502e156f3..0e4aaae5f 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: