import logging
import os
import time
from collections.abc import Callable, Sequence
from functools import partial
from itertools import count
from typing import Optional
import jax
import jax.numpy as jnp
import numpy as np
import optax
from tqdm.auto import tqdm, trange
from uncertainties import ufloat
from .conf.custom_resolvers import process_idx_suffix
from .ewm import init_multi_mol_multi_state_ewm
from .exceptions import NanError, TrainingBlowup, TrainingCrash
from .fit import fit_wf
from .hamil import MolecularHamiltonian
from .log import CheckpointStore, H5Logger, MetricLogger, TensorboardMetricLogger
from .loss import LossFunctionFactory, create_loss_fn
from .loss.clip import median_log_squeeze_and_mask
from .molecule import Molecule
from .observable import (
ObservableMonitor,
default_observable_monitors,
observable_monitor_from_name,
)
from .optimizer import NoOptimizer
from .parallel import pmap_pmean, split_on_devices, split_rng_key_to_devices
from .physics import pairwise_self_distance
from .pretrain.pretraining import pretrain
from .pretrain.pyscfext import compute_scf_solution
from .sampling import equilibrate, initialize_sampler_state
from .types import Ansatz, OptimizerFactory, SamplerFactory, TrainState
from .wf.base import init_wf_params
__all__ = ['train']
log = logging.getLogger(__name__)
[docs]
def train( # noqa: C901
hamil: MolecularHamiltonian,
ansatz: Ansatz,
opt: Optional[OptimizerFactory],
sampler_factory: SamplerFactory,
steps: int,
seed: Optional[int],
electron_batch_size: int,
molecule_batch_size: int = 1,
electronic_states: int = 1,
mols: Optional[list[Molecule]] = None,
workdir: Optional[str] = None,
train_state: Optional[TrainState] = None,
init_step: int = 0,
max_restarts: int = 3,
max_eq_steps: int = 1000,
eq_allow_early_stopping: bool = True,
pretrain_steps: Optional[int] = None,
pretrain_kwargs: Optional[dict] = None,
chkpt_constructor: Optional[Callable[..., CheckpointStore]] = None,
metric_logger_constructor: Optional[Callable[..., MetricLogger]] = None,
h5_logger_constructor: Optional[Callable[..., H5Logger]] = None,
merge_keys: Optional[list[str]] = None,
loss_function_factory: Optional[LossFunctionFactory] = None,
observable_monitors: Optional[list[ObservableMonitor | str]] = None,
):
r"""Train or evaluate a JAX wave function model.
It initializes and equilibrates the MCMC sampling of the wave function ansatz,
then optimizes or samples it using the variational principle. It optionally
saves checkpoints and rewinds the training/evaluation if an error is encountered.
If an optimizer factory is supplied, the Ansatz is optimized, otherwise the
Ansatz is only sampled.
Args:
hamil (~deepqmc.hamil.MolecularHamiltonian): the Hamiltonian of the physical
system.
ansatz (~deepqmc.types.Ansatz): the wave function Ansatz.
opt (Optional[~deepqmc.types.OptimizerFactory]): optional optimizer factory
or ``None``. Possible values include partially-initialized KFAC
optimizers, an ``optax`` optimizer instance, or ``None`` to run
evaluation-only.
sampler_factory (~deepqmc.types.SamplerFactory): callable that returns a
``(molecule_idx_sampler, sampler)`` pair used to create sampler state.
steps (int): number of optimization steps.
seed (Optional[int]): the seed used for PRNG; if omitted a random seed is
drawn using ``numpy.random.default_rng().integers(2**32)``.
electron_batch_size (int): the number of electron samples considered in a
batch.
molecule_batch_size (int): optional, the number of molecules considered in a
batch (used for transferable training).
electronic_states (int): optional, the number of electronic states to
consider.
mols (Optional[list[~deepqmc.molecule.Molecule]]): optional sequence of
molecules to consider for transferable training. If ``None``, the default
molecule in ``hamil`` is used.
workdir (Optional[str]): optional path where results should be saved.
train_state (Optional[~deepqmc.types.TrainState]): optional training
checkpoint to restore training or run evaluation.
init_step (int): optional initial step index, useful when restarting from a
checkpoint saved on disk.
max_restarts (int): optional maximum number of times the training is
retried before a :class:`NaNError` is raised.
max_eq_steps (int): optional maximum number of equilibration steps if not
detected earlier.
eq_allow_early_stopping (bool): whether to allow equilibration to stop early
when an equilibration criterion is met.
pretrain_steps (Optional[int]): optional number of pretraining steps with the
baseline wave function obtained from pyscf.
pretrain_kwargs (Optional[dict]): optional extra arguments for pretraining.
chkpt_constructor (Optional[~collections.abc.Callable[..., ~deepqmc.log.CheckpointStore]]): optional factory callable that returns a
:class:`~deepqmc.log.CheckpointStore` for saving training checkpoints to
``workdir``.
metric_logger_constructor (Optional[~collections.abc.Callable[..., ~deepqmc.log.MetricLogger]]): optional factory callable that returns a
:class:`~deepqmc.log.MetricLogger`; defaults to
:class:`~deepqmc.log.TensorboardMetricLogger` when omitted.
h5_logger_constructor (Optional[~collections.abc.Callable[..., ~deepqmc.log.H5Logger]]): optional factory callable that returns an
:class:`~deepqmc.log.H5Logger`; defaults to :class:`~deepqmc.log.H5Logger`
when omitted.
merge_keys (Optional[list[str]]): optional list of parameter-key substrings
that should be shared across electronic states.
loss_function_factory (Optional[~deepqmc.loss.LossFunctionFactory]): optional
callable returning a loss function for the optimization.
observable_monitors (Optional[list[~deepqmc.observable.ObservableMonitor | str]]): optional list of observable monitors (or monitor names) to evaluate
during training/evaluation.
"""
mode = 'evaluation' if opt is None else 'training'
if seed is None:
seed = int(np.random.default_rng().integers(2**32))
log.info('Seed not provided, using random seed %d', seed)
rng = jax.random.PRNGKey(seed + jax.process_index())
rng, rng_smpl = jax.random.split(rng)
mols = mols if isinstance(mols, Sequence) else [hamil.mol]
mols = mols * molecule_batch_size if len(mols) == 1 else mols
molecule_idx_sampler, sampler = sampler_factory(
rng_smpl,
hamil,
ansatz,
mols,
electronic_states,
molecule_batch_size,
)
optimizer_factory = opt or NoOptimizer
observable_monitors = observable_monitors or []
observables = [
(
observable
if isinstance(observable, ObservableMonitor)
else observable_monitor_from_name(observable)
)
for observable in observable_monitors
]
observables = default_observable_monitors() + observables
chkpts = None
if workdir:
workdir = os.path.join(workdir, mode + process_idx_suffix())
os.makedirs(workdir, exist_ok=True)
chkpts = (chkpt_constructor or CheckpointStore)(workdir)
log.debug('Setting up metric_logger...')
metric_logger = (metric_logger_constructor or TensorboardMetricLogger)(
workdir, molecule_batch_size
)
log.debug('Setting up h5_logger...')
h5_logger = (h5_logger_constructor or H5Logger)(
workdir,
init_step,
[observable.name for observable in observables],
aux_data={f'mol-{i}': m.coords for i, m in enumerate(mols)},
)
init_time = time.time()
else:
metric_logger = None
h5_logger = None
init_time = None
pbar = None
try:
if train_state:
log.info(
{
'training': f'Restart training from step {init_step}',
'evaluation': 'Start evaluation',
}[mode]
)
params = train_state.params
else:
rng, rng_init = jax.random.split(rng)
params = init_wf_params(
rng_init, hamil, ansatz, electronic_states, merge_keys=merge_keys
)
if pretrain_steps and mode == 'training':
log.info('Pretraining wrt. baseline wave function')
rng, rng_pretrain = jax.random.split(rng)
pretrain_kwargs = pretrain_kwargs or {}
pretrain_dataset = compute_scf_solution(
mols,
hamil,
electronic_states,
workdir=pretrain_kwargs.pop('pyscf_chkpt_path', None) or workdir,
**pretrain_kwargs.pop('scf_kwargs', {}),
)
opt_pretrain = getattr(optax, pretrain_kwargs.pop('opt', 'adam'))(
**pretrain_kwargs.pop('opt_kwargs', {'learning_rate': 3.0e-4})
)
ewm_state, update_ewm = init_multi_mol_multi_state_ewm(
shape=(len(mols), electronic_states), decay_alpha=1.0
)
mse_rep = None
_, rng_pretrain_smpl_init = split_on_devices(
split_rng_key_to_devices(rng), 2
)
pretrain_smpl_state = initialize_sampler_state(
rng_pretrain_smpl_init,
sampler,
params,
electron_batch_size,
jnp.stack([mol.coords for mol in mols]),
)
pbar = tqdm(range(pretrain_steps), desc='pretrain', disable=None)
for step, params, per_sample_losses, mol_idxs in pretrain( # noqa: B007
rng_pretrain,
hamil,
ansatz,
params,
opt_pretrain,
molecule_idx_sampler,
sampler,
pretrain_smpl_state,
pretrain_dataset,
merge_keys,
steps=pbar,
): # noqa: B007
per_mol_state_losses = per_sample_losses.mean(axis=-1)
ewm_state = update_ewm(per_mol_state_losses, ewm_state, mol_idxs)
pretrain_stats = {
'MSE': per_mol_state_losses,
'MSE/ewm': ewm_state.mean,
}
mse_rep = '|'.join(
'(' + '|'.join(f'{mses:0.2e}' for mses in msem) + ')'
for msem in ewm_state.mean
)
pbar.set_postfix(MSE=mse_rep)
if metric_logger:
metric_logger.update(
step, pretrain_stats, {}, mol_idxs, prefix='pretraining'
)
log.info(f'Pretraining completed with MSE = {mse_rep}')
rng = split_rng_key_to_devices(rng)
if train_state is None or train_state.sampler is None:
rng, rng_eq, rng_smpl_init = split_on_devices(rng, 3)
smpl_state = initialize_sampler_state(
rng_smpl_init,
sampler,
params,
electron_batch_size,
jnp.stack([mol.coords for mol in mols]),
)
log.info('Equilibrating sampler...')
pbar = tqdm(
count() if max_eq_steps is None else range(max_eq_steps),
desc='equilibrate sampler',
disable=None,
)
for step, smpl_state, mol_idxs, smpl_stats in equilibrate( # noqa: B007
rng_eq,
params,
molecule_idx_sampler,
sampler,
smpl_state,
lambda phys_conf: pmap_pmean(
pairwise_self_distance(phys_conf.r)
).mean(),
pbar,
block_size=10,
allow_early_stopping=eq_allow_early_stopping,
):
tau_rep = '|'.join(
'(' + '|'.join(f'{taus:.3f}' for taus in taum) + ')'
for taum in smpl_state['elec']['tau'].mean(axis=0)
)
pbar.set_postfix(tau=tau_rep)
if metric_logger:
metric_logger.update(
step, {}, smpl_stats, mol_idxs, prefix='equilibration'
)
pbar.close()
train_state = TrainState(smpl_state, params, None)
if workdir and mode == 'training':
assert chkpts
chkpts.update(init_step, train_state)
log.info(f'Start {mode}')
loss_function_factory = loss_function_factory or partial(
create_loss_fn, clip_mask_fn=median_log_squeeze_and_mask
)
best_ene = None
step = init_step
ewm_energies = len(mols) * [electronic_states * [ufloat(jnp.nan, 1)]]
for attempt in range(max_restarts):
try:
pbar = trange(
init_step,
steps,
initial=init_step,
total=steps,
desc=mode,
disable=None,
)
for (
step,
train_state,
mol_idxs,
stats,
observable_samples,
) in fit_wf( # noqa: B007
rng,
hamil,
ansatz,
optimizer_factory,
molecule_idx_sampler,
sampler,
pbar,
train_state,
merge_keys,
loss_function_factory,
observable_monitors=[
observable.finalize(hamil, ansatz.apply)
for observable in observables
],
):
ewm_energies, best_ene = update_progress(
pbar, best_ene, ewm_energies, mol_idxs, stats
)
if jnp.isnan(observable_samples['psi/samples']['log']).any():
raise NanError()
if workdir:
assert init_time is not None
assert h5_logger is not None
if mode == 'training':
assert chkpts
# the convention is that chkpt-i contains the step i-1 -> i
chkpts.update(step + 1, train_state)
if metric_logger:
metric_logger.update(step, stats, {}, mol_idxs)
observable_samples |= {
'mol_idxs': mol_idxs,
'step': step,
'time': time.time() - init_time,
**stats,
}
h5_logger.update(observable_samples)
log.info(f'The {mode} has been completed!')
return train_state
except (NanError, TrainingBlowup) as e:
if pbar:
pbar.close()
log.warning(f'Restarting due to {type(e).__name__}...')
if attempt < max_restarts and chkpts is not None:
init_step, train_state = chkpts.last
(rng,) = split_on_devices(rng, 1)
log.warning(
f'The {mode} has crashed before all steps were completed ({step}/{steps})!'
)
raise TrainingCrash(train_state)
finally:
if pbar:
pbar.close()
if chkpts:
chkpts.close()
if metric_logger:
metric_logger.close()
if h5_logger:
h5_logger.close()
def update_progress(pbar, best_ene, ewm_energies, mol_idxs, stats):
r"""Update the tqdm progress bar, maybe print progress message."""
if stats.get('energy/ewm') is not None:
for i, mol_idx in enumerate(mol_idxs):
ewm_energies[mol_idx] = [
ufloat(mean, sqerr)
for (mean, sqerr) in zip(
stats['energy/ewm'][i], stats['energy/ewm_error'][i]
)
]
energies = '|'.join(
'(' + '|'.join(f'{es:S}' for es in em) + ')' for em in ewm_energies
)
pbar.set_postfix(E=energies)
if best_ene is None or jnp.any(
jnp.array(
jax.tree.map(lambda x, y: x.s < 0.5 * y.s, ewm_energies, best_ene)
)
):
best_ene = ewm_energies
log.info(f'Progress: {pbar.n + 1}/{pbar.total}, energy = {energies}')
return ewm_energies, best_ene