larsupb-LoRA-Merger-ComfyUI/tests/test_decomposition.py

"""
Unit tests for decomposition module.

Tests tensor decomposition functionality including SVD, QR, and error handling.
"""

import pytest
import torch

# Imported as `src.*` so the package-relative imports inside src/ resolve.
# conftest.py puts the project root on sys.path and mocks the ComfyUI modules.
from src.decomposition import (
    SVDDecomposer,
    RandomizedSVDDecomposer,
    EnergyBasedRandomizedSVDDecomposer,
    SingularValueDistribution,
)


class TestSVDDecomposer:
    """Tests for standard SVD decomposer."""

    def test_basic_2d_decomposition(self):
        """Test basic 2D tensor decomposition."""
        # return_statistics is a constructor option, not a decompose() argument.
        decomposer = SVDDecomposer(return_statistics=True)
        weight = torch.randn(100, 50)
        target_rank = 10

        up, down, alpha, stats = decomposer.decompose(weight, target_rank)

        # Check shapes
        assert up.shape == (100, 10)
        assert down.shape == (10, 50)
        assert isinstance(alpha, float)

        # Reconstruction quality is bounded by the discarded singular values, not
        # by the decomposer: for a full-rank random Gaussian, a rank-10 truncation
        # of 50 singular values necessarily loses ~75% of the Frobenius norm.
        # Assert we land at that optimum rather than at an arbitrary threshold.
        reconstructed = up @ down
        reconstruction_error = torch.norm(weight - reconstructed) / torch.norm(weight)

        S = torch.linalg.svdvals(weight)
        optimal_error = (S[target_rank:].pow(2).sum() / S.pow(2).sum()).sqrt()
        assert reconstruction_error < optimal_error + 0.01

    def test_low_rank_input_reconstructs_accurately(self):
        """A genuinely low-rank matrix is recovered with little error."""
        decomposer = SVDDecomposer()
        weight = torch.randn(100, 10) @ torch.randn(10, 50)
        weight = weight + 0.01 * torch.randn(100, 50)

        up, down, alpha, _ = decomposer.decompose(weight, target_rank=10)

        reconstruction_error = torch.norm(weight - up @ down) / torch.norm(weight)
        assert reconstruction_error < 0.05

    def test_4d_conv_decomposition(self):
        """Test 4D convolutional tensor decomposition."""
        decomposer = SVDDecomposer()
        weight = torch.randn(64, 32, 3, 3)  # Conv layer
        target_rank = 16

        up, down, alpha, _ = decomposer.decompose(weight, target_rank)

        # Check shapes
        assert up.shape == (64, 16, 1, 1)
        assert down.shape == (16, 32, 3, 3)

    def test_symmetric_vs_asymmetric_distribution(self):
        """Test different singular value distributions."""
        weight = torch.randn(50, 30)
        target_rank = 10

        # Symmetric distribution
        decomposer_sym = SVDDecomposer(
            distribution=SingularValueDistribution.SYMMETRIC
        )
        up_sym, down_sym, _, _ = decomposer_sym.decompose(weight, target_rank)

        # Asymmetric distribution
        decomposer_asym = SVDDecomposer(
            distribution=SingularValueDistribution.ASYMMETRIC
        )
        up_asym, down_asym, _, _ = decomposer_asym.decompose(weight, target_rank)

        # Both should reconstruct similarly but with different scaling
        recon_sym = up_sym @ down_sym
        recon_asym = up_asym @ down_asym

        # The product is mathematically identical either way -- the distributions
        # only decide whether S goes into up, into down, or is split as sqrt(S)
        # across both. atol is needed because the default (1e-8) is below float32
        # rounding for entries near zero; observed difference is ~5e-7.
        assert torch.allclose(recon_sym, recon_asym, rtol=1e-4, atol=1e-5)

    def test_statistics_calculation(self):
        """Test that statistics are calculated correctly."""
        decomposer = SVDDecomposer(return_statistics=True)
        weight = torch.randn(80, 40)
        target_rank = 20

        _, _, _, stats = decomposer.decompose(weight, target_rank)

        assert stats is not None
        assert 'new_rank' in stats
        assert 'new_alpha' in stats
        assert 'sum_retained' in stats
        assert 'fro_retained' in stats
        assert 'max_ratio' in stats

        # Check statistics are reasonable
        assert stats['new_rank'] == 20
        assert 0.0 <= stats['sum_retained'] <= 1.0
        assert 0.0 <= stats['fro_retained'] <= 1.0

    def test_dynamic_rank_selection_ratio(self):
        """Test dynamic rank selection by singular value ratio."""
        decomposer = SVDDecomposer(return_statistics=True)
        weight = torch.randn(100, 50)

        _, _, _, stats = decomposer.decompose(
            weight,
            target_rank=50,
            dynamic_method="sv_ratio",
            dynamic_param=100.0  # Ratio threshold
        )

        # Rank should be selected based on ratio
        assert stats['new_rank'] <= 50
        assert stats['new_rank'] >= 1

    def test_dynamic_rank_selection_cumulative(self):
        """Test dynamic rank selection by cumulative singular values."""
        decomposer = SVDDecomposer(return_statistics=True)
        weight = torch.randn(100, 50)

        _, _, _, stats = decomposer.decompose(
            weight,
            target_rank=50,
            dynamic_method="sv_cumulative",
            dynamic_param=0.95  # 95% of cumulative sum
        )

        # Rank should be selected to capture 95% of singular values
        assert stats['sum_retained'] >= 0.90  # Allow some tolerance

    def test_dynamic_rank_selection_frobenius(self):
        """Test dynamic rank selection by Frobenius norm."""
        decomposer = SVDDecomposer(return_statistics=True)
        weight = torch.randn(100, 50)

        _, _, _, stats = decomposer.decompose(
            weight,
            target_rank=50,
            dynamic_method="sv_fro",
            dynamic_param=0.99  # 99% of Frobenius norm
        )

        # Rank should be selected to retain 99% of Frobenius norm
        assert stats['fro_retained'] >= 0.95  # Allow some tolerance


class TestRandomizedSVDDecomposer:
    """Tests for randomized SVD decomposer."""

    def test_randomized_svd_approximation(self):
        """Test that randomized SVD produces good approximation."""
        weight = torch.randn(200, 100)
        target_rank = 20

        # Standard SVD
        decomposer_std = SVDDecomposer()
        up_std, down_std, _, _ = decomposer_std.decompose(weight, target_rank)
        recon_std = up_std @ down_std

        # Randomized SVD
        decomposer_rand = RandomizedSVDDecomposer(n_oversamples=10, n_iter=2)
        up_rand, down_rand, _, _ = decomposer_rand.decompose(weight, target_rank)
        recon_rand = up_rand @ down_rand

        # Reconstructions should be similar
        error_std = torch.norm(weight - recon_std)
        error_rand = torch.norm(weight - recon_rand)

        # Randomized should be close to standard (within 50% relative error)
        assert abs(error_rand - error_std) / error_std < 0.5

    def test_randomized_svd_small_matrix(self):
        """Test that small matrices fall back to standard SVD."""
        decomposer = RandomizedSVDDecomposer()
        weight = torch.randn(50, 30)  # Small matrix
        target_rank = 10

        # Should not raise error
        up, down, alpha, _ = decomposer.decompose(weight, target_rank)

        assert up.shape == (50, 10)
        assert down.shape == (10, 30)


class TestEnergyBasedRandomizedSVDDecomposer:
    """Tests for energy-based randomized SVD."""

    def test_energy_based_rank_selection(self):
        """Test that energy threshold affects rank selection."""
        weight = torch.randn(100, 50)

        # Low energy threshold (fewer components)
        decomposer_low = EnergyBasedRandomizedSVDDecomposer(
            energy_threshold=0.8,
            return_statistics=True
        )
        _, _, _, stats_low = decomposer_low.decompose(weight, target_rank=50)

        # High energy threshold (more components)
        decomposer_high = EnergyBasedRandomizedSVDDecomposer(
            energy_threshold=0.99,
            return_statistics=True
        )
        _, _, _, stats_high = decomposer_high.decompose(weight, target_rank=50)

        # Higher threshold should generally use more components
        # (though not guaranteed due to randomness)
        assert stats_low is not None
        assert stats_high is not None


class TestErrorHandling:
    """Tests for error handling in decomposition."""

    def test_invalid_tensor_dimensions(self):
        """Test that invalid tensor dimensions raise errors."""
        decomposer = SVDDecomposer()
        weight = torch.randn(10)  # 1D tensor (invalid)

        with pytest.raises(ValueError, match="must be 2D, 3D, or 4D"):
            decomposer.decompose(weight, target_rank=5)

    def test_zero_matrix_handling(self):
        """Test handling of numerically zero matrices."""
        decomposer = SVDDecomposer(return_statistics=True)
        weight = torch.zeros(50, 30)

        # The zero-matrix guard lives in dynamic rank selection; without a
        # dynamic_method the caller has pinned the rank and target_rank is honoured.
        up, down, alpha, stats = decomposer.decompose(
            weight, target_rank=10, dynamic_method="sv_ratio", dynamic_param=2.0
        )

        # Rank should collapse to the minimum for a numerically zero matrix
        assert stats['new_rank'] == 1

    def test_zero_matrix_fixed_rank_is_honoured(self):
        """A pinned rank (no dynamic_method) is kept even for a zero matrix."""
        decomposer = SVDDecomposer(return_statistics=True)
        weight = torch.zeros(50, 30)

        up, down, alpha, stats = decomposer.decompose(weight, target_rank=10)

        assert stats['new_rank'] == 10
        assert torch.allclose(up @ down, torch.zeros(50, 30))

    def test_invalid_dynamic_method(self):
        """Test that invalid dynamic method raises error."""
        decomposer = SVDDecomposer()
        weight = torch.randn(50, 30)

        with pytest.raises(ValueError, match="Unknown dynamic method"):
            decomposer.decompose(
                weight,
                target_rank=10,
                dynamic_method="invalid_method"
            )


# Fixtures

@pytest.fixture
def sample_2d_weight():
    """Fixture providing sample 2D weight tensor."""
    return torch.randn(100, 50)


@pytest.fixture
def sample_4d_weight():
    """Fixture providing sample 4D convolutional weight."""
    return torch.randn(64, 32, 3, 3)


class TestIntegrationWithFixtures:
    """Integration tests using fixtures."""

    def test_all_decomposers_with_2d(self, sample_2d_weight):
        """Test all decomposers work with 2D tensors."""
        decomposers = [
            SVDDecomposer(),
            RandomizedSVDDecomposer(),
            EnergyBasedRandomizedSVDDecomposer(),
        ]

        for decomposer in decomposers:
            up, down, alpha, _ = decomposer.decompose(
                sample_2d_weight,
                target_rank=20
            )

            assert up.shape[0] == 100
            assert up.shape[1] == 20
            assert down.shape[0] == 20
            assert down.shape[1] == 50

    def test_all_decomposers_with_4d(self, sample_4d_weight):
        """Test all decomposers work with 4D conv tensors."""
        decomposers = [
            SVDDecomposer(),
            RandomizedSVDDecomposer(),
        ]

        for decomposer in decomposers:
            up, down, alpha, _ = decomposer.decompose(
                sample_4d_weight,
                target_rank=16
            )

            assert up.shape == (64, 16, 1, 1)
            assert down.shape == (16, 32, 3, 3)


if __name__ == "__main__":
    pytest.main([__file__, "-v"])