From a8fb4c2382ea28152d4ec27bacaedee4046cca6f Mon Sep 17 00:00:00 2001 From: maclariz Date: Thu, 30 Jul 2026 10:17:34 -0700 Subject: [PATCH 01/14] Digital Dark Field Digital Dark Field functions introduced --- src/quantem/core/datastructures/vector.py | 57 +++++++ src/quantem/diffraction/__init__.py | 3 +- src/quantem/diffraction/digital_dark_field.py | 159 ++++++++++++++++++ 3 files changed, 218 insertions(+), 1 deletion(-) create mode 100644 src/quantem/diffraction/digital_dark_field.py diff --git a/src/quantem/core/datastructures/vector.py b/src/quantem/core/datastructures/vector.py index aa2627d27..5f53fce84 100644 --- a/src/quantem/core/datastructures/vector.py +++ b/src/quantem/core/datastructures/vector.py @@ -7,6 +7,8 @@ import numpy as np from numpy.typing import NDArray +from tqdm import tqdm + from quantem.core.io.serialize import AutoSerialize from quantem.core.utils.validators import ( validate_fields, @@ -978,6 +980,61 @@ 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 + # ------------------------------------------------------------------ # + # Digital Dark Field Enabling + # ------------------------------------------------------------------ # + +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 _resolve_fields( fields: Sequence[str] | None, diff --git a/src/quantem/diffraction/__init__.py b/src/quantem/diffraction/__init__.py index 3cd0027f8..aa01328a8 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 import * \ No newline at end of file diff --git a/src/quantem/diffraction/digital_dark_field.py b/src/quantem/diffraction/digital_dark_field.py new file mode 100644 index 000000000..a3a6b680f --- /dev/null +++ b/src/quantem/diffraction/digital_dark_field.py @@ -0,0 +1,159 @@ +import numpy as np + +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 + ) Date: Thu, 30 Jul 2026 20:56:03 -0700 Subject: [PATCH 02/14] updating DDF functions Restructured and some cluster support --- src/quantem/core/datastructures/vector.py | 55 --- src/quantem/diffraction/digital_dark_field.py | 159 ------- .../diffraction/digital_dark_field_cluster.py | 413 ++++++++++++++++++ 3 files changed, 413 insertions(+), 214 deletions(-) delete mode 100644 src/quantem/diffraction/digital_dark_field.py create mode 100644 src/quantem/diffraction/digital_dark_field_cluster.py diff --git a/src/quantem/core/datastructures/vector.py b/src/quantem/core/datastructures/vector.py index 5f53fce84..9724113eb 100644 --- a/src/quantem/core/datastructures/vector.py +++ b/src/quantem/core/datastructures/vector.py @@ -980,61 +980,6 @@ 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 - # ------------------------------------------------------------------ # - # Digital Dark Field Enabling - # ------------------------------------------------------------------ # - -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 _resolve_fields( fields: Sequence[str] | None, diff --git a/src/quantem/diffraction/digital_dark_field.py b/src/quantem/diffraction/digital_dark_field.py deleted file mode 100644 index a3a6b680f..000000000 --- a/src/quantem/diffraction/digital_dark_field.py +++ /dev/null @@ -1,159 +0,0 @@ -import numpy as np - -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 - )=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 + ) Date: Thu, 30 Jul 2026 20:59:30 -0700 Subject: [PATCH 03/14] Update __init__.py --- src/quantem/diffraction/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/quantem/diffraction/__init__.py b/src/quantem/diffraction/__init__.py index aa01328a8..0eecd61e7 100644 --- a/src/quantem/diffraction/__init__.py +++ b/src/quantem/diffraction/__init__.py @@ -2,4 +2,4 @@ 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 -from quantem.diffraction.digital_dark_field import * \ No newline at end of file +from quantem.diffraction.digital_dark_field_cluster import * \ No newline at end of file From 24f1404b2d8332701ef4e0ae98c85a9bb81765cc Mon Sep 17 00:00:00 2001 From: maclariz Date: Thu, 30 Jul 2026 21:08:29 -0700 Subject: [PATCH 04/14] Update digital_dark_field_cluster.py missed import --- src/quantem/diffraction/digital_dark_field_cluster.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/quantem/diffraction/digital_dark_field_cluster.py b/src/quantem/diffraction/digital_dark_field_cluster.py index 30da424a5..de3791438 100644 --- a/src/quantem/diffraction/digital_dark_field_cluster.py +++ b/src/quantem/diffraction/digital_dark_field_cluster.py @@ -6,6 +6,8 @@ from sklearn.cluster import DBSCAN +from quantem.core.datastructures.vector import Vector + # ------------------------------------------------------------------ # # Digital Dark Field Basics # ------------------------------------------------------------------ # From 7fd35ba4bed093fbdedd475e5a03fda91cc27a53 Mon Sep 17 00:00:00 2001 From: maclariz Date: Thu, 30 Jul 2026 21:25:20 -0700 Subject: [PATCH 05/14] Update digital_dark_field_cluster.py missed import --- src/quantem/diffraction/digital_dark_field_cluster.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/quantem/diffraction/digital_dark_field_cluster.py b/src/quantem/diffraction/digital_dark_field_cluster.py index de3791438..f268a4aae 100644 --- a/src/quantem/diffraction/digital_dark_field_cluster.py +++ b/src/quantem/diffraction/digital_dark_field_cluster.py @@ -4,6 +4,8 @@ from matplotlib.colors import LinearSegmentedColormap, Normalize, PowerNorm import matplotlib.gridspec as GridSpec +from tqdm import tqdm + from sklearn.cluster import DBSCAN from quantem.core.datastructures.vector import Vector From e3f6bcd5a2c6a32764bb8525f32fdf3d997f2214 Mon Sep 17 00:00:00 2001 From: maclariz Date: Thu, 30 Jul 2026 21:31:54 -0700 Subject: [PATCH 06/14] Update digital_dark_field_cluster.py Double import --- src/quantem/diffraction/digital_dark_field_cluster.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/src/quantem/diffraction/digital_dark_field_cluster.py b/src/quantem/diffraction/digital_dark_field_cluster.py index f268a4aae..a5158738c 100644 --- a/src/quantem/diffraction/digital_dark_field_cluster.py +++ b/src/quantem/diffraction/digital_dark_field_cluster.py @@ -221,9 +221,6 @@ def DDFimage_from_maskstack(pointsvector,maskstack): im[rx,ry]+=I return im - import numpy as np -import matplotlib.pyplot as plt -from matplotlib.colors import LinearSegmentedColormap, Normalize, PowerNorm # ------------------------------------------------------------------ # # Clustering Functions From fb97b10887c1260fe3c9630752c2513ff6f5eef5 Mon Sep 17 00:00:00 2001 From: maclariz Date: Thu, 30 Jul 2026 21:41:39 -0700 Subject: [PATCH 07/14] Update digital_dark_field_cluster.py bugfix (relics of old code) --- .../diffraction/digital_dark_field_cluster.py | 14 ++++++-------- 1 file changed, 6 insertions(+), 8 deletions(-) diff --git a/src/quantem/diffraction/digital_dark_field_cluster.py b/src/quantem/diffraction/digital_dark_field_cluster.py index a5158738c..55703c41b 100644 --- a/src/quantem/diffraction/digital_dark_field_cluster.py +++ b/src/quantem/diffraction/digital_dark_field_cluster.py @@ -365,7 +365,7 @@ def plot_L1_clusters_kspace(L1labels, pointsvector, fields, max_kr, cmap=califor ) def show_L1_clusters_in_real_space( - L1labels, pointsvector, col=5, gamma=0.25, cmapname='inferno' + L1labels, pointsvector, ncols=5, gamma=0.25, cmapname='inferno' ): """ Function to show real space plots of L1 clustering outputs @@ -376,7 +376,7 @@ def show_L1_clusters_in_real_space( The labels list from a clustering algorithm pointsvector: Vector A points array, as defined in py4DSTEM.process.diffraction.digital_dark_field - cols: int + ncols: int number of columns to be used gamma: float Image gamma. <1 boosts lower intensities in display. @@ -390,17 +390,15 @@ def show_L1_clusters_in_real_space( l = cluster_list.shape[0] ar = shape[1] / shape[0] w = 10 - row = int(np.ceil(l / col)) - fig = plt.figure(figsize=(w, w * row / col / ar)) - gs = GridSpec.GridSpec(row, col) + row = int(np.ceil(l / ncols)) + fig = plt.figure(figsize=(w, w * row / ncols / ar)) + gs = GridSpec.GridSpec(row, ncols) for n, cluster_label in enumerate(cluster_list): - i, j = int(n / col), n % col + i, j = int(n / ncols), n % ncols ax = plt.subplot(gs[i, j]) ax.set_axis_off() mask = L1labels == cluster_label - DDFimage_from_maskstack - selpoints = pointsarray[L1labels == cluster_label] im = DDFimage_from_maskstack(pointsvector,mask[None,:]) ax.imshow(im, norm=colors.PowerNorm(gamma=gamma), cmap=cmapname) ax.text( From a84cd57e21202ec118163150184c7d620f3b1fcb Mon Sep 17 00:00:00 2001 From: maclariz Date: Thu, 30 Jul 2026 21:49:08 -0700 Subject: [PATCH 08/14] Update digital_dark_field_cluster.py Added the function to read pointsarrays from py4DSTEM and fixed a minor bug in cluster real space plot --- .../diffraction/digital_dark_field_cluster.py | 56 ++++++++++++++++++- 1 file changed, 54 insertions(+), 2 deletions(-) diff --git a/src/quantem/diffraction/digital_dark_field_cluster.py b/src/quantem/diffraction/digital_dark_field_cluster.py index 55703c41b..6ccc2031d 100644 --- a/src/quantem/diffraction/digital_dark_field_cluster.py +++ b/src/quantem/diffraction/digital_dark_field_cluster.py @@ -11,7 +11,7 @@ from quantem.core.datastructures.vector import Vector # ------------------------------------------------------------------ # - # Digital Dark Field Basics + # Create suitable Vector object # ------------------------------------------------------------------ # def make_FullPointsVector_centres(vecs,centers): @@ -65,6 +65,58 @@ def make_FullPointsVector_centres(vecs,centers): )) return pointsvector +def make_FullPointsVector_from_pointsarray(pointsarray): + ''' + For back compatibility, this reads in pointsarray objects made with py4DSTEM + digital dark field + + 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, @@ -400,7 +452,7 @@ def show_L1_clusters_in_real_space( mask = L1labels == cluster_label im = DDFimage_from_maskstack(pointsvector,mask[None,:]) - ax.imshow(im, norm=colors.PowerNorm(gamma=gamma), cmap=cmapname) + ax.imshow(im, norm=PowerNorm(gamma=gamma), cmap=cmapname) ax.text( 5, 5, From 7124dcc20aa142fac77aad2728599ed33320d27d Mon Sep 17 00:00:00 2001 From: maclariz Date: Sun, 2 Aug 2026 10:01:27 -0500 Subject: [PATCH 09/14] Update digital_dark_field_cluster.py Now writing cluster labels into the vector at L1 --- .../diffraction/digital_dark_field_cluster.py | 289 ++++++++++++++++-- 1 file changed, 263 insertions(+), 26 deletions(-) diff --git a/src/quantem/diffraction/digital_dark_field_cluster.py b/src/quantem/diffraction/digital_dark_field_cluster.py index 6ccc2031d..c8aa90013 100644 --- a/src/quantem/diffraction/digital_dark_field_cluster.py +++ b/src/quantem/diffraction/digital_dark_field_cluster.py @@ -69,7 +69,7 @@ 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 @@ -241,6 +241,49 @@ def DDFpointsmask(pointsvector,selectionpoints,tolerance): ) return maskstack +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 + ------- + 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. + ''' + 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 = pointsvector2.select_fields("kr").flatten() + radialmask = np.squeeze(np.logical_and(kr >= min_kr, kr <= max_kr)) + + # It's pretty easy to trim either the selected fields or a whole flattened array + pointsarray = (np.array(scaling)*pointsvector.select_fields(*fields).flatten())[radialmask] + + + db = DBSCAN(eps=eps, min_samples=min_samples) + db.fit(pointsarray) + + # but how do we rebuild the vector easily from the filtered flattened version + if 'L1labels' in pointsvector.fields: + pointsvector.remove_fields('L1labels') + pointsvector.add_fields('L1labels',db.labels_) + if plot: + plot_L1_clusters_kspace( + pointsvector, + fields, + max_kr_plot=int(pointsvector.select_fields('kx').flatten().max()*1.1) ) - return db.labels_ ''' A custom colormap for the k-space plots ''' - california = LinearSegmentedColormap.from_list( 'cali', [ @@ -341,7 +470,7 @@ def DBSCAN_pointsvector(pointsvector, fields=['kx','ky'], scaling = [1,1], eps=0 california.set_under('lightgrey') california.set_bad('red') -def plot_L1_clusters_kspace(L1labels, pointsvector, fields, max_kr, cmap=california, figax=None): +def plot_L1_clusters_kspace(pointsvector, fields, max_kr_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 @@ -371,7 +500,7 @@ def plot_L1_clusters_kspace(L1labels, pointsvector, fields, max_kr, cmap=califor """ if figax is None: - fig, ax = plt.subplots(1, 1, figsize=(6, 6)) + fig, ax = plt.subplots(1, 1, figsize=(6, 6)) else: fig, ax = figax assert isinstance(fig, Figure) @@ -380,16 +509,16 @@ def plot_L1_clusters_kspace(L1labels, pointsvector, fields, max_kr, cmap=califor ax.set_title("DBSCAN "+", ".join(fields)) ax.set_xlabel("kx", fontsize=24) ax.set_ylabel("ky", fontsize=24) - ax.set_ylim(max_kr, -max_kr) - ax.set_xlim(-max_kr, max_kr) - - uniquelabels = np.unique(L1labels) + ax.set_ylim(max_kr_plot, -max_kr_plot) + ax.set_xlim(-max_kr_plot, max_kr_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, @@ -400,16 +529,17 @@ def plot_L1_clusters_kspace(L1labels, pointsvector, fields, max_kr, cmap=califor rasterized=True ) - for label in np.unique(L1labels): - maxint = np.argmax(pointsvector.select_fields("intensity").flatten()[L1labels==label]) - r = kr[L1labels==label][maxint][0] + 6 - ang = np.radians(kphi[L1labels==label][maxint])[0] + for label in uniquelabels: + clustermask = L1labels==label + maxint = np.argmax(I[clustermask]) + r = kr[clustermask][maxint] + 6 + ang = np.radians(kphi[clustermask][maxint]) labx = np.sin(ang) * r laby = np.cos(ang) * r ax.annotate( label, - (ky[L1labels==label][maxint][0], kx[L1labels==label][maxint][0]), + (ky[clustermask][maxint], kx[clustermask][maxint]), (laby, -labx), horizontalalignment="center", verticalalignment="center", @@ -417,10 +547,14 @@ def plot_L1_clusters_kspace(L1labels, pointsvector, fields, max_kr, cmap=califor ) def show_L1_clusters_in_real_space( - L1labels, pointsvector, ncols=5, gamma=0.25, cmapname='inferno' + pointsvector, ncols=5, gamma=0.25, cmapname='inferno', ordering='sequential' ): """ - Function to show real space plots of L1 clustering outputs + 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, Parameters ---------- @@ -434,11 +568,20 @@ def show_L1_clusters_in_real_space( 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 Returns ------- """ + assert ordering in ["sequential", "size"], "ordering must be either sequential or size" + L1labels = pointsvector.select_fields("L1labels").flatten().astype(int) + L1_unique_labels, L1_all_cluster_sizes = np.unique(L1labels, return_counts=True) shape = pointsvector.shape - cluster_list = np.unique(L1labels)[1:] + if ordering == "sequential": + cluster_list = L1_unique_labels[1:] + elif ordering == "size": + cluster_list = L1_unique_labels[1:][np.argsort(L1_all_cluster_sizes[1:])[::-1]] l = cluster_list.shape[0] ar = shape[1] / shape[0] w = 10 @@ -461,4 +604,98 @@ def show_L1_clusters_in_real_space( size=14, fontweight="bold", verticalalignment="top", - ) \ No newline at end of file + ) + +def cluster_mask(cluster_labels, selected_cluster_labels): + """ + 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_maskstack". + + 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 + ------- + maskstack: np.ndarray + + """ + for cluster_label in selected_cluster_labels: + assert cluster_label in cluster_labels, f"{cluster_label} not in the cluster labels" + maskstack = (cluster_labels in selected_cluster_labels) + return maskstack + +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").flatten() + ry = pointsvector.select_fields("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, +# method = "COM", +# eps=3, +# min_samples=2, +# # plot=True +# ): + +# COMs = Cluster_COMs_R(pointsvector, weighted=True): +# db2 = DBSCAN(eps=eps, min_samples=ms) +# db2.fit(COMs) + From 0aa5de6ca6825bb36cbe53b9e3f4ce327c29b35b Mon Sep 17 00:00:00 2001 From: cophus Date: Sun, 2 Aug 2026 10:56:28 -0500 Subject: [PATCH 10/14] Vector additional methods --- src/quantem/core/datastructures/vector.py | 407 +++++++++++++++++++++- 1 file changed, 396 insertions(+), 11 deletions(-) diff --git a/src/quantem/core/datastructures/vector.py b/src/quantem/core/datastructures/vector.py index 9724113eb..e48549f79 100644 --- a/src/quantem/core/datastructures/vector.py +++ b/src/quantem/core/datastructures/vector.py @@ -2,13 +2,11 @@ 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 -from tqdm import tqdm - from quantem.core.io.serialize import AutoSerialize from quantem.core.utils.validators import ( validate_fields, @@ -17,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. @@ -90,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 @@ -313,6 +326,182 @@ 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") + # ------------------------------------------------------------------ # # Field management # ------------------------------------------------------------------ # @@ -485,6 +674,89 @@ 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 + -------- + 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, work through + ``flatten()`` and ``set_flattened()``. + + 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 compact(self) -> None: """Repack the backing row buffer to remove dead rows. @@ -535,14 +807,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() @@ -751,6 +1016,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: @@ -796,6 +1083,55 @@ 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_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. @@ -1080,6 +1416,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): From 2ce5b36e00bd1012fc09a9d875c311e81a6db9ab Mon Sep 17 00:00:00 2001 From: cophus Date: Sun, 2 Aug 2026 10:56:37 -0500 Subject: [PATCH 11/14] Vector tests --- tests/datastructures/test_vector.py | 223 ++++++++++++++++++++++++++++ 1 file changed, 223 insertions(+) diff --git a/tests/datastructures/test_vector.py b/tests/datastructures/test_vector.py index 954059d6b..e4fa07048 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,131 @@ 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_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() From a9b94b2535da3e891bbbbc2d9f7cde9b33824eab Mon Sep 17 00:00:00 2001 From: cophus Date: Sun, 2 Aug 2026 11:09:58 -0500 Subject: [PATCH 12/14] Adding filter_rows method --- src/quantem/core/datastructures/vector.py | 111 +++++++++++++++++++++- tests/datastructures/test_vector.py | 87 +++++++++++++++++ 2 files changed, 196 insertions(+), 2 deletions(-) diff --git a/src/quantem/core/datastructures/vector.py b/src/quantem/core/datastructures/vector.py index e48549f79..85f18e808 100644 --- a/src/quantem/core/datastructures/vector.py +++ b/src/quantem/core/datastructures/vector.py @@ -714,13 +714,13 @@ def mask(self, mask: Any, modify_in_place: bool = False) -> "Vector | None": 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, work through - ``flatten()`` and ``set_flattened()``. + field value, e.g. peaks below an intensity threshold, use ``filter_rows``. Examples -------- @@ -757,6 +757,81 @@ def mask(self, mask: Any, modify_in_place: bool = False) -> "Vector | None": 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. @@ -1119,6 +1194,38 @@ def _as_dataset(self, array: NDArray[Any], label: str, signal_units: str) -> "Da 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) diff --git a/tests/datastructures/test_vector.py b/tests/datastructures/test_vector.py index e4fa07048..ef51dfb26 100644 --- a/tests/datastructures/test_vector.py +++ b/tests/datastructures/test_vector.py @@ -572,6 +572,93 @@ def test_mask_in_place_empties_cells_across_all_fields(self): 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() From 9c1596416c7bdc2c2b2182d0c711b5436e74b7da Mon Sep 17 00:00:00 2001 From: maclariz Date: Mon, 3 Aug 2026 11:43:31 -0500 Subject: [PATCH 13/14] imgreduce Totally changes how we make images from the vector object. This really will speed up when we make complex selections in objects since we now just select with one mask, not a stack. --- src/quantem/core/datastructures/vector.py | 19 ++ .../diffraction/digital_dark_field_cluster.py | 234 ++++++++++-------- 2 files changed, 149 insertions(+), 104 deletions(-) diff --git a/src/quantem/core/datastructures/vector.py b/src/quantem/core/datastructures/vector.py index 85f18e808..8a82568bd 100644 --- a/src/quantem/core/datastructures/vector.py +++ b/src/quantem/core/datastructures/vector.py @@ -502,6 +502,25 @@ def count(self, per_cell: bool = False, as_dataset: bool = False) -> "int | NDAr 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 # ------------------------------------------------------------------ # diff --git a/src/quantem/diffraction/digital_dark_field_cluster.py b/src/quantem/diffraction/digital_dark_field_cluster.py index c8aa90013..391292d70 100644 --- a/src/quantem/diffraction/digital_dark_field_cluster.py +++ b/src/quantem/diffraction/digital_dark_field_cluster.py @@ -284,7 +284,7 @@ def DDFrphimask(pointsvector,r,rtol,phi=None,phitol=None): maskstack = radial_selection return maskstack.T -def DDFimage_from_maskstack(pointsvector,maskstack): +def DDFimage_from_mask(pointsvector,mask): ''' This calculates a DDF image from a Boolean maskstack object, which is a set of Boolean mask layers. Each mask layer may contribute to some of the same pixels, so calculating each separately is necessary. @@ -306,14 +306,8 @@ def DDFimage_from_maskstack(pointsvector,maskstack): im: np.ndarray The Digital Dark Field Image ''' - im = np.zeros(shape=pointsvector.shape) - if len(maskstack.shape) == 1: - maskstack = maskstack[None,:] - for mask in maskstack: - rx = pointsvector.select_fields("rx").flatten()[mask].astype(int) - ry = pointsvector.select_fields("ry").flatten()[mask].astype(int) - I = pointsvector.select_fields("intensity").flatten()[mask] - im[rx,ry]+=I + cut = pointsvector.filter_rows(mask) + im = cut.imgreduce() return im @@ -321,71 +315,15 @@ def DDFimage_from_maskstack(pointsvector,maskstack): # Clustering Functions # ------------------------------------------------------------------ # -def DBSCAN_pointsvector( - pointsvector, - fields=['kx','ky'], - scaling = [1,1], - eps=0.5, - min_samples=20, - plot=True -): - ''' - This is the working code that DBSCANs everything - - Runs DBSCAN on selected fields in a pointsvector - See scikit-learn documentation for general comments on their implementation of the DBSCAN function - Experience suggests about eps should be about 0.3-0.5 for detecting diffraction spots in kx,ky 2D - clustering, and about 1 will connect arcs/rings of spots for nanocrystalline / amorphous materials. - Too small and you see no clusters at all. - For 4D rx,ry,kx,ky clustering, eps needs to be larger, perhaps 3-10, depending on your scaling parameters. - Alter the relative weighting of real and reciprocal space depending on your dataset and the size of your - crystals in real space compared to the spacing of diffraction peaks in reciprocal space. - - Parameters - ---------- - pointsvector: Vector - Should contain any fields you are selecting to cluster on - fields: list of str - Strings in ["rx", "ry", 'kx', "ky", "kr", "kphi"] to cluster on - scaling: list of int, float - Relative scaling factors for different dimensions - eps: float - As defined in scikit-learn - min_samples: int - As defined in scikit-learn - 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" - pointsarray = (np.array(scaling)*pointsvector.select_fields(*fields).flatten()) - db = DBSCAN(eps=eps, min_samples=min_samples) - db.fit(pointsarray) - if 'L1labels' in pointsvector.fields: - pointsvector.remove_fields('L1labels') - pointsvector.add_fields('L1labels',db.labels_) - if plot: - plot_L1_clusters_kspace( - pointsvector, - fields, - max_kr_plot=int(pointsvector.select_fields('kx').flatten().max()*1.1) - ) -def incomplete_radial_filtering( +def DBSCAN_pointsvector( pointsvector, fields=['kx','ky'], scaling = [1,1], eps=0.5, min_samples=20, - min_kr = 0, - max_kr = 1000, + kr_min = 0, + kr_max = 1000, plot=True ): ''' @@ -410,7 +348,7 @@ def incomplete_radial_filtering( As defined in scikit-learn min_samples: int As defined in scikit-learn - min_kr, max_kr: int, float + kr_min, kr_max: int, float minimum and maximum peak radii to use for clustering. Setting min_kr>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 @@ -431,26 +369,20 @@ def incomplete_radial_filtering( pointsvector2 = pointsvector.copy() # making the mask is obvious - kr = pointsvector2.select_fields("kr").flatten() - radialmask = np.squeeze(np.logical_and(kr >= min_kr, kr <= max_kr)) - - # It's pretty easy to trim either the selected fields or a whole flattened array - pointsarray = (np.array(scaling)*pointsvector.select_fields(*fields).flatten())[radialmask] - + 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) - - # but how do we rebuild the vector easily from the filtered flattened version - if 'L1labels' in pointsvector.fields: - pointsvector.remove_fields('L1labels') - pointsvector.add_fields('L1labels',db.labels_) + pointsvector2.add_fields('L1labels',db.labels_) if plot: plot_L1_clusters_kspace( - pointsvector, + pointsvector2, fields, - max_kr_plot=int(pointsvector.select_fields('kx').flatten().max()*1.1) + kr_max_plot=int(pointsvector2.select_fields('kx').flatten().max()*1.05) ) + return pointsvector2 ''' A custom colormap for the k-space plots @@ -470,7 +402,7 @@ def incomplete_radial_filtering( california.set_under('lightgrey') california.set_bad('red') -def plot_L1_clusters_kspace(pointsvector, fields, max_kr_plot, cmap=california, figax=None): +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 @@ -488,7 +420,7 @@ def plot_L1_clusters_kspace(pointsvector, fields, max_kr_plot, cmap=california, pointsvector: Vector A Vector object from this repo, preferably constructed with ["rx", "ry", 'kx', "ky", "kr", "kphi", "I] as the fields - max_kr: int, float + 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 @@ -509,8 +441,8 @@ def plot_L1_clusters_kspace(pointsvector, fields, max_kr_plot, cmap=california, ax.set_title("DBSCAN "+", ".join(fields)) ax.set_xlabel("kx", fontsize=24) ax.set_ylabel("ky", fontsize=24) - ax.set_ylim(max_kr_plot, -max_kr_plot) - ax.set_xlim(-max_kr_plot, max_kr_plot) + 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() @@ -529,10 +461,10 @@ def plot_L1_clusters_kspace(pointsvector, fields, max_kr_plot, cmap=california, rasterized=True ) - for label in uniquelabels: + for label in uniquelabels[1:]: clustermask = L1labels==label maxint = np.argmax(I[clustermask]) - r = kr[clustermask][maxint] + 6 + r = kr[clustermask][maxint] + 3 ang = np.radians(kphi[clustermask][maxint]) labx = np.sin(ang) * r @@ -543,18 +475,19 @@ def plot_L1_clusters_kspace(pointsvector, fields, max_kr_plot, cmap=california, (laby, -labx), horizontalalignment="center", verticalalignment="center", - size=7, + size=8, ) def show_L1_clusters_in_real_space( - pointsvector, ncols=5, gamma=0.25, cmapname='inferno', ordering='sequential' + pointsvector, ncols=5, gamma=0.25, cmapname='inferno', ordering='sequential', save_ims=False ): """ 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, + 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 ---------- @@ -571,8 +504,12 @@ def show_L1_clusters_in_real_space( 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 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" L1labels = pointsvector.select_fields("L1labels").flatten().astype(int) @@ -582,12 +519,20 @@ def show_L1_clusters_in_real_space( cluster_list = L1_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]) @@ -605,6 +550,10 @@ def show_L1_clusters_in_real_space( fontweight="bold", verticalalignment="top", ) + if save_ims: + ims+=[im] + if save_ims: + return np.array(ims) def cluster_mask(cluster_labels, selected_cluster_labels): """ @@ -672,8 +621,7 @@ def Cluster_COMs_R(pointsvector, weighted=True): """ assert "L1labels" in pointsvector.fields, "This Vector does not appear to have been clustered" - rxy = pointsvector.select_fields("rx").flatten() - ry = pointsvector.select_fields("ry").flatten() + rxy = pointsvector.select_fields("rx","ry").flatten() I = pointsvector.select_fields("intensity").flatten() L1labels = pointsvector.select_fields("L1labels").flatten().astype(int) @@ -681,21 +629,99 @@ def Cluster_COMs_R(pointsvector, weighted=True): 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) + mask = np.squeeze(L1labels==label) if weighted: COMs[n] = (I * rxy)[mask].sum(axis=0) / I[mask].sum() - else: + else: COMs[n] = (rxy)[mask].sum(axis=0) / (rxy)[mask].shape[0] return COMs -# def DBSCAN_L2_(pointsvector, -# method = "COM", -# eps=3, -# min_samples=2, -# # plot=True -# ): +def DBSCAN_L2( + pointsvector, + eps=5, + min_samples=2, + plot=True, + method='COMs' +): + assert method in ['COMs','Jaccard'], 'method currently restricted to COMs or Jaccard' + + db2 = DBSCAN(eps=eps, min_samples=min_samples) + + if method == "COMs": + COMs = Cluster_COMs_R(pointsvector, weighted=True) + db2.fit(COMs) + + elif method == "Jaccard": + L1labels = pointsvector.select_fields("L1labels").flatten().astype(int) + L1_unique_labels = np.unique(L1labels) + ims = [] + for L1_label in L1_unique_labels[1:]: + mask = L1labels == L1_label + + corrs = jaccard_image_similarity(imsarray, plot=False) + + L2_unique_labels, L2_all_cluster_sizes = np.unique(db2.labels_, return_counts=True) + L2_unique_labels_proper = L2_unique_labels[1:] -# COMs = Cluster_COMs_R(pointsvector, weighted=True): -# db2 = DBSCAN(eps=eps, min_samples=ms) -# db2.fit(COMs) + L1labels = np.squeeze(pointsvector.select_fields("L1labels").flatten().astype(int)) + L1_unique_labels_proper = np.unique(L1labels)[1:] + Rshape = pointsvector.shape + + 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]) + + 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: + 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_similarity(imsarray, plot=False): + + imsmask[imsarray>1] = 1 + + corrs = 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)) + corrs[i] = both/either + + if plot: + plt.imshow(corrs) + return corrs \ No newline at end of file From 0596768e953677900a238852c5131b819fdd313b Mon Sep 17 00:00:00 2001 From: maclariz Date: Mon, 3 Aug 2026 21:09:46 -0500 Subject: [PATCH 14/14] Update digital_dark_field_cluster.py Complete set of functions using new img reduction method. Maybe I should also make a new method to reduce a selection to a diffraction pattern. --- .../diffraction/digital_dark_field_cluster.py | 126 ++++++++++++------ 1 file changed, 84 insertions(+), 42 deletions(-) diff --git a/src/quantem/diffraction/digital_dark_field_cluster.py b/src/quantem/diffraction/digital_dark_field_cluster.py index 391292d70..962e5c50e 100644 --- a/src/quantem/diffraction/digital_dark_field_cluster.py +++ b/src/quantem/diffraction/digital_dark_field_cluster.py @@ -3,6 +3,7 @@ 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 @@ -239,7 +240,9 @@ def DDFpointsmask(pointsvector,selectionpoints,tolerance): 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): ''' @@ -264,8 +267,8 @@ def DDFrphimask(pointsvector,r,rtol,phi=None,phitol=None): Returns ------- - maskstack: np.ndarray - A set of Boolean masks for selecting points. Each will have the same length as the flattened fields + 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)1] = 1 + imsmask = (imsarray>1).astype(int) - corrs = np.zeros(shape=(imsarray.shape[0],imsarray.shape[0])) + 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)) - corrs[i] = both/either + dists[i] = 1-both/either if plot: - plt.imshow(corrs) - return corrs \ No newline at end of file + plt.imshow(dists) + return dists \ No newline at end of file