diff --git a/pyproject.toml b/pyproject.toml index 7782c31..587d188 100644 --- a/pyproject.toml +++ b/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 = [ diff --git a/src/stamp/classify/features.py b/src/stamp/classify/features.py index c56759b..8455c2c 100644 --- a/src/stamp/classify/features.py +++ b/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) \ No newline at end of file + # 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 \ No newline at end of file diff --git a/src/stamp/cli/classify.py b/src/stamp/cli/classify.py index 5f0c53e..5895079 100644 --- a/src/stamp/cli/classify.py +++ b/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, diff --git a/src/stamp/commands/classify.py b/src/stamp/commands/classify.py index 8d57c82..31e8a90 100644 --- a/src/stamp/commands/classify.py +++ b/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), diff --git a/tests/unit/test_classify.py b/tests/unit/test_classify.py index 0d6e60c..4ee17e9 100644 --- a/tests/unit/test_classify.py +++ b/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: diff --git a/uv.lock b/uv.lock index 89f579d..71e4652 100644 --- a/uv.lock +++ b/uv.lock @@ -1,5 +1,5 @@ version = 1 -revision = 3 +revision = 5 requires-python = ">=3.14" resolution-markers = [ "sys_platform == 'win32'", @@ -850,7 +850,7 @@ wheels = [ [[package]] name = "stamp" -version = "0.13.2" +version = "0.14.0" source = { editable = "." } dependencies = [ { name = "fast-hdbscan" },