dubfi.linalg.mpi_worker_helpers

Helper functions for MPI linear algebra worker (private module).

This module contains alternative implementations of helper functions used by dubfi.linalg.mpi_worker.

Functions

_trace_product_self_parallel(s, e, hens, x, loc, chainT)

Compute part of trace product for MpiGradOpWorker (special case).

_trace_product_self_ensopt(s, e, hens, x, loc, chain)

Compute part of trace product for MpiGradOpWorker (special case).

_trace_product_self_ens(s, e, hens, x, loc, chain)

Compute part of trace product for MpiGradOpWorker (special case).

Module Contents

dubfi.linalg.mpi_worker_helpers._trace_product_self_parallel(s, e, hens, x, loc, chainT)

Compute part of trace product for MpiGradOpWorker (special case).

Helper for MpiGradOpWorker._trace_product() in commonly used case.

Runtime: O(n_state * n_obs**3) + O(n_state * n_ens * n_obs**2) + O(n_state**2 * n_ens * n_obs)

Required memory (theoretically): 16 * n_obs**2 Bytes per numba thread

Parameters:

  • s (int) – start index in reduced observation dimension

  • e (int) – end index in reduced observation dimension

  • hens (array, shape (n_ens, n_state, n_obs)) – centered ensemble observation operator

  • x (array, shape (n_ens, n_obs)) – ensemble of states

  • loc (array, shape (n_obs, n_obs)) – localization, real-symmetric matrix

  • chainT (array, shape (n_obs, n_obs)) – chained array, real-symmetric matrix

Return type:

numpy.ndarray

dubfi.linalg.mpi_worker_helpers._trace_product_self_ensopt(s, e, hens, x, loc, chain)

Compute part of trace product for MpiGradOpWorker (special case).

Helper for MpiGradOpWorker._trace_product() in commonly used case, assuming that loc and chain are symmetric. This function can be faster than _trace_product_self_parallel() if n_state > 2*n_obs**2.

Runtime: O(n_ens**2 * n_obs**3) + …

Required memory: 16 * (n_ens + 1) * n_obs**2 Bytes

Define: s[k, i1, j2] = (hens[m1, k, i1] x[m1, j1] + i1 <-> j1) loc[i1, j1] chain[j1, j2] (sum over m1, j1) Compute: result[k, l] = s[k, i1, j2] s[l, j2, i1] (sum over i1=s:e, j2=0:n_obs)

Consider fixed m1, m2. Then result consists of two parts: A: hens[m1, k, i1] x[m1, j1] loc[i1, j1] chain[j1, j2] hens[m2, l, j2] x[m2, i2] loc[i2, j2] chain[i2, i1] B: hens[m1, k, i1] x[m1, j1] loc[i1, j1] chain[j1, j2] hens[m2, l, i2] x[m2, j2] loc[i2, j2] chain[i2, i1]

Parameters:

  • s (int)

  • e (int)

  • hens (numpy.ndarray)

  • x (numpy.ndarray)

  • loc (numpy.ndarray)

  • chain (numpy.ndarray)

Return type:

numpy.ndarray

dubfi.linalg.mpi_worker_helpers._trace_product_self_ens(s, e, hens, x, loc, chain)

Compute part of trace product for MpiGradOpWorker (special case).

Helper for MpiGradOpWorker._trace_product() in commonly used case, assuming that loc and chain are symmetric.

Runtime: O(n_ens**2 * n_obs**3) + …

Define: s[k, i1, j2] = (hens[m1, k, i1] x[m1, j1] + i1 <-> j1) loc[i1, j1] chain[j1, j2] (sum over m1, j1) Compute: result[k, l] = s[k, i1, j2] s[l, j2, i1] (sum over i1=s:e, j2=0:n_obs)

Consider fixed m1, m2. Then result consists of two parts: A: hens[m1, k, i1] x[m1, j1] loc[i1, j1] chain[j1, j2] hens[m2, l, j2] x[m2, i2] loc[i2, j2] chain[i2, i1] B: hens[m1, k, i1] x[m1, j1] loc[i1, j1] chain[j1, j2] hens[m2, l, i2] x[m2, j2] loc[i2, j2] chain[i2, i1]

Parameters:

  • s (int)

  • e (int)

  • hens (numpy.ndarray)

  • x (numpy.ndarray)

  • loc (numpy.ndarray)

  • chain (numpy.ndarray)

Return type:

numpy.ndarray