"""
Module ``qlat_utils.jackknife_utils``
=====================================\n
Jackknife resampling and error estimation: the plain jackknife, the
super-jackknife, the randomized jackknife-bootstrap hybrid, and the unified
``g_*`` interface with its shared ``default_g_jk_kwargs``. Each family also has
collective MPI variants, which distribute the input and/or the output between
the nodes, implemented with ``mpi4py``.\n
The generic helpers (the type tuples, ``use_kwargs``, ``average``,
``block_data``, ``filter_np_results``, ``fsqr``, ``fsqrt``, ``qnorm`` and
``NewDictValues``) are imported from ``qlat_utils.data``.\n
Documentation: ``docs/qlat-utils/qlat_jackknife_utils.md``\n
.. note:: Update the documentation when updating this source file.
"""
import math
import numpy as np
class q:
from qlat_utils.data import (
int_types,
real_types,
number_types,
use_kwargs,
average,
block_data,
filter_np_results,
fsqr,
fsqrt,
qnorm,
NewDictValues,
)
from qlat_utils.timer import (
timer,
displayln_info,
)
from qlat_utils.rng_state import (
RngState,
)
# ----------
[docs]
def jackknife(data_list, *, eps=1):
r"""
Return jk[i] = avg - \frac{eps}{N} (v[i] - avg)
normal jackknife uses eps=1, scale the fluctuation by eps
"""
is_np_arr = isinstance(data_list, np.ndarray)
data_list_real = [d for d in data_list if d is not None]
n = len(data_list_real)
fac = eps / n
avg = q.average(data_list_real)
jks = [
avg,
]
for data in data_list:
if data is None:
jks.append(avg)
else:
jks.append(avg - fac * (data - avg))
if is_np_arr:
jks = np.array(jks, dtype=data_list.dtype)
return jks
# ----
[docs]
def jk_avg(jk_arr):
val = jk_arr[0]
return q.filter_np_results(val)
[docs]
def jk_err(jk_arr, *, eps=1, block_size=1):
"""
Return\n
.. math::\n
\\frac{1}{eps} \\sqrt{ N/(N-block_size) \\sum_{i=1}^N (jk[i] - jk_avg)^2 }.\n
when ``block_size=1``.
Note: ``len(jk_arr) = N + 1``.
Same ``eps`` as the ``eps`` used in the ``jackknife`` function.
Does not properly honor the :math:`(N-1)` formula in error calculation
if there were missing data in the original ``data_list`` in the ``jackknife`` function.
"""
assert block_size >= 1
avg = jk_avg(jk_arr)
n = len(jk_arr) - 1
if n <= 1:
fac = 1 / abs(eps)
val = fac * avg
val = q.filter_np_results(val)
return val
if n < 2 * block_size:
block_size = 1
assert n > block_size
blocks = q.block_data(jk_arr[1:], block_size)
diff_sqr = q.average([q.fsqr(jk - avg) for jk in blocks])
fac = math.sqrt(block_size / (n - block_size)) * n / abs(eps)
val = fac * q.fsqrt(diff_sqr)
val = q.filter_np_results(val)
return val
[docs]
def jk_avg_err(jk_arr, *, eps=1, block_size=1):
return jk_avg(jk_arr), jk_err(jk_arr, eps=eps, block_size=block_size)
[docs]
@q.timer
def sjackknife(
data_list,
jk_idx_list,
*,
avg=None,
is_hash_jk_idx=True,
jk_idx_hash_size=None,
rng_state=None,
all_jk_idx=None,
get_all_jk_idx=None,
jk_blocking_func=None,
eps=1,
):
"""
Super jackknife.
Return ``jk_arr``.
``len(jk_idx_list) == len(data_list)``
``len(jk_arr) == len(all_jk_idx)``
``jk_idx_list`` (after processed by ``jk_blocking_func``) should be contained in ``all_jk_idx``,
otherwise, if ``is_hash_jk_idx`` is true, then, a hash of ``jk_idx`` will be used instead.
Ideally, ``all_jk_idx`` should only contain distinct indices.
However, if there are repeatations, indices appear later take precedence.
if ``all_jk_idx`` and ``get_all_jk_idx`` are both ``None``, then a trivial
``all_jk_idx`` will be created based on ``jk_idx_hash_size``.
"""
if jk_idx_hash_size is None:
jk_idx_hash_size = 1024
if rng_state is None:
rng_state = q.RngState("rejk")
rs = rng_state
assert len(jk_idx_list) == len(data_list)
if isinstance(data_list, np.ndarray):
dtype = data_list.dtype
else:
dtype = None
data_list_real = [d for d in data_list if d is not None]
data_arr = np.array(data_list_real, dtype=dtype)
if avg is None:
avg = q.average(data_arr)
dtype = data_arr.dtype
jk_idx_list = [jk_idx for jk_idx, d in zip(jk_idx_list, data_list) if d is not None]
if jk_blocking_func is not None:
jk_idx_list = [jk_blocking_func(0, jk_idx) for jk_idx in jk_idx_list]
n = len(data_arr)
assert n == len(jk_idx_list)
if all_jk_idx is None:
if get_all_jk_idx is None:
assert is_hash_jk_idx
all_jk_idx = [
"avg",
] + list(range(jk_idx_hash_size))
else:
all_jk_idx = get_all_jk_idx()
assert all_jk_idx[0] == "avg"
n_super_sample = len(all_jk_idx) - 1
i_dict = dict()
for i, jk_idx in enumerate(all_jk_idx):
jk_idx_str = str(jk_idx)
i_dict[jk_idx_str] = i
i_arr = np.zeros(n, dtype=np.int32)
for j in range(n):
jk_idx = jk_idx_list[j]
jk_idx_str = str(jk_idx)
if jk_idx_str in i_dict:
i = i_dict[jk_idx_str]
else:
assert is_hash_jk_idx
rsi = rs.split(jk_idx_str)
i = 1 + int(rsi.rand_gen() % n_super_sample)
assert i > 0
i_arr[j] = i
count_dict = dict()
for j in range(n):
i = i_arr[j]
if i in count_dict:
count_dict[i] += 1
else:
count_dict[i] = 1
jk_arr = np.empty(
(
1 + n_super_sample,
*data_arr[0].shape,
),
dtype=dtype,
)
jk_arr[:] = avg
data_diff = data_arr - avg
for j in range(n):
i = i_arr[j]
assert i > 0
assert i in count_dict
if n > count_dict[i]:
n_b = n - count_dict[i]
fac = -eps * np.sqrt(1 / (n * n_b))
jk_arr[i] += fac * data_diff[j]
return jk_arr
@q.timer
def sjackknife_distributed(
data_list,
jk_idx_list,
*,
avg=None,
is_hash_jk_idx=True,
jk_idx_hash_size=None,
rng_state=None,
all_jk_idx=None,
get_all_jk_idx=None,
jk_blocking_func=None,
eps=1,
):
r"""
``sjackknife`` with distributed input and distributed output.\n
This is a collective MPI operation: every node must call it with the same
parameters and with its own disjoint part of the data set (``data_list``
and ``jk_idx_list`` are the local parts); every node returns its own part
of the ``jk_arr`` of ``sjackknife``. Node ``r`` out of ``num_node`` nodes
holds the samples in ``range(*get_distributed_range(len(all_jk_idx), r,
num_node))``, so that::
jk_arr = np.concatenate(q.get_comm().allgather(jk_local))
is the complete data set, in the same order as the one obtained with
``sjackknife``. No node needs to hold the whole data set or the whole
result; only the small ``jk_idx`` metadata is gathered on every node.\n
The result agrees with the one of ``sjackknife`` up to the floating-point
roundoff, but not bit-for-bit: the average and the sums over the data set
are reduced across the nodes, which changes the order of the
floating-point additions.\n
``sjackknife_sync_node`` performs the corresponding collective operation
for the case where every node has the whole input; ``g_mk_jk_distributed``
dispatches to this function.
"""
fname = "sjackknife_distributed"
comm, id_node, num_node = get_collective_comm(fname)
if jk_idx_hash_size is None:
jk_idx_hash_size = 1024
if rng_state is None:
rng_state = q.RngState("rejk")
(
data_arr,
jk_idx_list_local,
jk_idx_list_glb,
n,
elem_shape,
dtype,
) = get_distributed_jk_input(fname, comm, data_list, jk_idx_list)
avg = get_distributed_avg(comm, data_arr, n, avg)
if all_jk_idx is None:
if get_all_jk_idx is None:
assert is_hash_jk_idx
all_jk_idx = [
"avg",
] + list(range(jk_idx_hash_size))
else:
all_jk_idx = get_all_jk_idx()
assert all_jk_idx[0] == "avg"
total_size = len(all_jk_idx)
i_start, i_end = get_distributed_range(total_size, id_node, num_node)
# ``partial_arr`` holds the contribution of the local data to every sample
partial_arr = np.zeros((total_size,) + elem_shape, dtype=dtype)
if i_start == 0 and i_end > i_start:
# the sample 0 is the average and must be counted only once
partial_arr[0] = avg
#
n_super_sample = total_size - 1
if jk_blocking_func is None:
b_jk_idx_list_local = jk_idx_list_local
b_jk_idx_list_glb = jk_idx_list_glb
else:
b_jk_idx_list_local = [
jk_blocking_func(0, jk_idx) for jk_idx in jk_idx_list_local
]
b_jk_idx_list_glb = [jk_blocking_func(0, jk_idx) for jk_idx in jk_idx_list_glb]
i_dict = dict()
for i, jk_idx in enumerate(all_jk_idx):
i_dict[str(jk_idx)] = i
rs = rng_state
#
def get_i(jk_idx):
jk_idx_str = str(jk_idx)
if jk_idx_str in i_dict:
i = i_dict[jk_idx_str]
else:
assert is_hash_jk_idx
i = 1 + int(rs.split(jk_idx_str).rand_gen() % n_super_sample)
assert i > 0
return i
#
count_dict = dict()
for jk_idx in b_jk_idx_list_glb:
i = get_i(jk_idx)
count_dict[i] = count_dict.get(i, 0) + 1
data_diff = data_arr - avg
for j in range(len(data_arr)):
i = get_i(b_jk_idx_list_local[j])
count = count_dict[i]
if n > count:
fac = -eps * np.sqrt(1 / (n * (n - count)))
partial_arr[i] += fac * data_diff[j]
return get_reduce_scattered_jk_arr(partial_arr, comm, id_node, num_node, avg)
@q.timer
def sjackknife_sync_node(
data_list,
jk_idx_list,
*,
avg=None,
is_hash_jk_idx=True,
jk_idx_hash_size=None,
rng_state=None,
all_jk_idx=None,
get_all_jk_idx=None,
jk_blocking_func=None,
eps=1,
):
r"""
``sjackknife`` as a collective MPI operation where every node has the same
(whole) input and obtains the same complete result.\n
This implements the ``is_sync_node=True`` option of ``sjackknife``: the
input is split between the nodes in the order of the nodes and
``sjackknife_distributed`` is called on the local parts; the local results
are then gathered with ``Allgatherv``, so that every node obtains the
complete ``jk_arr`` of ``sjackknife``. The result agrees with the one of
``sjackknife`` up to the floating-point roundoff, but not bit-for-bit.\n
``g_mk_jk_sync_node`` dispatches to this function.
"""
comm, id_node, num_node = get_collective_comm("sjackknife_sync_node")
i_start, i_end = get_distributed_range(len(data_list), id_node, num_node)
jk_local = sjackknife_distributed(
data_list[i_start:i_end],
jk_idx_list[i_start:i_end],
avg=avg,
is_hash_jk_idx=is_hash_jk_idx,
jk_idx_hash_size=jk_idx_hash_size,
rng_state=rng_state,
all_jk_idx=all_jk_idx,
get_all_jk_idx=get_all_jk_idx,
jk_blocking_func=jk_blocking_func,
eps=eps,
)
return get_gathered_jk_arr(jk_local, comm, num_node)
[docs]
@q.timer
def sjk_mk_jk_val(
rs_tag,
val,
err,
*,
is_hash_jk_idx=True,
jk_idx_hash_size=None,
rng_state=None,
all_jk_idx=None,
get_all_jk_idx=None,
eps=1,
):
"""
return jk_arr
n = n_rand_sample
len(jk_arr) == 1 + n
jk_arr[i] = val + err * r[i] for i in 1..n
where r[i] ~ N(0, 1)
"""
if jk_idx_hash_size is None:
jk_idx_hash_size = 1024
assert jk_idx_hash_size >= 0
if rng_state is None:
rng_state = q.RngState("rejk")
rs = rng_state
if all_jk_idx is None:
if get_all_jk_idx is None:
assert is_hash_jk_idx
all_jk_idx = [
"avg",
] + list(range(jk_idx_hash_size))
else:
all_jk_idx = get_all_jk_idx()
assert all_jk_idx[0] == "avg"
n_super_sample = len(all_jk_idx) - 1
assert n_super_sample >= 0
assert isinstance(rng_state, q.RngState)
assert isinstance(val, q.real_types)
assert isinstance(err, q.real_types)
rs = rng_state.split(str(rs_tag))
jk_arr = np.zeros((n_super_sample + 1,), dtype=np.float64)
jk_arr[0] = val
r_arr = rs.g_rand_arr((n_super_sample,))
r_arr_qnorm = q.qnorm(r_arr)
r_arr = r_arr * np.sqrt(1 / r_arr_qnorm)
assert abs(q.qnorm(r_arr) - 1) < 1e-8
jk_arr[1:] = val + eps * r_arr * err
return jk_arr
[docs]
def sjk_avg(jk_arr):
return jk_avg(jk_arr)
[docs]
def sjk_err(jk_arr, *, eps=1):
"""
Return\n
.. math::\n
\\frac{1}{eps} \\sqrt{ \\sum_{i=1}^N (jk[i] - jk_avg)^2 }.\n
Note: ``len(jk_arr) = N + 1``.
Same ``eps`` as the ``eps`` used in the ``jackknife`` function.
"""
avg = jk_avg(jk_arr)
n = len(jk_arr) - 1
if n <= 1:
fac = 1 / abs(eps)
val = fac * avg
val = q.filter_np_results(val)
return val
diff_sqr = q.average([q.fsqr(jk - avg) for jk in jk_arr[1:]])
fac = math.sqrt(n) / abs(eps)
val = fac * q.fsqrt(diff_sqr)
val = q.filter_np_results(val)
return val
[docs]
def sjk_avg_err(jk_arr, *, eps=1):
return sjk_avg(jk_arr), sjk_err(jk_arr, eps=eps)
# ----------
@q.timer
def _fill_garr_dict(
rng_state,
garr_dict,
name_list,
*,
n_rand_sample,
is_normalizing_rand_sample,
):
r"""
Fill ``garr_dict`` with the random arrays of the jackknife streams.\n
For every name in ``name_list`` which is not in ``garr_dict`` yet, add
``garr_dict[name] = rng_state.split(name).g_rand_arr(n_rand_sample)``,
normalized to ``\sum_i r_i^2 = n_rand_sample`` when
``is_normalizing_rand_sample`` is True. A stream only depends on its
name, so the dict may be reused between calls (e.g. between the sample
blocks of a chunked computation) to generate every stream only once.\n
``name_list`` may contain ``None`` entries, which are skipped.
"""
for name in name_list:
if name is None or name in garr_dict:
continue
garr = rng_state.split(name).g_rand_arr(n_rand_sample)
if is_normalizing_rand_sample:
# garr_qnorm \approx n_rand_sample
garr_qnorm = q.qnorm(garr)
garr = garr * np.sqrt(n_rand_sample / garr_qnorm)
assert abs(q.qnorm(garr) / n_rand_sample - 1) < 1e-8
garr_dict[name] = garr
def _apply_jk_fac(r_arr, b_arr, n, eps):
r"""
Multiply ``r_arr`` in place by the jackknife factor.\n
The factor is ``-eps / sqrt(n * n_b)`` with ``n_b = n - b_arr``, where the
entries with ``n_b <= 0`` are replaced by 1 and the factor is set to 0
there (a block which holds the whole data set contributes nothing).
``r_arr`` and ``b_arr`` have shape ``(n_sample, n_data)`` and may be views
(e.g. one sample block of a chunked computation); ``n`` is the size of the
whole data set and ``eps`` the overall scale.\n
The factor is built in a single float64 buffer, keeping ``n * n_b`` in it,
so no extra ``(n_sample, n_data)`` array and no ``n_b`` int64 array are
allocated and the rounding is the same as with the separate arrays.
"""
fac_arr = np.asarray(n - b_arr, dtype=np.float64)
np.maximum(fac_arr, 1, out=fac_arr)
fac_arr *= n
np.sqrt(fac_arr, out=fac_arr)
fac_arr = -eps / fac_arr
fac_arr[b_arr >= n] = 0
r_arr *= fac_arr
@q.timer
def mk_r_i_j_mat(
n_rand_sample,
jk_idx_list,
rng_state,
*,
jk_blocking_func,
is_normalizing_rand_sample,
is_apply_rand_sample_jk_idx_blocking_shift,
is_use_old_rand_alg,
i_range=None,
jk_idx_list_for_count=None,
garr_dict=None,
):
"""
Return ``(r_arr, b_arr)``.\n
``r_arr`` and ``b_arr`` have shape ``(n_sample, n)`` with
``n_sample = i_end - i_start``, where ``i_range = (i_start, i_end,)``
(default ``None`` means the full range ``(0, n_rand_sample,)``).
Only the rows in ``i_range`` are computed; the per-block random streams
do not depend on ``i_range``, so every returned row is identical to the
corresponding row of a full-range call.\n
``jk_idx_list_for_count`` (default ``None`` means ``jk_idx_list``) is the
list over which the block sizes ``b_arr`` are counted; the columns of
``r_arr`` and ``b_arr`` always correspond to ``jk_idx_list``. This is used
by ``g_mk_jk_distributed`` where the columns are the locally held
``jk_idx`` while the block sizes must be counted over the whole data set.\n
``garr_dict`` (default ``None``) is the dict of the random streams keyed by
the ``rng_state.split`` name; when it is given, the streams already in it
are reused and only the missing names are generated, so a caller which
calls this function once per ``i_range`` block can pass the same dict and
generate every stream only once. ``None`` uses a private dict, which is
discarded together with the returned arrays.
"""
assert n_rand_sample >= 0
if i_range is None:
i_start, i_end = 0, n_rand_sample
else:
i_start, i_end = i_range
assert 0 <= i_start <= i_end <= n_rand_sample
n_sample = i_end - i_start
rs = rng_state
n = len(jk_idx_list)
r_arr = np.empty(
(
n_sample,
n,
),
dtype=np.float64,
)
b_arr = np.empty(
(
n_sample,
n,
),
dtype=np.int32,
)
codes = np.empty(
(
n_sample,
n,
),
dtype=np.int32,
)
if jk_blocking_func is None:
is_apply_rand_sample_jk_idx_blocking_shift = False
# ``b_arr`` counts the blocks of ``jk_idx_list_for_count`` (default
# ``jk_idx_list``); the columns of ``r_arr`` and ``b_arr`` always
# correspond to ``jk_idx_list``.
is_separate_count = jk_idx_list_for_count is not None
count_jk_idx_list = (
jk_idx_list if jk_idx_list_for_count is None else jk_idx_list_for_count
)
#
# ``codes[i, j]`` identifies the blocked ``jk_idx`` of the cell ``(i, j)``
# with a small int, and ``names[code] == str(blocked_jk_idx)`` is the name
# of its random stream. The name is what the random numbers are keyed by,
# but building a fresh ``str(blocked_jk_idx)`` for every cell used to
# dominate the runtime, so it is built once per distinct name and reused
# through ``code_dict``/``name_dict``.
names = list()
code_dict = dict()
name_dict = dict()
cell_codes = set()
#
def get_code(blocked_jk_idx, *, is_cell):
try:
code = code_dict.get(blocked_jk_idx)
except TypeError:
# an unhashable blocked jk_idx: fall back to keying on the name
code = None
if code is None:
name = str(blocked_jk_idx)
code = name_dict.get(name)
if code is None:
code = len(names)
name_dict[name] = code
names.append(name)
try:
code_dict[blocked_jk_idx] = code
except TypeError:
pass
if is_cell:
cell_codes.add(code)
return code
#
def is_plain_jk_idx(jk_idx):
if type(jk_idx) is int:
return True
return (
type(jk_idx) is tuple
and len(jk_idx) == 2
and type(jk_idx[0]) is str
and type(jk_idx[1]) is int
)
#
def try_set_jk_idx_default_shift():
# For the default blocking function, ``jk_blocking_func_default(i + 1,
# jk_idx)`` (with ``i`` the zero based sample) depends on ``i`` only
# through ``b_shift = shift % block_size``, where ``shift =
# jk_blocking_traj_shift_arr[i % len(jk_blocking_traj_shift_arr)]``.
# One call of the real blocking function per distinct ``(block_size,
# b_shift)`` pair therefore reproduces every cell exactly, so the
# ``n_sample * n`` per-cell calls are replaced by a few such calls and
# vectorized fills.
if jk_blocking_func is not jk_blocking_func_default:
return False
for jk_idx in jk_idx_list:
if not is_plain_jk_idx(jk_idx):
return False
for jk_idx in count_jk_idx_list:
if not is_plain_jk_idx(jk_idx):
return False
block_size = default_g_jk_kwargs.get("block_size")
if not isinstance(block_size, q.int_types):
return False
block_size_dict = default_g_jk_kwargs.get("block_size_dict")
if block_size_dict is None:
block_size_dict = dict()
all_jk_idx_set = default_g_jk_kwargs.get("all_jk_idx_set")
if all_jk_idx_set is not None:
# ``jk_blocking_func_default`` adds every ``jk_idx`` it sees to
# this set; the content is the same when each one is added once.
for jk_idx in jk_idx_list:
all_jk_idx_set.add(jk_idx)
for jk_idx in count_jk_idx_list:
all_jk_idx_set.add(jk_idx)
#
def get_block_size(jk_idx):
if type(jk_idx) is int:
return block_size
return block_size_dict.get(jk_idx[0], block_size)
#
block_size_list = sorted({get_block_size(jk_idx) for jk_idx in jk_idx_list})
count_block_size_list = sorted(
{get_block_size(jk_idx) for jk_idx in count_jk_idx_list}
)
all_block_size_list = sorted(set(block_size_list).union(count_block_size_list))
i_arr = np.arange(i_start, i_end)
shift_arr = jk_blocking_traj_shift_arr[i_arr % len(jk_blocking_traj_shift_arr)]
# The rows of ``codes`` and ``b_arr`` only depend on the tuple of the
# ``b_shift`` values, so group them and call the blocking function once
# per group (the group representative) instead of once per cell.
rows_of_key = dict()
rep_i_of_key = dict()
for i_local in range(n_sample):
key = tuple(int(shift_arr[i_local]) % bs for bs in all_block_size_list)
rows_of_key.setdefault(key, list()).append(i_local)
if key not in rep_i_of_key:
rep_i_of_key[key] = i_start + i_local
cells_of_block_size = dict()
for bs in block_size_list:
cells_of_block_size[bs] = np.asarray(
[
j
for j, jk_idx in enumerate(jk_idx_list)
if get_block_size(jk_idx) == bs
],
dtype=np.int64,
)
code_arr_of = dict()
for key, i_rep in rep_i_of_key.items():
row = np.empty(n, dtype=np.int32)
for bs, b_shift in zip(all_block_size_list, key):
cols = cells_of_block_size.get(bs)
if cols is None:
continue
code_arr = code_arr_of.get((bs, b_shift))
if code_arr is None:
code_arr = np.empty(len(cols), dtype=np.int32)
for pos, j in enumerate(cols.tolist()):
code_arr[pos] = get_code(
jk_blocking_func(i_rep + 1, jk_idx_list[j]),
is_cell=True,
)
code_arr_of[(bs, b_shift)] = code_arr
row[cols] = code_arr
for i_local in rows_of_key[key]:
codes[i_local, :] = row
for key, i_rep in rep_i_of_key.items():
if count_jk_idx_list:
count_codes = np.empty(len(count_jk_idx_list), dtype=np.int32)
for pos, jk_idx in enumerate(count_jk_idx_list):
count_codes[pos] = get_code(
jk_blocking_func(i_rep + 1, jk_idx), is_cell=False
)
counts = np.bincount(count_codes, minlength=len(names))
else:
counts = np.zeros(0, dtype=np.int64)
row = codes[rows_of_key[key][0], :]
rows = np.asarray(rows_of_key[key], dtype=np.int64)
b_arr[rows, :] = counts[row].astype(np.int32)
return True
#
def set_jk_idx_generic_shift():
for i_local in range(n_sample):
i_arg = i_start + i_local + 1
row = codes[i_local]
blocked_cache = dict()
for j in range(n):
jk_idx = jk_idx_list[j]
blocked_jk_idx = jk_blocking_func(i_arg, jk_idx)
try:
blocked_cache[jk_idx] = blocked_jk_idx
except TypeError:
pass
row[j] = get_code(blocked_jk_idx, is_cell=True)
count_dict = dict()
if is_separate_count:
for jk_idx in count_jk_idx_list:
try:
blocked_jk_idx = blocked_cache[jk_idx]
except (KeyError, TypeError):
blocked_jk_idx = jk_blocking_func(i_arg, jk_idx)
code = get_code(blocked_jk_idx, is_cell=False)
count_dict[code] = count_dict.get(code, 0) + 1
else:
for j in range(n):
code = row[j]
count_dict[code] = count_dict.get(code, 0) + 1
for j in range(n):
b_arr[i_local, j] = count_dict[row[j]]
#
def set_jk_idx_no_shift():
row = np.empty(n, dtype=np.int32)
for j in range(n):
jk_idx = jk_idx_list[j]
if jk_blocking_func is not None:
jk_idx = jk_blocking_func(0, jk_idx)
row[j] = get_code(jk_idx, is_cell=True)
count_dict = dict()
if is_separate_count:
for jk_idx in count_jk_idx_list:
if jk_blocking_func is not None:
jk_idx = jk_blocking_func(0, jk_idx)
code = get_code(jk_idx, is_cell=False)
count_dict[code] = count_dict.get(code, 0) + 1
else:
for j in range(n):
code = row[j]
count_dict[code] = count_dict.get(code, 0) + 1
if n_sample > 0:
# the blocked jk_idx does not depend on the sample: every row of
# ``codes`` and ``b_arr`` is the same
codes[:, :] = row
for j in range(n):
b_arr[:, j] = count_dict[row[j]]
#
@q.timer
def set_jk_idx():
if is_apply_rand_sample_jk_idx_blocking_shift:
if try_set_jk_idx_default_shift():
return
set_jk_idx_generic_shift()
else:
set_jk_idx_no_shift()
#
set_jk_idx()
if is_use_old_rand_alg == "v1":
assert not is_normalizing_rand_sample
for i_local in range(n_sample):
i = i_start + i_local
rsi = rs.split(str(i))
row = codes[i_local]
for j in range(n):
r_arr[i_local, j] = rsi.split(names[row[j]]).g_rand_gen()
return r_arr, b_arr
assert not is_use_old_rand_alg
#
# The random stream of a blocked ``jk_idx`` only depends on its name; the
# caller may pass ``garr_dict`` to keep the streams between calls (e.g.
# between the sample blocks of a chunked computation).
if garr_dict is None:
garr_dict = dict()
_fill_garr_dict(
rs,
garr_dict,
[names[code] for code in cell_codes],
n_rand_sample=n_rand_sample,
is_normalizing_rand_sample=is_normalizing_rand_sample,
)
garr_of_code = dict()
for code in cell_codes:
garr_of_code[code] = garr_dict[names[code]]
#
@q.timer
def set_r_arr():
if not is_apply_rand_sample_jk_idx_blocking_shift:
if n_sample > 0:
for j in range(n):
r_arr[:, j] = garr_of_code[int(codes[0, j])][i_start:i_end]
return
# Group the rows which share the same blocked codes, so that every
# column is filled with one slice/gather per group instead of one
# element per cell.
rows_of_key = dict()
for i_local in range(n_sample):
rows_of_key.setdefault(codes[i_local].tobytes(), list()).append(i_local)
for rows in rows_of_key.values():
row = codes[rows[0]]
if len(rows) == n_sample:
for j in range(n):
r_arr[:, j] = garr_of_code[int(row[j])][i_start:i_end]
else:
rows_arr = np.asarray(rows, dtype=np.int64)
glb_rows = rows_arr + i_start
for j in range(n):
r_arr[rows_arr, j] = garr_of_code[int(row[j])][glb_rows]
#
set_r_arr()
return r_arr, b_arr
def get_distributed_range(total_size, id_node, num_node):
"""
Return ``(i_start, i_end)``, the range of the ``total_size`` indices owned
by node ``id_node`` out of ``num_node`` nodes.\n
The ranges are contiguous, ordered by ``id_node``, cover
``range(total_size)`` exactly once and differ in size by at most one::
i_start = (total_size * id_node) // num_node
i_end = (total_size * (id_node + 1)) // num_node
Concatenating the parts in the order of the nodes reproduces the whole
range, so ``np.concatenate(comm.allgather(x_local))`` is the complete
array when ``x_local = x[i_start:i_end]``.
"""
assert isinstance(total_size, q.int_types)
assert 0 <= total_size
assert isinstance(id_node, q.int_types)
assert isinstance(num_node, q.int_types)
assert 0 <= id_node < num_node
i_start = (total_size * id_node) // num_node
i_end = (total_size * (id_node + 1)) // num_node
return i_start, i_end
def is_mpi_dtype(dtype):
"""
Return whether ``dtype`` is a numpy dtype that MPI can use to sum buffers,
i.e. a dtype supported by ``MPI.SUM``, like the ones used by
``g_mk_jk_distributed``.
"""
dtype = np.dtype(dtype)
if dtype.kind == "f":
return dtype.itemsize in (4, 8)
if dtype.kind == "c":
return dtype.itemsize in (8, 16)
if dtype.kind in ("i", "u"):
return dtype.itemsize in (1, 2, 4, 8)
return False
def get_collective_comm(tag):
"""
Return ``(comm, id_node, num_node)`` for a collective MPI operation.\n
``tag`` names the caller for the error messages, e.g.
``"g_mk_jk_sync_node"``.\n
``qlat`` (used for ``q.get_comm()``) is imported only here. The
communicator must be initialized on the whole MPI communicator, i.e. with
``q.begin_with_mpi()``, ``q.begin_with_gpt()`` or ``q.begin_with_grid()``
(or set with ``q.set_comm(...)``), so that
``comm.rank == q.get_id_node()`` and ``comm.size == q.get_num_node()``.
"""
import qlat
#
comm = qlat.get_comm()
if comm is None:
raise Exception(
f"{tag} requires the qlat communicator;"
" use q.begin_with_mpi(), q.begin_with_gpt() or"
" q.begin_with_grid() (or set it with q.set_comm(...))"
)
if comm.size != qlat.get_num_node():
raise Exception(
f"{tag} requires qlat to be initialized"
f" on the whole MPI communicator, but comm.size={comm.size} and"
f" q.get_num_node()={qlat.get_num_node()}"
)
assert comm.rank == qlat.get_id_node()
return comm, comm.rank, comm.size
def get_distributed_jk_input(fname, comm, data_list, jk_idx_list):
r"""
Return ``(data_arr, jk_idx_list_local, jk_idx_list_glb, n, elem_shape, dtype)``
for the distributed jackknife functions.\n
``data_list`` and ``jk_idx_list`` are the local parts of the data set and
``comm`` is the communicator from ``get_collective_comm``. ``data_arr`` is
the local data, with shape ``(n_local, *elem_shape)`` (``n_local == 0``
when the node holds no data); ``jk_idx_list_local`` are the local
``jk_idx`` and ``jk_idx_list_glb`` is the whole list of ``jk_idx`` gathered
on every node; ``n`` is the total number of data points; ``elem_shape`` and
``dtype`` describe a single data point.\n
``fname`` names the caller for the error messages. The data must have a
dtype supported by MPI, so that the results can be summed and gathered.
"""
from mpi4py import MPI
#
assert len(data_list) == len(jk_idx_list)
if isinstance(data_list, np.ndarray):
dtype = data_list.dtype
else:
dtype = None
data_list_local = [d for d in data_list if d is not None]
jk_idx_list_local = [
jk_idx for jk_idx, d in zip(jk_idx_list, data_list) if d is not None
]
data_arr = np.array(data_list_local, dtype=dtype)
n_local = len(data_arr)
# the dtype and the element shape of the data; the lowest rank with data
# provides them, so that a node without any data still knows them
if n_local > 0:
elem_info = (data_arr.dtype.str, tuple(data_arr.shape[1:]))
else:
elem_info = None
elem_info = next(
(info for info in comm.allgather(elem_info) if info is not None), None
)
if elem_info is None:
raise Exception(f"{fname}: the distributed data set is empty")
dtype = np.dtype(elem_info[0])
elem_shape = tuple(elem_info[1])
if not is_mpi_dtype(dtype):
raise Exception(
f"{fname}: the data must have a dtype supported by MPI so that the"
f" results can be summed, but dtype={dtype}"
)
if n_local == 0:
data_arr = np.zeros((0,) + elem_shape, dtype=dtype)
else:
data_arr = np.asarray(data_arr, dtype=dtype)
n_arr = np.array([n_local], dtype=np.int64)
comm.Allreduce(MPI.IN_PLACE, n_arr, op=MPI.SUM)
n = int(n_arr[0])
# only the small ``jk_idx`` metadata is gathered on every node
jk_idx_list_glb = [
jk_idx
for jk_idx_list_node in comm.allgather(jk_idx_list_local)
for jk_idx in jk_idx_list_node
]
return data_arr, jk_idx_list_local, jk_idx_list_glb, n, elem_shape, dtype
def get_distributed_avg(comm, data_arr, n, avg=None):
r"""
Return the average of the whole distributed data set.\n
When ``avg`` is not ``None``, it is returned as it is; otherwise the local
sums of ``data_arr`` are summed over the nodes with ``Allreduce`` and
divided by ``n``. This is the mean over the whole data set, which is
needed by the (randomized) Super-Jackknife, and it agrees with the one of
the sequential functions up to the floating-point roundoff.
"""
from mpi4py import MPI
#
if avg is not None:
return avg
local_sum = np.asarray(np.sum(data_arr, axis=0))
glb_sum = np.zeros_like(local_sum)
comm.Allreduce(local_sum, glb_sum, op=MPI.SUM)
return q.filter_np_results(glb_sum / n)
def get_reduce_scattered_jk_arr(partial_arr, comm, id_node, num_node, avg):
"""
Return the local part of the (randomized) Super-Jackknife data set.\n
``partial_arr`` has shape ``(total_size, *elem_shape)`` and holds the
contribution of the local data to every sample, with the sample 0 equal to
``avg`` on the node which owns it and 0 elsewhere. The contributions of
the nodes are summed with ``Reduce_scatter``, so that every node obtains
its own samples; ``avg`` is then added to the samples, since the samples
are ``avg + sum_j (...)`` while the sample 0 is ``avg`` itself.\n
``partial_arr`` may be any array, in particular a non-contiguous view such
as a column of a 2-D array; a contiguous copy is made when needed.
"""
from mpi4py import MPI
#
# The buffers of the collectives must be contiguous; ``reshape(-1)`` is
# only a view when the array is already 1-D, so copy a strided array here.
partial_arr = np.ascontiguousarray(partial_arr)
total_size = partial_arr.shape[0]
elem_shape = partial_arr.shape[1:]
elem_size = 1
for x in elem_shape:
elem_size *= x
recvcounts = []
for r in range(num_node):
r_start, r_end = get_distributed_range(total_size, r, num_node)
recvcounts.append((r_end - r_start) * elem_size)
i_start, i_end = get_distributed_range(total_size, id_node, num_node)
jk_arr = np.empty((i_end - i_start,) + elem_shape, dtype=partial_arr.dtype)
comm.Reduce_scatter(
partial_arr.reshape(-1),
jk_arr.reshape(-1),
recvcounts,
op=MPI.SUM,
)
if i_end - i_start > 0 and i_end > 1:
jk_arr[max(1 - i_start, 0) :] += avg
return jk_arr
def get_gathered_jk_arr(jk_local, comm, num_node):
"""
Return the complete (randomized) Super-Jackknife data set.\n
``jk_local`` is the local part of the result of a distributed jackknife
function; the parts are gathered with ``Allgatherv``, in the order of the
nodes, so that every node obtains the complete ``jk_arr``.\n
``jk_local`` may be any array, in particular a non-contiguous view such as
a column of a 2-D array; a contiguous copy is made when needed.
"""
from mpi4py import MPI
#
# The send buffer of the collective must be contiguous; ``reshape(-1)`` is
# only a view when the array is already 1-D, so copy a strided array here.
jk_local = np.ascontiguousarray(jk_local)
n_arr = np.array([len(jk_local)], dtype=np.int64)
comm.Allreduce(MPI.IN_PLACE, n_arr, op=MPI.SUM)
total_size = int(n_arr[0])
elem_size = 1
for x in jk_local.shape[1:]:
elem_size *= x
recvcounts = []
displs = []
displ = 0
for r in range(num_node):
r_start, r_end = get_distributed_range(total_size, r, num_node)
recvcounts.append((r_end - r_start) * elem_size)
displs.append(displ)
displ += recvcounts[-1]
jk_arr = np.empty((total_size,) + jk_local.shape[1:], dtype=jk_local.dtype)
comm.Allgatherv(
jk_local.reshape(-1),
(jk_arr.reshape(-1), (recvcounts, displs)),
)
return jk_arr
def get_jk_rows(fac_r_arr, data_diff, out=None):
r"""
Return the contribution of the local data to every sample, i.e.
``sum_j fac_r_arr[i, j] data_diff[j]``, with shape ``(n_sample,
*data_diff.shape[1:])`` and ``n_sample = fac_r_arr.shape[0]``.\n
``fac_r_arr`` has shape ``(n_sample, n_data)`` and ``data_diff`` has shape
``(n_data_diff, *elem_shape)``; the contraction sums over the
``n_data_diff`` data points of ``data_diff`` (``n_data`` and
``n_data_diff`` usually agree, and a caller may restrict both to the same
subset of the data).\n
The contraction is a matrix product, which is evaluated by BLAS without
materializing the broadcast temporary of
``fac_r_arr[:, :, None, ...] * data_diff[None]``: that temporary has
``n_sample * n_data * prod(elem_shape)`` elements, while the result has
only ``n_sample * prod(elem_shape)``. The matrix product is mathematically
identical to the broadcast contraction; it only changes the order of the
floating-point additions. The node which holds no data (``n_data == 0``)
returns zeros.\n
``out`` (default ``None``) is an optional array with the shape of the
result and C-contiguous storage, which is written in place and returned, so
a caller which already owns a buffer (e.g. the contribution buffer of the
(randomized) Super-Jackknife, which passes a slice of it so that its first
row stays untouched) does not pay for the extra ``n_sample *
prod(elem_shape)`` result array. Its dtype is the dtype in which the
result is returned: when it is the dtype of ``data_diff`` (the usual case,
equal to ``np.result_type(fac_r_arr.dtype, data_diff.dtype)`` there) the
contraction is written into it in place, exactly as before; when it is
narrower than the contraction (e.g. float32 data with a float64 ``avg``,
which widens ``data_diff``) the contraction is still computed in its
natural dtype and cast on the assignment into ``out``, which is what a
caller doing ``out[...] = get_jk_rows(...)`` used to get. With
``out=None`` a new array of the natural dtype is returned, exactly as
before.
"""
n_sample, n_data = fac_r_arr.shape
n_data_diff = data_diff.shape[0]
elem_shape = tuple(data_diff.shape[1:])
if out is not None:
assert out.shape == (n_sample,) + elem_shape
assert out.flags.c_contiguous
if n_data == 0 or n_data_diff == 0:
if out is None:
return np.zeros((n_sample,) + elem_shape, dtype=data_diff.dtype)
out[...] = 0
return out
if out is None:
return np.matmul(
fac_r_arr,
data_diff.reshape(n_data_diff, -1),
).reshape((n_sample,) + elem_shape)
if n_sample == 0:
# A node which owns no sample has an empty result (``n_rand_sample ==
# 0`` is the extreme case); the allocating path returns an empty array
# and the caller copies nothing, so keep the empty buffer zeroed.
out[...] = 0
return out
fac_r_mat = fac_r_arr
data_mat = data_diff.reshape(n_data_diff, -1)
result_dtype = np.result_type(fac_r_arr.dtype, data_diff.dtype)
if out.dtype == data_diff.dtype or out.dtype == result_dtype:
# ``out`` receives the contraction without a cast (the usual case, and
# the one which was already supported): keep the in-place matmul.
np.matmul(fac_r_mat, data_mat, out=out.reshape(n_sample, -1))
return out
# ``out`` is narrower than the contraction (e.g. float32 data with a
# float64 ``avg``, which widens ``data_diff``): compute the contraction in
# its natural dtype and cast on the assignment, which is exactly what the
# pre-``out=`` callers got from ``out[...] = get_jk_rows(...)``. Casting
# the inputs first would change the rounding.
out[...] = np.matmul(fac_r_mat, data_mat).reshape((n_sample,) + elem_shape)
return out
def _get_jk_sample_blocks(total_size, num_node):
r"""
Return ``(R, K, blocks)`` for the block-wise random-sample loop.\n
``R = ceil(total_size / num_node)`` rows per block, ``K = ceil(total_size /
R)`` blocks, and ``blocks`` is the list of ``(b0, b1)`` row ranges, which
are consecutive and cover ``range(total_size)`` exactly once.\n
The result only depends on the job-wide sizes (``total_size = 1 +
n_rand_sample`` and ``num_node``), so every node of a collective call
derives the same list without any communication. ``K <= num_node`` and
``total_size <= num_node`` gives one row per block; the first ``K - 1``
blocks have exactly ``R`` rows and the last has the remainder in
``[1, R]``.\n
One block per node: the sample axis is distributed with the same ``ceil``
arithmetic that already distributes the output, so the send buffer of one
block is of the order of the local output instead of the whole sample
axis, and the peak memory of the contribution buffer drops by about
``num_node``.
"""
assert total_size >= 1
assert num_node >= 1
R = -(-int(total_size) // int(num_node))
K = -(-int(total_size) // R)
blocks = [(b * R, min((b + 1) * R, total_size)) for b in range(K)]
assert blocks[0][0] == 0
assert blocks[-1][1] == total_size
assert all(b1 > b0 for b0, b1 in blocks)
assert all(blocks[b][1] == blocks[b + 1][0] for b in range(K - 1))
assert K <= num_node
return R, K, blocks
def _get_reduce_scattered_jk_block(
partial_blk, jk_view, block_start, block_end, total_size, comm
):
r"""
Sum the local contributions of the samples in ``range(block_start,
block_end)`` over the nodes and write this node's part into ``jk_view``.\n
``partial_blk`` is the ``(block_end - block_start, *elem_shape)``
C-contiguous contribution of the local data (the send buffer) and
``jk_view`` is the matching C-contiguous slice ``jk_arr[off:off +
n_rows]`` of the local output (possibly empty). The average is added by
the caller after the loop.\n
``Reduce_scatter`` has no ``displs``: the send buffer is the concatenation
of the node segments in node order, which matches the row order because
``get_distributed_range`` returns contiguous, node-ordered ranges. A node
with no row in the block gets ``recvcounts[node] == 0`` and passes a
0-length array, which is a legal ``Reduce_scatter`` argument.
"""
from mpi4py import MPI
#
partial_blk = np.ascontiguousarray(partial_blk)
elem_size = 1
for x in partial_blk.shape[1:]:
elem_size *= x
recvcounts = []
for r in range(comm.size):
r_start, r_end = get_distributed_range(total_size, r, comm.size)
n_rows = min(block_end, r_end) - max(block_start, r_start)
recvcounts.append(max(0, n_rows) * elem_size)
assert sum(recvcounts) == partial_blk.size
assert jk_view.size == recvcounts[comm.rank]
comm.Reduce_scatter(
partial_blk.reshape(-1),
jk_view.reshape(-1),
recvcounts,
op=MPI.SUM,
)
[docs]
@q.timer
def rjackknife(
data_list,
jk_idx_list,
*,
avg=None,
rng_state=None,
n_rand_sample=None,
jk_blocking_func=None,
is_normalizing_rand_sample=False,
is_apply_rand_sample_jk_idx_blocking_shift=True,
is_use_old_rand_alg=False,
eps=1,
is_sync_node=False,
):
r"""
Jackknife-bootstrap hybrid resampling.
Return ``jk_arr``.
``len(jk_arr) == 1 + n_rand_sample``
distribution of ``jk_arr`` should be similar as the distribution of ``avg``.
``r_{i,j} ~ N(0, 1)``\n
::\n
if is_normalizing_rand_sample:
n_j = \sum_i r_{i,j}^2
r_{i,j} <- \sqrt{n_rand_sample / n_j} r_{i,j}
data_list_real = [d for d in data_list if d is not None]
data_arr = np.array(data_list_real, dtype=dtype)
avg = average(data_arr)
len(data_list_real) = n
jk_arr[0] = avg
jk_arr[i] = avg + \sum_{j=1}^{n} (-eps/\sqrt{n (n - b(i,j))}) r_{i,j} (data_list_real[j] - avg)\n
where ``b(i,j)`` represent the ``block_size``.\n
if ``jk_blocking_func`` is provided::\n
``jk_blocking_func(i, jk_idx) => blocked jk_idx``\n
::\n
jk_arr[i] = avg + \sum_{j=1}^{n} r_{i,jk_block_func(j)} (jk_arr[j] - avg)\n
If ``is_sync_node`` is True:\n
Assume this is a collective operation in a MPI program where every
node have the same input. The operation is performed by
``rjackknife_sync_node``, which splits the input between the nodes and
calls ``rjackknife_distributed``; every node obtains the complete
``jk_arr``, which agrees with the one obtained with
``is_sync_node=False`` up to the floating-point roundoff (not
bit-for-bit, because the average and the sums over the data set are
reduced across the nodes). ``qlat`` (used for ``q.get_comm()``) and
``mpi4py`` are imported only when ``is_sync_node`` is True.
"""
if is_sync_node:
# Collective MPI operation: every node has the same input and obtains
# the complete result, which is computed distributedly.
return rjackknife_sync_node(
data_list,
jk_idx_list,
avg=avg,
rng_state=rng_state,
n_rand_sample=n_rand_sample,
jk_blocking_func=jk_blocking_func,
is_normalizing_rand_sample=is_normalizing_rand_sample,
is_apply_rand_sample_jk_idx_blocking_shift=is_apply_rand_sample_jk_idx_blocking_shift,
is_use_old_rand_alg=is_use_old_rand_alg,
eps=eps,
)
if n_rand_sample is None:
n_rand_sample = 1024
if rng_state is None:
rng_state = q.RngState("rejk")
assert len(data_list) == len(jk_idx_list)
assert isinstance(n_rand_sample, q.int_types)
assert n_rand_sample >= 0
assert isinstance(rng_state, q.RngState)
if isinstance(data_list, np.ndarray):
dtype = data_list.dtype
else:
dtype = None
data_list_real = [d for d in data_list if d is not None]
data_arr = np.array(data_list_real, dtype=dtype)
if avg is None:
avg = q.average(data_arr)
dtype = data_arr.dtype
jk_idx_list = [jk_idx for jk_idx, d in zip(jk_idx_list, data_list) if d is not None]
n = len(data_arr)
#
r_arr, b_arr = mk_r_i_j_mat(
n_rand_sample,
jk_idx_list,
rng_state,
jk_blocking_func=jk_blocking_func,
is_normalizing_rand_sample=is_normalizing_rand_sample,
is_apply_rand_sample_jk_idx_blocking_shift=is_apply_rand_sample_jk_idx_blocking_shift,
is_use_old_rand_alg=is_use_old_rand_alg,
)
_apply_jk_fac(r_arr, b_arr, n, eps)
data_diff = data_arr - avg
jk_arr = np.empty(
(
1 + n_rand_sample,
*data_arr[0].shape,
),
dtype=dtype,
)
jk_arr[0] = avg
if np.result_type(r_arr.dtype, data_diff.dtype) == jk_arr.dtype:
# Write the contraction straight into the result array; adding ``avg``
# afterwards is elementwise the same as ``avg + get_jk_rows(...)``.
get_jk_rows(r_arr, data_diff, out=jk_arr[1:])
jk_arr[1:] += avg
else:
# ``jk_arr`` has a narrower dtype than the contraction (e.g. float32
# data): keep the two step computation so that the rounding does not
# change.
jk_arr[1:] = avg + get_jk_rows(r_arr, data_diff)
return jk_arr
@q.timer
def rjackknife_distributed(
data_list,
jk_idx_list,
*,
avg=None,
rng_state=None,
n_rand_sample=None,
jk_blocking_func=None,
is_normalizing_rand_sample=False,
is_apply_rand_sample_jk_idx_blocking_shift=True,
is_use_old_rand_alg=False,
eps=1,
):
r"""
``rjackknife`` with distributed input and distributed output.\n
This is a collective MPI operation: every node must call it with the same
parameters and with its own disjoint part of the data set (``data_list``
and ``jk_idx_list`` are the local parts); every node returns its own part
of the ``jk_arr`` of ``rjackknife``. Node ``r`` out of ``num_node`` nodes
holds the samples in
``range(*get_distributed_range(1 + n_rand_sample, r, num_node))``, so
that::
jk_arr = np.concatenate(q.get_comm().allgather(jk_local))
is the complete data set, in the same order as the one obtained with
``rjackknife``. No node needs to hold the whole data set, the whole
random matrix or the whole result; only the small ``jk_idx`` metadata is
gathered on every node.\n
The result agrees with the one of ``rjackknife`` up to the floating-point
roundoff, but not bit-for-bit: the average and the sums over the data set
are reduced across the nodes, which changes the order of the
floating-point additions.\n
The samples are built one block of rows at a time, so the contribution
buffer of one call is ``R`` rows instead of the whole sample axis ``T = 1 +
n_rand_sample``: with ``N = num_node`` rows per block ``R = ceil(T / N)``
and ``K = ceil(T / R) <= N`` blocks, i.e. about one block per node (the
rule is automatic, there is no option to change it). The blocks are a
pure function of ``T`` and ``N``, so every node issues the same ``K``
``Reduce_scatter`` calls in the same order, and the peak memory of the
contribution buffer drops by about ``N``. The random numbers are
bit-identical to the ones of a full-range call; only the order of the
floating-point sums over the nodes changes, at the roundoff level.\n
``rjackknife_sync_node`` performs the corresponding collective operation
for the case where every node has the whole input and inherits the same
sample blocks; ``g_mk_jk_distributed`` dispatches to this function.
"""
fname = "rjackknife_distributed"
comm, id_node, num_node = get_collective_comm(fname)
if n_rand_sample is None:
n_rand_sample = 1024
if rng_state is None:
rng_state = q.RngState("rejk")
(
data_arr,
jk_idx_list_local,
jk_idx_list_glb,
n,
elem_shape,
dtype,
) = get_distributed_jk_input(fname, comm, data_list, jk_idx_list)
avg = get_distributed_avg(comm, data_arr, n, avg)
#
total_size = 1 + n_rand_sample
i_start, i_end = get_distributed_range(total_size, id_node, num_node)
jk_arr = np.empty((i_end - i_start,) + elem_shape, dtype=dtype)
# The sample blocks are a pure function of the job-wide sizes: every node
# derives the same list and the loop issues the same ``K`` Reduce_scatter
# calls in the same order on every node (no collective is needed to agree
# on ``R``, and there is no per-node decision).
R, n_block, block_list = _get_jk_sample_blocks(total_size, num_node)
# One stream cache for the whole loop, so that every random stream is
# generated once instead of once per block; the cache is filled lazily
# (``_fill_garr_dict``, called by ``mk_r_i_j_mat``) with the streams of the
# names it needs. With a single block the mk call keeps its private dict.
garr_dict = dict() if n_block > 1 else None
# ``data_arr`` is a private copy made by ``get_distributed_jk_input``, so
# the average can be subtracted in place (same elementwise result as the
# separate ``data_diff`` array, without the extra buffer). The branch is
# rank independent: the dtype comes from the allgathered element info and
# ``avg`` is collective input or an ``Allreduce`` result.
if data_arr.dtype == np.asarray(avg).dtype:
data_arr -= avg
data_diff = data_arr
else:
data_diff = data_arr - avg
#
for b0, b1 in block_list:
# samples 1..n_rand_sample are the mk rows 0..n_rand_sample-1
mk_i0 = max(b0 - 1, 0)
mk_i1 = max(b1 - 1, 0)
if b0 == 0:
partial_blk = np.zeros((b1,) + elem_shape, dtype=dtype)
if i_start == 0 and i_end > i_start:
# the sample 0 is the average and must be counted only once
partial_blk[0] = avg
out_blk = partial_blk[1:]
else:
partial_blk = np.empty((b1 - b0,) + elem_shape, dtype=dtype)
out_blk = partial_blk
r_blk, b_blk = mk_r_i_j_mat(
n_rand_sample,
jk_idx_list_local,
rng_state,
jk_blocking_func=jk_blocking_func,
is_normalizing_rand_sample=is_normalizing_rand_sample,
is_apply_rand_sample_jk_idx_blocking_shift=is_apply_rand_sample_jk_idx_blocking_shift,
is_use_old_rand_alg=is_use_old_rand_alg,
i_range=(mk_i0, mk_i1),
jk_idx_list_for_count=jk_idx_list_glb,
garr_dict=garr_dict,
)
_apply_jk_fac(r_blk, b_blk, n, eps)
# The contraction is a matrix product (``get_jk_rows``), which does not
# materialize the broadcast temporary of
# ``r_blk[:, :, None, ...] * data_diff[None]``, whose size
# ``n_sample * n_local * prod(elem_shape)`` used to be the peak memory
# of the operation; ``out=out_blk`` writes the result straight into the
# send buffer of the ``Reduce_scatter``, so the result array is not
# materialized either and the peak is one block of element shapes.
get_jk_rows(r_blk, data_diff, out=out_blk)
off = max(b0, i_start) - i_start
n_mine = max(0, min(b1, i_end) - max(b0, i_start))
_get_reduce_scattered_jk_block(
partial_blk,
jk_arr[off : off + n_mine],
b0,
b1,
total_size,
comm,
)
#
if i_end - i_start > 0 and i_end > 1:
jk_arr[max(1 - i_start, 0) :] += avg
return jk_arr
@q.timer
def rjackknife_sync_node(
data_list,
jk_idx_list,
*,
avg=None,
rng_state=None,
n_rand_sample=None,
jk_blocking_func=None,
is_normalizing_rand_sample=False,
is_apply_rand_sample_jk_idx_blocking_shift=True,
is_use_old_rand_alg=False,
eps=1,
):
r"""
``rjackknife`` as a collective MPI operation where every node has the same
(whole) input and obtains the same complete result.\n
This implements the ``is_sync_node=True`` option of ``rjackknife``: the
input is split between the nodes in the order of the nodes and
``rjackknife_distributed`` is called on the local parts; the local results
are then gathered with ``Allgatherv``, so that every node obtains the
complete ``jk_arr`` of ``rjackknife``. The result agrees with the one of
``rjackknife`` up to the floating-point roundoff, but not bit-for-bit.\n
The local part is computed with the sample blocks of
``rjackknife_distributed`` (``R = ceil(T / num_node)`` rows per block), so
the memory is one block of samples instead of the whole sample axis.\n
``g_mk_jk_sync_node`` dispatches to this function.
"""
comm, id_node, num_node = get_collective_comm("rjackknife_sync_node")
i_start, i_end = get_distributed_range(len(data_list), id_node, num_node)
jk_local = rjackknife_distributed(
data_list[i_start:i_end],
jk_idx_list[i_start:i_end],
avg=avg,
rng_state=rng_state,
n_rand_sample=n_rand_sample,
jk_blocking_func=jk_blocking_func,
is_normalizing_rand_sample=is_normalizing_rand_sample,
is_apply_rand_sample_jk_idx_blocking_shift=is_apply_rand_sample_jk_idx_blocking_shift,
is_use_old_rand_alg=is_use_old_rand_alg,
eps=eps,
)
return get_gathered_jk_arr(jk_local, comm, num_node)
[docs]
@q.timer
def rjk_mk_jk_val(
rs_tag,
val,
err,
*,
n_rand_sample=None,
rng_state=None,
eps=1,
):
"""
return jk_arr
n = n_rand_sample
len(jk_arr) == 1 + n
jk_arr[i] = val + err * r[i] for i in 1..n
where r[i] ~ N(0, 1)
"""
if n_rand_sample is None:
n_rand_sample = 1024
if rng_state is None:
rng_state = q.RngState("rejk")
assert n_rand_sample >= 0
assert isinstance(rng_state, q.RngState)
assert isinstance(val, q.real_types)
assert isinstance(err, q.real_types)
rs = rng_state.split(str(rs_tag))
jk_arr = np.zeros((n_rand_sample + 1,), dtype=np.float64)
jk_arr[0] = val
r_arr = rs.g_rand_arr((n_rand_sample,))
r_arr_qnorm = q.qnorm(r_arr)
r_arr = r_arr * np.sqrt(n_rand_sample / r_arr_qnorm)
assert abs(q.qnorm(r_arr) / n_rand_sample - 1) < 1e-8
jk_arr[1:] = val + eps * r_arr * err
return jk_arr
[docs]
def rjk_avg(jk_arr):
return jk_avg(jk_arr)
[docs]
def rjk_err(jk_arr, eps=1):
"""
Return\n
.. math::\n
\\frac{1}{eps} \\sqrt{ 1/N \\sum_{i=1}^N (jk[i] - jk_avg)^2 }.\n
Note: ``
len(jk_arr) = N + 1.
jk_avg = jk_arr[0]
``
Same ``eps`` as the ``eps`` used in the ``jackknife`` function.
"""
avg = jk_avg(jk_arr)
n = len(jk_arr) - 1
if n <= 0:
fac = 1 / abs(eps)
val = fac * avg
val = q.filter_np_results(val)
return val
diff_sqr = q.average([q.fsqr(jk - avg) for jk in jk_arr[1:]])
fac = 1 / abs(eps)
val = fac * q.fsqrt(diff_sqr)
val = q.filter_np_results(val)
return val
[docs]
def rjk_avg_err(rjk_list, eps=1):
return rjk_avg(rjk_list), rjk_err(rjk_list, eps)
# ----------
default_g_jk_kwargs = dict()
def mk_g_jk_kwargs():
"""
Return the predefined ``default_g_jk_kwargs``.
"""
g_jk_kwargs = dict()
#
g_jk_kwargs["jk_type"] = "rjk" # choices: "rjk", "super"
g_jk_kwargs["eps"] = 1
#
# for jk_type = "rjk"
g_jk_kwargs["n_rand_sample"] = 1024
g_jk_kwargs["is_normalizing_rand_sample"] = False
g_jk_kwargs["is_apply_rand_sample_jk_idx_blocking_shift"] = True
#
# for jk_type = "super"
g_jk_kwargs["is_hash_jk_idx"] = True
g_jk_kwargs["jk_idx_hash_size"] = 1024
#
# Is only needed to reproduce old results
# Possible choice: "v1" (also need default_g_jk_kwargs["is_normalizing_rand_sample"] == False)
g_jk_kwargs["is_use_old_rand_alg"] = False
#
# these parameters are used in jk_blocking_func_default
g_jk_kwargs["block_size"] = 1
g_jk_kwargs["block_size_dict"] = {
"job_tag": 1,
}
#
# Below are items which are not touched in
# ``get_jk_state`` or ``set_jk_state``
#
g_jk_kwargs["rng_state"] = q.RngState("rejk")
#
g_jk_kwargs["all_jk_idx"] = None
g_jk_kwargs["get_all_jk_idx"] = None
#
g_jk_kwargs["all_jk_idx_set"] = set()
#
# ``is_sync_node`` runs ``g_mk_jk`` as a collective MPI operation. It only
# changes how the result is computed (the result agrees up to the
# floating-point roundoff), so it is deliberately not touched in
# ``get_jk_state`` or ``set_jk_state`` (and hence not part of the cache key).
g_jk_kwargs["is_sync_node"] = False
#
# jk_blocking_func(i, jk_idx) => blocked_jk_idx
g_jk_kwargs["jk_blocking_func"] = jk_blocking_func_default
#
return g_jk_kwargs
def reset_default_g_jk_kwargs():
default_g_jk_kwargs.clear()
default_g_jk_kwargs.update(mk_g_jk_kwargs())
[docs]
@q.use_kwargs(default_g_jk_kwargs)
def get_jk_state(
*,
jk_type,
eps,
n_rand_sample,
is_normalizing_rand_sample,
is_apply_rand_sample_jk_idx_blocking_shift,
is_hash_jk_idx,
jk_idx_hash_size,
is_use_old_rand_alg,
block_size,
block_size_dict,
**_kwargs,
):
"""
Currently only useful if we set::\n
q.default_g_jk_kwargs["jk_type"] = "rjk" # this is the default now\n
and::\n
q.default_g_jk_kwargs["jk_blocking_func"] = jk_blocking_func_default\n
Used for ``q.cache_call``.\n
Example::\n
@cache_call(get_state=q.get_jk_state)
def func(...):
...
"""
return (
jk_type,
eps,
n_rand_sample,
is_normalizing_rand_sample,
is_apply_rand_sample_jk_idx_blocking_shift,
is_hash_jk_idx,
jk_idx_hash_size,
is_use_old_rand_alg,
block_size,
block_size_dict,
)
[docs]
def set_jk_state(state):
(
jk_type,
eps,
n_rand_sample,
is_normalizing_rand_sample,
is_apply_rand_sample_jk_idx_blocking_shift,
is_hash_jk_idx,
jk_idx_hash_size,
is_use_old_rand_alg,
block_size,
block_size_dict,
) = state
g_dict = default_g_jk_kwargs
g_dict["jk_type"] = jk_type
g_dict["eps"] = eps
g_dict["n_rand_sample"] = n_rand_sample
g_dict["is_normalizing_rand_sample"] = is_normalizing_rand_sample
g_dict["is_apply_rand_sample_jk_idx_blocking_shift"] = (
is_apply_rand_sample_jk_idx_blocking_shift
)
g_dict["is_hash_jk_idx"] = is_hash_jk_idx
g_dict["jk_idx_hash_size"] = jk_idx_hash_size
g_dict["is_use_old_rand_alg"] = is_use_old_rand_alg
g_dict["block_size"] = block_size
g_dict["block_size_dict"] = block_size_dict
jk_blocking_traj_shift_arr = q.RngState("jk_blocking_traj_shift_arr").rand_arr(
16 * 1024
) % (1024 * 1024 * 1024 * 1024)
@q.use_kwargs(default_g_jk_kwargs)
def jk_blocking_func_default(
i,
jk_idx,
*,
block_size,
block_size_dict,
all_jk_idx_set,
**_kwargs,
):
"""
return ``blocked_jk_idx``.
``blocked_jk_idx`` should uniquely identify the block
that configuration identified by ``jk_idx`` belongs to.
The block scheme can be different for different J-B hybrid sample.
The J-B hybrid sample is indexed by ``i`` (``1 <= i <= n_rand_sample``).
``
block_size_for_this_job_tag = block_size_dict.get(job_tag, block_size)
``
use default_g_jk_kwargs for
block_size, block_size_dict, all_jk_idx_set
"""
if i == 0:
shift = 0
else:
assert i >= 1
shift = int(
jk_blocking_traj_shift_arr[(i - 1) % len(jk_blocking_traj_shift_arr)]
)
if block_size_dict is None:
block_size_dict = dict()
if all_jk_idx_set is not None:
all_jk_idx_set.add(jk_idx)
if isinstance(jk_idx, q.int_types):
traj = jk_idx
b_shift = shift % block_size
return (traj + b_shift) // block_size
elif (
isinstance(jk_idx, tuple)
and len(jk_idx) == 2
and isinstance(jk_idx[1], q.int_types)
):
job_tag, traj = jk_idx
assert isinstance(job_tag, str)
assert isinstance(traj, q.int_types)
block_size_for_this_job_tag = block_size_dict.get(job_tag, block_size)
assert isinstance(block_size_for_this_job_tag, q.int_types)
b_shift = shift % block_size_for_this_job_tag
return (
job_tag,
(traj + b_shift) // block_size_for_this_job_tag,
)
else:
return jk_idx
assert False
[docs]
@q.use_kwargs(default_g_jk_kwargs)
@q.timer
def g_mk_jk(
data_list,
jk_idx_list,
*,
avg=None,
jk_type,
all_jk_idx,
get_all_jk_idx,
n_rand_sample,
rng_state,
jk_blocking_func,
is_normalizing_rand_sample,
is_apply_rand_sample_jk_idx_blocking_shift,
is_use_old_rand_alg,
is_hash_jk_idx,
jk_idx_hash_size,
eps,
is_sync_node=False,
**_kwargs,
):
"""
Create a (randomized) Super-Jackknife data set from un-jackknifed data.\n
``jk_arr[0]`` is the average of the data and ``jk_arr[1:]`` are the
resampled values, from which the error is estimated with ``g_jk_avg_err``.
The data set has ``g_jk_size()`` samples, i.e. ``1 + n_rand_sample`` for
``jk_type == "rjk"`` and ``1 + len(all_jk_idx)`` for
``jk_type == "super"``, and its dtype is the dtype of the data.\n
:param data_list: the un-jackknifed data, a list or ``np.ndarray`` of
values (each value being a ``float``, a ``complex`` or an
``np.ndarray``); ``None`` entries are ignored. For the collective MPI
variants below the data must have a dtype supported by MPI.
:param jk_idx_list: the indices that name the entries of ``data_list``,
with ``len(jk_idx_list) == len(data_list)``, usually
``jk_idx_list = [(job_tag, traj,) for traj in traj_list]``. The
indices are mapped to the jackknife blocks by ``jk_blocking_func``
(see ``jk_blocking_func_default``, ``block_size`` and
``block_size_dict``).
:param avg: the average of the data; when ``None`` (the default) it is
computed from ``data_list``. Pass a precomputed value to reuse it (it
must be the average of the whole data set).
:param is_sync_node: when ``True``, the operation is performed by
``g_mk_jk_sync_node``, i.e. as a collective MPI operation in which
every node holds the whole data set, every node must call this function
with the same parameters, and every node obtains the complete
``jk_arr``. Use ``g_mk_jk_distributed`` instead when the data set
itself is split between the nodes. Both ``jk_type`` values are
supported; the result agrees with the ``is_sync_node=False`` result up
to the floating-point roundoff.
:return: the (randomized) Super-Jackknife data set ``jk_arr``.\n
The other keyword parameters are the entries of ``default_g_jk_kwargs``,
which supplies their defaults; set them there, pass them explicitly or use
the ``q.JkKwargs(...)`` context manager. The most commonly used are:\n
- ``jk_type``: ``"rjk"`` (the default) or ``"super"``.
- ``eps`` (default ``1``): the overall scale of the fluctuations; when the
data is already jackknifed, multiply it by ``len(data_list)``.
- ``n_rand_sample`` (default ``1024``): the number of random samples of
``"rjk"``.
- ``is_normalizing_rand_sample``,
``is_apply_rand_sample_jk_idx_blocking_shift`` and
``is_use_old_rand_alg``: options of the random numbers of ``"rjk"``.
- ``is_hash_jk_idx``, ``jk_idx_hash_size``, ``all_jk_idx`` and
``get_all_jk_idx``: the samples of ``"super"``.
- ``block_size``, ``block_size_dict`` and ``jk_blocking_func``: the
jackknife blocks.
- ``rng_state``: the random numbers of ``"rjk"``.\n
See ``rjackknife`` and ``sjackknife`` for the formulas, and
``docs/qlat-utils/qlat_data.md`` for the full documentation.\n
Example::
jk_arr = q.g_mk_jk(data_list, jk_idx_list)
avg, err = q.g_jk_avg_err(jk_arr)
"""
if is_sync_node:
# Collective MPI operation where every node has the same input: only a
# part of the result is computed on each node and the parts are then
# gathered, so that every node obtains the complete result.
jk_arr = g_mk_jk_sync_node(
data_list,
jk_idx_list,
avg=avg,
jk_type=jk_type,
all_jk_idx=all_jk_idx,
get_all_jk_idx=get_all_jk_idx,
n_rand_sample=n_rand_sample,
rng_state=rng_state,
jk_blocking_func=jk_blocking_func,
is_normalizing_rand_sample=is_normalizing_rand_sample,
is_apply_rand_sample_jk_idx_blocking_shift=is_apply_rand_sample_jk_idx_blocking_shift,
is_use_old_rand_alg=is_use_old_rand_alg,
is_hash_jk_idx=is_hash_jk_idx,
jk_idx_hash_size=jk_idx_hash_size,
eps=eps,
)
return jk_arr
if jk_type == "super":
jk_arr = sjackknife(
data_list,
jk_idx_list,
avg=avg,
is_hash_jk_idx=is_hash_jk_idx,
jk_idx_hash_size=jk_idx_hash_size,
rng_state=rng_state,
all_jk_idx=all_jk_idx,
get_all_jk_idx=get_all_jk_idx,
jk_blocking_func=jk_blocking_func,
eps=eps,
)
elif jk_type == "rjk":
jk_arr = rjackknife(
data_list,
jk_idx_list,
avg=avg,
n_rand_sample=n_rand_sample,
rng_state=rng_state,
jk_blocking_func=jk_blocking_func,
is_normalizing_rand_sample=is_normalizing_rand_sample,
is_apply_rand_sample_jk_idx_blocking_shift=is_apply_rand_sample_jk_idx_blocking_shift,
is_use_old_rand_alg=is_use_old_rand_alg,
eps=eps,
is_sync_node=is_sync_node,
)
else:
assert False
return jk_arr
@q.use_kwargs(default_g_jk_kwargs)
@q.timer
def g_mk_jk_distributed(
data_list,
jk_idx_list,
*,
avg=None,
jk_type,
all_jk_idx,
get_all_jk_idx,
n_rand_sample,
rng_state,
jk_blocking_func,
is_normalizing_rand_sample,
is_apply_rand_sample_jk_idx_blocking_shift,
is_use_old_rand_alg,
is_hash_jk_idx,
jk_idx_hash_size,
eps,
**_kwargs,
):
"""
Create a (randomized) Super-Jackknife data set when the data set itself is
split between the MPI nodes.\n
This is a collective MPI operation: every node must call it with the same
parameters and with its own disjoint part of the data set, and every node
returns its own part of the result. The split of the data set between the
nodes is free as long as the local parts are disjoint and cover the whole
data set exactly once (``get_distributed_range(len(data_list), ...)`` is a
convenient balanced split); no node needs to hold the whole data set, the
whole random matrix or the whole result, and only the small ``jk_idx``
metadata is gathered on every node.\n
The returned ``jk_arr`` is the local part of the data set, which has
``g_jk_size()`` samples in total. The samples are split between the nodes
in the order of the nodes, independently of the split of the input: node
``r`` out of ``num_node`` nodes holds the samples in
``range(*get_distributed_range(g_jk_size(...), r, num_node))`` (a node
holds 0 samples when there are more nodes than samples), so that::
jk_arr = np.concatenate(q.get_comm().allgather(jk_local))
is the complete data set, in the same order as the one obtained with
``g_mk_jk``.\n
``jk_type == "rjk"`` calls ``rjackknife_distributed`` (the random samples
are split between the nodes and every node computes the contribution of
its own data to all of them); ``jk_type == "super"`` calls
``sjackknife_distributed`` (the samples are the ``all_jk_idx`` entries, or
the hash based samples).\n
Requires qlat to be initialized on the whole MPI communicator, i.e. with
``q.begin_with_mpi()``, ``q.begin_with_gpt()`` or ``q.begin_with_grid()``
(or ``q.set_comm(...)``), and the data must have a dtype supported by MPI.\n
The result agrees with the one of ``g_mk_jk`` (and of
``g_mk_jk(..., is_sync_node=True)``) up to the floating-point roundoff, but
not bit-for-bit: the average and the sums over the data set are reduced
across the nodes, which changes the order of the floating-point additions.
The ``is_sync_node`` entry of ``default_g_jk_kwargs`` is ignored, since the
output of this function is always distributed.\n
:param data_list: the local part of the un-jackknifed data.
:param jk_idx_list: the indices that name the local ``data_list``, with
``len(jk_idx_list) == len(data_list)``.
:param avg: the average of the whole data set, the same on every node; when
``None`` (the default) it is computed from the data set by summing the
local contributions over the nodes with ``Allreduce``.
:return: the local part of the (randomized) Super-Jackknife data set (an
array with 0 samples on a node which owns no sample).\n
The other keyword parameters, including ``jk_type`` and ``eps``, are the
same as for ``g_mk_jk``; their defaults are the entries of
``default_g_jk_kwargs``. See ``g_mk_jk`` for the description of the
entries.\n
Example::
comm = q.get_comm()
i_start, i_end = q.get_distributed_range(len(data_list), comm.rank, comm.size)
jk_local = q.g_mk_jk_distributed(
data_list[i_start:i_end], jk_idx_list[i_start:i_end],
)
jk_arr = np.concatenate(comm.allgather(jk_local))
"""
if jk_type == "super":
jk_arr = sjackknife_distributed(
data_list,
jk_idx_list,
avg=avg,
is_hash_jk_idx=is_hash_jk_idx,
jk_idx_hash_size=jk_idx_hash_size,
rng_state=rng_state,
all_jk_idx=all_jk_idx,
get_all_jk_idx=get_all_jk_idx,
jk_blocking_func=jk_blocking_func,
eps=eps,
)
elif jk_type == "rjk":
jk_arr = rjackknife_distributed(
data_list,
jk_idx_list,
avg=avg,
rng_state=rng_state,
n_rand_sample=n_rand_sample,
jk_blocking_func=jk_blocking_func,
is_normalizing_rand_sample=is_normalizing_rand_sample,
is_apply_rand_sample_jk_idx_blocking_shift=is_apply_rand_sample_jk_idx_blocking_shift,
is_use_old_rand_alg=is_use_old_rand_alg,
eps=eps,
)
else:
assert False
return jk_arr
@q.use_kwargs(default_g_jk_kwargs)
@q.timer
def g_mk_jk_sync_node(
data_list,
jk_idx_list,
*,
avg=None,
jk_type,
all_jk_idx,
get_all_jk_idx,
n_rand_sample,
rng_state,
jk_blocking_func,
is_normalizing_rand_sample,
is_apply_rand_sample_jk_idx_blocking_shift,
is_use_old_rand_alg,
is_hash_jk_idx,
jk_idx_hash_size,
eps,
**_kwargs,
):
"""
Create a (randomized) Super-Jackknife data set as a collective MPI
operation in which every node holds the whole input and obtains the whole
result. This is what ``g_mk_jk(..., is_sync_node=True)`` calls.\n
Every node must call this function with the same parameters and with the
whole data set (``data_list`` and ``jk_idx_list`` are the complete lists);
every node then obtains the complete ``jk_arr``, in the same order as the
one obtained with ``g_mk_jk`` on a single node. The input is split between
the nodes in the order of the nodes with ``get_distributed_range`` and the
local parts are computed by ``rjackknife_sync_node`` or
``sjackknife_sync_node`` (which call ``rjackknife_distributed`` or
``sjackknife_distributed``) and then gathered, so the work is parallelized
over the nodes even though every node starts with the whole data set.\n
Both ``jk_type == "rjk"`` and ``jk_type == "super"`` are supported. The
result agrees with the one of ``g_mk_jk`` up to the floating-point
roundoff, but not bit-for-bit: the average and the sums over the data set
are reduced across the nodes, which changes the order of the floating-point
additions.\n
Requires qlat to be initialized on the whole MPI communicator, i.e. with
``q.begin_with_mpi()``, ``q.begin_with_gpt()`` or ``q.begin_with_grid()``
(or ``q.set_comm(...)``), and the data must have a dtype supported by MPI.\n
:param data_list: the whole un-jackknifed data, the same on every node.
:param jk_idx_list: the indices that name the whole ``data_list``, the same
on every node.
:param avg: the average of the whole data set, the same on every node; when
``None`` (the default) it is computed from the data set by summing the
local contributions over the nodes with ``Allreduce``.
:return: the complete (randomized) Super-Jackknife data set ``jk_arr``, the
same on every node.\n
The other keyword parameters, including ``jk_type`` and ``eps``, are the
same as for ``g_mk_jk``; their defaults are the entries of
``default_g_jk_kwargs``. See ``g_mk_jk`` for the description of the
entries, and ``g_mk_jk_distributed`` for the case in which the data set is
split between the nodes instead.\n
Example::
jk_arr = q.g_mk_jk_sync_node(data_list, jk_idx_list)
# is equivalent to
jk_arr = q.g_mk_jk(data_list, jk_idx_list, is_sync_node=True)
"""
if jk_type == "super":
jk_arr = sjackknife_sync_node(
data_list,
jk_idx_list,
avg=avg,
is_hash_jk_idx=is_hash_jk_idx,
jk_idx_hash_size=jk_idx_hash_size,
rng_state=rng_state,
all_jk_idx=all_jk_idx,
get_all_jk_idx=get_all_jk_idx,
jk_blocking_func=jk_blocking_func,
eps=eps,
)
elif jk_type == "rjk":
jk_arr = rjackknife_sync_node(
data_list,
jk_idx_list,
avg=avg,
rng_state=rng_state,
n_rand_sample=n_rand_sample,
jk_blocking_func=jk_blocking_func,
is_normalizing_rand_sample=is_normalizing_rand_sample,
is_apply_rand_sample_jk_idx_blocking_shift=is_apply_rand_sample_jk_idx_blocking_shift,
is_use_old_rand_alg=is_use_old_rand_alg,
eps=eps,
)
else:
assert False
return jk_arr
[docs]
@q.use_kwargs(default_g_jk_kwargs)
@q.timer
def g_mk_jk_val(
rs_tag,
val,
err,
*,
jk_type,
all_jk_idx,
get_all_jk_idx,
n_rand_sample,
rng_state,
is_hash_jk_idx,
jk_idx_hash_size,
eps,
**_kwargs,
):
"""
Create a jackknife sample with random numbers based on central value ``val`` and error ``err``.\n
Need::\n
default_g_jk_kwargs["jk_type"] = "rjk"
default_g_jk_kwargs["n_rand_sample"] = n_rand_sample
# e.g. n_rand_sample = 1024
default_g_jk_kwargs["rng_state"] = rng_state
# e.g. rng_state = q.RngState("rejk")
"""
if jk_type == "super":
jk_val = sjk_mk_jk_val(
rs_tag,
val,
err,
is_hash_jk_idx=is_hash_jk_idx,
jk_idx_hash_size=jk_idx_hash_size,
rng_state=rng_state,
all_jk_idx=all_jk_idx,
get_all_jk_idx=get_all_jk_idx,
eps=eps,
)
elif jk_type == "rjk":
jk_val = rjk_mk_jk_val(
rs_tag,
val,
err,
n_rand_sample=n_rand_sample,
rng_state=rng_state,
eps=eps,
)
else:
assert False
return jk_val
[docs]
def g_jk_avg(jk_arr, **_kwargs):
"""
Return ``avg`` of the ``jk_arr``.
"""
if isinstance(jk_arr, q.number_types):
return jk_arr
return jk_avg(jk_arr)
[docs]
@q.use_kwargs(default_g_jk_kwargs)
def g_jk_err(jk_arr, *, eps, jk_type, **_kwargs):
"""
Return ``err`` of the ``jk_arr``.
"""
if isinstance(jk_arr, q.number_types):
return 0
if jk_type == "super":
return sjk_err(jk_arr, eps=eps)
elif jk_type == "rjk":
return rjk_err(jk_arr, eps=eps)
else:
assert False
return None
[docs]
@q.timer
def g_jk_avg_err(jk_arr, **kwargs):
"""
Return ``(avg, err,)`` of the ``jk_arr``.
"""
return g_jk_avg(jk_arr), g_jk_err(jk_arr, **kwargs)
[docs]
@q.timer
def g_jk_avg_err_arr(jk_arr, **kwargs):
"""
Return ``avg_err_arr`` of the ``jk_arr``.
``
avg_err_arr.shape = jk_arr[0].shape + (2,)
``
"""
avg, err = g_jk_avg_err(jk_arr, **kwargs)
avg_err_arr = np.stack(
[
avg,
err,
]
)
avg_err_arr = np.moveaxis(avg_err_arr, 0, -1).copy()
return avg_err_arr
[docs]
@q.use_kwargs(default_g_jk_kwargs)
def g_jk_size(
*,
jk_type,
all_jk_idx,
get_all_jk_idx,
n_rand_sample,
is_hash_jk_idx,
jk_idx_hash_size,
**_kwargs,
):
"""
Return number of samples for the (randomized) Super-Jackknife data set.
"""
if jk_type == "super":
if all_jk_idx is None:
if get_all_jk_idx is None:
assert is_hash_jk_idx
all_jk_idx = [
"avg",
] + list(range(jk_idx_hash_size))
else:
all_jk_idx = get_all_jk_idx()
assert all_jk_idx[0] == "avg"
n_super_sample = len(all_jk_idx) - 1
assert n_super_sample >= 0
return 1 + n_super_sample
elif jk_type == "rjk":
return 1 + n_rand_sample
else:
assert False
return None
[docs]
@q.use_kwargs(default_g_jk_kwargs)
def g_jk_blocking_func(
i,
jk_idx,
*,
jk_blocking_func,
**_kwargs,
):
"""
Return ``jk_blocking_func(jk_idx)``.
"""
if jk_blocking_func is None:
return jk_idx
else:
return jk_blocking_func(i, jk_idx)
@q.use_kwargs(default_g_jk_kwargs)
def g_jk_sample_size(
job_tag,
traj_list,
**_kwargs,
):
jk_idx_list = [
(
job_tag,
traj,
)
for traj in traj_list
]
b_jk_idx_set = set(
g_jk_blocking_func(0, jk_idx, **kwargs) for jk_idx in jk_idx_list
)
return len(b_jk_idx_set)
reset_default_g_jk_kwargs()
# ----
class JkKwargs(q.NewDictValues):
"""
Example:
#
with q.JkKwargs(n_rand_sample=1024, block_size=10, block_size_dict={ "48I": 20, }):
...
#
"""
def __init__(self, **kwargs):
super().__init__(default_g_jk_kwargs, **kwargs)
# ----
# ---- old funcs
def merge_jk_idx(*jk_idx_list_list):
for jk_idx_list in jk_idx_list_list:
assert jk_idx_list[0] == "avg"
return [
"avg",
] + [jk_idx for jk_idx_list in jk_idx_list_list for jk_idx in jk_idx_list[1:]]
@q.timer
def rejk_list(jk_list, jk_idx_list, all_jk_idx):
"""
Super jackknife
``jk_idx_list`` should be contained in ``all_jk_idx`` and have the same order.
Does not properly honor the (N-1) formula in error calculation.
"""
assert jk_idx_list[0] == "avg"
assert all_jk_idx[0] == "avg"
assert len(jk_idx_list) == len(jk_list)
assert len(jk_idx_list) <= len(all_jk_idx)
is_np_arr = isinstance(jk_list, np.ndarray)
jk_avg = jk_list[0]
size_new = len(all_jk_idx)
i_new = 0
jk_list_new = []
for i, idx in enumerate(jk_idx_list):
while all_jk_idx[i_new] != idx:
jk_list_new.append(jk_avg)
i_new += 1
assert i_new < size_new
jk_list_new.append(jk_list[i])
i_new += 1
while i_new < size_new:
jk_list_new.append(jk_avg)
i_new += 1
assert i_new == size_new
assert size_new == len(jk_list_new)
if is_np_arr:
jk_list_new = np.array(jk_list_new, dtype=jk_list.dtype)
return jk_list_new
@q.timer
def rjk_jk_list(
jk_list,
jk_idx_list,
n_rand_sample,
rng_state,
jk_blocking_func=None,
is_normalizing_rand_sample=False,
is_apply_rand_sample_jk_idx_blocking_shift=True,
is_use_old_rand_alg=False,
):
r"""
return jk_list
len(jk_list) == 1 + n_rand_sample
distribution of jk_list should be similar as the distribution of avg
r_{i,j} ~ N(0, 1)
if is_normalizing_rand_sample:
n_j = \sum_i r_{i,j}^2
r_{i,j} <- \sqrt{n_rand_sample / n_j} r_{i,j}
avg = jk_list[0]
len(jk_list) = n + 1
jk_list[i] = avg + \sum_{j=1}^{n} r_{i,j} (jk_list[j] - avg)
#
if ``jk_blocking_func`` is provided:
``
jk_blocking_func(i, jk_idx) => blocked jk_idx
``
Note that: ``1 <= i <= n_rand_sample``
``
jk_list[i] = avg + \sum_{j=1}^{n} r_{i,jk_block_func(i, j)} (jk_list[j] - avg)
``
"""
assert jk_idx_list[0] == "avg"
assert isinstance(n_rand_sample, q.int_types)
assert n_rand_sample >= 0
assert isinstance(rng_state, q.RngState)
is_np_arr = isinstance(jk_list, np.ndarray)
n = len(jk_list) - 1
r_arr, b_arr = mk_r_i_j_mat(
n_rand_sample,
jk_idx_list[1:],
rng_state,
jk_blocking_func=jk_blocking_func,
is_normalizing_rand_sample=is_normalizing_rand_sample,
is_apply_rand_sample_jk_idx_blocking_shift=is_apply_rand_sample_jk_idx_blocking_shift,
is_use_old_rand_alg=is_use_old_rand_alg,
)
avg = jk_list[0]
if is_np_arr:
jk_arr = jk_list
jk_diff = jk_arr[1:] - avg
rjk_arr = np.empty(
(
1 + n_rand_sample,
*avg.shape,
),
dtype=jk_arr.dtype,
)
rjk_arr[:] = avg
for j in range(n):
for i in range(n_rand_sample):
rjk_arr[i + 1] += r_arr[i, j] * jk_diff[j]
return rjk_arr
else:
rjk_list = [
avg,
]
jk_diff = [jk_list[j] - avg for j in range(1, n + 1)]
for i in range(n_rand_sample):
rjk_list.append(avg + sum([r_arr[i, j] * jk_diff[j] for j in range(n)]))
return rjk_list
@q.use_kwargs(default_g_jk_kwargs)
@q.timer
def g_jk(data_list, *, eps, **_kwargs):
"""
Obsolete, call ``g_mk_jk`` instead.
--
Perform initial Jackknife for the original data set.\n
"""
return jackknife(data_list, eps=eps)
@q.use_kwargs(default_g_jk_kwargs)
@q.timer
def g_rejk(
jk_list,
jk_idx_list,
*,
jk_type,
all_jk_idx,
get_all_jk_idx,
n_rand_sample,
rng_state,
jk_blocking_func,
is_normalizing_rand_sample,
is_apply_rand_sample_jk_idx_blocking_shift,
is_use_old_rand_alg,
**_kwargs,
):
"""
Obsolete, call ``g_mk_jk`` instead.
--
Perform (randomized) Super-Jackknife for the Jackknife data set.
--
:jk_list: usually the Jackknife data set obtained with ``g_jk(data_list)``.
:jk_idx_list: should be list of indices that names the ``jk_list``.
:jk_type: ``[ "rjk", "super", ]``
:returns: (randomized) Super-Jackknife data set.
Note that::
len(jk_list) == len(jk_idx_list)
jk_idx_list[0] == "avg"
"""
if jk_type == "super":
if jk_blocking_func is not None:
q.displayln_info(
f"g_rejk: jk_type={jk_type} does not support jk_blocking_func={jk_blocking_func}"
)
if all_jk_idx is None:
assert get_all_jk_idx is not None
all_jk_idx = get_all_jk_idx()
return rejk_list(
jk_list,
jk_idx_list,
all_jk_idx,
)
elif jk_type == "rjk":
return rjk_jk_list(
jk_list,
jk_idx_list,
n_rand_sample,
rng_state,
jk_blocking_func,
is_normalizing_rand_sample,
is_apply_rand_sample_jk_idx_blocking_shift,
is_use_old_rand_alg,
)
else:
assert False
return None
# ----
def mk_jk_blocking_func(block_size=1, block_size_dict=None, all_jk_idx_set=None):
"""
Recommend to use ``jk_blocking_func_default`` instead.
#
block_size_for_this_job_tag = block_size_dict.get(job_tag, block_size)
"""
if block_size_dict is None:
block_size_dict = dict()
#
def jk_blocking_func(jk_idx):
if all_jk_idx_set is not None:
all_jk_idx_set.add(jk_idx)
if isinstance(jk_idx, q.int_types):
traj = jk_idx
return traj // block_size
elif (
isinstance(jk_idx, tuple)
and len(jk_idx) == 2
and isinstance(jk_idx[1], q.int_types)
):
job_tag, traj = jk_idx
assert isinstance(job_tag, str)
assert isinstance(traj, q.int_types)
block_size_for_this_job_tag = block_size_dict.get(job_tag, block_size)
assert isinstance(block_size_for_this_job_tag, q.int_types)
return (
job_tag,
traj // block_size_for_this_job_tag,
)
else:
return jk_idx
#
return jk_blocking_func
# ----
def add_jk_idx(arr):
"""
arr: no jk index
return: add trivial jk index in the LAST axis
"""
return arr.reshape(arr.shape + (1,))
def jk_transpose(arr):
"""
arr: jk index is the 0th axis
return: jk index is the last axis
"""
shape = arr.shape
ndim = len(shape)
if ndim <= 1:
return arr
axes = list(range(1, ndim)) + [
0,
]
return arr.transpose(axes)
def jk_transpose_back(arr):
"""
jk_transpose_back(jk_transpose(arr)) == arr
"""
shape = arr.shape
ndim = len(shape)
if ndim <= 1:
return arr
axes = [
ndim - 1,
] + list(range(0, ndim - 1))
return arr.transpose(axes)