import re
from pathlib import Path
from typing import Optional
import h5py
import jax
import numpy as np
[docs]
def gather_electron_axis(pytree, electron_batch_axis=4):
r"""Merge the per-device electron-sample axis of arrays read from a result file.
Local energies, wave function values and other per-sample quantities are stored
in the ``result.h5`` file of a deepQMC workdir with the same
``[n_iterations, n_device, ..., electron_batch_size / n_device, ...]`` layout
they have during training, see
:func:`~deepqmc.parallel.gather_electrons_on_one_device`. This function merges
the device axis (axis 1) into the electron batch axis, to arrive at the more
convenient ``[n_iterations, ..., electron_batch_size, ...]`` layout.
Args:
pytree: a pytree of arrays all with shape ``[n_iterations, n_device, ...,
electron_batch_size / n_device, ...]``.
electron_batch_axis (int): optional, the axis carrying the electron batch
before merging, e.g. 4 for arrays of shape ``[n_iterations, n_device,
molecule_batch_size, electronic_states, electron_batch_size / n_device,
...]``, the most common case.
Returns:
a pytree of arrays all with shape ``[n_iterations, ..., electron_batch_size,
...]``.
"""
return jax.tree.map(
lambda x: np.moveaxis(x, 1, electron_batch_axis - 1).reshape(
x.shape[0],
*x.shape[2:electron_batch_axis],
-1,
*x.shape[electron_batch_axis + 1 :],
),
pytree,
)
def subscript_sorting_key(string_with_subscript: str):
r"""Extracts the integer subscript from strings such as foo_2."""
re_match = re.search(r'.+_(\d+)', string_with_subscript)
assert re_match, f'Invalid string with substring {string_with_subscript}'
return int(re_match.group(1))
def is_multi_node_subdir(subdir_name: str):
r"""Checks if a subdir name is of the form training_0, or evaluation_1, etc."""
assert subdir_name.startswith('training') or subdir_name.startswith(
'evaluation'
), f'Invalid subdir name {subdir_name}'
return re.search(r'.+_\d+', subdir_name) is not None
def sorted_subdirs(subdirs: list[str]) -> list[str]:
r"""Sorts subdirs with potential integer subscripts."""
are_multi_node_subdir = [is_multi_node_subdir(subdir) for subdir in subdirs]
if any(are_multi_node_subdir):
assert all(are_multi_node_subdir), 'Mix of single and multi node subdirs'
assert sorted([subscript_sorting_key(subdir) for subdir in subdirs]) == list(
range(len(subdirs))
), 'Invalid subscripts for multi node subdirs'
return sorted(subdirs, key=subscript_sorting_key)
else:
assert len(subdirs) == 1, 'Multiple single node subdirs found'
return subdirs
def chkpt_file_iteration(chkpt_file_name: str):
r"""Extract the iteration count from the name of a checkpoint file."""
re_match = re.search(r'chkpt-(\d+).pt', chkpt_file_name)
assert re_match, f'Invalid checkpoint file name: {chkpt_file_name}'
return int(re_match.group(1))
[docs]
def last_checkpoint_iteration(path: Path) -> Optional[int]:
r"""Return the iteration of the last checkpoint file in a deepQMC subdir."""
chkpt_iterations = sorted(
[chkpt_file_iteration(file.name) for file in path.glob('chkpt-*.pt')]
)
if len(chkpt_iterations) > 0:
return chkpt_iterations[-1]
return None
def concatenate_subdir_results(
subdir_results: list[tuple[dict, Optional[int]]],
) -> tuple[dict, Optional[int]]:
r"""Concatenate results from multiple deepQMC subdirs."""
if len(subdir_results) == 1:
return subdir_results[0]
results, last_chkpt_iters = zip(*subdir_results)
assert all(
last_chkpt_iter == last_chkpt_iters[0]
for last_chkpt_iter in last_chkpt_iters[1:]
), 'Mismatching last checkpoint iterations between subdirs'
assert all(
result.keys() == results[0].keys() for result in results[1:]
), 'Mismatching keys between subdirs'
min_lengths = { # pyright: ignore
key: min(len(result[key]) for result in results) for key in results[0].keys()
}
results = {
key: (
results[0][key]
if 'samples' not in key
else np.concatenate(
[result[key][: min_lengths[key]] for result in results], axis=1
)
)
for key in results[0].keys()
}
return results, last_chkpt_iters[0]
def read_subdir(path: Path, keys: list[str]) -> tuple[dict, Optional[int]]:
r"""Read values of given keys from a result.h5 file in a deepQMC subdir."""
last_chkpt_iter = last_checkpoint_iteration(path)
result_file = path / 'result.h5'
if not result_file.exists():
return {}, None
with h5py.File(result_file, swmr=True, libver='v110') as f:
results = {key: np.array(f[key]) for key in keys if key in f.keys()}
return results, last_chkpt_iter
[docs]
def read_workdir(path: Path, keys: list[str]) -> tuple[dict, Optional[int]]:
r"""Read values of given keys from result.h5 files in a deepQMC workdir.
Automatically detects whether ``path`` contains a single-node ``training``/
``evaluation`` subdir, or multiple ``training_0``, ``training_1``, ... /
``evaluation_0``, ``evaluation_1``, ... multi-node subdirs, and concatenates the
results of the latter along the electron batch axis.
Args:
path (~pathlib.Path): the deepQMC workdir, e.g. as passed to
``hydra.run.dir``.
keys (list[str]): the names of the datasets to read from the ``result.h5``
files, e.g. ``['local_energy', 'mol_idxs']``.
Returns:
tuple[dict, Optional[int]]: a dictionary mapping each of the requested
``keys`` present in the result files to the corresponding array, and the
iteration of the last checkpoint file found in the workdir, or :data:`None`
if no ``training``/``evaluation`` subdir or checkpoint was found.
"""
eval_subdirs = [subdir.name for subdir in path.glob('evaluation*')]
train_subdirs = [subdir.name for subdir in path.glob('training*')]
if not eval_subdirs and not train_subdirs:
return {}, None
if eval_subdirs and train_subdirs:
raise ValueError(
f'workdir {path} contains both evaluation and training subdirs:'
f' {eval_subdirs + train_subdirs}'
)
subdirs = eval_subdirs if not train_subdirs else train_subdirs
subdir_results = [
read_subdir(path / subdir, keys) for subdir in sorted_subdirs(subdirs)
]
workdir_result, last_chkpt_iter = concatenate_subdir_results(subdir_results)
return workdir_result, last_chkpt_iter
[docs]
def read_and_reshape_result(
path, *keys, read_workdir=read_workdir, gather_electrons=True
):
r"""Read and reshape results from a deepQMC workdir.
Use when retrieving all individual samples from a result file that
contains results for a single molecule, and no ``mol_idxs`` entry,
e.g. for a standard (non-transferable) training or evaluation run.
Args:
path (~pathlib.Path): the deepQMC workdir, e.g. as passed to
``hydra.run.dir``.
*keys (str): the names of the datasets to read from the ``result.h5`` files,
e.g. ``'local_energy'``.
read_workdir (~collections.abc.Callable): optional, the function used to
read the raw results from ``path``, defaults to
:func:`~deepqmc.postprocess.workdir.read_workdir`.
gather_electrons (bool): optional, whether to merge the per-device electron
batch axis of the results via
:func:`~deepqmc.postprocess.workdir.gather_electron_axis`. If
:data:`False`, only the results of the first device are kept.
Returns:
the array for the requested key, or a dictionary mapping each requested key
to its array if more than one key was requested.
"""
results, _ = read_workdir(path, keys)
min_idxs = {
key: min(len(result) for result in results.values()) for key in results.keys()
}
results = {
key: (
gather_electron_axis(results[key][: min_idxs[key]])
if gather_electrons
else results[key][: min_idxs[key], 0]
)
for key in keys
}
return list(results.values())[0] if len(results.keys()) == 1 else results
[docs]
def read_and_convert_result(
path, *keys, read_workdir=read_workdir, gather_electrons=True
):
r"""Read and convert results from a deepQMC workdir to per molecule format.
Use when the result file contains results for multiple molecules, and the
``mol_idxs`` entry, e.g. for a transferable training or evaluation run. See also
:func:`~deepqmc.postprocess.workdir.convert_to_per_molecule_format`.
Args:
path (~pathlib.Path): the deepQMC workdir, e.g. as passed to
``hydra.run.dir``.
*keys (str): the names of the datasets to read from the ``result.h5`` files,
e.g. ``'local_energy'``.
read_workdir (~collections.abc.Callable): optional, the function used to
read the raw results from ``path``, defaults to
:func:`~deepqmc.postprocess.workdir.read_workdir`.
gather_electrons (bool): optional, whether to merge the per-device electron
batch axis of the results via
:func:`~deepqmc.postprocess.workdir.gather_electron_axis`. If
:data:`False`, only the results of the first device are kept.
Returns:
the array for the requested key, or a dictionary mapping each requested key
to its array if more than one key was requested, rearranged into per
molecule format.
"""
results, _ = read_workdir(path, [*keys, 'mol_idxs'])
min_idxs = {
key: min(len(result) for result in results.values()) for key in results.keys()
}
electrons_gathered = {
key: (
gather_electron_axis(results[key][: min_idxs[key]])
if gather_electrons
else results[key][: min_idxs[key], 0]
)
for key in keys
}
results = {
k: convert_to_per_molecule_format(
electrons_gathered[k],
results['mol_idxs'][: min_idxs[k]],
)
for k in keys
}
return list(results.values())[0] if len(results.keys()) == 1 else results
[docs]
def read_average_iteration_time(path, discard_first=10):
r"""Read the average iteration time from a deepQMC workdir.
Needs the ``time`` entry in the result.h5 file. Computes the average wall-clock
time per training/evaluation iteration, discarding the first few iterations
(which include JIT compilation overhead) from the estimate.
Args:
path (~pathlib.Path): the deepQMC workdir, e.g. as passed to
``hydra.run.dir``.
discard_first (int): optional, the number of initial iterations to discard
from the estimate.
Returns:
float: the average wall-clock time per iteration, in seconds.
"""
times = read_workdir(path, ['time'])[0]['time']
average_time = (times[-1] - times[discard_first - 1]) / (len(times) - discard_first)
return average_time