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`(→ numpy.ndarray)	Compute part of trace product for `MpiGradOpWorker` (special case).
`_trace_product_self_ensopt`(→ numpy.ndarray)	Compute part of trace product for `MpiGradOpWorker` (special case).
`_trace_product_self_ens`(→ numpy.ndarray)	Compute part of trace product for `MpiGradOpWorker` (special case).

Module Contents¶

dubfi.linalg.mpi_worker_helpers._trace_product_self_parallel(s: int, e: int, hens: numpy.ndarray, x: numpy.ndarray, loc: numpy.ndarray, chainT: numpy.ndarray) → numpy.ndarray¶

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) – centered ensemble observation operator, shape (n_ens, n_state, n_obs)
x (array) – ensemble of states, shape (n_ens, n_obs)
loc (array) – localization, real-symmetric matrix, shape (n_obs, n_obs)
chainT (array) – chained array, real-symmetric matrix, shape (n_obs, n_obs)

dubfi.linalg.mpi_worker_helpers._trace_product_self_ensopt(s: int, e: int, hens: numpy.ndarray, x: numpy.ndarray, loc: numpy.ndarray, chain: numpy.ndarray) → numpy.ndarray¶

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]

dubfi.linalg.mpi_worker_helpers._trace_product_self_ens(s: int, e: int, hens: numpy.ndarray, x: numpy.ndarray, loc: numpy.ndarray, chain: numpy.ndarray) → numpy.ndarray¶

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)