Permute B's eigen-state axis alongside its components

[aarmey]

Problem

parafac2.utils.standardize_pf2 finishes by permuting the rows of B (the eigen-state axis) via linear_sum_assignment to maximize |diag(B)|, permuting the columns of every projection P_k to match, and flipping signs so diag(B) is non-negative. That pairs eigen-state i with component i.

order_components_by_energy then applied its energy ordering to the columns of A, C, and B. Column-only permutation of B moves each diagonal entry from (i, i) to (i, order⁻¹(i)), so the maximal-diagonal form is destroyed:

# B = eye(rank) + 0.1 * noise, then reordered by energy
diag dominant before: True
diag dominant after : False
[[ 0.01 -0.01  0.06  1.01]
 [ 0.09  1.04  0.13 -0.05]
 [ 0.   -0.13  0.94 -0.07]
 [ 0.93 -0.02 -0.12 -0.23]]

The model itself was never wrong — weighted_projections was recomputed as projections @ B[:, order], so A, C, cell scores, and the reconstruction were all consistent. What broke was the interpretive structure: plot_eigenstate_factors renders Pf2_B as a rank x rank heatmap and it came out scrambled, and the eigen-state row labels no longer corresponded to the component numbers used in the A/C plots.

Fix

Relabel the eigen-state axis by the same permutation (B[np.ix_(order, order)]) 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 are bit-for-bit what the previous code produced, while diag(B) and its non-negativity survive the ordering.

Verified end to end on a synthetic fit: the diagonal of B is now exactly unchanged by the ordering step (sum|diag| identical before and after), and weighted_projections == projections @ B still holds.

Tests

• test_order_components_by_energy_preserves_diagonal_B — a diagonally dominant B, as standardize_pf2 produces, stays diagonally dominant with a positive diagonal after ordering. This fails on main.
• test_order_components_by_energy_relabels_projection_columns — projection columns are permuted to match, and P_k B is unchanged up to the component permutation.
• test_order_components_by_energy_reorders_weights_and_B and the test_order_components_by_energy_is_a_permutation hypothesis invariant updated to assert both axes of B are permuted.

Full suite: 137 passed, 2 skipped.

Note for reviewers

obsm["projections"] is now written in the relabeled basis. It is consistent everywhere it is used within a run (PaCMAP input, factor_io round-trip, weighted_projections reconstruction), but projections saved by an older version are in the old eigen-state basis and should not be mixed with a Pf2_B from this version.

🤖 Generated with Claude Code