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..104ff3a6 100644 --- a/scrise/tests/test_component_ordering.py +++ b/scrise/tests/test_component_ordering.py @@ -126,7 +126,71 @@ 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( + np.array(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( + 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"]), + (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)