From 905b460f698d902a48c5c7fddb4d51647fe93479 Mon Sep 17 00:00:00 2001 From: Kayahan Saritas Date: Fri, 21 Aug 2026 12:28:52 -0400 Subject: [PATCH 1/2] dynamic walk workflows for ecut and k-grid --- nexus/nexus/dynamic_workflows/__init__.py | 18 ++ nexus/nexus/dynamic_workflows/base.py | 167 ++++++++++++++ .../nexus/dynamic_workflows/pwscf/__init__.py | 13 ++ nexus/nexus/dynamic_workflows/pwscf/ecut.py | 115 +++++++++ nexus/nexus/dynamic_workflows/pwscf/kgrid.py | 169 ++++++++++++++ nexus/nexus/dynamic_workflows/walk.py | 218 ++++++++++++++++++ .../dynamic_workflows/03_walk/ecut_walk.py | 62 +++++ .../dynamic_workflows/03_walk/kgrid_walk.py | 70 ++++++ 8 files changed, 832 insertions(+) create mode 100644 nexus/nexus/dynamic_workflows/__init__.py create mode 100644 nexus/nexus/dynamic_workflows/base.py create mode 100644 nexus/nexus/dynamic_workflows/pwscf/__init__.py create mode 100644 nexus/nexus/dynamic_workflows/pwscf/ecut.py create mode 100644 nexus/nexus/dynamic_workflows/pwscf/kgrid.py create mode 100644 nexus/nexus/dynamic_workflows/walk.py create mode 100644 nexus/nexus/examples/dynamic_workflows/03_walk/ecut_walk.py create mode 100644 nexus/nexus/examples/dynamic_workflows/03_walk/kgrid_walk.py diff --git a/nexus/nexus/dynamic_workflows/__init__.py b/nexus/nexus/dynamic_workflows/__init__.py new file mode 100644 index 0000000000..0b3afd0b0c --- /dev/null +++ b/nexus/nexus/dynamic_workflows/__init__.py @@ -0,0 +1,18 @@ +################################################################## +## (c) Copyright 2015- by Jaron T. Krogel ## +################################################################## + + +#====================================================================# +# dynamic_workflows # +# Modes for dynamic (poll-loop) Nexus workflows. # +# # +# Walk sequential parameter walk (implemented). # +# Spawn fan-out / pick-first / pick-min (not yet implemented). # +#====================================================================# + + +from .base import DynamicDecision, DynamicMode, Target +from .walk import DynamicWalk, WalkDecision, SuccessiveChange +from .pwscf.ecut import Ecut, next_ecutwfc +from .pwscf.kgrid import Kgrid, kgrid_from_kspacing, next_kgrid diff --git a/nexus/nexus/dynamic_workflows/base.py b/nexus/nexus/dynamic_workflows/base.py new file mode 100644 index 0000000000..ed6875c3e0 --- /dev/null +++ b/nexus/nexus/dynamic_workflows/base.py @@ -0,0 +1,167 @@ +################################################################## +## (c) Copyright 2015- by Jaron T. Krogel ## +################################################################## + + +#====================================================================# +# base.py # +# Shared types for dynamic-workflow modes (walk, later spawn). # +# # +# Content summary: # +# Target # +# Sequential stop rule: reached(history). Used by walk. # +# Spawn pick-rules should not inherit Target. # +# DynamicDecision # +# status / products / completed / max_runs. # +# WalkDecision and spawn decisions subclass this. # +# DynamicMode # +# Shared protocol: initial / propose / stop / observe / drive. # +# history, max_runs. # +#====================================================================# + + +from ..developer import DevBase + + +def _iter_work(params): + """Yield parameter dicts. Walk: one dict. Spawn: a sequence of dicts.""" + if params is None: + return + #end if + if isinstance(params, dict): + yield params + return + #end if + for item in params: + yield item + #end for +#end def _iter_work + + +def _as_products(spec, sim): + if callable(spec): + return dict(spec(sim)) + #end if + if isinstance(spec, str): + spec = (spec,) + #end if + products = {} + for key in spec: + products[key] = getattr(sim.products, key) + #end for + return products +#end def _as_products + + +class Target(DevBase): + """Sequential stop rule. Subclasses implement ``reached(history)``. + + ``history`` is a list of ``{params, products}`` records in time order. + This protocol is sequential; spawn pick-rules do not inherit Target. + """ + + def reached(self, history): + self.error('Target.reached() must be implemented in a subclass') + #end def reached +#end class Target + + +class DynamicDecision(DevBase): + """Outcome of one dynamic-mode observation. + + ``status`` and ``products`` are shared with future spawn decisions. + + status + ``continue`` keep polling this step + ``completed`` finished successfully (walk: target reached) + ``max_runs`` stopped without completing (cap, or propose() is None) + """ + + def __init__(self, status, products=None): + self.status = status + self.products = None if products is None else dict(products) + #end def __init__ + + @property + def completed(self): + return self.status == 'completed' + #end def completed + + @property + def max_runs(self): + return self.status == 'max_runs' + #end def max_runs +#end class DynamicDecision + + +class DynamicMode(DevBase): + """Shared protocol for dynamic workflows: initial/propose/drive/observe/stop. + + initial: First work to launch. + propose(history): Next work(s) to launch, or None if there is no further work. + drive(make, wm, products): Poll until observe is not continue. + observe(params, products): Record a finished job and return a ``DynamicDecision``. + stop(history): True when the mode has finished successfully. Default is False. + """ + + def __init__(self, max_runs=10): + max_runs = int(max_runs) + if max_runs < 1: + self.error(f'max_runs must be >= 1, received {max_runs}') + #end if + self.max_runs = max_runs + self.history = [] + #end def __init__ + + def initial(self): + self.error('initial() must be implemented in a subclass') + #end def initial + + def propose(self, history): + self.error('propose() must be implemented in a subclass') + #end def propose + + def drive(self, make, wm, products, poll=1): + """Return a DynamicDecision after polling until the mode is finished.""" + jobs = [] + for params in _iter_work(self.initial()): + jobs.append([params, make(params)]) + #end for + if len(jobs) == 0: + self.error('initial() produced no work to launch') + #end if + while True: + for job in list(jobs): + params, sim = job + if sim.succ: + decision = self.observe(params, _as_products(products, sim)) + jobs.remove(job) + if decision.status != 'continue': + return decision + #end if + nxt = list(_iter_work(decision.next_params)) + if len(nxt) == 0: + self.error('continue decision has no next_params') + #end if + for p in nxt: + jobs.append([p, make(p)]) + #end for + elif sim.fail: + self.error('simulation failed') + #end if + #end for + if len(jobs) == 0: + self.error('no remaining jobs') + #end if + wm.poll(poll) + #end while + #end def drive + + def observe(self, params, products): + self.error('observe() must be implemented in a subclass') + #end def observe + + def stop(self, history): + return False + #end def stop +#end class DynamicMode diff --git a/nexus/nexus/dynamic_workflows/pwscf/__init__.py b/nexus/nexus/dynamic_workflows/pwscf/__init__.py new file mode 100644 index 0000000000..b003cf38a7 --- /dev/null +++ b/nexus/nexus/dynamic_workflows/pwscf/__init__.py @@ -0,0 +1,13 @@ +################################################################## +## (c) Copyright 2015- by Jaron T. Krogel ## +################################################################## + + +#====================================================================# +# pwscf # +# PWscf-specific dynamic walks (Ecut, Kgrid). # +#====================================================================# + + +from .ecut import Ecut, next_ecutwfc +from .kgrid import Kgrid, kgrid_from_kspacing, next_kgrid diff --git a/nexus/nexus/dynamic_workflows/pwscf/ecut.py b/nexus/nexus/dynamic_workflows/pwscf/ecut.py new file mode 100644 index 0000000000..8c884d4859 --- /dev/null +++ b/nexus/nexus/dynamic_workflows/pwscf/ecut.py @@ -0,0 +1,115 @@ +################################################################## +## (c) Copyright 2015- by Jaron T. Krogel ## +################################################################## + + +#====================================================================# +# ecut.py # +# Walk-mode PWscf planewave-cutoff (example 02 repeat). # +#====================================================================# + + +from ..walk import DynamicWalk, SuccessiveChange +from ...developer import error + + +def next_ecutwfc(ecut, factor=1.5, round_ry=10): + """Increase cutoff as in the diamond energy-convergence example. + + ``int(((factor * ecut) // round_ry) * round_ry)``, then bump by + ``round_ry`` if the rounded value did not increase. + """ + ecut = float(ecut) + factor = float(factor) + round_ry = int(round_ry) + if ecut <= 0.0: + error( + f'ecutwfc must be positive, received {ecut}', + header='Ecut', + ) + #end if + if factor <= 1.0: + error( + f'ecut factor must be > 1, received {factor}', + header='Ecut', + ) + #end if + if round_ry < 1: + error( + f'round_ry must be >= 1, received {round_ry}', + header='Ecut', + ) + #end if + nxt = int(((factor * ecut) // round_ry) * round_ry) + if nxt <= int(ecut): + nxt = int(ecut) + round_ry + #end if + return nxt +#end def next_ecutwfc + + +class Ecut(DynamicWalk): + """Walk ``ecutwfc`` until a successive-change target is reached. + + Parameter dict: ``{'ecutwfc': int}``. Default stepping matches + ``examples/dynamic_workflows/02_energy_convergence`` (×1.5, rounded + down to 10 Ry). The poll loop and ``generate_pwscf`` stay in the + user script. + + ``target`` may be a target object (``reached(history)``), or the + name ``'consecutive'`` (default), which builds + ``SuccessiveChange('energy', atol=1e-4, ...)``. ``tol=`` is accepted + as an alias for ``atol=``. + """ + + def __init__( + self, + start=50, + end=None, + factor=1.5, + round_ry=10, + max_runs=10, + target='consecutive', + **target_kw + ): + start = int(start) + if start < 1: + self.error(f'start ecutwfc must be >= 1, received {start}') + #end if + if end is not None: + end = int(end) + if end < start: + self.error(f'end must be >= start, received {end} < {start}') + #end if + #end if + if target is None or target == 'consecutive': + if 'tol' in target_kw and 'atol' not in target_kw: + target_kw['atol'] = target_kw.pop('tol') + #end if + if 'atol' not in target_kw and 'rtol' not in target_kw: + target_kw['atol'] = 1e-4 + #end if + target = SuccessiveChange('energy', **target_kw) + elif not hasattr(target, 'reached'): + self.error(f'invalid target: {target}') + #end if + DynamicWalk.__init__(self, target=target, max_runs=max_runs) + self.start = start + self.end = end + self.factor = float(factor) + self.round_ry = int(round_ry) + #end def __init__ + + def initial(self): + return dict(ecutwfc=self.start) + #end def initial + + def propose(self, history): + ecut = int(history[-1]['params']['ecutwfc']) + nxt = next_ecutwfc(ecut, factor=self.factor, round_ry=self.round_ry) + if self.end is not None and nxt > self.end: + return None + #end if + return dict(ecutwfc=nxt) + #end def propose +#end class Ecut diff --git a/nexus/nexus/dynamic_workflows/pwscf/kgrid.py b/nexus/nexus/dynamic_workflows/pwscf/kgrid.py new file mode 100644 index 0000000000..03668bf9bd --- /dev/null +++ b/nexus/nexus/dynamic_workflows/pwscf/kgrid.py @@ -0,0 +1,169 @@ +################################################################## +## (c) Copyright 2015- by Jaron T. Krogel ## +################################################################## + + +#====================================================================# +# kgrid.py # +# Walk-mode PWscf k-point mesh (example 02 policy). # +# # +# increment add to the integer Monkhorst-Pack grid each step. # +# spacing initial mesh from Structure.kgrid_from_kspacing. # +#====================================================================# + + +from ..walk import DynamicWalk, SuccessiveChange +from ...developer import error + + +def _as_kgrid(kgrid, name='kgrid'): + if isinstance(kgrid, (int, float)): + n = int(kgrid) + kgrid = (n, n, n) + else: + kgrid = tuple(int(n) for n in kgrid) + #end if + if len(kgrid) != 3 or min(kgrid) < 1: + error( + f'{name} must be a positive int or 3-tuple, received {kgrid}', + header='Kgrid', + ) + #end if + return kgrid +#end def _as_kgrid + + +def _as_increment(increment): + if isinstance(increment, (int, float)): + d = int(increment) + increment = (d, d, d) + else: + increment = tuple(int(n) for n in increment) + #end if + if len(increment) != 3 or min(increment) < 0 or max(increment) < 1: + error( + f'increment must be a positive int or 3-tuple with at least one > 0, received {increment}', + header='Kgrid', + ) + #end if + return increment +#end def _as_increment + + +def _structure(system): + if system is None: + error( + 'structure is required when spacing is set', + header='Kgrid', + ) + #end if + struct = system.structure if hasattr(system, 'structure') else system + if not hasattr(struct, 'kgrid_from_kspacing'): + tname = type(system).__name__ + error( + f'structure must provide kgrid_from_kspacing, received {tname}', + header='Kgrid', + ) + #end if + return struct +#end def _structure + + +def kgrid_from_kspacing(structure, kspacing): + """Integer mesh from ``Structure.kgrid_from_kspacing``.""" + kspacing = float(kspacing) + if kspacing <= 0.0: + error( + f'kspacing must be positive, received {kspacing}', + header='Kgrid', + ) + #end if + return tuple(int(n) for n in _structure(structure).kgrid_from_kspacing(kspacing)) +#end def kgrid_from_kspacing + + +def next_kgrid(kgrid, increment=1): + """Increase a Monkhorst-Pack grid by ``increment`` (example 02: +1). + + ``kgrid`` and ``increment`` may be a positive int (applied on all + three axes) or a 3-tuple. + """ + kgrid = _as_kgrid(kgrid) + increment = _as_increment(increment) + return tuple(n + d for n, d in zip(kgrid, increment)) +#end def next_kgrid + + +class Kgrid(DynamicWalk): + """Walk ``kgrid`` until a successive-change target is reached. + + Parameter dict: ``{'kgrid': (nx, ny, nz)}``. Default stepping + matches ``examples/dynamic_workflows/02_energy_convergence`` + (start ``(1,1,1)``, add 1 to each axis). The poll loop and + ``generate_pwscf`` stay in the user script. + + ``spacing`` + If set, the first mesh is ``structure.kgrid_from_kspacing(spacing)``. + Later points still use ``increment``. ``structure`` may be a + ``Structure`` or a ``PhysicalSystem``. + ``increment`` + Added to the integer grid each step (int or 3-tuple). + ``end`` + Optional cap (int or 3-tuple). ``propose`` returns ``None`` if + the next mesh would exceed it on any axis. + """ + + def __init__( + self, + start=(1, 1, 1), + end=None, + increment=1, + spacing=None, + structure=None, + max_runs=10, + target='successive_change', + **target_kw + ): + increment = _as_increment(increment) + if spacing is not None: + start = kgrid_from_kspacing(structure, spacing) + else: + start = _as_kgrid(start, 'start') + #end if + if end is not None: + end = _as_kgrid(end, 'end') + if any(s > e for s, e in zip(start, end)): + self.error(f'end must be >= start, received {end} < {start}') + #end if + #end if + if target is None or target == 'successive_change': + if 'tol' in target_kw and 'atol' not in target_kw: + target_kw['atol'] = target_kw.pop('tol') + #end if + if 'atol' not in target_kw and 'rtol' not in target_kw: + target_kw['atol'] = 1e-3 + #end if + target = SuccessiveChange('energy', **target_kw) + elif not hasattr(target, 'reached'): + self.error(f'invalid target: {target}') + #end if + DynamicWalk.__init__(self, target=target, max_runs=max_runs) + self.start = start + self.end = end + self.increment = increment + self.spacing = None if spacing is None else float(spacing) + #end def __init__ + + def initial(self): + return dict(kgrid=self.start) + #end def initial + + def propose(self, history): + kgrid = _as_kgrid(history[-1]['params']['kgrid']) + nxt = next_kgrid(kgrid, increment=self.increment) + if self.end is not None and any(n > e for n, e in zip(nxt, self.end)): + return None + #end if + return dict(kgrid=nxt) + #end def propose +#end class Kgrid diff --git a/nexus/nexus/dynamic_workflows/walk.py b/nexus/nexus/dynamic_workflows/walk.py new file mode 100644 index 0000000000..5701368f21 --- /dev/null +++ b/nexus/nexus/dynamic_workflows/walk.py @@ -0,0 +1,218 @@ +################################################################## +## (c) Copyright 2015- by Jaron T. Krogel ## +################################################################## + + +#====================================================================# +# walk.py # +# Walk mode: sequential parameter walk (initial / propose / stop).# +# # +# Content summary: # +# SuccessiveChange(Target) # +# Successive |Δ| of one or more products within atol/rtol. # +# WalkDecision(DynamicDecision) # +# Outcome of one walk observation. # +# DynamicWalk(DynamicMode) # +# Sequential walk. stop is Target.reached; observe returns # +# WalkDecision. Subclasses implement initial / propose. # +#====================================================================# + + +from ..developer import error +from .base import DynamicDecision, DynamicMode, Target + + +def _as_params(params): + if not isinstance(params, dict) or len(params) == 0: + error( + f'walk parameters must be a non-empty dict, received {params}', + header='DynamicWalk', + ) + #end if + return dict(params) +#end def _as_params + + +class WalkDecision(DynamicDecision): + """Outcome of ``DynamicWalk.observe``. + + status + ``continue`` propose a new parameter point + ``completed`` target reached + ``max_runs`` history length reached max_runs, or propose() returned None + """ + + def __init__(self, status, params, next_params=None, products=None): + DynamicDecision.__init__(self, status, products=products) + self.params = None if params is None else dict(params) + self.next_params = None if next_params is None else dict(next_params) + #end def __init__ +#end class WalkDecision + + +class SuccessiveChange(Target): + """Target: successive values of one or more products have settled. + + A pair ``(old, new)`` of a product is within tolerance if:: + + abs(new - old) <= atol + rtol * abs(old) + + (same combination as ``numpy.isclose``). Every listed key must pass + for the last ``consecutive`` successive pairs. + + Parameters + ---------- + keys : str or sequence of str + Product names, e.g. ``'energy'`` or ``('energy', 'variance')``. + atol : float or dict, optional + Absolute tolerance. A scalar applies to every key; a dict maps + key → atol. Default is ``0.0`` when ``rtol`` is set. + rtol : float or dict, optional + Relative tolerance against the previous value. Default is ``0.0`` + when ``atol`` is set. + consecutive : int, optional + Number of successive pairs that must all pass. Default is ``1``. + + Examples + -------- + abs_energy = SuccessiveChange('energy', atol=1e-4) + rel_energy = SuccessiveChange('energy', rtol=1e-3) + mixed_target = SuccessiveChange( + ('energy', 'variance'), + atol={'energy': 1e-4, 'variance': 0.01}, + ) + """ + + def __init__(self, keys='energy', atol=None, rtol=None, consecutive=1): + if isinstance(keys, str): + keys = (keys,) + else: + keys = tuple(keys) + #end if + if len(keys) == 0: + self.error('SuccessiveChange requires at least one product key') + #end if + if atol is None and rtol is None: + self.error('SuccessiveChange requires atol and/or rtol') + #end if + consecutive = int(consecutive) + if consecutive < 1: + self.error(f'consecutive must be >= 1, received {consecutive}') + #end if + self.keys = keys + self.atol = _tol_map(atol, keys, 0.0) + self.rtol = _tol_map(rtol, keys, 0.0) + self.consecutive = consecutive + #end def __init__ + + def _value(self, record, key): + products = record.get('products') or {} + if key not in products: + self.error(f'SuccessiveChange target requires products[{key!r}]') + #end if + value = float(products[key]) + return value + #end def _value + + def _within(self, key, old, new): + return abs(new - old) <= self.atol[key] + self.rtol[key] * abs(old) + #end def _within + + def reached(self, history): + needed = self.consecutive + 1 + if len(history) < needed: + return False + #end if + for offset in range(self.consecutive): + newer = history[-(offset + 1)] + older = history[-(offset + 2)] + for key in self.keys: + if not self._within(key, self._value(older, key), self._value(newer, key)): + return False + #end if + #end for + #end for + return True + #end def reached +#end class SuccessiveChange + + +available_targets = { + 'successive_change': SuccessiveChange, + } + + +def _tol_map(spec, keys, default): + if spec is None: + return {key: float(default) for key in keys} + #end if + if isinstance(spec, dict): + missing = [key for key in keys if key not in spec] + if missing and default is None: + error( + f'missing tolerances for product keys: {missing}', + header='SuccessiveChange', + ) + #end if + return {key: float(spec[key]) if key in spec else float(default) for key in keys} + #end if + val = float(spec) + return {key: val for key in keys} +#end def _tol_map + + +class DynamicWalk(DynamicMode): + """Sequential parameter walk (walk mode). + + ``stop`` is ``self.target.reached(history)``. ``observe`` returns + ``WalkDecision``. + """ + + def __init__(self, target=None, max_runs=10): + DynamicMode.__init__(self, max_runs=max_runs) + if target is not None and not isinstance(target, Target): + tname = type(target).__name__ + self.error( + f'target must be a Target instance with reached(history), received {tname}' + ) + #end if + self.target = target + #end def __init__ + + def stop(self, history): + if self.target is None: + return DynamicMode.stop(self, history) + #end if + return self.target.reached(history) + #end def stop + + def observe(self, params, products): + params = _as_params(params) + products = dict(products) + self.history.append(dict(params=params, products=products)) + if self.stop(self.history): + return WalkDecision( + 'completed', + params = params, + products = products, + ) + #end if + nxt = None + if len(self.history) < self.max_runs: + nxt = self.propose(self.history) + #end if + if nxt is None: + return WalkDecision( + 'max_runs', + params = params, + products = products, + ) + #end if + return WalkDecision( + 'continue', + params = params, + next_params = _as_params(nxt), + products = products, + ) + #end def observe +#end class DynamicWalk diff --git a/nexus/nexus/examples/dynamic_workflows/03_walk/ecut_walk.py b/nexus/nexus/examples/dynamic_workflows/03_walk/ecut_walk.py new file mode 100644 index 0000000000..7300474a05 --- /dev/null +++ b/nexus/nexus/examples/dynamic_workflows/03_walk/ecut_walk.py @@ -0,0 +1,62 @@ +#! /usr/bin/env python3 + +''' +Planewave cutoff walk using ``Ecut`` (walk mode). + +Same diamond / BFD setup and stepping as the ecut half of +``02_energy_convergence`` (start 50 Ry, ×1.5 rounded to 10 Ry, +successive |ΔE| ≤ 1e-4 Ry). ``drive`` owns the poll loop; +``Ecut`` only decides the next ``ecutwfc`` and when the target +is reached. ``generate_pwscf`` stays in this script. +''' + +import sys +from nexus import settings, job, workflow_manager +from nexus import generate_physical_system +from nexus import generate_pwscf +from nexus.dynamic_workflows.pwscf import Ecut + + +settings( + results = '', + pseudo_dir = '../../qmcpack/pseudopotentials', + machine = 'ws8', + dynamic = True, + ) + + +system = generate_physical_system( + structure = 'diamond', + cell = 'prim', + kgrid = (1, 1, 1), + C = 4, + ) + +walk = Ecut(start=50, tol=1e-4, max_runs=10) + + +def make_scf(params): + ecut = params['ecutwfc'] + return generate_pwscf( + identifier = 'scf', + path = f'ecut_{ecut}', + job = job(cores=4), + system = system, + pseudos = ['C.BFD.upf'], + input_type = 'generic', + ecutwfc = ecut, + kgrid = (1, 1, 1), + dynamic_id = f'ecut{ecut}', + requires = 'none', + ) +#end def make_scf + +wm = workflow_manager() +decision = walk.drive(make_scf, wm, products='energy') +print() +print('status :', decision.status) +print('ecut :', decision.params['ecutwfc']) +print('products:', decision.products) +if not decision.completed: + sys.exit(1) +#end if diff --git a/nexus/nexus/examples/dynamic_workflows/03_walk/kgrid_walk.py b/nexus/nexus/examples/dynamic_workflows/03_walk/kgrid_walk.py new file mode 100644 index 0000000000..17317f8d8b --- /dev/null +++ b/nexus/nexus/examples/dynamic_workflows/03_walk/kgrid_walk.py @@ -0,0 +1,70 @@ +#! /usr/bin/env python3 + +''' +k-point mesh walk using ``Kgrid`` (walk mode). + +Same diamond / BFD setup and stepping as the k-grid half of +``02_energy_convergence`` (start 1x1x1, +1 per axis, successive +|ΔE| ≤ 1e-3 Ry). ``drive`` owns the poll loop; ``Kgrid`` only +decides the next ``kgrid`` and when the target is reached. +``generate_pwscf`` stays in this script. + +Set ``ecutwfc`` to the cutoff from ``ecut_walk.py`` (example 02 +uses 330 Ry). 50 Ry is used here so this script can run on its +own. For a spacing-based start instead of ``start=1``:: + + walk = Kgrid(spacing=0.5, structure=system, increment=1, tol=1e-3) +''' + +import sys +from nexus import settings, job, workflow_manager +from nexus import generate_physical_system +from nexus import generate_pwscf +from nexus.dynamic_workflows.pwscf import Kgrid + + +settings( + results = '', + pseudo_dir = '../../qmcpack/pseudopotentials', + machine = 'ws8', + dynamic = True, + ) + + +system = generate_physical_system( + structure = 'diamond', + cell = 'prim', + kgrid = (1, 1, 1), + C = 4, + ) + +ecutwfc = 50 +walk = Kgrid(start=1, increment=1, tol=1e-3, max_runs=10) + + +def make_scf(params): + kgrid = params['kgrid'] + n1, n2, n3 = kgrid + return generate_pwscf( + identifier = 'scf', + path = f'kgrid_{n1}{n2}{n3}', + job = job(cores=4), + system = system, + pseudos = ['C.BFD.upf'], + input_type = 'generic', + ecutwfc = ecutwfc, + kgrid = kgrid, + dynamic_id = f'kgrid{n1}{n2}{n3}', + requires = 'none', + ) +#end def make_scf + +wm = workflow_manager() +decision = walk.drive(make_scf, wm, products='energy') +print() +print('status :', decision.status) +print('kgrid :', decision.params['kgrid']) +print('products:', decision.products) +if not decision.completed: + sys.exit(1) +#end if From d27f54ed24f24d94ca6c641962ba3edc8c880a57 Mon Sep 17 00:00:00 2001 From: Kayahan Saritas Date: Fri, 21 Aug 2026 15:42:59 -0400 Subject: [PATCH 2/2] error handling initial --- nexus/nexus/dynamic_workflows/__init__.py | 9 +- nexus/nexus/dynamic_workflows/base.py | 60 +-- .../dynamic_workflows/{walk.py => chain.py} | 96 +++-- .../dynamic_workflows/chain_error_handler.py | 199 ++++++++++ .../nexus/dynamic_workflows/pwscf/__init__.py | 4 +- nexus/nexus/dynamic_workflows/pwscf/ecut.py | 26 +- .../dynamic_workflows/pwscf/error_handler.py | 352 ++++++++++++++++++ nexus/nexus/dynamic_workflows/pwscf/kgrid.py | 10 +- .../ecut_walk.py => 03_chain/ecut_chain.py} | 12 +- .../kgrid_walk.py => 03_chain/kgrid_chain.py} | 18 +- .../04_pwscf_recover/recover_chain.py | 58 +++ nexus/nexus/pwscf.py | 68 +++- nexus/nexus/simulation.py | 3 + 13 files changed, 789 insertions(+), 126 deletions(-) rename nexus/nexus/dynamic_workflows/{walk.py => chain.py} (66%) create mode 100644 nexus/nexus/dynamic_workflows/chain_error_handler.py create mode 100644 nexus/nexus/dynamic_workflows/pwscf/error_handler.py rename nexus/nexus/examples/dynamic_workflows/{03_walk/ecut_walk.py => 03_chain/ecut_chain.py} (71%) rename nexus/nexus/examples/dynamic_workflows/{03_walk/kgrid_walk.py => 03_chain/kgrid_chain.py} (62%) create mode 100644 nexus/nexus/examples/dynamic_workflows/04_pwscf_recover/recover_chain.py diff --git a/nexus/nexus/dynamic_workflows/__init__.py b/nexus/nexus/dynamic_workflows/__init__.py index 0b3afd0b0c..cdc70ad12f 100644 --- a/nexus/nexus/dynamic_workflows/__init__.py +++ b/nexus/nexus/dynamic_workflows/__init__.py @@ -7,12 +7,15 @@ # dynamic_workflows # # Modes for dynamic (poll-loop) Nexus workflows. # # # -# Walk sequential parameter walk (implemented). # -# Spawn fan-out / pick-first / pick-min (not yet implemented). # +# Chain sequential parameter chain (implemented). # +# Error ChainErrorHandler; PWscf via error_handling=True. # +# Spawn fan-out / pick-first / pick-min (not yet implemented). # #====================================================================# from .base import DynamicDecision, DynamicMode, Target -from .walk import DynamicWalk, WalkDecision, SuccessiveChange +from .chain import DynamicChain, ChainDecision, SuccessiveChange +from .chain_error_handler import ChainErrorHandler from .pwscf.ecut import Ecut, next_ecutwfc from .pwscf.kgrid import Kgrid, kgrid_from_kspacing, next_kgrid +from .pwscf.error_handler import PwscfErrorHandler, parse_pwscf_text diff --git a/nexus/nexus/dynamic_workflows/base.py b/nexus/nexus/dynamic_workflows/base.py index ed6875c3e0..1ad95915a1 100644 --- a/nexus/nexus/dynamic_workflows/base.py +++ b/nexus/nexus/dynamic_workflows/base.py @@ -5,15 +5,15 @@ #====================================================================# # base.py # -# Shared types for dynamic-workflow modes (walk, later spawn). # +# Shared types for dynamic-workflow modes (chain, later spawn). # # # # Content summary: # # Target # -# Sequential stop rule: reached(history). Used by walk. # +# Sequential stop rule: reached(history). Used by chain. # # Spawn pick-rules should not inherit Target. # # DynamicDecision # -# status / products / completed / max_runs. # -# WalkDecision and spawn decisions subclass this. # +# status / products / completed / max_runs / failed. # +# ChainDecision and spawn decisions subclass this. # # DynamicMode # # Shared protocol: initial / propose / stop / observe / drive. # # history, max_runs. # @@ -24,7 +24,7 @@ def _iter_work(params): - """Yield parameter dicts. Walk: one dict. Spawn: a sequence of dicts.""" + """Yield parameter dicts. Chain: one dict. Spawn: a sequence of dicts.""" if params is None: return #end if @@ -73,8 +73,9 @@ class DynamicDecision(DevBase): status ``continue`` keep polling this step - ``completed`` finished successfully (walk: target reached) + ``completed`` finished successfully (chain: target reached) ``max_runs`` stopped without completing (cap, or propose() is None) + ``failed`` a simulation failed (chain: stop; spawn can ignore) """ def __init__(self, status, products=None): @@ -91,6 +92,11 @@ def completed(self): def max_runs(self): return self.status == 'max_runs' #end def max_runs + + @property + def failed(self): + return self.status == 'failed' + #end def failed #end class DynamicDecision @@ -99,7 +105,7 @@ class DynamicMode(DevBase): initial: First work to launch. propose(history): Next work(s) to launch, or None if there is no further work. - drive(make, wm, products): Poll until observe is not continue. + drive(sim_generator, wm, products): Implemented by chain / spawn. observe(params, products): Record a finished job and return a ``DynamicDecision``. stop(history): True when the mode has finished successfully. Default is False. """ @@ -113,6 +119,10 @@ def __init__(self, max_runs=10): self.history = [] #end def __init__ + def reset(self): + self.history = [] + #end def reset + def initial(self): self.error('initial() must be implemented in a subclass') #end def initial @@ -121,40 +131,8 @@ def propose(self, history): self.error('propose() must be implemented in a subclass') #end def propose - def drive(self, make, wm, products, poll=1): - """Return a DynamicDecision after polling until the mode is finished.""" - jobs = [] - for params in _iter_work(self.initial()): - jobs.append([params, make(params)]) - #end for - if len(jobs) == 0: - self.error('initial() produced no work to launch') - #end if - while True: - for job in list(jobs): - params, sim = job - if sim.succ: - decision = self.observe(params, _as_products(products, sim)) - jobs.remove(job) - if decision.status != 'continue': - return decision - #end if - nxt = list(_iter_work(decision.next_params)) - if len(nxt) == 0: - self.error('continue decision has no next_params') - #end if - for p in nxt: - jobs.append([p, make(p)]) - #end for - elif sim.fail: - self.error('simulation failed') - #end if - #end for - if len(jobs) == 0: - self.error('no remaining jobs') - #end if - wm.poll(poll) - #end while + def drive(self, sim_generator, wm, products, poll=1): + self.error('drive() must be implemented in a subclass') #end def drive def observe(self, params, products): diff --git a/nexus/nexus/dynamic_workflows/walk.py b/nexus/nexus/dynamic_workflows/chain.py similarity index 66% rename from nexus/nexus/dynamic_workflows/walk.py rename to nexus/nexus/dynamic_workflows/chain.py index 5701368f21..5521f43b0c 100644 --- a/nexus/nexus/dynamic_workflows/walk.py +++ b/nexus/nexus/dynamic_workflows/chain.py @@ -4,42 +4,44 @@ #====================================================================# -# walk.py # -# Walk mode: sequential parameter walk (initial / propose / stop).# +# chain.py # +# Chain mode: sequential parameter chain (initial / propose / stop).# # # # Content summary: # # SuccessiveChange(Target) # # Successive |Δ| of one or more products within atol/rtol. # -# WalkDecision(DynamicDecision) # -# Outcome of one walk observation. # -# DynamicWalk(DynamicMode) # -# Sequential walk. stop is Target.reached; observe returns # -# WalkDecision. Subclasses implement initial / propose. # +# ChainDecision(DynamicDecision) # +# Outcome of one chain observation. # +# DynamicChain(DynamicMode) # +# Sequential chain. stop is Target.reached; observe returns # +# ChainDecision. drive polls until completed / max_runs / # +# failed. Subclasses implement initial / propose. # #====================================================================# from ..developer import error -from .base import DynamicDecision, DynamicMode, Target +from .base import DynamicDecision, DynamicMode, Target, _as_products, _iter_work def _as_params(params): if not isinstance(params, dict) or len(params) == 0: error( - f'walk parameters must be a non-empty dict, received {params}', - header='DynamicWalk', + f'chain parameters must be a non-empty dict, received {params}', + header='DynamicChain', ) #end if return dict(params) #end def _as_params -class WalkDecision(DynamicDecision): - """Outcome of ``DynamicWalk.observe``. +class ChainDecision(DynamicDecision): + """Outcome of ``DynamicChain.observe`` or ``drive``. status ``continue`` propose a new parameter point ``completed`` target reached - ``max_runs`` history length reached max_runs, or propose() returned None + ``max_runs`` history length reached max_runs + ``failed`` no next work (propose returned None), or a sim failed """ def __init__(self, status, params, next_params=None, products=None): @@ -47,7 +49,7 @@ def __init__(self, status, params, next_params=None, products=None): self.params = None if params is None else dict(params) self.next_params = None if next_params is None else dict(next_params) #end def __init__ -#end class WalkDecision +#end class ChainDecision class SuccessiveChange(Target): @@ -93,7 +95,7 @@ def __init__(self, keys='energy', atol=None, rtol=None, consecutive=1): self.error('SuccessiveChange requires at least one product key') #end if if atol is None and rtol is None: - self.error('SuccessiveChange requires atol and/or rtol') + self.error('SuccessiveChange requires atol or rtol') #end if consecutive = int(consecutive) if consecutive < 1: @@ -115,7 +117,9 @@ def _value(self, record, key): #end def _value def _within(self, key, old, new): - return abs(new - old) <= self.atol[key] + self.rtol[key] * abs(old) + cond_atol = abs(new - old) <= self.atol[key] + cond_rtol = abs(new - old) <= self.rtol[key] * abs(old) + return cond_atol and cond_rtol #end def _within def reached(self, history): @@ -161,11 +165,11 @@ def _tol_map(spec, keys, default): #end def _tol_map -class DynamicWalk(DynamicMode): - """Sequential parameter walk (walk mode). +class DynamicChain(DynamicMode): + """Sequential parameter chain (chain mode). ``stop`` is ``self.target.reached(history)``. ``observe`` returns - ``WalkDecision``. + ``ChainDecision``. ``drive`` polls until the chain is not continue. """ def __init__(self, target=None, max_runs=10): @@ -179,6 +183,45 @@ def __init__(self, target=None, max_runs=10): self.target = target #end def __init__ + def drive(self, sim_generator, wm, products, poll=1): + """Poll until the chain completes, hits max_runs, or a sim fails.""" + sims = [] + for params in _iter_work(self.initial()): + sims.append([params, sim_generator(params)]) + #end for + if len(sims) == 0: + self.error('initial() produced no work to launch') + #end if + while True: + for sim_info in list(sims): + params, sim = sim_info + if sim.succ: + decision = self.observe(params, _as_products(products, sim)) + sims.remove(sim_info) + if decision.status != 'continue': + return decision + #end if + nxt = list(_iter_work(decision.next_params)) + if len(nxt) == 0: + self.error('continue decision has no next_params') + #end if + for p in nxt: + sims.append([p, sim_generator(p)]) + #end for + elif sim.fail: + return ChainDecision( + 'failed', + params = params, + ) + #end if + #end for + if len(sims) == 0: + self.error('no remaining jobs') + #end if + wm.poll(poll) + #end while + #end def drive + def stop(self, history): if self.target is None: return DynamicMode.stop(self, history) @@ -191,7 +234,7 @@ def observe(self, params, products): products = dict(products) self.history.append(dict(params=params, products=products)) if self.stop(self.history): - return WalkDecision( + return ChainDecision( 'completed', params = params, products = products, @@ -202,17 +245,22 @@ def observe(self, params, products): nxt = self.propose(self.history) #end if if nxt is None: - return WalkDecision( - 'max_runs', + if len(self.history) >= self.max_runs: + status = 'max_runs' + else: + status = 'failed' + #end if + return ChainDecision( + status, params = params, products = products, ) #end if - return WalkDecision( + return ChainDecision( 'continue', params = params, next_params = _as_params(nxt), products = products, ) #end def observe -#end class DynamicWalk +#end class DynamicChain diff --git a/nexus/nexus/dynamic_workflows/chain_error_handler.py b/nexus/nexus/dynamic_workflows/chain_error_handler.py new file mode 100644 index 0000000000..63c768bc78 --- /dev/null +++ b/nexus/nexus/dynamic_workflows/chain_error_handler.py @@ -0,0 +1,199 @@ +################################################################## +## (c) Copyright 2015- by Jaron T. Krogel ## +################################################################## + + +#====================================================================# +# chain_error_handler.py # +# Chain recovery for a failed simulation. # +# # +# Content summary: # +# ChainErrorHandler(DynamicChain) # +# Parse a failed job, propose one input patch per attempt. # +# First matching handler wins. # +# Subclasses implement parse / apply_patch / params_from_input. # +# Spawn can later apply every matching handler in parallel. # +#====================================================================# + + +import os + +from .base import DynamicMode +from .chain import DynamicChain + + +BOOKKEEPING = frozenset(('attempt',)) +CHUNK_BYTES = 64 * 1024 + + +def as_simulation(sim): + inner = getattr(sim, 'sim', None) + if inner is not None and hasattr(sim, 'dpid'): + return inner + #end if + return sim +#end def as_simulation + + +def scan_file(path, markers, extra_re=None): + """Return (matched marker strings, extra_re match or None). Chunked.""" + found = [] + extra = None + if not path or not os.path.isfile(path): + return found, extra + #end if + remaining = [] + overlap = 80 + for marker in markers: + raw = marker.encode('utf-8') if isinstance(marker, str) else marker + remaining.append((marker, raw)) + overlap = max(overlap, len(raw)) + #end for + try: + with open(path, 'rb') as handle: + prev = b'' + while True: + chunk = handle.read(CHUNK_BYTES) + if not chunk: + break + #end if + window = prev + chunk + if extra_re is not None and extra is None: + match = extra_re.search(window) + if match: + extra = match + #end if + #end if + still = [] + for marker, raw in remaining: + if raw in window: + if isinstance(marker, str): + found.append(marker) + else: + found.append(marker.decode('utf-8', 'replace')) + #end if + else: + still.append((marker, raw)) + #end if + #end for + remaining = still + if not remaining and (extra_re is None or extra is not None): + break + #end if + prev = window[-overlap:] + #end while + except (OSError, MemoryError): + return found, extra + #end try + return found, extra +#end def scan_file + + +def drop_bookkeeping(params): + out = {} + for key, value in params.items(): + if key in BOOKKEEPING or str(key).startswith('_'): + continue + #end if + out[key] = value + #end for + return out +#end def drop_bookkeeping + + +class ChainErrorHandler(DynamicChain): + """Chain recovery: parse a failure and propose one input patch. + """ + + unrecoverable = frozenset() + + def __init__(self, start=None, max_runs=3, handlers=None): + DynamicChain.__init__(self, max_runs=max_runs) + self.start = dict(start or {}) + self.handlers = tuple(self.default_handlers() if handlers is None else handlers) + #end def __init__ + + def default_handlers(self): + return () + #end def default_handlers + + def parse(self, sim): + self.error('parse() must be implemented in a subclass') + #end def parse + + def apply_patch(self, inp, patch): + self.error('apply_patch() must be implemented in a subclass') + #end def apply_patch + + def params_from_input(self, inp): + return {} + #end def params_from_input + + def recover_failed(self, sim): + """If a patch exists, archive the attempt, apply it, and resubmit. + + Returns True when the simulation should be retried. + """ + sim.log('recovering failed run'+sim.idstr(), n=3) + products = self.parse(sim) + sim.log('parsed products'+str(products), n=3) + params = self.params_from_input(sim.input) + params['attempt'] = len(self.history) + decision = self.observe(params, products) + sim.log('observed decision'+str(decision), n=3) + if decision.status != 'continue': + return False + #end if + sim.save_attempt() + sim.log('saved attempt', n=3) + self.apply_patch(sim.input, decision.next_params) + sim.log('applied patch', n=3) + sim.input.write(os.path.join(sim.locdir, sim.infile)) + sim.log('wrote input files'+sim.idstr(), n=3) + sim.reset_indicators() + sim.log('reset indicators', n=3) + return True + #end def recover_failed + + def reset(self): + DynamicMode.reset(self) + #end def reset + + def initial(self): + if not self.start: + self.error(f'{type(self).__name__}.initial needs start params') + #end if + params = dict(self.start) + params.setdefault('attempt', 0) + return params + #end def initial + + def propose(self, history): + last = history[-1] + params = dict(last['params']) + products = last.get('products') or {} + errors = list(products.get('errors') or []) + if any(tag in self.unrecoverable for tag in errors): + return None + #end if + patch = self._patch(params, products) + if patch is None: + return None + #end if + nxt = dict(params) + nxt.update(patch) + nxt['attempt'] = int(params.get('attempt', 0)) + 1 + return nxt + #end def propose + + def _patch(self, params, products): + """Chain: first matching handler.""" + for handler in self.handlers: + patch = handler(params, products) + if patch is not None: + return patch + #end if + #end for + return None + #end def _patch +#end class ChainErrorHandler diff --git a/nexus/nexus/dynamic_workflows/pwscf/__init__.py b/nexus/nexus/dynamic_workflows/pwscf/__init__.py index b003cf38a7..8e32b3c064 100644 --- a/nexus/nexus/dynamic_workflows/pwscf/__init__.py +++ b/nexus/nexus/dynamic_workflows/pwscf/__init__.py @@ -5,9 +5,11 @@ #====================================================================# # pwscf # -# PWscf-specific dynamic walks (Ecut, Kgrid). # +# PWscf-specific dynamic chains (Ecut, Kgrid). # +# PwscfErrorHandler is the error_handling=True option. # #====================================================================# from .ecut import Ecut, next_ecutwfc from .kgrid import Kgrid, kgrid_from_kspacing, next_kgrid +from .error_handler import PwscfErrorHandler, parse_pwscf_text diff --git a/nexus/nexus/dynamic_workflows/pwscf/ecut.py b/nexus/nexus/dynamic_workflows/pwscf/ecut.py index 8c884d4859..a3a3257aa2 100644 --- a/nexus/nexus/dynamic_workflows/pwscf/ecut.py +++ b/nexus/nexus/dynamic_workflows/pwscf/ecut.py @@ -5,19 +5,17 @@ #====================================================================# # ecut.py # -# Walk-mode PWscf planewave-cutoff (example 02 repeat). # +# Chain-mode PWscf planewave-cutoff (example 02 repeat). # #====================================================================# -from ..walk import DynamicWalk, SuccessiveChange +from ..chain import DynamicChain, SuccessiveChange from ...developer import error def next_ecutwfc(ecut, factor=1.5, round_ry=10): - """Increase cutoff as in the diamond energy-convergence example. - - ``int(((factor * ecut) // round_ry) * round_ry)``, then bump by - ``round_ry`` if the rounded value did not increase. + """Increase cutoff for convergence example. + Uses ecut_nxt = factor * ecut, but rounded using round_ry. """ ecut = float(ecut) factor = float(factor) @@ -48,18 +46,8 @@ def next_ecutwfc(ecut, factor=1.5, round_ry=10): #end def next_ecutwfc -class Ecut(DynamicWalk): - """Walk ``ecutwfc`` until a successive-change target is reached. - - Parameter dict: ``{'ecutwfc': int}``. Default stepping matches - ``examples/dynamic_workflows/02_energy_convergence`` (×1.5, rounded - down to 10 Ry). The poll loop and ``generate_pwscf`` stay in the - user script. - - ``target`` may be a target object (``reached(history)``), or the - name ``'consecutive'`` (default), which builds - ``SuccessiveChange('energy', atol=1e-4, ...)``. ``tol=`` is accepted - as an alias for ``atol=``. +class Ecut(DynamicChain): + """1D Ecut parameter scan until target is reached. """ def __init__( @@ -93,7 +81,7 @@ def __init__( elif not hasattr(target, 'reached'): self.error(f'invalid target: {target}') #end if - DynamicWalk.__init__(self, target=target, max_runs=max_runs) + DynamicChain.__init__(self, target=target, max_runs=max_runs) self.start = start self.end = end self.factor = float(factor) diff --git a/nexus/nexus/dynamic_workflows/pwscf/error_handler.py b/nexus/nexus/dynamic_workflows/pwscf/error_handler.py new file mode 100644 index 0000000000..ee3294e889 --- /dev/null +++ b/nexus/nexus/dynamic_workflows/pwscf/error_handler.py @@ -0,0 +1,352 @@ +################################################################## +## (c) Copyright 2015- by Jaron T. Krogel ## +################################################################## + + +#====================================================================# +# error_handler.py # +# PWscf parse + recovery # +# # +# generate_pwscf(..., error_handling=True) attaches this chain. # +# Spawn can later fan the same handlers out and pick first # +# success. # +#====================================================================# + + +import os +import re + +from ...pwscf_input import PwscfInput +from ..chain_error_handler import ChainErrorHandler, as_simulation, drop_bookkeeping, scan_file + + +# Tags mainly obtained from AiiDA PwBaseWorkChain +PWSCF_ERROR_MARKERS = ( + ('Maximum CPU time exceeded', 'walltime'), + ('Program stopped by user request', 'user_stop'), + ('convergence NOT achieved after', 'electronic_convergence'), + ('history already reset at previous step: stopping', 'bfgs_history'), + ('charge is wrong', 'wrong_charge'), + ('not orthogonal operation', 'symmetry_not_orthogonal'), + ('Cholesky failed in invchol', 'cholesky'), + ('problems computing cholesky', 'cholesky'), + ('too many bands are not converged', 'unconverged_bands'), + ('S matrix not positive definite', 's_matrix_not_positive_definite'), + ('zhegvd failed', 'zhegvd_failed'), + ('[Q, R] = qr(X, 0) failed', 'qr_failed'), + ('eigenvectors failed to converge', 'unconverged_eigenvectors'), + ('dexx is negative', 'negative_dexx'), + ('Error in routine broyden', 'broyden_failure'), + ( + 'Not enough space allocated for radial FFT: try restarting with a larger cell_factor', + 'not_enough_fft_space', + ), + ('some nodes have no k-points', 'npools_too_high'), + ('probably because G_par is NOT a reciprocal lattice vector', 'gpar_error'), + ) + +PWSCF_ABORT_MARKERS = ( + 'slurmstepd: error', + 'CANCELLED AT', + 'Job step aborted', + 'srun: error', + 'srun: Job step aborted', + 'DUE TO TIME LIMIT', + 'oom-kill', + 'Out Of Memory', + 'Caught signal', + 'Segmentation fault', + 'Fatal Error', + 'mpirun: command not found', + 'mpiexec: command not found', + ) + +UNRECOVERABLE_TAGS = frozenset(( + 'wrong_charge', + 'symmetry_not_orthogonal', + 'gpar_error', + 'negative_dexx', + 'npools_too_high', + 'command_not_found', + )) + +DIAGO_TAGS = frozenset(( + 'cholesky', + 'unconverged_bands', + 's_matrix_not_positive_definite', + 'zhegvd_failed', + 'qr_failed', + 'unconverged_eigenvectors', + 'broyden_failure', + )) + +QE_MIXING_DEFAULT = 0.7 +QE_CELL_FACTOR_DEFAULT = 2.0 +DIAGO_ORDER = ('david', 'cg') + +NBND_RE = re.compile(br'number of Kohn-Sham states\s*=\s*(\d+)') + + +def parse_pwscf_text(text): + """Return a unique list of error tags found in a stdout string.""" + tags = [] + seen = set() + if not text: + return tags + #end if + for marker, tag in PWSCF_ERROR_MARKERS: + if marker in text and tag not in seen: + tags.append(tag) + seen.add(tag) + #end if + #end for + return tags +#end def parse_pwscf_text + + +class PwscfErrorHandler(ChainErrorHandler): + """PWscf recovery for ``generate_pwscf(..., error_handling=True)``. + + Chain: one handler patch per attempt. Pass this object (or an + int for ``max_runs``) as ``error_handling``. + """ + + unrecoverable = UNRECOVERABLE_TAGS + + def __init__( + self, + start=None, + max_runs=3, + handlers=None, + mixing_factor=0.7, + mixing_min=0.05, + mixing_default=QE_MIXING_DEFAULT, + diago_order=DIAGO_ORDER, + nbnd_factor=0.05, + nbnd_minimum=4, + ): + ChainErrorHandler.__init__(self, start=start, max_runs=max_runs, handlers=handlers) + self.mixing_factor = float(mixing_factor) + self.mixing_min = float(mixing_min) + self.mixing_default = float(mixing_default) + self.diago_order = tuple(diago_order) + self.nbnd_factor = float(nbnd_factor) + self.nbnd_minimum = int(nbnd_minimum) + self.tried_diago = set() + #end def __init__ + + def default_handlers(self): + return ( + self.handle_walltime, + self.handle_electronic_convergence, + self.handle_diagonalization, + self.handle_nbnd, + self.handle_fft_space, + self.handle_retry, + ) + #end def default_handlers + + def reset(self): + ChainErrorHandler.reset(self) + self.tried_diago = set() + #end def reset + + def parse(self, sim): + '''Overlapping with parsing in PWscf.check_sim_status. Can be combined in the future. + In short term, check_sim_status can reverse read last N lines from stdout to reduce load. + It currently parses for errors from stdout and stderr, plus number of bands. + Number of bands is useful for error handling where nbnd is increased for convergence.''' + sim = as_simulation(sim) + locdir = getattr(sim, 'locdir', None) + outfile = getattr(sim, 'outfile', None) + errfile = getattr(sim, 'errfile', None) + out_path = None + err_path = None + if locdir and outfile: + out_path = os.path.join(locdir, outfile) + #end if + if locdir and errfile: + err_path = os.path.join(locdir, errfile) + #end if + + markers = [m for m, _ in PWSCF_ERROR_MARKERS] + found, nbnd_match = scan_file(out_path, markers, extra_re=NBND_RE) + tags = [] + seen = set() + marker_to_tag = {m: t for m, t in PWSCF_ERROR_MARKERS} + for marker in found: + tag = marker_to_tag.get(marker) + if tag and tag not in seen: + tags.append(tag) + seen.add(tag) + #end if + #end for + abort_found, _ = scan_file(err_path, PWSCF_ABORT_MARKERS) + for marker in abort_found: + tag = 'command_not_found' if 'command not found' in marker else 'abort' + if tag not in seen: + tags.append(tag) + seen.add(tag) + #end if + #end for + + job_done = False + if out_path and os.path.isfile(out_path): + done_found, _ = scan_file(out_path, ('JOB DONE',)) + job_done = bool(done_found) + #end if + + nbnd = None + if nbnd_match is not None: + nbnd = int(nbnd_match.group(1)) + #end if + return dict(errors=tags, job_done=job_done, nbnd=nbnd) + #end def parse + + def params_from_input(self, inp): + params = {} + electrons = getattr(inp, 'electrons', None) + if electrons is not None: + if 'mixing_beta' in electrons: + params['mixing_beta'] = electrons.mixing_beta + #end if + if 'diagonalization' in electrons: + params['diagonalization'] = electrons.diagonalization + #end if + if 'electron_maxstep' in electrons: + params['electron_maxstep'] = electrons.electron_maxstep + #end if + #end if + control = getattr(inp, 'control', None) + if control is not None and 'restart_mode' in control: + params['restart_mode'] = control.restart_mode + #end if + system = getattr(inp, 'system', None) + if system is not None and 'nbnd' in system: + params['nbnd'] = system.nbnd + #end if + cell = getattr(inp, 'cell', None) + if cell is not None and 'cell_factor' in cell: + params['cell_factor'] = cell.cell_factor + #end if + return params + #end def params_from_input + + def apply_patch(self, inp, patch): + '''Apply the patch to the input file. + Ideally these are tightly organized with error markers''' + patch = drop_bookkeeping(patch) + if not patch: + return + #end if + if 'mixing_beta' in patch or 'diagonalization' in patch or 'electron_maxstep' in patch: + electrons = getattr(inp, 'electrons', None) + if electrons is None: + electrons = PwscfInput.section_types['electrons']() + inp.electrons = electrons + #end if + if 'mixing_beta' in patch: + electrons.mixing_beta = patch['mixing_beta'] + #end if + if 'diagonalization' in patch: + electrons.diagonalization = patch['diagonalization'] + #end if + if 'electron_maxstep' in patch: + electrons.electron_maxstep = patch['electron_maxstep'] + #end if + #end if + if 'restart_mode' in patch: + inp.control.restart_mode = patch['restart_mode'] + #end if + if 'nbnd' in patch: + inp.system.nbnd = patch['nbnd'] + #end if + if 'cell_factor' in patch: + cell = getattr(inp, 'cell', None) + if cell is None: + cell = PwscfInput.section_types['cell']() + inp.cell = cell + #end if + cell.cell_factor = patch['cell_factor'] + #end if + #end def apply_patch + + def handle_walltime(self, params, products): + errors = products.get('errors') or [] + if 'walltime' not in errors and 'user_stop' not in errors: + return None + #end if + return dict(restart_mode='restart') + #end def handle_walltime + + def handle_electronic_convergence(self, params, products): + errors = products.get('errors') or [] + if 'electronic_convergence' not in errors: + return None + #end if + beta = float(params.get('mixing_beta', self.mixing_default)) + nxt = max(self.mixing_min, beta * self.mixing_factor) + patch = dict(restart_mode='from_scratch') + if nxt < beta: + patch['mixing_beta'] = nxt + #end if + maxstep = params.get('electron_maxstep') + if maxstep is not None and int(maxstep) < 80: + patch['electron_maxstep'] = 80 + #end if + if 'mixing_beta' not in patch and 'electron_maxstep' not in patch: + return None + #end if + return patch + #end def handle_electronic_convergence + + def handle_diagonalization(self, params, products): + errors = products.get('errors') or [] + if not any(tag in DIAGO_TAGS for tag in errors): + return None + #end if + current = params.get('diagonalization', DIAGO_ORDER[0]) + self.tried_diago.add(current) + for algo in self.diago_order: + if algo not in self.tried_diago: + self.tried_diago.add(algo) + return dict(diagonalization=algo, restart_mode='from_scratch') + #end if + #end for + return None + #end def handle_diagonalization + + def handle_nbnd(self, params, products): + errors = products.get('errors') or [] + if 'unconverged_bands' not in errors: + return None + #end if + nbnd = params.get('nbnd') + if nbnd is None: + nbnd = products.get('nbnd') + #end if + if nbnd is None: + return None + #end if + nbnd = int(nbnd) + bump = max(int(nbnd * self.nbnd_factor), self.nbnd_minimum) + return dict(nbnd=nbnd + bump, restart_mode='from_scratch') + #end def handle_nbnd + + def handle_fft_space(self, params, products): + errors = products.get('errors') or [] + if 'not_enough_fft_space' not in errors: + return None + #end if + factor = float(params.get('cell_factor', QE_CELL_FACTOR_DEFAULT)) + return dict(cell_factor=2.0 * factor, restart_mode='from_scratch') + #end def handle_fft_space + + def handle_retry(self, params, products): + errors = products.get('errors') or [] + if not errors or 'abort' in errors: + return dict() + #end if + return None + #end def handle_retry +#end class PwscfErrorHandler diff --git a/nexus/nexus/dynamic_workflows/pwscf/kgrid.py b/nexus/nexus/dynamic_workflows/pwscf/kgrid.py index 03668bf9bd..beff5a5562 100644 --- a/nexus/nexus/dynamic_workflows/pwscf/kgrid.py +++ b/nexus/nexus/dynamic_workflows/pwscf/kgrid.py @@ -5,14 +5,14 @@ #====================================================================# # kgrid.py # -# Walk-mode PWscf k-point mesh (example 02 policy). # +# Chain-mode PWscf k-point mesh (example 02). # # # # increment add to the integer Monkhorst-Pack grid each step. # # spacing initial mesh from Structure.kgrid_from_kspacing. # #====================================================================# -from ..walk import DynamicWalk, SuccessiveChange +from ..chain import DynamicChain, SuccessiveChange from ...developer import error @@ -94,8 +94,8 @@ def next_kgrid(kgrid, increment=1): #end def next_kgrid -class Kgrid(DynamicWalk): - """Walk ``kgrid`` until a successive-change target is reached. +class Kgrid(DynamicChain): + """Chain ``kgrid`` until a successive-change target is reached. Parameter dict: ``{'kgrid': (nx, ny, nz)}``. Default stepping matches ``examples/dynamic_workflows/02_energy_convergence`` @@ -147,7 +147,7 @@ def __init__( elif not hasattr(target, 'reached'): self.error(f'invalid target: {target}') #end if - DynamicWalk.__init__(self, target=target, max_runs=max_runs) + DynamicChain.__init__(self, target=target, max_runs=max_runs) self.start = start self.end = end self.increment = increment diff --git a/nexus/nexus/examples/dynamic_workflows/03_walk/ecut_walk.py b/nexus/nexus/examples/dynamic_workflows/03_chain/ecut_chain.py similarity index 71% rename from nexus/nexus/examples/dynamic_workflows/03_walk/ecut_walk.py rename to nexus/nexus/examples/dynamic_workflows/03_chain/ecut_chain.py index 7300474a05..c173bd51f5 100644 --- a/nexus/nexus/examples/dynamic_workflows/03_walk/ecut_walk.py +++ b/nexus/nexus/examples/dynamic_workflows/03_chain/ecut_chain.py @@ -1,13 +1,7 @@ #! /usr/bin/env python3 ''' -Planewave cutoff walk using ``Ecut`` (walk mode). - -Same diamond / BFD setup and stepping as the ecut half of -``02_energy_convergence`` (start 50 Ry, ×1.5 rounded to 10 Ry, -successive |ΔE| ≤ 1e-4 Ry). ``drive`` owns the poll loop; -``Ecut`` only decides the next ``ecutwfc`` and when the target -is reached. ``generate_pwscf`` stays in this script. +Same diamond / BFD setup and stepping as the ecut part of example 02. ''' import sys @@ -32,7 +26,7 @@ C = 4, ) -walk = Ecut(start=50, tol=1e-4, max_runs=10) +chain = Ecut(start=50, tol=1e-4, max_runs=10) def make_scf(params): @@ -52,7 +46,7 @@ def make_scf(params): #end def make_scf wm = workflow_manager() -decision = walk.drive(make_scf, wm, products='energy') +decision = chain.drive(make_scf, wm, products='energy') print() print('status :', decision.status) print('ecut :', decision.params['ecutwfc']) diff --git a/nexus/nexus/examples/dynamic_workflows/03_walk/kgrid_walk.py b/nexus/nexus/examples/dynamic_workflows/03_chain/kgrid_chain.py similarity index 62% rename from nexus/nexus/examples/dynamic_workflows/03_walk/kgrid_walk.py rename to nexus/nexus/examples/dynamic_workflows/03_chain/kgrid_chain.py index 17317f8d8b..1a0217be46 100644 --- a/nexus/nexus/examples/dynamic_workflows/03_walk/kgrid_walk.py +++ b/nexus/nexus/examples/dynamic_workflows/03_chain/kgrid_chain.py @@ -1,19 +1,7 @@ #! /usr/bin/env python3 ''' -k-point mesh walk using ``Kgrid`` (walk mode). - -Same diamond / BFD setup and stepping as the k-grid half of -``02_energy_convergence`` (start 1x1x1, +1 per axis, successive -|ΔE| ≤ 1e-3 Ry). ``drive`` owns the poll loop; ``Kgrid`` only -decides the next ``kgrid`` and when the target is reached. -``generate_pwscf`` stays in this script. - -Set ``ecutwfc`` to the cutoff from ``ecut_walk.py`` (example 02 -uses 330 Ry). 50 Ry is used here so this script can run on its -own. For a spacing-based start instead of ``start=1``:: - - walk = Kgrid(spacing=0.5, structure=system, increment=1, tol=1e-3) +Same diamond / BFD setup and stepping as the k-grid part of example 02. ''' import sys @@ -39,7 +27,7 @@ ) ecutwfc = 50 -walk = Kgrid(start=1, increment=1, tol=1e-3, max_runs=10) +chain = Kgrid(start=1, increment=1, tol=1e-3, max_runs=10) def make_scf(params): @@ -60,7 +48,7 @@ def make_scf(params): #end def make_scf wm = workflow_manager() -decision = walk.drive(make_scf, wm, products='energy') +decision = chain.drive(make_scf, wm, products='energy') print() print('status :', decision.status) print('kgrid :', decision.params['kgrid']) diff --git a/nexus/nexus/examples/dynamic_workflows/04_pwscf_recover/recover_chain.py b/nexus/nexus/examples/dynamic_workflows/04_pwscf_recover/recover_chain.py new file mode 100644 index 0000000000..c3d6002eff --- /dev/null +++ b/nexus/nexus/examples/dynamic_workflows/04_pwscf_recover/recover_chain.py @@ -0,0 +1,58 @@ +#! /usr/bin/env python3 + +''' +Planewave cutoff chain with recoverable PWscf jobs. +''' + +import sys +from nexus import settings, job, workflow_manager +from nexus import generate_physical_system +from nexus import generate_pwscf +from nexus.dynamic_workflows.pwscf import Ecut +from nexus.dynamic_workflows.pwscf.error_handler import PwscfErrorHandler + +settings( + results = '', + pseudo_dir = '../../qmcpack/pseudopotentials', + machine = 'ws8', + dynamic = True, + verbose = True, + ) + + +system = generate_physical_system( + structure = 'diamond', + cell = 'prim', + kgrid = (1, 1, 1), + C = 4, + ) + +chain = Ecut(start=50, tol=1e-4, max_runs=10) + + +def make_scf(params): + ecut = params['ecutwfc'] + return generate_pwscf( + identifier = 'scf', + path = f'ecut_{ecut}', + job = job(cores=4), + system = system, + pseudos = ['C.BFD.upf'], + input_type = 'generic', + ecutwfc = ecut, + kgrid = (1, 1, 1), + electron_maxstep = 2, # Added to simulate failure + mixing_beta = 0.7, + dynamic_id = f'ecut{ecut}', + requires = 'none', + error_handling = True, + # error_handling = 3, + # error_handling = PwscfErrorHandler(max_runs=10, mixing_factor=0.3) + ) +#end def make_scf + +wm = workflow_manager() +decision = chain.drive(make_scf, wm, products='energy') +print('status :', decision.status) +print('ecut :', decision.params['ecutwfc']) +print('products:', decision.products) diff --git a/nexus/nexus/pwscf.py b/nexus/nexus/pwscf.py index e336120348..e71966614d 100644 --- a/nexus/nexus/pwscf.py +++ b/nexus/nexus/pwscf.py @@ -120,7 +120,9 @@ def restore_default_settings(): def __init__(self,**sim_args): group_atoms = sim_args.pop('group_atoms',False) sync_from_scf = sim_args.pop('sync_from_scf',True) + error_handling = sim_args.pop('error_handling',False) Simulation.__init__(self,**sim_args) + self.error_handler = self._make_error_handler(error_handling) self.sync_from_scf = False calc = None cont = self.input.control @@ -135,6 +137,24 @@ def __init__(self,**sim_args): #end if #end def __init__ + def _make_error_handler(self, error_handling): + from .dynamic_workflows.pwscf.error_handler import PwscfErrorHandler + from .dynamic_workflows.chain_error_handler import ChainErrorHandler + if error_handling is True: + return PwscfErrorHandler() + #end if + if isinstance(error_handling, int): + return PwscfErrorHandler(max_runs=int(error_handling)) + #end if + if isinstance(error_handling, ChainErrorHandler): + return error_handling + #end if + tname = type(error_handling).__name__ + self.error( + f'error_handling must be True, False, an int (max_runs), or a ChainErrorHandler, received {tname}' + ) + #end def _make_error_handler + def write_prep(self): #make sure the output directory exists @@ -342,22 +362,52 @@ def check_sim_status(self): time_exceeded = 'Maximum CPU time exceeded' in output user_stop = 'Program stopped by user request' in output run_finished = 'JOB DONE' in output - restartable = not_converged or time_exceeded or user_stop - restart = run_finished and self.restartable and restartable - if restart: + error_in_routine = 'Error in routine' in output + failed = ( + not run_finished + or not_converged + or time_exceeded + or user_stop + or error_in_routine + ) + + if not failed: + self.finished = True + self.failed = False + return + #end if + + if self.restartable and run_finished and (not_converged or time_exceeded or user_stop): self.save_attempt() self.input.control.restart_mode = 'restart' self.reset_indicators() - else: - error_in_routine = 'Error in routine' in output - failed = not_converged or time_exceeded or user_stop - failed |= error_in_routine + return + #end if + + # Error handler works last + if self.error_handler is not None: + if self._recover_failed_run(): + return + #end if self.finished = run_finished - self.failed = failed + self.failed = True + return #end if + + self.finished = run_finished + self.failed = True #end def check_sim_status + def _recover_failed_run(self): + handler = getattr(self, 'error_handler', None) + if handler is None: + return False + #end if + return handler.recover_failed(self) + #end def _recover_failed_run + + def get_output_files(self): output_files = [] return output_files @@ -502,7 +552,7 @@ def generate_pwscf(**kwargs): kwargs['files'] = list(kwargs.get('files',[])) + list(pseudos.values()) #end if - sim_args,inp_args = Pwscf.separate_inputs(kwargs) + sim_args,inp_args = Pwscf.separate_inputs(kwargs, sim_kw={'error_handling'}) if 'input' not in sim_args: input_type = inp_args.pop('input_type','generic') diff --git a/nexus/nexus/simulation.py b/nexus/nexus/simulation.py index 76519d04b7..181f6f53af 100644 --- a/nexus/nexus/simulation.py +++ b/nexus/nexus/simulation.py @@ -1095,6 +1095,9 @@ def submit(self): if (self.job.batch_mode or not nexus_core.monitor) and not nexus_core.generate_only: self.save_image() #end if + if self.job is not None and self.job.finished: + self.check_status() + #end if elif not self.finished: self.check_status() #end if