assystant 0.1.0__py3-none-any.whl

assystant/__init__.py

@@ -0,0 +1,38 @@
+"""assystant -- run an ASSYST campaign end to end from one YAML file.
+
+``gen`` generates PyXtal structures, ``run`` submits one SLURM job per
+structure, ``status`` reports progress, ``collect`` concatenates the results
+into a single training dataset.
+"""
+
+from .config import (
+    CampaignConfig,
+    ConfigError,
+    StaleConfigError,
+    system_name,
+)
+from .generate import (
+    enumerate_compositions,
+    parse_struct_dir,
+    poscar_hash,
+    struct_dir_name,
+)
+from .incar import get_incar_magmom_string, stack_element_string
+from .submit import RESULT_FILENAME
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "CampaignConfig",
+    "ConfigError",
+    "StaleConfigError",
+    "system_name",
+    "enumerate_compositions",
+    "parse_struct_dir",
+    "poscar_hash",
+    "struct_dir_name",
+    "get_incar_magmom_string",
+    "stack_element_string",
+    "RESULT_FILENAME",
+    "__version__",
+]

assystant/cli.py

@@ -0,0 +1,194 @@
+"""``assystant`` command line: ``gen`` | ``run`` | ``status`` | ``collect``."""
+
+from __future__ import annotations
+
+import argparse
+import sys
+from pathlib import Path
+
+import yaml
+
+from .config import CampaignConfig, ConfigError, StaleConfigError
+
+__all__ = ["build_parser", "main", "OVERRIDE_FLAGS"]
+
+#: CLI flag -> dotted config path.  Any scalar in the YAML is reachable this
+#: way; ``--set a.b=c`` covers the ones without a dedicated flag.
+OVERRIDE_FLAGS = {
+    "root": "root",
+    "ntasks": "slurm.ntasks",
+    "time": "slurm.time",
+    "partition": "slurm.partition",
+    "env": "slurm.env",
+    "vasp_binary": "vasp.binary",
+    "ionic_steps": "assyst.ionic_steps",
+    "n_rattle": "assyst.n_rattle",
+    "n_stretch": "assyst.n_stretch",
+    "min_dist_backend": "filters.min_dist_backend",
+    "seed": "system.seed",
+    "n_jobs": "system.n_jobs",
+}
+
+
+def _add_override_flags(parser: argparse.ArgumentParser) -> None:
+    parser.add_argument("--root", help="Override `root` (campaign directory)")
+    parser.add_argument("--ntasks", type=int, help="Override `slurm.ntasks`")
+    parser.add_argument("--time", help="Override `slurm.time` (HH:MM:SS)")
+    parser.add_argument("--partition", help="Override `slurm.partition`")
+    parser.add_argument("--env", help="Override `slurm.env` (conda prefix for the runner)")
+    parser.add_argument("--vasp-binary", dest="vasp_binary", help="Override `vasp.binary`")
+    parser.add_argument("--ionic-steps", dest="ionic_steps", type=int,
+                        help="Override `assyst.ionic_steps`")
+    parser.add_argument("--n-rattle", dest="n_rattle", type=int, help="Override `assyst.n_rattle`")
+    parser.add_argument("--n-stretch", dest="n_stretch", type=int,
+                        help="Override `assyst.n_stretch`")
+    parser.add_argument("--min-dist-backend", dest="min_dist_backend",
+                        choices=["neighbor_list", "mic"], help="Override `filters.min_dist_backend`")
+    parser.add_argument("--seed", type=int, help="Override `system.seed`")
+    parser.add_argument("--n-jobs", dest="n_jobs", type=int, help="Override `system.n_jobs`")
+    parser.add_argument(
+        "--set", dest="set_", action="append", default=[], metavar="PATH=VALUE",
+        help="Override any config scalar by dotted path, e.g. --set assyst.rattle_strain=0.02",
+    )
+
+
+def overrides_from_args(args) -> dict:
+    """Collect dotted-path overrides from parsed CLI arguments.
+
+    Only flags the user actually passed appear, so precedence is
+    CLI flag > YAML > dataclass default.
+    """
+    overrides = {}
+    for dest, dotted in OVERRIDE_FLAGS.items():
+        value = getattr(args, dest, None)
+        if value is not None:
+            overrides[dotted] = value
+    for item in getattr(args, "set_", []) or []:
+        if "=" not in item:
+            raise ConfigError(f"--set expects PATH=VALUE, got {item!r}")
+        path, _, raw = item.partition("=")
+        overrides[path.strip()] = yaml.safe_load(raw)
+    return overrides
+
+
+def load_config(args) -> CampaignConfig:
+    return CampaignConfig.from_yaml(args.config).resolve(overrides_from_args(args))
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        prog="assystant",
+        description=(
+            "Run an ASSYST campaign end to end: PyXtal structure generation -> "
+            "per-structure DFT via SLURM -> collated training dataset."
+        ),
+    )
+    sub = parser.add_subparsers(dest="command", required=True)
+
+    gen = sub.add_parser("gen", help="PyXtal spacegroup x stoichiometry scan -> POSCARs + manifest")
+    gen.add_argument("config", help="Campaign YAML")
+    gen.add_argument("--no-progress", action="store_true", help="Suppress the progress bar")
+    _add_override_flags(gen)
+
+    run = sub.add_parser("run", help="Render and submit one SLURM job per structure")
+    run.add_argument("config", help="Campaign YAML")
+    run.add_argument("--limit", type=int, help="Process at most N structures")
+    run.add_argument("--dry-run", action="store_true",
+                     help="Render sbatch scripts without submitting or querying the scheduler")
+    run.add_argument(
+        "--from-scratch", action="store_true",
+        help="Accept a changed config and DELETE each structure's run_dir in full, "
+             "including executorlib_cache",
+    )
+    _add_override_flags(run)
+
+    status = sub.add_parser("status", help="Per-system progress table")
+    status.add_argument("config", help="Campaign YAML")
+    status.add_argument("--csv", help="Append the table to a job_progress-shaped CSV")
+    status.add_argument("--watch", nargs="?", type=int, const=60, default=None,
+                        metavar="SECONDS", help="Refresh every SECONDS (default 60)")
+    _add_override_flags(status)
+
+    collect = sub.add_parser("collect", help="Concatenate every df_ASSYST_jobs.pkl into one dataset")
+    collect.add_argument("config", help="Campaign YAML")
+    collect.add_argument("-o", "--output", required=True, help="Output .pkl (or .csv)")
+    collect.add_argument("--all-systems", action="store_true",
+                         help="Scan every system directory under root, not only those in the config")
+    _add_override_flags(collect)
+
+    return parser
+
+
+def _cmd_gen(args) -> int:
+    from .generate import run_generation
+
+    cfg = load_config(args)
+    manifests = run_generation(cfg, progress=not args.no_progress)
+    for name, manifest in manifests.items():
+        print(f"{name}: {len(manifest)} structures -> {Path(cfg.root) / name}")
+    return 0
+
+
+def _cmd_run(args) -> int:
+    from .submit import submit_campaign
+
+    cfg = load_config(args)
+    stats = submit_campaign(
+        cfg, dry_run=args.dry_run, limit=args.limit, from_scratch=args.from_scratch
+    )
+    verb = "rendered" if args.dry_run else "submitted"
+    message = (
+        f"{verb} {len(stats['rendered']) if args.dry_run else stats['submitted']}, "
+        f"skipped {stats['skipped_queued']} queued / {stats['skipped_done']} finished"
+    )
+    if args.dry_run and args.from_scratch:
+        # A dry run never deletes: report what --from-scratch would wipe
+        # instead of silently omitting the most destructive part of what the
+        # real run would do.
+        message += (
+            f", would delete {stats['would_delete']} existing run_dir(s) "
+            "(not deleted -- dry run)"
+        )
+    print(message)
+    return 0
+
+
+def _cmd_status(args) -> int:
+    from .status import append_csv, campaign_status, format_table, watch
+
+    cfg = load_config(args)
+    if args.watch:
+        watch(cfg, interval=args.watch, csv_path=args.csv)
+        return 0
+    rows = campaign_status(cfg)
+    print(format_table(rows))
+    if args.csv:
+        append_csv(rows, args.csv)
+        print(f"appended to {args.csv}")
+    return 0
+
+
+def _cmd_collect(args) -> int:
+    from .collect import collect_campaign
+
+    cfg = load_config(args)
+    systems = None if args.all_systems else cfg.system.systems
+    frame = collect_campaign(cfg.root, systems, args.output)
+    print(f"collected {len(frame)} rows -> {args.output}")
+    return 0
+
+
+_COMMANDS = {"gen": _cmd_gen, "run": _cmd_run, "status": _cmd_status, "collect": _cmd_collect}
+
+
+def main(argv=None) -> int:
+    args = build_parser().parse_args(argv)
+    try:
+        return _COMMANDS[args.command](args)
+    except (ConfigError, StaleConfigError) as exc:
+        print(f"error: {exc}", file=sys.stderr)
+        return 2
+
+
+if __name__ == "__main__":  # pragma: no cover
+    sys.exit(main())

assystant/collect.py

@@ -0,0 +1,182 @@
+"""Walk a campaign root and concatenate every ``df_ASSYST_jobs.pkl`` into one dataset.
+
+Two on-disk layouts are read:
+
+``struct_<idx:05d>_<sha8>/``
+    written by this package; the hash column of ``manifest.csv`` ties each
+    result back to the structure that generated it.
+
+``struct_pyxtal_<idx>/``
+    the pre-2026 layout.  Existing campaign data is readable without migration,
+    which is why ``collect`` matches on layout rather than assuming one.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+
+import pandas as pd
+
+from .config import CampaignConfig, system_name
+from .generate import MANIFEST_DTYPES, MANIFEST_NAME, parse_struct_dir
+from .submit import RESULT_FILENAME
+
+__all__ = ["ResultDir", "iter_result_dirs", "collect_campaign", "load_manifests"]
+
+
+@dataclass(frozen=True)
+class ResultDir:
+    path: Path
+    system: str
+    index: int
+    structure_hash: str | None
+    layout: str
+
+
+def _scan_system_dir(sys_dir: Path, system: str):
+    if not sys_dir.is_dir():
+        return
+    for entry in sorted(sys_dir.iterdir()):
+        if not entry.is_dir():
+            continue
+        parsed = parse_struct_dir(entry.name)
+        if parsed is None:
+            continue
+        if not (entry / RESULT_FILENAME).is_file():
+            continue
+        index, sha8, layout = parsed
+        yield ResultDir(entry, system, index, sha8, layout)
+
+
+def iter_result_dirs(root, systems=None):
+    """Yield every finished structure directory under ``root``.
+
+    Scans one level below ``root`` (``<root>/<system>/struct_*``) plus ``root``
+    itself (``<root>/struct_*``, for a campaign root that *is* one system).
+    Bounded on purpose: a full ``rglob`` over a campaign with tens of thousands
+    of structure directories, each with calculation subdirectories, is far
+    slower for no additional coverage.
+    """
+    root = Path(root)
+    if not root.is_dir():
+        return
+    yield from _scan_system_dir(root, root.name)
+    names = (
+        [system_name(list(s)) for s in systems]
+        if systems
+        else sorted(e.name for e in root.iterdir() if e.is_dir())
+    )
+    for name in names:
+        yield from _scan_system_dir(root / name, name)
+
+
+def load_manifests(root, systems=None) -> pd.DataFrame:
+    """Concatenate every available ``manifest.csv`` under ``root``."""
+    root = Path(root)
+    frames = []
+    candidates = []
+    if (root / MANIFEST_NAME).is_file():
+        candidates.append((root.name, root / MANIFEST_NAME))
+    names = (
+        [system_name(list(s)) for s in systems]
+        if systems
+        else sorted(e.name for e in root.iterdir() if e.is_dir())
+        if root.is_dir()
+        else []
+    )
+    for name in names:
+        path = root / name / MANIFEST_NAME
+        if path.is_file():
+            candidates.append((name, path))
+    for name, path in candidates:
+        frame = pd.read_csv(path, dtype=MANIFEST_DTYPES)
+        if "system" not in frame.columns:
+            frame["system"] = name
+        frames.append(frame)
+    if not frames:
+        return pd.DataFrame(columns=["index", "system", "dirname", "hash"])
+    return pd.concat(frames, ignore_index=True)
+
+
+def collect_campaign(root, systems=None, output=None) -> pd.DataFrame:
+    """Concatenate all results under ``root``, joined to the manifest.
+
+    Hashed directories join on ``(system, structure_hash)``; legacy
+    ``struct_pyxtal_<idx>`` directories have no hash on disk, so they join on
+    ``(system, struct_index)`` when a manifest exists and simply carry no
+    provenance columns when it does not.
+    """
+    root = Path(root)
+    results = list(iter_result_dirs(root, systems))
+    frames = []
+    for result in results:
+        try:
+            frame = pd.read_pickle(result.path / RESULT_FILENAME)
+        except Exception as exc:
+            print(f"Warning: could not read {result.path / RESULT_FILENAME}: {exc}")
+            continue
+        frame = frame.copy()
+        frame["system"] = result.system
+        frame["struct_index"] = result.index
+        frame["structure_hash"] = result.structure_hash
+        frame["layout"] = result.layout
+        frame["run_dir"] = str(result.path)
+        frames.append(frame)
+
+    if not frames:
+        combined = pd.DataFrame(
+            columns=["system", "struct_index", "structure_hash", "layout", "run_dir"]
+        )
+    else:
+        combined = pd.concat(frames, ignore_index=True)
+
+    manifest = load_manifests(root, systems)
+    if len(manifest) and len(combined):
+        combined = _join_manifest(combined, manifest)
+
+    if output is not None:
+        output = Path(output)
+        output.parent.mkdir(parents=True, exist_ok=True)
+        if output.suffix in (".csv",):
+            combined.to_csv(output, index=False)
+        else:
+            combined.to_pickle(output)
+    return combined
+
+
+_MANIFEST_COLUMNS = ["spacegroup", "composition", "formula", "n_atoms", "volume_per_atom"]
+
+
+def _join_manifest(combined: pd.DataFrame, manifest: pd.DataFrame) -> pd.DataFrame:
+    keep = [c for c in _MANIFEST_COLUMNS if c in manifest.columns]
+    hashed = combined["layout"] == "hashed"
+
+    out = combined.copy()
+    for column in keep:
+        out[column] = None
+
+    if "hash" in manifest.columns:
+        by_hash = manifest.drop_duplicates(subset=["system", "hash"]).set_index(
+            ["system", "hash"]
+        )[keep]
+        keys = list(zip(out.loc[hashed, "system"], out.loc[hashed, "structure_hash"]))
+        if keys:
+            joined = by_hash.reindex(keys)
+            for column in keep:
+                out.loc[hashed, column] = joined[column].to_numpy()
+
+    if "index" in manifest.columns and (~hashed).any():
+        by_index = manifest.drop_duplicates(subset=["system", "index"]).set_index(
+            ["system", "index"]
+        )[keep]
+        keys = list(zip(out.loc[~hashed, "system"], out.loc[~hashed, "struct_index"]))
+        if keys:
+            joined = by_index.reindex(keys)
+            for column in keep:
+                out.loc[~hashed, column] = joined[column].to_numpy()
+    return out
+
+
+def collect_from_config(cfg: CampaignConfig, output=None) -> pd.DataFrame:
+    return collect_campaign(cfg.root, cfg.system.systems, output)