STAMP v0.14.0 release

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "stamp"
-version = "0.13.2"
+version = "0.14.0"
 description = "A workflow for Sub-Tomogram Averaging Membrane Proteins"
 readme = "README.md"
 authors = [

src/stamp/classify/features.py

@@ -16,6 +16,12 @@
 from stamp.utils.log import get_worker_log_config, init_worker_logging, log
 from stamp.utils.parallel import run_parallel_ordered
 
+# FEATURE_NORMALISATIONS: how 1+ modes are made comparable across particles and tomograms
+FEATURE_NORMALISATIONS = ('none', 'particle', 'tomogram', 'both')
+
+# _MIN_TOMOGRAM_PARTICLES: smallest tomogram standardised on its own, smaller ones fall back to population statistics
+_MIN_TOMOGRAM_PARTICLES = 50
+
 # cylindrical_bins: precompute per-voxel radial bin index and validity mask
 def cylindrical_bins(box_voxels: int, n_radial_bins: int) -> tuple[np.ndarray, np.ndarray]:
     half = (box_voxels - 1) / 2.0
@@ -100,12 +106,26 @@ def _standardise_rows(matrix: np.ndarray) -> np.ndarray:
     stds[stds == 0.0] = 1.0
     return (matrix - means) / stds
 
-# _standardise_columns: z-score each column across the particle population
-def _standardise_columns(matrix: np.ndarray) -> np.ndarray:
-    means = matrix.mean(axis=0, keepdims=True)
-    stds = matrix.std(axis=0, keepdims=True)
-    stds[stds == 0.0] = 1.0
-    return (matrix - means) / stds
+# _standardise_columns: z-score each column across the particle population, or within each tomogram of at least _MIN_TOMOGRAM_PARTICLES particles when tomogram_of specified
+def _standardise_columns(matrix: np.ndarray, tomogram_of: np.ndarray | None = None) -> np.ndarray:
+    # _zscore: standardise block using the statistics of reference
+    def _zscore(block: np.ndarray, reference: np.ndarray) -> np.ndarray:
+        stds = reference.std(axis=0, keepdims=True)
+        stds[stds == 0.0] = 1.0
+        return (block - reference.mean(axis=0, keepdims=True)) / stds
+    result = _zscore(matrix, matrix)
+    if tomogram_of is not None:
+        for tomogram in np.unique(tomogram_of):
+            members = tomogram_of == tomogram
+            if members.sum() >= _MIN_TOMOGRAM_PARTICLES:
+                result[members] = _zscore(matrix[members], matrix[members])
+    return result
+
+# _normalise_power: divide a particle's 1+ modes by joint L2 norm (so overall non-axisymmetric amplitude doesn't dominate)
+def _normalise_power(stacked_modes: np.ndarray) -> np.ndarray:
+    norms = np.linalg.norm(stacked_modes.reshape(stacked_modes.shape[0], -1), axis=1)
+    norms[norms == 0.0] = 1.0
+    return stacked_modes / norms[:, None, None, None]
 
 # _load_cached_magnitudes: one tomogram's cached magnitudes, or None when absent or built for different rows
 def _load_cached_magnitudes(cache_dir: Path | None, tomogram_id: str, rows: list[int]) -> np.ndarray | None:
@@ -130,12 +150,15 @@ def build_feature_matrix(
     *,
     batch_size: int = 512,
     cache_dir: Path | None = None,
+    normalisation: str = 'both',
 ) -> np.ndarray:
     log.debug(f'Building features for {len(subvols)} subvolumes, max_mode={max_azimuthal_mode}')
     if len(subvols) == 0:
         return np.empty((0, 0))
     if not 0.0 <= min_radius_fraction < 1.0:
         raise ValueError('min_radius_fraction must be between [0, 1)')
+    if normalisation not in FEATURE_NORMALISATIONS:
+        raise ValueError(f'normalisation must be one of {FEATURE_NORMALISATIONS}, got {normalisation!r}')
 
     worker = partial(azimuthal_magnitudes, n_radial_bins=n_radial_bins, n_azimuthal_samples=n_azimuthal_samples, max_mode=max_azimuthal_mode)
     magnitudes: list[np.ndarray | None] = [None] * len(subvols)
@@ -186,9 +209,18 @@ def build_feature_matrix(
     # Mode 0: all radial bins, standardised per particle
     blocks = [_standardise_rows(stacked[:, 0].reshape(n_particles, -1))]
 
-    # Modes 1+: outer bins only, standardised per column across particle so weak-but-consistent symmetry signal survives into the PCA
-    for mode in range(1, max_azimuthal_mode + 1):
-        block = stacked[:, mode, first_outer_bin:, :].reshape(n_particles, -1)
-        blocks.append(_standardise_columns(block))
-
-    return np.hstack(blocks)
+    # Modes 1+: outer bins only, standardised per column so weak-but-consistent symmetry signal survives into the PCA
+    # 'particle' first divides each particle by its own total 1+ mode power, 'tomogram' standardises columns within each tomogram
+    tomogram_of = np.empty(n_particles, dtype=np.int64)
+    for tomogram_number, local_indices in enumerate(by_tomogram.values()):
+        tomogram_of[local_indices] = tomogram_number
+    outer = stacked[:, 1:, first_outer_bin:, :]
+    if normalisation in ('particle', 'both') and outer.shape[1]:
+        outer = _normalise_power(outer)
+    for mode in range(outer.shape[1]):
+        block = outer[:, mode].reshape(n_particles, -1)
+        blocks.append(_standardise_columns(block, tomogram_of if normalisation in ('tomogram', 'both') else None))
+
+    features = np.hstack(blocks)
+    features[~stacked.reshape(n_particles, -1).any(axis=1)] = 0.0  # empty subvolumes stay all-zero (column standardisation gives non-zero rows that aren't recognised as degenerate)
+    return features

src/stamp/cli/classify.py

@@ -98,6 +98,10 @@ def classify(
         int,
         typer.Option(help='Azimuthal sampling points (must be at least 2*(modes+1)).', rich_help_panel = 'Subvolume extraction'),
     ] = 64,
+    feature_normalisation: Annotated[
+        Literal['none', 'particle', 'tomogram', 'both'],
+        typer.Option('--feature-normalisation', help='Method for azimuthal mode (above 0) normalisation.', rich_help_panel = 'Subvolume extraction'),
+    ] = 'both',
     n_workers: Annotated[
         int,
         typer.Option('-n', '--n-processes', help='Number of processes to use for per-cluster in-plane alignment (1 = sequential).'),
@@ -139,6 +143,7 @@ def classify(
         azimuthal_modes=azimuthal_modes,
         min_radius_fraction=min_radius_fraction,
         n_azimuthal_samples=n_azimuthal_samples,
+        feature_normalisation=feature_normalisation,
         make_plots=make_plots,
         plot_format=plot_format,
         keep_raw=keep_raw,

src/stamp/commands/classify.py

@@ -53,6 +53,7 @@ def run_classify(
     azimuthal_modes = 4,
     min_radius_fraction = 0.25,
     n_azimuthal_samples = 64,
+    feature_normalisation = 'both',
     n_workers: int = 1,
     make_plots: bool = True,
     plot_format: str = 'tiff',
@@ -92,6 +93,7 @@ def run_classify(
         'azimuthal_modes': azimuthal_modes,
         'n_azimuthal_samples': n_azimuthal_samples,
         'min_radius_fraction': min_radius_fraction,
+        'feature_normalisation': feature_normalisation,
         'method': method,
         'min_cluster_size': min_cluster_size,
         'n_clusters': n_clusters,
@@ -124,6 +126,7 @@ def run_classify(
         min_radius_fraction=min_radius_fraction,
         n_workers=n_workers,
         cache_dir=output_dir / 'raw' / 'features',
+        normalisation=feature_normalisation,
     )
     log.progress(f'Feature matrix complete: {features.shape[0]} particles x {features.shape[1]} features (modes 0-{azimuthal_modes})')
 
@@ -199,6 +202,7 @@ def run_classify(
             'azimuthal_modes': azimuthal_modes,
             'n_azimuthal_samples': n_azimuthal_samples,
             'min_radius_fraction': min_radius_fraction,
+            'feature_normalisation': feature_normalisation,
             'n_workers': n_workers,
             'is_decoy': is_decoy,
             'membrane_subtracted': bool(segmentation_paths),

tests/unit/test_classify.py

@@ -321,6 +321,14 @@ def test_degenerate_rows_are_all_zero(self, tmp_path: Path):
         features = build_feature_matrix(subvols, n_radial_bins=4, max_azimuthal_mode=0)
         assert not features[0].any() and features[1].any()
 
+    def test_degenerate_rows_stay_zero_with_azimuthal_modes(self, tmp_path: Path):
+        '''An empty subvolume is all-zero in every block, including column-standardised modes 1+.'''
+        rng = np.random.default_rng(0)
+        stack = np.concatenate([np.zeros((1, 15, 15, 15)), rng.random((6, 15, 15, 15))])
+        for normalisation in ('none', 'particle', 'tomogram', 'both'):
+            features = build_feature_matrix(_subvols_from_array(tmp_path / normalisation, stack), n_radial_bins=6, max_azimuthal_mode=3, normalisation=normalisation)
+            assert not features[0].any() and features[1:].any(axis=1).all()
+
     def test_feature_checkpoint_skips_recomputation(self, tmp_path: Path, monkeypatch) -> None:
         from stamp.classify import features as features_module
         subvols = _subvols_from_array(tmp_path / 'sv', np.random.default_rng(0).random((6, 11, 11, 11)))
@@ -337,6 +345,41 @@ def _fail(*args, **kwargs):
         with pytest.raises(AssertionError):
             build_feature_matrix(subvols.take([0, 1, 2]), n_radial_bins=4, cache_dir=cache)
 
+    def test_particle_normalisation_removes_amplitude_scale(self, tmp_path: Path) -> None:
+        '''Amplifying one particle's signal disturbs other particles' features if normalisation isn't applied.'''
+        rng = np.random.default_rng(0)
+        base = np.stack([_rotate_in_plane(_c_n_particle(n_fold=4), rng.uniform(0, 360)) + rng.normal(0, 0.05, (25, 25, 25)) for _ in range(8)])
+        loud = base.copy()
+        loud[0] = base[0].mean() + 20.0 * (base[0] - base[0].mean())
+        kwargs = dict(n_radial_bins=8, max_azimuthal_mode=3)
+        shifts = {}
+        for mode in ('none', 'particle'):
+            before = build_feature_matrix(_subvols_from_array(tmp_path / f'{mode}_a', base), normalisation=mode, **kwargs)
+            after = build_feature_matrix(_subvols_from_array(tmp_path / f'{mode}_b', loud), normalisation=mode, **kwargs)
+            shifts[mode] = np.abs(before[1:] - after[1:]).mean()  # the untouched particles
+        assert shifts['particle'] < 0.2 * shifts['none']
+
+    def test_tomogram_normalisation_removes_tomogram_offsets(self, tmp_path: Path) -> None:
+        '''Two tomograms with different mode 1+ levels are centred on each other post-tomogram standardisation.'''
+        rng = np.random.default_rng(0)
+        n, box = 60, 15  # n is at least the per-tomogram minimum, so each tomogram is standardised on its own
+        tmp_path.mkdir(parents=True, exist_ok=True)
+        np.save(tmp_path / 'quiet.npy', rng.normal(0, 0.05, (n, box, box, box)).astype(np.float32))
+        np.save(tmp_path / 'loud.npy', rng.normal(0, 0.5, (n, box, box, box)).astype(np.float32))
+        subvols = Subvolumes(tmp_path, [('quiet', i) for i in range(n)] + [('loud', i) for i in range(n)], box)
+        kwargs = dict(n_radial_bins=6, max_azimuthal_mode=2, min_radius_fraction=0.0)
+        mode0_width = 6 * box
+        population = build_feature_matrix(subvols, normalisation='none', **kwargs)[:, mode0_width:]
+        per_tomogram = build_feature_matrix(subvols, normalisation='tomogram', **kwargs)[:, mode0_width:]
+        gap_population = np.abs(population[:n].mean() - population[n:].mean())
+        gap_tomogram = np.abs(per_tomogram[:n].mean() - per_tomogram[n:].mean())
+        assert gap_tomogram < 0.1 * gap_population
+
+    def test_unknown_normalisation_is_rejected(self, tmp_path: Path) -> None:
+        subvols = _subvols_from_array(tmp_path, np.random.default_rng(0).random((4, 11, 11, 11)))
+        with pytest.raises(ValueError, match='normalisation'):
+            build_feature_matrix(subvols, n_radial_bins=4, normalisation='bogus')
+
 # TestCluster: class containing unit tests for src/stamp/classify/cluster.py
 class TestCluster:
     def test_hdbscan_recovers_three_groups(self) -> None: