"""Command-line interface for Atom-HiFi.
Two subcommands:
* ``atom-hifi init [path] [--force]`` — write a commented ``config.yaml`` starter.
* ``atom-hifi run <config.yaml>`` — execute the full L/R workflow.
"""
from __future__ import annotations
import argparse
import logging
import os
from pathlib import Path
import shutil
import sys
_LEVELS = {
'debug': logging.DEBUG,
'info': logging.INFO,
'warning': logging.WARNING,
'error': logging.ERROR,
}
def _configure_logging(level: int) -> None:
root = logging.getLogger('atom_hifi')
for previous in root.handlers[:]:
if getattr(previous, '_atom_hifi_cli', False):
root.removeHandler(previous)
previous.close()
handler = logging.StreamHandler(sys.stdout)
handler._atom_hifi_cli = True
handler.setLevel(level)
handler.setFormatter(logging.Formatter('%(message)s'))
root.setLevel(level)
root.addHandler(handler)
def _configure_runtime_threads(n_threads: int) -> None:
"""Set BLAS/OpenMP thread counts before NumPy/BLAS-heavy modules are imported.
Single-process stages (e.g. the microstate KMeans bisection) have no
subprocess pool of their own and only get parallelism from these thread
counts, so they're sized to the run's own worker counts rather than a
fixed guess — the workflow's parallel stages run one at a time, so there
is no oversubscription risk in letting a single-process stage use as many
threads as the process-pool stages are configured for.
"""
n_threads = str(int(n_threads))
os.environ.setdefault('OPENBLAS_NUM_THREADS', n_threads)
os.environ.setdefault('OMP_NUM_THREADS', n_threads)
os.environ.setdefault('MKL_NUM_THREADS', n_threads)
def _cmd_init(args: argparse.Namespace) -> int:
from atom_hifi.config import default_config_yaml
path = args.path
if os.path.exists(path) and not args.force:
print(f"error: {path} already exists (use --force to overwrite)",
file=sys.stderr)
return 1
with open(path, 'w') as fh:
fh.write(default_config_yaml())
print(f"Wrote {path}. Edit it, then run: atom-hifi run {path}")
return 0
def _stage_config(config_path: str, output_dir: str) -> Path:
"""Copy the input YAML to the results directory before workflow work starts."""
source = Path(config_path).resolve()
destination = Path(output_dir).resolve() / 'config.yaml'
destination.parent.mkdir(parents=True, exist_ok=True)
if source != destination:
shutil.copy2(source, destination)
return destination
def _cmd_run(args: argparse.Namespace) -> int:
# config.py and PyYAML are scientific-stack-free. Read only the two worker
# counts first so BLAS is capped before load_config validates cluster.backend
# and imports NumPy/sklearn/decaf. Full parsing/validation remains below.
import yaml
from atom_hifi.config import _load_runtime_defaults, load_config
defaults = _load_runtime_defaults()
with open(args.config) as fh:
raw_config = yaml.safe_load(fh) or {}
if not isinstance(raw_config, dict):
raw_config = {}
raw_compute = raw_config.get('compute')
if not isinstance(raw_compute, dict):
raw_compute = {}
n_jobs = int(raw_config.get('n_jobs', defaults['n_jobs']))
n_jobs_desc = int(raw_compute.get(
'n_jobs_desc', defaults['compute']['n_jobs_desc']))
_configure_runtime_threads(max(n_jobs, n_jobs_desc))
config = load_config(args.config)
_stage_config(args.config, config['paths']['output_dir'])
if getattr(args, 'f_tol', None) is not None:
if args.f_tol <= 0:
raise ValueError(
f"eps_noise.force.f_tol must be positive, got {args.f_tol!r}.")
config.setdefault('eps_noise', {}).setdefault('force', {})['f_tol'] = args.f_tol
# --f-tol materialises an explicit force block, so an r<NUM> reference
# resolves against ε_F (mirrors the raw-config _force_active rule).
config.setdefault('microstate', {})['_force_active'] = True
import matplotlib
matplotlib.use('Agg') # headless; must precede any pyplot import
# Imported lazily so `atom-hifi init` stays fast and dependency-light.
from atom_hifi.runner import run
# CLI flags override the config's log_level; default falls back to INFO.
if args.verbose:
level = logging.DEBUG
elif args.quiet:
level = logging.WARNING
else:
level = _LEVELS.get(str(config.get('log_level', 'info')).lower(), logging.INFO)
_configure_logging(level)
run(config)
return 0
[docs]
def build_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(
prog='atom-hifi',
description='Atom-HiFi: high-fidelity representative-set selection.')
sub = parser.add_subparsers(dest='command', required=True)
p_init = sub.add_parser('init', help='write a commented config.yaml starter')
p_init.add_argument('path', nargs='?', default='config.yaml',
help='output path (default: config.yaml)')
p_init.add_argument('--force', action='store_true',
help='overwrite an existing file')
p_init.set_defaults(func=_cmd_init)
p_run = sub.add_parser('run', help='run the full L/R workflow from a config file')
p_run.add_argument('config', help='path to config.yaml')
p_run.add_argument('--verbose', '-v', action='store_true',
help='show debug-level output (overrides config log_level)')
p_run.add_argument('--quiet', '-q', action='store_true',
help='suppress info output, warnings only (overrides config log_level)')
p_run.add_argument('--f-tol', type=float, default=None,
help='override eps_noise.force.f_tol for f_tol or '
'force-resolution r<NUM> references')
p_run.set_defaults(func=_cmd_run)
return parser
[docs]
def main(argv=None) -> int:
parser = build_parser()
args = parser.parse_args(argv)
return args.func(args)
if __name__ == '__main__':
sys.exit(main())