diff --git a/src/quantem/core/datastructures/vector.py b/src/quantem/core/datastructures/vector.py index aa2627d27..8a82568bd 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: + from quantem.core.datastructures.dataset import Dataset + class Vector(AutoSerialize): """Ragged cell data on a fixed grid. @@ -88,6 +91,18 @@ class Vector(AutoSerialize): ... kx.flatten(), ... ) ... ) + + Empty the fixed-grid cells that a boolean grid mask deselects: + + >>> kept = v.mask(np.array([[True, False], [False, True]])) + >>> kept.row_counts() + [2, 0, 0, 0] + + Reduce the ragged rows of each cell down to a fixed-grid image: + + >>> total = v.select_fields("intensity").sum(per_cell=True, as_dataset=True) + >>> total.shape + (2, 2) """ __array_priority__ = 1000 @@ -311,6 +326,201 @@ 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()] + # ------------------------------------------------------------------ # + # Reductions + # ------------------------------------------------------------------ # + + def sum(self, per_cell: bool = False, as_dataset: bool = False) -> "NDArray[Any] | Dataset": + """Sum the ragged rows of the current selection, per field. + + Parameters + ---------- + per_cell : bool, optional + If False (default), reduce every selected row into one value per + field, giving shape ``(num_fields,)``. If True, reduce within each + fixed-grid cell instead, giving shape ``shape + (num_fields,)``. + as_dataset : bool, optional + If True, return the per-cell result as a fixed-grid ``Dataset`` + instead of an array. Requires ``per_cell=True`` and exactly one + selected field. + + Returns + ------- + numpy.ndarray or Dataset + Summed values. Cells with no rows sum to ``0.0``. + + See Also + -------- + mean : Average instead of total. + count : Number of rows, rather than a reduction over field values. + + Examples + -------- + Total intensity recorded at each scan position, as an image: + + >>> total = v.select_fields("intensity").sum(per_cell=True, as_dataset=True) + >>> total.shape + (128, 128) + """ + return self._reduce("sum", per_cell, as_dataset) + + def mean(self, per_cell: bool = False, as_dataset: bool = False) -> "NDArray[Any] | Dataset": + """Average the ragged rows of the current selection, per field. + + Parameters + ---------- + per_cell : bool, optional + If False (default), reduce every selected row into one value per + field, giving shape ``(num_fields,)``. If True, reduce within each + fixed-grid cell instead, giving shape ``shape + (num_fields,)``. + as_dataset : bool, optional + If True, return the per-cell result as a fixed-grid ``Dataset`` + instead of an array. Requires ``per_cell=True`` and exactly one + selected field. + + Returns + ------- + numpy.ndarray or Dataset + Mean values. Cells with no rows are ``np.nan``. + + Examples + -------- + Mean peak position over the whole scan: + + >>> v.select_fields("q_row", "q_col").mean() + array([63.8, 64.1]) + """ + return self._reduce("mean", per_cell, as_dataset) + + def min(self, per_cell: bool = False, as_dataset: bool = False) -> "NDArray[Any] | Dataset": + """Reduce the ragged rows of the current selection to their minimum, per field. + + Parameters + ---------- + per_cell : bool, optional + If False (default), reduce every selected row into one value per + field, giving shape ``(num_fields,)``. If True, reduce within each + fixed-grid cell instead, giving shape ``shape + (num_fields,)``. + as_dataset : bool, optional + If True, return the per-cell result as a fixed-grid ``Dataset`` + instead of an array. Requires ``per_cell=True`` and exactly one + selected field. + + Returns + ------- + numpy.ndarray or Dataset + Minimum values. Cells with no rows are ``np.nan``. + """ + return self._reduce("min", per_cell, as_dataset) + + def max(self, per_cell: bool = False, as_dataset: bool = False) -> "NDArray[Any] | Dataset": + """Reduce the ragged rows of the current selection to their maximum, per field. + + Parameters + ---------- + per_cell : bool, optional + If False (default), reduce every selected row into one value per + field, giving shape ``(num_fields,)``. If True, reduce within each + fixed-grid cell instead, giving shape ``shape + (num_fields,)``. + as_dataset : bool, optional + If True, return the per-cell result as a fixed-grid ``Dataset`` + instead of an array. Requires ``per_cell=True`` and exactly one + selected field. + + Returns + ------- + numpy.ndarray or Dataset + Maximum values. Cells with no rows are ``np.nan``. + + Examples + -------- + Brightest peak found at each scan position: + + >>> brightest = v.select_fields("intensity").max(per_cell=True, as_dataset=True) + """ + return self._reduce("max", per_cell, as_dataset) + + def std(self, per_cell: bool = False, as_dataset: bool = False) -> "NDArray[Any] | Dataset": + """Standard deviation of the ragged rows of the current selection, per field. + + The population standard deviation is used, matching ``numpy.std`` + defaults. + + Parameters + ---------- + per_cell : bool, optional + If False (default), reduce every selected row into one value per + field, giving shape ``(num_fields,)``. If True, reduce within each + fixed-grid cell instead, giving shape ``shape + (num_fields,)``. + as_dataset : bool, optional + If True, return the per-cell result as a fixed-grid ``Dataset`` + instead of an array. Requires ``per_cell=True`` and exactly one + selected field. + + Returns + ------- + numpy.ndarray or Dataset + Standard deviations. Cells with no rows are ``np.nan``. + """ + return self._reduce("std", per_cell, as_dataset) + + def count(self, per_cell: bool = False, as_dataset: bool = False) -> "int | NDArray | Dataset": + """Count the ragged rows of the current selection. + + Counts are a property of rows rather than of field values, so the result + carries no field axis. + + Parameters + ---------- + per_cell : bool, optional + If False (default), return the total row count as an int. If True, + return one count per fixed-grid cell, with shape ``shape``. + as_dataset : bool, optional + If True, return the per-cell result as a fixed-grid ``Dataset`` + instead of an array. Requires ``per_cell=True``. + + Returns + ------- + int, numpy.ndarray or Dataset + Row counts, as ``np.int64`` when ``per_cell`` is True. + + Examples + -------- + Number of Bragg peaks detected at each scan position: + + >>> num_peaks = v.count(per_cell=True, as_dataset=True) + >>> num_peaks.shape + (128, 128) + """ + if not per_cell: + if as_dataset: + raise ValueError("as_dataset=True requires per_cell=True.") + return self.total_rows + + counts = np.asarray(self.row_counts(), dtype=np.int64).reshape(self.shape) + if not as_dataset: + return counts + return self._as_dataset(counts, "count", signal_units="counts") + + def imgreduce(self): + ''' + This method reduces a Vector to an image by summing all the intensity + in each pixel + + Return + ------ + im: np.ndarray + the final image + ''' + intensity = self.select_fields("intensity") + + flat = intensity.flatten()[:, 0] # (total_rows,) — flatten() is always 2D + counts = np.asarray(intensity.row_counts()) # row-major over vec.shape + cells = np.repeat(np.arange(counts.size), counts) + + img = np.bincount(cells, weights=flat, minlength=counts.size).reshape(self.shape) + return img + # ------------------------------------------------------------------ # # Field management # ------------------------------------------------------------------ # @@ -483,6 +693,164 @@ def set_flattened(self, values: Any) -> None: cell[:, field_indices] = flat_values[cursor : cursor + rows] cursor += rows + def mask(self, mask: Any, modify_in_place: bool = False) -> "Vector | None": + """Keep only the fixed-grid cells selected by a boolean mask. + + The mask is a boolean array with the same shape as this Vector, holding + one entry per fixed-grid cell, so a Vector of Bragg peaks over a + ``(N_row, N_col)`` scan takes a ``(N_row, N_col)`` mask. Any fixed-grid + dimensionality is supported, from 0D up. + + Cells marked False are emptied: they keep their place in the fixed grid + and simply hold zero ragged rows. The fixed-grid shape and the field + schema are always preserved, so a masked ``(256, 256)`` scan is still + ``(256, 256)``. + + Parameters + ---------- + mask : array-like + Boolean array of shape ``self.shape``, or a flat boolean array with + one entry per cell in row-major order. Integer masks are accepted and + read as nonzero-means-keep. + modify_in_place : bool, optional + If True, empty the deselected cells in this Vector's backing storage + and return None. The change is visible to every view sharing that + storage. If False (default), return a new Vector holding only the + selected cells and leave this one untouched. + + Returns + ------- + Vector or None + Masked copy of the current selection if ``modify_in_place`` is False, + otherwise None. + + Raises + ------ + ValueError + If the mask shape does not match the selected fixed-grid cells. + TypeError + If the mask is neither boolean nor integer typed. + + See Also + -------- + filter_rows : Rowwise counterpart, dropping ragged rows instead of cells. + select_fields : Field-wise counterpart, selecting fields instead of cells. + + Notes + ----- + Masking selects whole cells, never individual ragged rows. To drop rows by + field value, e.g. peaks below an intensity threshold, use ``filter_rows``. + + Examples + -------- + Keep the scan positions inside a region of interest: + + >>> roi = (scan_row > 32) & (scan_row < 96) # shape == v.shape == (128, 128) + >>> region = v.mask(roi) + >>> region.shape + (128, 128) + + Empty a few cells of a 1D Vector, in place: + + >>> v.mask(np.array([True, False, True, True]), modify_in_place=True) + + Mask a single cell of a 0D selection: + + >>> kept = v[3, 7].mask(np.True_) + """ + keep = self._resolve_cell_mask(mask) + targets = self._selected_cell_indices() + + if modify_in_place: + dropped = targets[~keep] + empty = np.empty((0, self._full_num_fields), dtype=self.dtype) + self._replace_cells(dropped, [empty] * dropped.size) + return None + + empty = np.empty((0, self.num_fields), dtype=self.dtype) + kept = [ + self._selected_cell_matrix(int(index)) if flag else empty + for index, flag in zip(targets, keep) + ] + result = self._empty_like() + result._replace_cells(result._selected_cell_indices(), kept) + return result + + def filter_rows(self, mask: Any, modify_in_place: bool = False) -> "Vector | None": + """Keep only the ragged rows selected by a rowwise boolean mask. + + The mask holds one entry per ragged row, in the row-major order produced + by ``flatten()``, which is how a mask built from field values arrives: + + >>> kr = v.select_fields("kr").flatten()[:, 0] + >>> annulus = v.filter_rows((kr > k_min) & (kr < k_max)) + + Rows are kept or dropped in full, across every field, so only per-cell row + counts change. The fixed-grid shape and the field schema are preserved, and + cells that lose all their rows simply become empty. + + Parameters + ---------- + mask : array-like or Vector + Rows to keep, as a 1D boolean array of ``total_rows`` entries or a 2D + array of shape ``(total_rows, 1)``, i.e. the direct result of comparing + a single-field ``flatten()`` against a value. A single-field ``Vector`` + with matching per-cell row counts also works, with nonzero meaning + keep. Integer masks are read as nonzero-means-keep; masks with more + than one column must be reduced first, e.g. with ``.any(axis=1)``. + modify_in_place : bool, optional + If True, drop the rows from this Vector's backing storage and return + None. Rows are removed across *all* fields, even when the mask was + built from a field-selected view, and the change is visible to every + view sharing that storage. If False (default), return a new Vector + holding only the kept rows and leave this one untouched. + + Returns + ------- + Vector or None + Filtered copy of the current selection if ``modify_in_place`` is False, + otherwise None. + + Raises + ------ + ValueError + If the mask length does not match the number of selected rows. + TypeError + If the mask is neither boolean nor integer typed. + + See Also + -------- + mask : Cellwise counterpart, emptying fixed-grid cells instead of rows. + + Examples + -------- + Keep the peaks inside a reciprocal-space annulus: + + >>> kr = v.select_fields("kr").flatten()[:, 0] + >>> annulus = v.filter_rows((kr > k_min) & (kr < k_max)) + >>> annulus.shape == v.shape + True + + Discard weak peaks from the Vector itself: + + >>> intensity = v.select_fields("intensity").flatten() + >>> v.filter_rows(intensity > 0.1, modify_in_place=True) + """ + row_masks = self._resolve_row_mask(mask) + targets = self._selected_cell_indices() + + if modify_in_place: + kept = [self._cell_matrix(int(index))[keep] for index, keep in zip(targets, row_masks)] + self._replace_cells(targets, kept) + return None + + kept = [ + self._selected_cell_matrix(int(index))[keep] for index, keep in zip(targets, row_masks) + ] + result = self._empty_like() + result._replace_cells(result._selected_cell_indices(), kept) + return result + def compact(self) -> None: """Repack the backing row buffer to remove dead rows. @@ -533,14 +901,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, - ) + copied = self._empty_like() target_cells = copied._selected_cell_indices() source_arrays = [ self._selected_cell_matrix(index).copy() for index in self._selected_cell_indices() @@ -749,6 +1110,28 @@ def save( def _full_num_fields(self) -> int: return len(self._state["fields"]) + def _empty_like(self) -> "Vector": + """Return an empty root Vector matching this selection's shape and schema. + + Unlike the public constructor this accepts zero-length fixed-grid axes, so + selections such as ``vector[[]]`` can still be copied or masked. + """ + obj = self.__class__.__new__(self.__class__) + obj._state = { + "shape": self.shape, + "fields": list(self.fields), + "units": list(self.units), + "name": self.name, + "metadata": copy.deepcopy(self.metadata), + "data": np.empty((0, self.num_fields), dtype=self.dtype), + "cell_starts": np.zeros(_cell_count(self.shape), dtype=np.int64), + "cell_lengths": np.zeros(_cell_count(self.shape), dtype=np.int64), + } + obj._selection_shape = self.shape + obj._selection_indices = None + obj._selected_fields = None + return obj + def _field_indices(self) -> NDArray[np.int64]: """Map selected field names to column indices in the backing buffer.""" if self._selected_fields is None: @@ -794,6 +1177,87 @@ def _selected_cell_matrix(self, linear_index: int) -> NDArray[Any]: return cell[:, int(cols[0]) : int(cols[-1]) + 1] return cell[:, cols].copy() + def _reduce(self, op: str, per_cell: bool, as_dataset: bool) -> "NDArray[Any] | Dataset": + """Reduce the selected rows over one field column at a time.""" + values = self.flatten() + if not per_cell: + if as_dataset: + raise ValueError("as_dataset=True requires per_cell=True.") + return _reduce_rows(values, op) + + lengths = np.asarray(self.row_counts(), dtype=np.int64) + reduced = _reduce_segments(values, lengths, op) + reduced = reduced.reshape(self.shape + (self.num_fields,)) + if not as_dataset: + return reduced + if self.num_fields != 1: + raise ValueError( + f"as_dataset=True requires exactly one selected field, got {self.num_fields}. " + "Narrow the selection with select_fields(...) first." + ) + return self._as_dataset(reduced[..., 0], op, signal_units=self.units[0]) + + def _as_dataset(self, array: NDArray[Any], label: str, signal_units: str) -> "Dataset": + """Wrap a fixed-grid result array in the Dataset subclass for its dimensionality.""" + from quantem.core.datastructures import Dataset + + if self.shape == (): + raise ValueError( + "as_dataset=True requires a Vector with at least one fixed-grid axis." + ) + cls = Dataset._registry.get(len(self.shape), Dataset) + fields = ", ".join(self.fields) + return cls.from_array( + array=array, + name=f"{self.name} {label}({fields})", + signal_units=signal_units, + ) + + def _resolve_row_mask(self, mask: Any) -> list[NDArray[np.bool_]]: + """Validate a rowwise mask and split it into one boolean array per selected cell.""" + row_counts = self.row_counts() + if isinstance(mask, Vector): + if mask.num_fields != 1: + raise ValueError( + f"A Vector mask must have exactly one field, got {mask.num_fields}." + ) + if mask.row_counts() != row_counts: + raise ValueError("A Vector mask must have matching per-cell row counts.") + mask = mask.flatten()[:, 0] != 0 + + array = np.asarray(mask) + if array.ndim == 2 and array.shape[1] == 1: + array = array[:, 0] + if array.ndim != 1: + raise ValueError( + f"Mask must be 1D or of shape (n_rows, 1), got shape {array.shape}. " + "Reduce multi-column masks first, e.g. with .any(axis=1)." + ) + if array.size and array.dtype != bool and not np.issubdtype(array.dtype, np.integer): + raise TypeError(f"Mask must be boolean or integer typed, got dtype {array.dtype}.") + if array.shape[0] != sum(row_counts): + raise ValueError( + f"Mask has {array.shape[0]} entries, expected {sum(row_counts)} rows." + ) + + if not row_counts: + return [] + bounds = np.cumsum(row_counts[:-1], dtype=np.int64) + return list(np.split(array.astype(bool, copy=False), bounds)) + + def _resolve_cell_mask(self, mask: Any) -> NDArray[np.bool_]: + """Validate a fixed-grid mask and flatten it to one boolean per selected cell.""" + array = np.asarray(mask) + num_cells = self.num_cells + if array.shape != self.shape and not (array.ndim == 1 and array.shape[0] == num_cells): + raise ValueError( + f"Mask has shape {array.shape}, expected {self.shape} or a flat mask " + f"with {num_cells} entries." + ) + if array.size and array.dtype != bool and not np.issubdtype(array.dtype, np.integer): + raise TypeError(f"Mask must be boolean or integer typed, got dtype {array.dtype}.") + return array.astype(bool, copy=False).reshape(-1) + def _replace_cells(self, targets: NDArray[np.int64], arrays: Sequence[NDArray[Any]]) -> None: """Replace complete cells in the compact row buffer. @@ -1078,6 +1542,55 @@ def _coerce_cell_array(value: Any, num_fields: int) -> NDArray[Any]: return array +def _reduce_rows(values: NDArray[Any], op: str) -> NDArray[Any]: + """Reduce a flattened ``(n_rows, num_fields)`` array down to one value per field.""" + if values.shape[0] == 0: + fill = 0.0 if op == "sum" else np.nan + return np.full(values.shape[1], fill, dtype=float) + if op == "sum": + return values.sum(axis=0) + if op == "mean": + return values.mean(axis=0) + if op == "min": + return values.min(axis=0) + if op == "max": + return values.max(axis=0) + if op == "std": + return values.std(axis=0) + raise ValueError(f"Unknown reduction {op!r}.") + + +def _reduce_segments(values: NDArray[Any], lengths: NDArray[np.int64], op: str) -> NDArray[Any]: + """Reduce contiguous row segments of ``values``, one segment per fixed-grid cell. + + ``values`` holds the selected rows in cell order and ``lengths`` their per-cell + row counts, so each cell owns one contiguous slice. Empty cells have no rows to + reduce and are filled with ``0.0`` for sums and ``np.nan`` otherwise. + """ + out = np.full((lengths.size, values.shape[1]), 0.0 if op == "sum" else np.nan, dtype=float) + nonempty = lengths > 0 + if not nonempty.any(): + return out + + starts = (np.cumsum(lengths) - lengths)[nonempty] + counts = lengths[nonempty].astype(float)[:, None] + if op == "sum": + out[nonempty] = np.add.reduceat(values, starts, axis=0) + elif op == "mean": + out[nonempty] = np.add.reduceat(values, starts, axis=0) / counts + elif op == "min": + out[nonempty] = np.minimum.reduceat(values, starts, axis=0) + elif op == "max": + out[nonempty] = np.maximum.reduceat(values, starts, axis=0) + elif op == "std": + means = np.add.reduceat(values, starts, axis=0) / counts + deviations = (values - np.repeat(means, lengths[nonempty], axis=0)) ** 2 + out[nonempty] = np.sqrt(np.add.reduceat(deviations, starts, axis=0) / counts) + else: + raise ValueError(f"Unknown reduction {op!r}.") + return out + + def _flatten_fixed_grid(node: Any) -> tuple[tuple[int, ...], list[NDArray[Any]]]: """Recursively flatten nested fixed-grid input into row-major cell order.""" if isinstance(node, np.ndarray): diff --git a/src/quantem/diffraction/__init__.py b/src/quantem/diffraction/__init__.py index 3cd0027f8..0eecd61e7 100644 --- a/src/quantem/diffraction/__init__.py +++ b/src/quantem/diffraction/__init__.py @@ -1,4 +1,5 @@ from quantem.diffraction.bragg_vectors import BraggVectors as BraggVectors from quantem.diffraction.strain import StrainMap as StrainMap from quantem.diffraction.strain_autocorrelation import StrainMapAutocorrelation as StrainMapAutocorrelation -from quantem.diffraction.model_fitting import ModelDiffraction as ModelDiffraction \ No newline at end of file +from quantem.diffraction.model_fitting import ModelDiffraction as ModelDiffraction +from quantem.diffraction.digital_dark_field_cluster import * \ No newline at end of file diff --git a/src/quantem/diffraction/digital_dark_field_cluster.py b/src/quantem/diffraction/digital_dark_field_cluster.py new file mode 100644 index 000000000..962e5c50e --- /dev/null +++ b/src/quantem/diffraction/digital_dark_field_cluster.py @@ -0,0 +1,769 @@ +import numpy as np + +import matplotlib.pyplot as plt +from matplotlib.colors import LinearSegmentedColormap, Normalize, PowerNorm +import matplotlib.gridspec as GridSpec +import matplotlib.patheffects as path_effects + +from tqdm import tqdm + +from sklearn.cluster import DBSCAN + +from quantem.core.datastructures.vector import Vector + + # ------------------------------------------------------------------ # + # Create suitable Vector object + # ------------------------------------------------------------------ # + +def make_FullPointsVector_centres(vecs,centers): + ''' + This may be a bit wasteful but it builds a new vector object that contains everything you need for + DDF imaging. Maybe you could do this instead by augmenting the existing object + from disk detection, but I couldn't figure out how + azimuthal angle is measured anticlockwise from horizontal right + + Parameters + ---------- + vecs: Vector + Currently must contain fields for kx, ky and intensity + centers: np.ndarray + A (2,Rx,Ry) array of kx and ky centres + + Returns + ------- + pointsvector: Vector + Containing fields ["rx", "ry", "kx", "ky", "kr", "kphi", "intensity"] + + ''' + Rshape = centers.shape[1:] + if 'q_row' in vecs.fields: + fields = ["q_row","q_col"] + elif 'kx' in vecs.fields: + fields = ["kx","ky"] + pointsvector = Vector.from_shape( + shape=Rshape, + fields=("rx", "ry", "kx", "ky", "kr", "kphi", "intensity"), + units=("pixels", "pixels", "pixels", "pixels", "pixels", "degrees", "counts"), + name="diffraction_vectors", + ) + + for rx in tqdm(range(Rshape[0])): + for ry in range(Rshape[1]): + kx = vecs[rx,ry].select_fields(fields[0]).flatten()-centers[0,rx,ry] + ky = vecs[rx,ry].select_fields(fields[1]).flatten()-centers[1,rx,ry] + kr = (kx**2+ky**2)**.5 + kphi = np.degrees(np.arctan2(-kx, ky)) + I = vecs[rx,ry].select_fields("intensity").flatten() + + pointsvector[rx, ry] = np.column_stack(( + rx * np.ones_like(kx), + ry * np.ones_like(kx), + kx, + ky, + kr, + kphi, + I + )) + return pointsvector + +def make_FullPointsVector_from_pointsarray(pointsarray): + ''' + For back compatibility, this reads in pointsarray objects made with py4DSTEM + digital dark field + x + Parameters + ---------- + pointsarray: np.ndarray + Nx7 array + + Returns + ------- + pointsvector: Vector + Containing fields ["rx", "ry", "kx", "ky", "kr", "kphi", "intensity"] + + ''' + Rshape = (int(pointsarray.T[3].max()+1),int(pointsarray.T[4].max()+1)) + print(Rshape) + pointsvector = Vector.from_shape( + shape=Rshape, + fields=("rx", "ry", "kx", "ky", "kr", "kphi", "intensity"), + units=("pixels", "pixels", "pixels", "pixels", "pixels", "degrees", "counts"), + name="diffraction_vectors", + ) + + for rx in tqdm(range(Rshape[0])): + for ry in range(Rshape[1]): + mask = np.logical_and( + pointsarray.T[3]==rx, + pointsarray.T[4]==ry, + ) + kx = pointsarray.T[0][mask] + ky = pointsarray.T[1][mask] + kr = pointsarray.T[5][mask] + kphi = pointsarray.T[6][mask] + I = pointsarray.T[2][mask] + + pointsvector[rx, ry] = np.column_stack(( + rx * np.ones_like(kx), + ry * np.ones_like(kx), + kx, + ky, + kr, + kphi, + I + )) + return pointsvector + + # ------------------------------------------------------------------ # + # Digital Dark Field Basics + # ------------------------------------------------------------------ # + +def generate_DDF_pointselect_array( + Qshape, + g1=None, + g2=None, + g1min=-1, + g1max=1, + g2min=-1, + g2max=1, + arrayorigin=np.array([0,0]), + rmin=0, + rmax=100 +): + ''' + Drop in replacement for earlier functions for creating a list selection points for forming + Digital Dark Field images. The function is more compact in construction, however. This is only + for spots in regular arrangements: single spots, lines (2-beam conditions) or arrays (zone axes). + + If you specify neither basis vector, g1 or g2, it just produces one point at the array origin, + i.e. classic bright or dark field with one aperture. + + If you specify a g1, then it will make a line of spots along this. Default is that this will be + -g, 0 and g. + + If you specify both g1 and g2, you get a grid, currently 3x3 by default. You adjust this by changing + g1min, g1max, g2min, and g2max, which are the maximum multipliers for g1 and g2 in negative and positive + senses. + + An array need not be centered on 0,0, if you move array origin (e.g. to g1 / 2 for a half RL cell shift) + + It is convenient to get g1 and g2 from the strain module. + + If you want a grid but to skip the central beam, then just set rmin as something larger than 0. 1 pixel will + usually work with aligned data (if working in uncalibrated pixels). + + You can set a maximum radius cutoff too, if required. rmin and rmax measure from (0,0), regardless of what you + set for an arrayorigin. + + Parameters + ---------- + Qshape: tuple + Shape of the diffraction pattern + g1: np.ndarray + A [kx,ky] vector + g2: np.ndarray + A [kx,ky] vector + g1min, g1max, g2min, g2max: int + maximum multiples of each g-vector in either direction + arrayorigin: np.ndarray + A [kx,ky] vector, which sets where either a single aperture or the centre of some line or grid + will go + rmin, rmax: int, float + min and max radii from [0,0] within which points will be selected + + Returns + ------- + selected_points: np.ndarray + A Nx2 vector which lists a number of kx,ky points chosen as selection positions for DDF imaging + + ''' + if isinstance(g1, np.ndarray): + if isinstance(g2, np.ndarray): + # Compute an array of points + grids = np.mgrid[ + g1min:g1max+1, + g2min:g2max+1 + ] + selected_points = np.outer(grids[0].flatten(),g1)+np.outer(grids[1].flatten(),g2)+arrayorigin + else: + # Compute a line of points + grids = np.mgrid[ + g1min:g1max+1, + ] + selected_points = np.outer(grids,g1)+arrayorigin + else: + selected_points = np.array([[arrayorigin[0],arrayorigin[1]]]) + radii = (selected_points**2).sum(axis=1)**.5 + selected_points = selected_points[ + np.logical_and( + radii>=rmin, + radii<=rmax + ) + ] + return selected_points + +def DDFpointsmask(pointsvector,selectionpoints,tolerance): + ''' + This makes a Boolean mask for selection of diffraction peaks for DDF imaging from a set of selected + positions in the reciprocal space plane. This will work with regular arrangements from + generate_DDF_pointselect_array, as well as lists of points from other sources, such as the diffraction points + extracted from some particular pixel in the dataset. + + If there are multiple points, then this will generate + multiple masks and the object will be MxN in size, where N is the length of the flattened pointsvector and + M is the number of masks. Each mask needs to be separate since multiple diffraction spots may contribute to + total intensity in a pixel, so all need counting separately and adding and there are multiple contributions + to the bright pixels + + Parameters + ---------- + pointsvector: Vector + Currently must contain fields for rx, ry, kx, ky and intensity + selectionpoints: np.ndarray + This will have shape (M,2) and will contain M pairs of kx,ky coordinates + tolerance: int, float + This is the tolerance for selection of a peak near any of the selectionpoints + in whatever units are used for the selectionpoints (will work in pixels or calibrated units) + + Returns + ------- + maskstack: np.ndarray + A set of Boolean masks for selecting points. Each will have the same length as the flattened fields + in the pointsvector it is to be used on. + ''' + if 'q_row' in pointsvector.fields: + fields = ["q_row","q_col"] + elif 'kx' in pointsvector.fields: + fields = ["kx","ky"] + maskstack = np.transpose( + np.linalg.norm( + pointsvector.select_fields(*fields).flatten()[:,None,:]-selectionpoints,axis=2 + )1: + mask = maskstack.sum(axis=0) + return mask + +def DDFrphimask(pointsvector,r,rtol,phi=None,phitol=None): + ''' + This selects points that fit within a certain radial range, and optionally, + within a certain azimuthal angle range. + + In general, the azimuthal angle is defined in the range -180 - 180, so + selections are recommended in this range. + + Parameters + ---------- + pointsvector: Vector + Currently must contain fields for rx, ry, kr, kphi and intensity + r: int, float + The reciprocal space radius chosen + rtol: int, float + The tolerance on the reciprocal space radius chosen + phi: None, int, float + The azimuthal angle chosen (in degrees) + phitol: int, float + The tolerance on the azimuthal angle radius chosen (in degrees) + + Returns + ------- + mask: np.ndarray + A Boolean mask for selecting points with the same length as the flattened fields + in the pointsvector it is to be used on. + ''' + radial_selection = np.abs(pointsvector.select_fields('kr').flatten()-r) 180: + additional_phi_selection = np.abs(pointsvector.select_fields('kphi').flatten()-phi+360)0 blocks the primary beam, which + may be sensible. Values will need adjusting for your data and detector, and whether you are working in + calibrated units or raw pixels + plot: bool + Turns plotting on or off + Returns + ------- + pointsvector2: Vector + A copy of the original Vector, with an additional field for L1labels. It may be shorter than + pointsvector if radial filtering has been applied + + ''' + for item in fields: + assert item in ["rx", "ry", 'kx', "ky", "kr", "kphi"], "field not found in [rx, ry, kx, ky, kr, kphi]" + assert len(scaling)==len(fields), "the scalings and fields must have the same number of entries" + + # We need to return a new Vector as it is changing length once we select only part of the data + pointsvector2 = pointsvector.copy() + + # making the mask is obvious + kr = pointsvector.select_fields("kr").flatten() + pointsvector2 = pointsvector.filter_rows((kr > kr_min) & (kr < kr_max)) + pointsarray = (np.array(scaling)*pointsvector2.select_fields(*fields).flatten()) + + db = DBSCAN(eps=eps, min_samples=min_samples) + db.fit(pointsarray) + pointsvector2.add_fields('L1labels',db.labels_) + if plot: + plot_L1_clusters_kspace( + pointsvector2, + fields, + kr_max_plot=int(pointsvector2.select_fields('kx').flatten().max()*1.05) + ) + return pointsvector2 + +''' +A custom colormap for the k-space plots +''' +california = LinearSegmentedColormap.from_list( + 'cali', + [ + (1,0.5,0), + (.9,.9,0), + (0.5,.9,0), + (0,.9,.9), + (0,.5,1) + ], + + # bad='gray' +) +california.set_under('lightgrey') +california.set_bad('red') + +def plot_L1_clusters_kspace(pointsvector, fields, kr_max_plot, cmap=california, figax=None): + """ + Takes a L1 cluster result of running some cluster algorithm in Scikit-Learn (e.g. DBSCAN) + on 4D data in a points array and plots the results in reciprocal and real space. Everything + is plotted in uncalibrated pixels, since this is just about seeing the results. + It expects you have applied some radial filtering (although this is not necessary) + and sets a maximum radius in reciprocal space, purely for visualisation (not calculation) + purposes. It also plots the unclustered points in pale grey. + You can use it for simple inline visualisation in a notebook, or can return a figure for + saving. + + Parameters + ---------- + L1labels: np.ndarray + The labels list from a clustering algorithm + pointsvector: Vector + A Vector object from this repo, preferably constructed with ["rx", "ry", 'kx', "ky", "kr", "kphi", "I] + as the fields + kr_max_plot: int, float + maximum radius for the reciprocal space plot + cmap: colormap + Either use the default one or provide your own. Note it needs to be a colormap (not + just a name of a colormap, such as matplotlib.colormaps["viridis"]) + figax: tuple or None + If None, a fig and ax are provided. But you can plot in your own defined axes. + Returns + ------- + """ + + if figax is None: + fig, ax = plt.subplots(1, 1, figsize=(6, 6)) + else: + fig, ax = figax + assert isinstance(fig, Figure) + assert isinstance(ax, Axes) + + ax.set_title("DBSCAN "+", ".join(fields)) + ax.set_xlabel("kx", fontsize=24) + ax.set_ylabel("ky", fontsize=24) + ax.set_ylim(kr_max_plot, -kr_max_plot) + ax.set_xlim(-kr_max_plot, kr_max_plot) + + kx = pointsvector.select_fields("kx").flatten() + ky = pointsvector.select_fields("ky").flatten() + I = pointsvector.select_fields("intensity").flatten() + kr = pointsvector.select_fields("kr").flatten() + kphi = pointsvector.select_fields("kphi").flatten() + L1labels = pointsvector.select_fields("L1labels").flatten().astype(int) + uniquelabels = np.unique(L1labels) + + ax.scatter( + ky,kx, + s=0.1, alpha=0.2, + cmap = cmap, + c=L1labels, + norm=Normalize(vmin=0, vmax=L1labels.max(), clip=False), + rasterized=True + ) + + for label in uniquelabels[1:]: + clustermask = L1labels==label + maxint = np.argmax(I[clustermask]) + r = kr[clustermask][maxint] + 3 + ang = np.radians(kphi[clustermask][maxint]) + + labx = np.sin(ang) * r + laby = np.cos(ang) * r + ax.annotate( + label, + (ky[clustermask][maxint], kx[clustermask][maxint]), + (laby, -labx), + horizontalalignment="center", + verticalalignment="center", + size=8, + ) + +def show_clusters_in_real_space( + pointsvector, ncols=5, gamma=0.25, cmapname='inferno', ordering='sequential', save_ims=False, level=1 +): + """ + Function to show real space plots of all L1 clustering outputs. This is designed purely + for in-line sanity checking, and not for publication quality output so there is no + savefig option. It is likely that in many cases, the output will be verbose and need + scrolling through. + There is an option to return the images themselves as a dict, which is especially useful + for image similarity based computation of L2 clusters. + + Parameters + ---------- + L1labels: np.ndarray + The labels list from a clustering algorithm + pointsvector: Vector + A points array, as defined in py4DSTEM.process.diffraction.digital_dark_field + ncols: int + number of columns to be used + gamma: float + Image gamma. <1 boosts lower intensities in display. + cmapname: str + Must be a valid name for a colormap in matplotlib + ordering: str + Either "sequential" for the ordering from the cluster output or "size" for ordering + by cluster size + save_ims: bool + Can turn on return of an image + level: int + Selects which level to pull from the pointsvector + Returns + ------- + imdict: dict + dictionary with cluster indices as keys and images as np.ndarray + """ + assert ordering in ["sequential", "size"], "ordering must be either sequential or size" + assert level in [1,2], "level must be either 1 or 2" + level_label = f"L{level}labels" + + labels = pointsvector.select_fields(level_label).flatten().astype(int) + unique_labels, all_cluster_sizes = np.unique(labels, return_counts=True) + shape = pointsvector.shape + if ordering == "sequential": + cluster_list = unique_labels[1:] + elif ordering == "size": + cluster_list = L1_unique_labels[1:][np.argsort(L1_all_cluster_sizes[1:])[::-1]] + + # Set up aspect ration for plotting + l = cluster_list.shape[0] + ar = shape[1] / shape[0] + w = 10 + row = int(np.ceil(l / ncols)) + + # Set up plot + fig = plt.figure(figsize=(w, w * row / ncols / ar)) + gs = GridSpec.GridSpec(row, ncols) + + # Do the plotting (and maybe save the images) + if save_ims: + ims = [] + for n, cluster_label in enumerate(cluster_list): + i, j = int(n / ncols), n % ncols + ax = plt.subplot(gs[i, j]) + ax.set_axis_off() + + mask = labels == cluster_label + im = DDFimage_from_mask(pointsvector,mask) + ax.imshow(im, norm=PowerNorm(gamma=gamma), cmap=cmapname) + ax.text( + 5, + 5, + cluster_label, + color="w", + size=14, + fontweight="bold", + verticalalignment="top", + ) + if save_ims: + ims+=[im] + if save_ims: + return np.array(ims) + +def cluster_mask(pointsvector, selected_cluster_labels, labelstitle): + """ + Makes a mask that selects only the points in a particular cluster. If applied on an output + from clustering directly on a Vector object, then it can be used for Digital Dark Field imaging + with that Vector using "DDFimage_from_mask". + + Parameters + ---------- + cluster_labels: np.ndarray + The labels list from a clustering algorithm + selected: int, list of int + An integer specifying one of the cluster labels in cluster_labels or a list of ints selecting + more than one cluster + Returns + ------- + mask: np.ndarray + A single mask that is True wherever any row matches one of the selected cluster labels + """ + maskstack = [] + assert labelstitle in ['L1labels','L2labels'], "You need to choose either L1labels or L2labels" + assert labelstitle in pointsvector.fields, "This clustering has not yet been done on this Vector" + cluster_labels = pointsvector.select_fields(labelstitle).flatten().astype(int) + + for cluster_label in selected_cluster_labels: + assert cluster_label in cluster_labels, f"{cluster_label} not in the cluster labels" + maskstack += [cluster_labels == cluster_label] + mask = np.squeeze(np.array(maskstack).sum(axis=0)).astype(bool) + return mask + +def apply_maskstack_to_Vector(pointsvector,maskstack): + """ + Applies a mask or stack of masks to a Vector to select one or more cluster components for further + analysis (e.g. plotting or statistical analysis). You could apply this to a Vector sampled from the + original with just some of the fields selected if you do not need the whole thing. + + Parameters + ---------- + pointsvector: Vector + The raw Vector that was run through clustering + maskstack: np.ndarray + A single mask or stack of masks selecting one or more clusters + Returns + ------- + maskstack: no.ndarray + + """ + assert isinstance(maskstack, np.ndarray), "the maskstack must be a numpy array" + assert maskstack.shape[-1] == pointsvector.flatten.shape[1], "the mask size does not match the Vector size" + if len(maskstack.shape) == 1: + return pointsvector.flatten()[maskstack] + else: + mask = maskstack.sum(axis=0).astype(bool) + return pointsvector.flatten()[mask] + +def Cluster_COMs_R(pointsvector, weighted=True): + """ + Calculates either real space centre of mass (weighted by intensity) or a simplified version with + no intensity from a specific cluster after running cluster analysis + with scikit.learn on a pointsarray + + Parameters + ---------- + pointsvector: Vector + The raw Vector that was run through L1 clustering. Must have a column giving the L1labels. + + Returns + ------- + COMs: np.ndarray + [COMx,COMy]xNclusters, shape=(N,2) + """ + assert "L1labels" in pointsvector.fields, "This Vector does not appear to have been clustered" + + rxy = pointsvector.select_fields("rx","ry").flatten() + I = pointsvector.select_fields("intensity").flatten() + L1labels = pointsvector.select_fields("L1labels").flatten().astype(int) + + L1_unique_labels = np.unique(L1labels)[1:] + + COMs = np.zeros_like(np.vstack((L1_unique_labels,L1_unique_labels)).T) + for n, label in enumerate(L1_unique_labels): + mask = np.squeeze(L1labels==label) + if weighted: + COMs[n] = (I * rxy)[mask].sum(axis=0) / I[mask].sum() + else: + COMs[n] = (rxy)[mask].sum(axis=0) / (rxy)[mask].shape[0] + return COMs + +def DBSCAN_L2( + pointsvector, + eps=5, + min_samples=2, + plotCOMs=True, + method='COMs', + printclustersizes=False +): + ''' + Groups L1 clusters into clusters via two different methods and then writes the Labels + into a new field ("L2labels") in the Vector, which can then be used to make images or + diffraction patterns + + The two methods are clustering real space centres of mass for L1 clusters, or clustering + image distances (1-image similarity) via the simple Jaccard metric. + + Parameters + ---------- + pointsvector: Vector + The raw Vector that was run through L1 clustering. Must have a column giving the L1labels. + eps: int, float + As defined by scikit-learn + min_samples: int + As defined by scikit-learn, wants at least 2 L1 clusters to correlate + plotCOMs: bool + True makes a plot if method is "COMs" + method: str + Either COMs or Jaccard + printclustersizes: bool + If True, then you see the cluster labels and sizes, which may help in tuning eps + Returns + ------- + ''' + assert method in ['COMs','Jaccard'], 'method currently restricted to COMs or Jaccard' + + if method == "COMs": + db2 = DBSCAN(eps=eps, min_samples=min_samples) + COMs = Cluster_COMs_R(pointsvector, weighted=True) + db2.fit(COMs) + + if method == "Jaccard": + L1labels = pointsvector.select_fields("L1labels").flatten().astype(int) + ims = [] + for cluster_label in np.unique(L1labels)[1:]: + mask = L1labels == cluster_label + ims += [DDFimage_from_mask(pointsvector,mask)] + imstack = np.array(ims) + corrs = jaccard_image_dist(imstack, plot=False) + db2 = DBSCAN(eps=eps, min_samples=min_samples,metric='precomputed') + db2.fit(corrs) + + L2_unique_labels, L2_all_cluster_sizes = np.unique(db2.labels_, return_counts=True) + L2_unique_labels_proper = L2_unique_labels[1:] + if printclustersizes: + print(L2_unique_labels, L2_all_cluster_sizes) + + L1labels = np.squeeze(pointsvector.select_fields("L1labels").flatten().astype(int)) + L1_unique_labels_proper = np.unique(L1labels)[1:] + + Rshape = pointsvector.shape + + L1toL2mapping = {-1:-2} + L2toL1mapping = {} + for L2cluster in L2_unique_labels: + L1labels_in_L2cluster = L1_unique_labels_proper[db2.labels_==L2cluster] + [L1toL2mapping.update({L1label: L2cluster}) for L1label in L1labels_in_L2cluster] + L2toL1mapping.update({L2cluster:L1labels_in_L2cluster}) + + L2labels = [L1toL2mapping[L1label] for L1label in L1labels] + if 'L2labels' in pointsvector.fields: + pointsvector.remove_fields('L2labels') + pointsvector.add_fields('L2labels',L2labels) + + if plotCOMs and method=='COMs': + fig,ax = plt.subplots(1,1, figsize=(12,12*Rshape[0]/Rshape[1])) + ax.set_ylim(Rshape[0],0) + ax.set_xlim(0,Rshape[1]) + for L2cluster in L2_unique_labels: + L1labels_in_L2cluster = L2toL1mapping[L2cluster] + chosenCOMs = COMs[L1labels_in_L2cluster] + ax.scatter( + chosenCOMs.T[1], + chosenCOMs.T[0], + cmap = california, + c=[L2cluster]*chosenCOMs.T[0].shape[0], + norm=Normalize(vmin=0, vmax=L2_unique_labels_proper.max(), clip=False), + rasterized=True + ) + if L2cluster!=-1: + ax.text( + chosenCOMs.T[1].mean(), + chosenCOMs.T[0].mean(), + str(L2cluster), + horizontalalignment='center', + verticalalignment='center', + fontsize=14, + path_effects = [ + path_effects.Stroke(linewidth=3, foreground='w'), + path_effects.Normal() + ] + ) + +def jaccard_image_dist(imsarray, plot=False): + + imsmask = (imsarray>1).astype(int) + + dists = np.zeros(shape=(imsarray.shape[0],imsarray.shape[0])) + masks = imsarray > 1 + + for i, mask in enumerate(masks): + either = (np.logical_or(masks,mask[np.newaxis,:,:])).sum(axis=(1,2)) + both = (masks*mask[np.newaxis,:,:]).sum(axis=(1,2)) + dists[i] = 1-both/either + + if plot: + plt.imshow(dists) + return dists \ No newline at end of file diff --git a/tests/datastructures/test_vector.py b/tests/datastructures/test_vector.py index 954059d6b..ef51dfb26 100644 --- a/tests/datastructures/test_vector.py +++ b/tests/datastructures/test_vector.py @@ -3,6 +3,7 @@ import numpy as np import pytest +from quantem.core.datastructures.dataset2d import Dataset2d from quantem.core.datastructures.vector import Vector from quantem.core.io.serialize import load @@ -111,6 +112,102 @@ def test_select_fields_and_chaining_equivalence(self): assert multi.total_rows == 6 assert multi.row_counts() == [2, 1, 2, 1] + def test_reductions_over_all_rows(self): + v = make_line_vector() + + np.testing.assert_allclose(v.sum(), np.array([21.0, 210.0, 2100.0])) + np.testing.assert_allclose(v.mean(), np.array([3.5, 35.0, 350.0])) + np.testing.assert_allclose(v.min(), np.array([1.0, 10.0, 100.0])) + np.testing.assert_allclose(v.max(), np.array([6.0, 60.0, 600.0])) + np.testing.assert_allclose(v.std(), np.std(v.flatten(), axis=0)) + assert v.count() == 6 + + # Field and fixed-grid selections narrow what is reduced + np.testing.assert_allclose(v.select_fields("kx").mean(), np.array([35.0])) + np.testing.assert_allclose(v[:2].sum(), np.array([6.0, 60.0, 600.0])) + assert v[:2].count() == 3 + + def test_reductions_per_cell(self): + v = make_line_vector() + + np.testing.assert_allclose( + v.sum(per_cell=True), + np.array( + [[3.0, 30.0, 300.0], [3.0, 30.0, 300.0], [9.0, 90.0, 900.0], [6.0, 60.0, 600.0]] + ), + ) + np.testing.assert_allclose( + v.mean(per_cell=True), + np.array( + [[1.5, 15.0, 150.0], [3.0, 30.0, 300.0], [4.5, 45.0, 450.0], [6.0, 60.0, 600.0]] + ), + ) + np.testing.assert_allclose(v.min(per_cell=True)[0], np.array([1.0, 10.0, 100.0])) + np.testing.assert_allclose(v.max(per_cell=True)[0], np.array([2.0, 20.0, 200.0])) + np.testing.assert_allclose( + v.select_fields("intensity").std(per_cell=True)[:, 0], + np.array([0.5, 0.0, 0.5, 0.0]), + ) + np.testing.assert_array_equal(v.count(per_cell=True), np.array([2, 1, 2, 1])) + + # Per-cell results keep the fixed-grid shape plus a trailing field axis + grid = make_grid_vector() + assert grid.sum(per_cell=True).shape == (3, 2, 3) + assert grid.count(per_cell=True).shape == (3, 2) + np.testing.assert_allclose(grid.max(per_cell=True)[2, 1], np.array([21.0, 121.0, 221.0])) + + def test_reductions_handle_empty_cells_and_selections(self): + v = Vector.from_shape(shape=(3,), fields=["intensity"]) + v[0] = np.array([[2.0], [4.0]]) + v[2] = np.array([[9.0]]) + + np.testing.assert_allclose(v.sum(per_cell=True)[:, 0], np.array([6.0, 0.0, 9.0])) + per_cell_mean = v.mean(per_cell=True)[:, 0] + np.testing.assert_allclose(per_cell_mean[[0, 2]], np.array([3.0, 9.0])) + assert np.isnan(per_cell_mean[1]) + assert np.isnan(v.min(per_cell=True)[1, 0]) + assert np.isnan(v.max(per_cell=True)[1, 0]) + assert np.isnan(v.std(per_cell=True)[1, 0]) + np.testing.assert_array_equal(v.count(per_cell=True), np.array([2, 0, 1])) + + # A selection with no rows at all + empty = v[1] + np.testing.assert_allclose(empty.sum(), np.array([0.0])) + assert np.isnan(empty.mean()).all() + assert empty.count() == 0 + + def test_reductions_as_dataset(self): + v = make_grid_vector() + + image = v.select_fields("intensity").max(per_cell=True, as_dataset=True) + assert isinstance(image, Dataset2d) + assert image.shape == (3, 2) + assert image.signal_units == "none" + assert "max" in image.name + np.testing.assert_allclose(image.array, np.array([[0.0, 1.0], [10.0, 11.0], [20.0, 21.0]])) + + counts = v.count(per_cell=True, as_dataset=True) + assert isinstance(counts, Dataset2d) + assert counts.signal_units == "counts" + np.testing.assert_array_equal(counts.array, np.ones((3, 2))) + + line = make_line_vector() + line_sum = line.select_fields("kx").sum(per_cell=True, as_dataset=True) + assert line_sum.shape == (4,) + assert line_sum.signal_units == "px" + + with pytest.raises(ValueError, match="exactly one selected field"): + v.max(per_cell=True, as_dataset=True) + + with pytest.raises(ValueError, match="requires per_cell=True"): + v.select_fields("intensity").max(as_dataset=True) + + with pytest.raises(ValueError, match="requires per_cell=True"): + v.count(as_dataset=True) + + with pytest.raises(ValueError, match="at least one fixed-grid axis"): + v[0, 0].select_fields("intensity").max(per_cell=True, as_dataset=True) + def test_array_mutation_writes_through_for_single_field(self): v = make_line_vector() cell = v.select_fields("kx")[1].array @@ -311,6 +408,7 @@ def test_empty_selection_is_valid_and_no_op_for_scalar_math(self): empty = v[[], :] assert empty.shape == (0, 2) assert empty.flatten().shape == (0, 3) + assert empty.copy().shape == (0, 2) empty.select_fields("kx")[...] += 1 np.testing.assert_array_equal(v.flatten(), before) @@ -371,6 +469,218 @@ def test_remove_fields_preserves_remaining_data(self): np.array([[1.0, 100.0], [2.0, 200.0]]), ) + def test_mask_empties_deselected_cells(self): + v = make_grid_vector() + + grid_mask = np.array([[True, False], [False, True], [True, True]]) + masked = v.mask(grid_mask) + + assert isinstance(masked, Vector) + assert masked.shape == v.shape + assert masked.fields == v.fields + assert masked.units == v.units + assert masked.name == v.name + assert masked.row_counts() == [1, 0, 0, 1, 1, 1] + np.testing.assert_array_equal(masked[0, 0].array, v[0, 0].array) + assert masked[0, 1].array.shape == (0, 3) + np.testing.assert_array_equal(masked[1, 1].array, v[1, 1].array) + + # The source Vector is untouched + assert v.row_counts() == [1] * 6 + + def test_mask_accepts_flat_and_integer_masks(self): + v = make_grid_vector() + grid_mask = np.array([[True, False], [False, True], [True, True]]) + + # A flat mask in row-major cell order matches the grid-shaped mask + np.testing.assert_array_equal( + v.mask(grid_mask.reshape(-1)).flatten(), + v.mask(grid_mask).flatten(), + ) + + # Integer masks are read as nonzero-means-keep + np.testing.assert_array_equal( + v.mask(grid_mask.astype(int)).flatten(), + v.mask(grid_mask).flatten(), + ) + + def test_mask_over_fixed_grid_dimensions(self): + # 1D + line = make_line_vector() + line_masked = line.mask(np.array([False, True, False, True])) + assert line_masked.shape == (4,) + assert line_masked.row_counts() == [0, 1, 0, 1] + np.testing.assert_array_equal( + line_masked.flatten(), + np.array([[3.0, 30.0, 300.0], [6.0, 60.0, 600.0]]), + ) + + # 0D, where the mask is a single boolean + assert line[0].mask(np.True_).array.shape == (2, 3) + assert line[0].mask(np.False_).array.shape == (0, 3) + + # 3D + cube = Vector.from_shape(shape=(2, 2, 2), fields=["kx", "ky"]) + for i in range(2): + for j in range(2): + for k in range(2): + cube[i, j, k] = np.array([[float(i), float(j + k)]]) + cube_mask = np.zeros((2, 2, 2), dtype=bool) + cube_mask[1, 0, 1] = True + cube_masked = cube.mask(cube_mask) + assert cube_masked.shape == (2, 2, 2) + assert cube_masked.total_rows == 1 + np.testing.assert_array_equal(cube_masked[1, 0, 1].array, np.array([[1.0, 1.0]])) + + def test_mask_on_field_and_grid_selections(self): + v = make_grid_vector() + + # Masking a field-selected view keeps only that field, like copy() + kx_masked = v.select_fields("kx").mask(np.array([[True, False]] * 3)) + assert kx_masked.fields == ["kx"] + np.testing.assert_array_equal(kx_masked.flatten(), np.array([[100.0], [110.0], [120.0]])) + + # Masking a fixed-grid selection is relative to that selection's shape + sub = v[:2] + sub_masked = sub.mask(np.array([[True, True], [False, False]])) + assert sub_masked.shape == (2, 2) + assert sub_masked.row_counts() == [1, 1, 0, 0] + + def test_mask_in_place_empties_cells_across_all_fields(self): + v = make_grid_vector() + + assert v.mask(np.array([[True, False], [True, False], [True, False]])) is not None + assert ( + v.mask(np.array([[True, False], [True, False], [True, False]]), modify_in_place=True) + is None + ) + assert v.shape == (3, 2) + assert v.row_counts() == [1, 0, 1, 0, 1, 0] + np.testing.assert_array_equal( + v.flatten(), + np.array([[0.0, 100.0, 200.0], [10.0, 110.0, 210.0], [20.0, 120.0, 220.0]]), + ) + + # Cells are emptied across every field, even through a field-selected view + v2 = make_grid_vector() + v2.select_fields("kx").mask(np.zeros((3, 2), dtype=bool), modify_in_place=True) + assert v2.fields == ["intensity", "kx", "ky"] + assert v2.row_counts() == [0] * 6 + + # In-place masking of a grid selection leaves unselected cells alone + v3 = make_grid_vector() + v3[0].mask(np.array([False, True]), modify_in_place=True) + assert v3.row_counts() == [0, 1, 1, 1, 1, 1] + + def test_filter_rows_keeps_selected_rows(self): + v = make_line_vector() + + intensity = v.select_fields("intensity").flatten() + filtered = v.filter_rows(intensity > 3.0) + + assert isinstance(filtered, Vector) + assert filtered.shape == v.shape + assert filtered.fields == v.fields + assert filtered.units == v.units + assert filtered.row_counts() == [0, 0, 2, 1] + np.testing.assert_array_equal( + filtered.flatten(), + np.array([[4.0, 40.0, 400.0], [5.0, 50.0, 500.0], [6.0, 60.0, 600.0]]), + ) + # The source Vector is untouched + assert v.row_counts() == [2, 1, 2, 1] + + # (n_rows, 1) and 1D masks are equivalent, as are integer masks + np.testing.assert_array_equal( + v.filter_rows((intensity > 3.0)[:, 0]).flatten(), filtered.flatten() + ) + np.testing.assert_array_equal( + v.filter_rows(np.array([0, 0, 0, 1, 1, 1])).flatten(), filtered.flatten() + ) + + # A single-field Vector mask works too + np.testing.assert_array_equal( + v.filter_rows(np.greater(v.select_fields("intensity"), 3.0)).flatten(), + filtered.flatten(), + ) + + def test_filter_rows_in_place_and_on_selections(self): + v = make_line_vector() + + kr = v.select_fields("ky").flatten()[:, 0] + assert v.filter_rows((kr > 150.0) & (kr < 550.0), modify_in_place=True) is None + assert v.row_counts() == [1, 1, 2, 0] + np.testing.assert_array_equal(v[0].array, np.array([[2.0, 20.0, 200.0]])) + + # Rows drop across all fields even when the mask came from a field view + v2 = make_line_vector() + kx = v2.select_fields("kx") + kx.filter_rows(kx.flatten() < 45.0, modify_in_place=True) + assert v2.fields == ["intensity", "kx", "ky"] + assert v2.row_counts() == [2, 1, 1, 0] + np.testing.assert_array_equal(v2[2].array, np.array([[4.0, 40.0, 400.0]])) + + # Filtering a field-selected view returns only that field, like copy() + kx_only = make_line_vector().select_fields("kx") + kx_filtered = kx_only.filter_rows(kx_only.flatten() >= 40.0) + assert kx_filtered.fields == ["kx"] + np.testing.assert_array_equal(kx_filtered.flatten(), np.array([[40.0], [50.0], [60.0]])) + + # A fixed-grid selection only sees its own rows, and leaves the rest alone + v3 = make_line_vector() + v3[:2].filter_rows(np.array([False, True, True]), modify_in_place=True) + assert v3.row_counts() == [1, 1, 2, 1] + np.testing.assert_array_equal(v3[0].array, np.array([[2.0, 20.0, 200.0]])) + + def test_filter_rows_edge_cases_and_validation(self): + v = make_line_vector() + + np.testing.assert_array_equal(v.filter_rows(np.ones(6, dtype=bool)).flatten(), v.flatten()) + + drop_all = v.filter_rows(np.zeros(6, dtype=bool)) + assert drop_all.row_counts() == [0, 0, 0, 0] + assert drop_all.flatten().shape == (0, 3) + + empty = v[[]] + assert empty.filter_rows(np.array([], dtype=bool)).flatten().shape == (0, 3) + + with pytest.raises(ValueError, match="expected 6 rows"): + v.filter_rows(np.ones(5, dtype=bool)) + + with pytest.raises(TypeError, match="boolean or integer"): + v.filter_rows(np.ones(6, dtype=float)) + + with pytest.raises(ValueError, match="Reduce multi-column masks"): + v.filter_rows(np.ones((6, 3), dtype=bool)) + + with pytest.raises(ValueError, match="exactly one field"): + v.filter_rows(np.greater(v.select_fields("intensity", "kx"), 3.0)) + + with pytest.raises(ValueError, match="matching per-cell row counts"): + v.filter_rows(np.greater(v[:2].select_fields("intensity"), 3.0)) + + def test_mask_edge_cases_and_validation(self): + v = make_grid_vector() + + keep_all = v.mask(np.ones((3, 2), dtype=bool)) + np.testing.assert_array_equal(keep_all.flatten(), v.flatten()) + + drop_all = v.mask(np.zeros((3, 2), dtype=bool)) + assert drop_all.row_counts() == [0] * 6 + assert drop_all.flatten().shape == (0, 3) + + empty = v[[], :] + assert empty.mask(np.zeros((0, 2), dtype=bool)).flatten().shape == (0, 3) + + with pytest.raises(ValueError, match=r"expected \(3, 2\)"): + v.mask(np.ones((2, 3), dtype=bool)) + + with pytest.raises(ValueError, match="flat mask with 6 entries"): + v.mask(np.ones(5, dtype=bool)) + + with pytest.raises(TypeError, match="boolean or integer"): + v.mask(np.ones((3, 2), dtype=float)) + def test_copy_is_deep(self): v = make_line_vector() v_copy = v.select_fields(["intensity", "kx"]).copy()