From 274c1ae662ea13e6576eab705096762d8b2b72c7 Mon Sep 17 00:00:00 2001 From: Aaron Meyer Date: Sat, 19 Sep 2026 06:30:27 -0700 Subject: [PATCH 1/2] Permute B's eigen-state axis alongside its components `parafac2.utils.standardize_pf2` permutes the rows of B (the eigen-state axis) via linear-sum assignment so that B's diagonal is maximal, pairing eigen-state i with component i, and permutes the columns of each projection to match. `order_components_by_energy` then reordered only B's columns, which moves each diagonal entry from (i, i) to (i, order^-1(i)) and destroys that pairing -- so the eigen-state heatmap comes out scrambled and the eigen-state row labels no longer line up with the component numbers used in the A/C plots. Relabel the eigen-state axis by the same permutation, and permute the columns of obsm["projections"] to match. This is an exact relabeling of the eigen-state basis: P_k B is unchanged up to the component permutation, so weighted_projections and the reconstruction come out identical to before. Co-Authored-By: Claude Opus 5 --- scrise/factorization.py | 19 ++++++-- scrise/tests/test_component_ordering.py | 62 ++++++++++++++++++++++++- scrise/tests/test_invariants.py | 6 +++ 3 files changed, 83 insertions(+), 4 deletions(-) diff --git a/scrise/factorization.py b/scrise/factorization.py index ea17061f..3ae2c582 100644 --- a/scrise/factorization.py +++ b/scrise/factorization.py @@ -121,8 +121,12 @@ def order_components_by_energy(X: anndata.AnnData) -> anndata.AnnData: ------- anndata.AnnData The same AnnData object, with Pf2_A, Pf2_B, Pf2_C, Pf2_weights - (and, if present, obsm["weighted_projections"]) reordered/updated - in place. + (and, if present, obsm["projections"] and + obsm["weighted_projections"]) reordered/updated in place. The + eigen-state axis of Pf2_B, and the matching columns of + obsm["projections"], are permuted alongside the components so that + the maximal-diagonal form of B established by + ``parafac2.utils.standardize_pf2`` is preserved. """ A = np.array(X.uns["Pf2_A"]) B = np.array(X.uns["Pf2_B"]) @@ -141,10 +145,19 @@ def order_components_by_energy(X: anndata.AnnData) -> anndata.AnnData: X.uns["Pf2_A"] = A[:, order] X.varm["Pf2_C"] = C[:, order] - X.uns["Pf2_B"] = B[:, order] X.uns["Pf2_weights"] = weights[order] + # B is indexed by (eigen-state, component). ``parafac2.utils.standardize_pf2`` + # permutes B's *rows* so that its diagonal is maximal, pairing eigen-state i + # with component i, and permutes the columns of each projection to match. + # Permuting only B's columns here would move each diagonal entry off the + # diagonal and destroy that pairing, so the eigen-state axis is relabeled by + # the same permutation. Since P_k B is left unchanged by this relabeling + # (up to the component permutation), the reconstruction is untouched. + X.uns["Pf2_B"] = B[np.ix_(order, order)] + if "projections" in X.obsm: + X.obsm["projections"] = np.asarray(X.obsm["projections"])[:, order] X.obsm["weighted_projections"] = ( X.obsm["projections"] @ X.uns["Pf2_B"] ).astype(np.float32, copy=False) diff --git a/scrise/tests/test_component_ordering.py b/scrise/tests/test_component_ordering.py index c9335158..a9bf7b4a 100644 --- a/scrise/tests/test_component_ordering.py +++ b/scrise/tests/test_component_ordering.py @@ -126,7 +126,67 @@ def test_order_components_by_energy_reorders_weights_and_B(): ordered = order_components_by_energy(adata) np.testing.assert_allclose(ordered.uns["Pf2_weights"], weights[order]) - np.testing.assert_allclose(ordered.uns["Pf2_B"], B[:, order]) + # Both axes of B are permuted: the component axis by the energy ordering, + # and the eigen-state axis by the same permutation (see below). + np.testing.assert_allclose(ordered.uns["Pf2_B"], B[np.ix_(order, order)]) + np.testing.assert_allclose(ordered.obsm["projections"], projections[:, order]) + + +def test_order_components_by_energy_preserves_diagonal_B(): + """``parafac2.utils.standardize_pf2`` leaves B with a maximal diagonal; + reordering components must not scramble it. Permuting only B's columns + would move entry (i, i) to (i, order^-1(i)) and destroy the pairing.""" + rng = np.random.default_rng(4) + n_cells, n_genes, n_conditions, rank = 30, 20, 5, 4 + + A = rng.normal(size=(n_conditions, rank)) + C = rng.normal(size=(n_genes, rank)) + # A diagonally dominant B, as standardize_pf2 produces. + B = np.eye(rank) + 0.1 * rng.normal(size=(rank, rank)) + weights = np.array([1.0, 5.0, 2.0, 9.0]) + projections, _ = np.linalg.qr(rng.normal(size=(n_cells, rank))) + + assert np.array_equal(np.argmax(np.abs(B), axis=0), np.arange(rank)) + + adata = make_mock_adata_from_factors( + A.copy(), B.copy(), C.copy(), weights.copy(), projections + ) + ordered = order_components_by_energy(adata) + + B_new = np.array(ordered.uns["Pf2_B"]) + assert np.array_equal(np.argmax(np.abs(B_new), axis=0), np.arange(rank)) + # The non-negative diagonal convention survives too. + assert np.all(np.diag(B_new) > 0) + + +def test_order_components_by_energy_relabels_projection_columns(): + """Permuting B's eigen-state axis is only valid if the projections are + permuted the same way, so that P_k B is unchanged up to the component + permutation.""" + rng = np.random.default_rng(5) + n_cells, n_genes, n_conditions, rank = 35, 22, 6, 4 + + A = rng.normal(size=(n_conditions, rank)) + C = rng.normal(size=(n_genes, rank)) + B = rng.normal(size=(rank, rank)) + weights = np.array([3.0, 1.0, 7.0, 2.0]) + projections, _ = np.linalg.qr(rng.normal(size=(n_cells, rank))) + + energy = np.abs(weights) * np.linalg.norm(A, axis=0) * np.linalg.norm(C, axis=0) + order = np.argsort(energy)[::-1] + + adata = make_mock_adata_from_factors( + A.copy(), B.copy(), C.copy(), weights.copy(), projections + ) + ordered = order_components_by_energy(adata) + + np.testing.assert_allclose(ordered.obsm["projections"], projections[:, order]) + # P_k B is unchanged, up to the component permutation. + np.testing.assert_allclose( + np.array(ordered.obsm["weighted_projections"]), + (projections @ B)[:, order], + atol=1e-5, + ) def test_match_components_across_ranks_identifies_new_component(): diff --git a/scrise/tests/test_invariants.py b/scrise/tests/test_invariants.py index ffc910cb..54c4a6a6 100644 --- a/scrise/tests/test_invariants.py +++ b/scrise/tests/test_invariants.py @@ -134,6 +134,12 @@ def test_order_components_by_energy_is_a_permutation( np.testing.assert_allclose( np.array(ordered.obsm["weighted_projections"]), before_wp[:, order], atol=1e-5 ) + # Both axes of B, and the projection columns, are relabeled by the same + # permutation, which is what keeps B's maximal-diagonal form intact. + np.testing.assert_allclose(np.array(ordered.uns["Pf2_B"]), B[np.ix_(order, order)]) + np.testing.assert_allclose( + np.array(ordered.obsm["projections"]), projections[:, order] + ) np.testing.assert_allclose(new_A.shape, A.shape) From 37b55b7daf796be129ae403abffcb7c9e9dc67fd Mon Sep 17 00:00:00 2001 From: Aaron Meyer Date: Sat, 19 Sep 2026 06:33:23 -0700 Subject: [PATCH 2/2] Wrap obsm accesses in np.array for the type checker Co-Authored-By: Claude Opus 5 --- scrise/tests/test_component_ordering.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/scrise/tests/test_component_ordering.py b/scrise/tests/test_component_ordering.py index a9bf7b4a..104ff3a6 100644 --- a/scrise/tests/test_component_ordering.py +++ b/scrise/tests/test_component_ordering.py @@ -129,7 +129,9 @@ def test_order_components_by_energy_reorders_weights_and_B(): # Both axes of B are permuted: the component axis by the energy ordering, # and the eigen-state axis by the same permutation (see below). np.testing.assert_allclose(ordered.uns["Pf2_B"], B[np.ix_(order, order)]) - np.testing.assert_allclose(ordered.obsm["projections"], projections[:, order]) + np.testing.assert_allclose( + np.array(ordered.obsm["projections"]), projections[:, order] + ) def test_order_components_by_energy_preserves_diagonal_B(): @@ -180,7 +182,9 @@ def test_order_components_by_energy_relabels_projection_columns(): ) ordered = order_components_by_energy(adata) - np.testing.assert_allclose(ordered.obsm["projections"], projections[:, order]) + np.testing.assert_allclose( + np.array(ordered.obsm["projections"]), projections[:, order] + ) # P_k B is unchanged, up to the component permutation. np.testing.assert_allclose( np.array(ordered.obsm["weighted_projections"]),