Group in nmf results and better coloring in plotting

[KondratievaOlesya]

Added

• Persistent sample grouping in NMFResult

  • Introduced a new groups attribute in NMFResult:

    • Stored as a pandas.DataFrame with:

      • index = sample IDs
      • column = "group"

    • Always present (at minimum, all samples assigned to group "1").

• Optional clustering at NMF runtime

  • run_nmf() now supports optional sample clustering via max_clusters:

    • If max_clusters > 1, samples are clustered based on signature exposures using KMeans.
    • Optimal number of clusters is selected using silhouette score.
    • Cluster assignments are stored in NMFResult.groups.
    • If clustering is disabled or not possible, all samples are assigned to group "1".

• Group-aware PCA plotting

  • plot_pca_samples() now:

    • Colors samples exclusively using NMFResult.groups

    • Automatically aligns PCA coordinates and group labels by sample ID

    • Automatically detects whether group is continuous vs categorical:

      • Treats group as continuous only if it is numeric and has sufficient diversity
        (more than 10 unique non-null values or > 30% unique fraction)
      • Continuous groups are shown with a color gradient + colorbar
      • Categorical groups are shown with separate legend entries

    • Does not perform clustering internally

    • Returns:

      • PCA coordinates
      • explained variance ratio
      • matplotlib axes

• Group-aware exposure plotting

  • plot_exposures() now:

    • Uses NMFResult.groups["group"] for:

      • sample ordering
      • visual separation between groups

    • No longer performs clustering internally

    • Ensures consistent grouping across all downstream plots

• Groups persisted in I/O

  • save_nmf_result() now writes:

    • groups.tsv

  • load_nmf_result() restores groups alongside signatures and exposures.

Changed

• Plotting API semantics

  • Clustering logic has been fully decoupled from plotting:

    • Clustering is now a modeling decision (run_nmf)
    • Plotting functions are purely representational

  • Removed cluster / color_by logic from PCA and exposure plotting.

• End-to-end pipeline consistency

  • Integration tests updated to:

    • Validate presence and correctness of groups
    • Use group-aware plotting APIs
    • Ensure consistent sample alignment across NMF, PCA, and exposure plots

Removed

• Internal clustering from plotting functions

  • plot_pca_samples() and plot_exposures() no longer:

    • run KMeans
    • select k
    • return cluster labels

  • All grouping information must come from NMFResult.groups.

[Copilot]

Pull request overview

This PR introduces persistent sample grouping as a core component of NMFResult and refactors the plotting API to decouple clustering from visualization. Clustering is now an optional modeling decision made during run_nmf(), and plotting functions consume the stored groups for consistent coloring and ordering.

Key Changes:

• Added groups DataFrame attribute to NMFResult with optional KMeans clustering during NMF runtime via max_clusters parameter
• Refactored plot_pca_samples() and plot_exposures() to use pre-computed groups instead of performing clustering internally
• Implemented automatic detection of continuous vs categorical groups for appropriate PCA coloring (gradient vs discrete legend)

Reviewed changes

Copilot reviewed 9 out of 9 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
src/str_mut_signatures/nmf/nmf.py	Added groups attribute to NMFResult, implemented cluster_samples() and helper functions for group management, updated I/O to persist groups
src/str_mut_signatures/nmf/plot.py	Removed internal clustering from plotting functions, added group-aware coloring logic for PCA with continuous/categorical detection, refactored exposure plotting for consistent sample ordering
tests/test_nmf.py	Added tests for cluster_samples() function and groups persistence, updated existing tests to validate groups attribute
tests/integration/test_pipeline.py	Updated integration tests to use new plotting API without clustering parameters
tests/cli/test_cli_commands.py	Updated version assertion from 0.3.0 to 1.0.0
examples/basic_usage.py	Updated example to use new max_clusters parameter and handle dict return from plot_exposures()
README.rst	Updated documentation to reflect new API, removed clustering parameters from plotting examples, added section on custom groups
CHANGELOG.md	Added comprehensive changelog for version 1.0.0 with all breaking changes and new features
src/str_mut_signatures/__init__.py	Bumped version from 0.3.0 to 1.0.0

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Missing trailing comma after the last argument in the NMFResult constructor call. Python style guides recommend trailing commas in multi-line function calls for better diff readability and consistency.

Suggested change

	model_params=model_params
	model_params=model_params,

[Copilot]

Missing space around assignment operator. According to PEP 8, there should be spaces around the '=' operator in keyword arguments. Should be 'max_clusters = 6' instead of 'max_clusters = 6' (note the extra space before the equals).

[Copilot]

The sample ordering differs between absolute and proportion plots. For absolute exposures, samples are sorted by ["group", "_total"], but for proportion exposures, they're sorted by ["group"] + sig_cols. This creates inconsistent sample ordering across the two plot types, making it difficult to compare the same samples between absolute and proportion views. Consider using the same sorting logic for both plot types.

Suggested change

	df_prop = df_prop.sort_values(
	by=["group"] + sig_cols, ascending=[True] + [False] * len(sig_cols), kind="mergesort"
	)
	# Keep the same row order as `df` to ensure consistent sample ordering
	# between absolute and proportion plots.

[KondratievaOlesya]

This one is intentional