Vector updates

[arthurmccray]

What problem this PR addreseses

Move Vector to torch, add to_polars()

Vector is only used by a few folks, so it seemed easiest to just break things by moving to torch rather than doing a slow migration.

Breaking changes

• .array → .tensor. The 0D cell accessor returns a torch.Tensor. Note this is one cell, unlike Dataset.tensor which is the whole payload.
• flatten() and all arithmetic return tensors. NumPy arrays, lists and scalars are still accepted as input everywhere; .numpy() converts back (read-only, matching Dataset.numpy()).
• Default dtype is torch.float32 (was float64). Pass dtype=torch.float64 to from_shape/from_data if you need double.
• NumPy ufuncs are disabled. np.sin(v) raises TypeError; use torch.sin(v). torch_function replaces array_ufunc, gated on an explicit pointwise-op allowlist — so torch.sin(v) returns a Vector while torch.sum(v) and torch.t(v) return plain tensors. Reductions now work, which they didn't under the ufunc protocol.

New

• to_polars(dim_names=None) — flattens to a DataFrame with the fixed-grid location as leading dim_0…dim_n integer columns. Honors the current view (both [] and select_fields), and reports root-grid coordinates so a row always says which scan position it came from. polars is an optional dep (quantem[dataframe]), lazily imported, not in dev/test.
• .device, .to(device), .numpy(), dtype=/device= on the factories. Offset bookkeeping stays on CPU by design — it's read one scalar at a time, so keeping it on-device would sync per cell access.

Bug fix

select_fields("c","b","a") returned data in storage order while .fields reported the requested order. Pre-existing (same logic in the NumPy version); fixed and tested.

Serialization

save() writes the row buffer as compressed Zarr rather than an uncompressed torch.save blob, and a new _post_load hook on AutoSerialize rehydrates to CPU tensors. Consequences: pre-migration files still load (keeping their original dtype), and a GPU-saved Vector loads onto CPU.

Known limitation: this only applies when the Vector is the save root. AutoSerialize walks nested objects via _recursive_save, so a Vector held as an attribute of another serializable object still round-trips correctly but writes uncompressed (~48% larger). The real fix is teaching _serialize_value to store plain non-grad tensors as Zarr arrays — which would also fix tensor-backed Dataset, currently saving fully uncompressed. Documented in save(), deliberately out of scope here.

What should the reviewer(s) do

Make sure it works for your workflows, make fixes.
Take a look at the tutorial: https://github.com/electronmicroscopy/quantem-tutorials/blob/main/tutorials/core/vector.ipynb

• This PR introduces a public-facing change (e.g., figures, CLI input/output, API).
  • For functional and algorithmic changes, tests are written or updated.
  • Documentation (e.g., tutorials, examples, README) has been updated.

[arthurmccray]

@cophus just pinging you as discussed so we can get this into dev :)

[cophus]

I pushed two commits to this branch.

41aeffc (fixes and docstrings)

• set_flattened, field assignment and in-place arithmetic now write all selected rows with a single index_put_ instead of looping over cells. On a 256×256 grid with 5–30 rows per cell, set_flattened drops from 291 ms to 2 ms, and += from 503 ms to 6 ms.
• In-place operations that change dtype kind, such as /= on an integer Vector, now raise a clear TypeError instead of the torch index_put dtype error.
• tensor + v (and t.add(v), etc.) now returns a Vector, matching v + tensor.
• from_data accepts [] for an empty cell, and [[], []] is now read as two empty cells rather than one cell with two zero-length rows. This bug predates the torch migration, and it matters for Bragg peak data, where many cells are empty.
• Public constructors and mutators now have Parameters and Returns sections, and the RST list formatting is fixed.

a78e0e5 merges dev and resolves the uv.lock conflict. I regenerated the lock from dev's version, so it differs from dev only by polars.

I also left five inline comments with specific changes. I tested each one against a78e0e5, and all Vector tests pass:

1. NumPy operands promote to float64 (_binary_op): cast non-tensor operands to the Vector dtype, as the in-place path already does.
2. Single-cell indexing (_selected_cell_indices): cache the cell index range and add an all-integer fast path. 1000 calls of v[i, j] drop from 55 ms to 2.8 ms (NumPy version: 13 ms).
3. from_data (from_data): build the row buffer with one concatenation. 284 ms drops to 36 ms (NumPy version: 52 ms).
4. 1D operands on single-field Vectors (__torch_function__): make x + v apply one value per row, matching v + x.
5. Stale views after rename_fields: document the behavior and add a test for now, and decide on a fix separately.

Items 1 to 4 are ready as code. If you agree with them, I'll push them with tests, or you can apply the suggestion blocks directly.

[cophus]

Views created before a rename keep the old field names, so kx = v.select_fields("a"); v.rename_fields({"a": "x"}); kx.flatten() raises KeyError. Fixing this properly means views would store column indices, which would change how add_fields and remove_fields update views. For this PR, please document it in the rename_fields docstring:

Notes

Views created before the rename keep the old field names and raise
KeyError; create new views with select_fields after renaming.

and add a test that asserts the KeyError, so a later fix has to update it.

[cophus]

v * np.ones(3) returns a float64 Vector, because NumPy arrays default to float64 and torch promotes float32 * float64 to float64. In a NumPy workflow almost every operand array is float64, so arithmetic silently doubles memory use and changes dtype. The in-place path (v *= np.ones(3)) keeps float32, so the two paths currently disagree. Please cast non-tensor operands (NumPy arrays, lists) to the Vector's dtype, and keep torch promotion for tensor operands:

            dtype=None if isinstance(other, torch.Tensor) else lhs.dtype,

Tested: (v * np.ones(3)).dtype is float32, (v * torch.ones(3, dtype=torch.float64)).dtype is still float64, and all tests pass. Please also add a test asserting both dtypes.

[cophus]

v[i, j] on a root Vector rebuilds torch.arange(num_cells) on every call. On a 256×256 grid this accounts for 87% of the time per access (55 µs, compared with 13 µs for the NumPy version). Please cache it:

        return _cell_arange(_cell_count(self._state["shape"]))

and add next to _cell_count (with import functools):

@functools.lru_cache(maxsize=8)
def _cell_arange(num_cells: int) -> torch.Tensor:
"""Cached arange(num_cells), shared read-only by every root Vector of that size."""
return torch.arange(num_cells, dtype=torch.int64)

Please also add an all-integer fast path at the top of _select_linear_indices. It should .clone() the result so views never share memory with the cache:

key = idx if isinstance(idx, tuple) else (idx,)
if len(key) == len(shape) and all(type(i) is int for i in key):
linear = 0
for i, size in zip(key, shape):
if not -size <= i < size:
raise IndexError("Vector index out of range")
linear = linear * size + (i % size)
return (), current_indices[linear : linear + 1].clone()

Tested: 1000 accesses take 2.8 ms instead of 55 ms. Negative and out-of-range indices behave as before.

[cophus]

from_data takes 284 ms on a 256×256 grid, compared with 52 ms for the NumPy version, because each cell goes through its own as_tensor, reshape and dtype cast. Please build the buffer directly:

    num_fields = vector._full_num_fields
    lengths = torch.tensor([array.shape[0] for array in cell_arrays], dtype=torch.int64)
    rows = [array.reshape(-1, num_fields) for array in cell_arrays if array.shape[0] > 0]
    if rows and all(isinstance(array, np.ndarray) for array in rows):
        vector._state["data"] = vector._to_buffer(_as_tensor(np.concatenate(rows)))
    elif rows:
        vector._state["data"] = vector._to_buffer(
            torch.cat([_as_tensor(array) for array in rows], dim=0)
        )
    vector._state["cell_lengths"] = lengths
    vector._state["cell_starts"] = torch.cumsum(lengths, 0) - lengths

For this to help, _coerce_inferred_cell_array also needs to return NumPy input unchanged (reshaped to 2D) instead of calling _as_tensor. Tested: 36 ms. The output buffer is identical to the current version, and mixed NumPy/torch/list leaves still work.

[cophus]

For a single-field Vector, v + x with a 1D x of length total_rows adds one value per row. x + v instead uses torch broadcasting and returns a (total_rows, total_rows) tensor. Please apply the same per-row rule in torch dispatch:

    # Single-field Vectors read a 1D tensor as one value per row, as in v + x.
    per_row = sum(row_counts) if template.num_fields == 1 else None
    flat_args = tuple(_flatten_torch_input(value, per_row) for value in args)
    flat_kwargs = {key: _flatten_torch_input(value, per_row) for key, value in kwargs.items()}

and in _flatten_torch_input(value, per_row=None), add
if isinstance(value, torch.Tensor) and value.ndim == 1 and value.shape[0] == per_row: return value.reshape(-1, 1).
Restricting this to single-field Vectors matters: for a multi-field Vector with total_rows == num_fields, the reshape would make x + v and v + x disagree. Tested: both orders now give the same result for single-field and multi-field Vectors. Please add a test with x = torch.arange(total_rows).

[cophus]

Looks good!