Compare inversion methods for ESMDA

[tommyod]

When I wrote ESMDA I implemented several inversion methods, that is, ways of computing

$$X Y^T /(N-1) (C_{DD} + \alpha C_D)^{-1} (D - Y).$$

We should benchmark these and bring the best ones into the public ESMDA API. More specifically:

• Benchmark practical running times on the typical case ($n \gg m \gg N$)
• Benchmark memory usage on the typical case ($n \gg m \gg N$)
• Clean up documentation on the best methods
• Bring the best methods into the ESMDA api

[tommyod]

Some results on memory

Ran only this test in the test suite, using --memray to log memory.

def test_esmda_inversion_memory():
    
    num_outputs=1000
    num_inputs=100_000
    ensemble_members=100
    
    C_D_diag = np.exp(np.random.randn(num_outputs))
    assert C_D_diag.ndim == 1

    # Set alpha to something other than 1 to test that it works
    alpha = 3

    # Create observations
    D = np.random.randn(num_outputs, ensemble_members)
    Y = np.random.randn(num_outputs, ensemble_members)
    X = np.random.randn(num_inputs, ensemble_members)
    
    # Base memory: 
    # Total memory allocated: 9.2MiB when 
    # num_outputs=1000
    # num_inputs=10000
    # ensemble_members=100
    
    # num_outputs=1000
    # num_inputs=10000
    # ensemble_members=100
    # inversion_exact_cholesky              91.3MiB
    # inversion_subspace                    86.2MiB
    # inversion_exact_subspace_woodbury     92.1MiB
    # inversion_rescaled_subspace           87.0MiB
    
    # num_outputs=10000
    # num_inputs=1000
    # ensemble_members=100
    # inversion_exact_cholesky              914.0MiB 
    # inversion_subspace                    128.2MiB
    # inversion_exact_subspace_woodbury       1.6GiB
    # inversion_rescaled_subspace           127.2MiB
    
    # num_outputs=1000
    # num_inputs=100_000
    # ensemble_members=100
    # inversion_exact_cholesky              229.6MiB
    # inversion_subspace                    223.5MiB
    # inversion_exact_subspace_woodbury     229.5MiB
    # inversion_rescaled_subspace           223.4MiB
    func = inversion_rescaled_subspace
    result_matrix = func(alpha=alpha, C_D=C_D_diag, D=D, Y=Y, X=X)

Some results on time

test_timing(num_outputs=10_000, num_inputs=1000, ensemble_members=100)
--------------------------------
Function: inversion_exact_cholesky on dense covariance: 3.6645 s
Function: inversion_exact_cholesky on diagonal covariance: 3.2296 s
--------------------------------
Function: inversion_subspace on dense covariance: 0.4175 s
Function: inversion_subspace on diagonal covariance: 0.0669 s
--------------------------------
Function: inversion_exact_subspace_woodbury on dense covariance: 20.8711 s
Function: inversion_exact_subspace_woodbury on diagonal covariance: 1.0465 s
--------------------------------
Function: inversion_rescaled_subspace on dense covariance: 5.9943 s
Function: inversion_rescaled_subspace on diagonal covariance: 0.0795 s
--------------------------------


test_timing(num_outputs=10_000, num_inputs=100_000, ensemble_members=100)
--------------------------------
Function: inversion_exact_cholesky on dense covariance: 3.5683 s
Function: inversion_exact_cholesky on diagonal covariance: 3.1073 s
--------------------------------
Function: inversion_subspace on dense covariance: 0.5029 s
Function: inversion_subspace on diagonal covariance: 0.1507 s
--------------------------------
Function: inversion_exact_subspace_woodbury on dense covariance: 22.2824 s
Function: inversion_exact_subspace_woodbury on diagonal covariance: 1.4583 s
--------------------------------
Function: inversion_rescaled_subspace on dense covariance: 5.9192 s
Function: inversion_rescaled_subspace on diagonal covariance: 0.2002 s
--------------------------------

[dafeda]

I am surprised by these results as I thought exact cholesky with diagonal covariance would be the fastest.