Bug fixes for the optimization and scan modules of MAPLE#46
Open
AxiEj wants to merge 19 commits into
Open
Conversation
Use UMA_DEFAULT_SIZE for omitted UMA size handling, set that default to uma-s-1p1, and route the UMA model alias plus constructor fallback through the same constant so SetClaculator and direct UMACalculator construction agree. Constraint: Omitted UMA size must keep resolving to the historical uma-s-1p1 local checkpoint. Rejected: Adding UMA_DEFAULT_LOCAL_SIZE | a separate constant is unnecessary when the intended default is uma-s-1p1. Rejected: Only changing SetClaculator | direct UMACalculator still resolved model='uma' to uma-s-1p2. Confidence: high Scope-risk: narrow Directive: Keep UMA_DEFAULT_SIZE, UMA_MODELS_MAP['uma'], and the constructor fallback aligned unless a default model change is intentional. Tested: python -m compileall -q maple/function/calculator/set_calculator.py maple/function/calculator/uma/_uma_calculator.py Tested: targeted Python mock checks for SetClaculator omitted/explicit UMA size and direct UMACalculator omitted size. Tested: local CUDA smoke using maple/function/calculator/model/uma-s-1p1.pt on H2 returned energy and forces without network. Not-tested: CPU energy smoke with UMA turbo; fairchem/Triton expects GPU kernels and failed on CPU tensors.
Select FAIR Chemistry's default inference backend for CPU while preserving turbo for CUDA, so local UMA checkpoints can run on both devices without changing model selection. Normalize CUDA-like device requests to FAIR Chemistry's supported cuda device token; fall back to CPU/default when CUDA is unavailable. Constraint: Turbo currently routes UMA through GPU/Triton execution paths that fail on CPU tensors, and FAIR Chemistry's MLIP unit accepts only cpu or cuda device tokens. Rejected: Disable turbo globally | CUDA smoke works and turbo is the intended fast path for repeated MAPLE workloads. Rejected: Add direct gpu0/gpu1 handling in UMACalculator | historical UMA behavior did not support direct gpu* aliases or preserve specific GPU indices; leave per-card support for a dedicated change. Confidence: high Scope-risk: narrow Directive: Keep CPU on default inference unless FAIR Chemistry provides a CPU-safe turbo backend. Tested: python -m compileall -q maple/function/calculator/uma/_uma_calculator.py maple/function/calculator/set_calculator.py Tested: targeted Python check that CUDA-like inputs select turbo when CUDA is available and CPU/default otherwise, while direct gpu* remains CPU like historical UMA behavior. Tested: targeted Python check that _build_predictor normalizes direct gpu0/cuda0/cpu inputs before calling FAIR Chemistry. Tested: real CPU H2 energy/forces smoke using local maple/function/calculator/model/uma-s-1p1.pt without network. Tested: real CUDA H2 energy/forces smoke using local maple/function/calculator/model/uma-s-1p1.pt without network. Not-tested: Full workflow/regression suite; no project test suite was available for UMA.
Keep repository status clean by ignoring local OMX/OMC session artifacts and generated docs workspace files that should not be versioned. Constraint: Local agent/session files are machine-specific and were appearing as untracked workspace noise. Rejected: Commit generated docs and session state | they are local artifacts, not source changes. Confidence: high Scope-risk: narrow Directive: Do not remove these ignores unless docs/ or .omc/ becomes an intentional tracked source tree. Tested: git check-ignore -v .omc docs docs/.omc/state/subagent-tracking.json Not-tested: N/A; ignore-only change.
Constraint: AGENTS.md and CLAUDE.md are local workspace guidance files in this checkout. Rejected: Commit generated local guide files | they are environment-specific and should not affect OPT review. Confidence: high Scope-risk: narrow Directive: Keep root-local assistant artifacts ignored unless the project decides to version them deliberately. Tested: git status no longer reports AGENTS.md or CLAUDE.md as untracked on debug/OPT. Not-tested: none.
Constraint: Preserve the historical Optmization spelling as a compatibility alias while exposing Optimization as the canonical class. Rejected: Leave unknown methods as silent no-ops | silent fallthrough hides invalid user input and makes failures hard to diagnose. Confidence: high Scope-risk: narrow Directive: Route new OPT methods through one resolved method string and return the selected optimizer's run result. Tested: python -m compileall -q maple; python -m pytest -q; python example/test.py -k opt --no-color. Not-tested: isolated unknown-method CLI fixture beyond static review of NotImplementedError path.
Constraint: Keep BatchLBFGS and deprecated DIIS behavior out of this cleanup. Rejected: Preserve both _traj.xyz and _opt_traj.xyz | duplicate trajectory sources made OPT output ambiguous. Rejected: Add verbose per-frame trajectory metadata | compact Image/Energy comments are the intended output. Confidence: high Scope-risk: moderate Directive: Keep OPT trajectory append-only in _opt_traj.xyz and do not reintroduce _traj.xyz without a format migration. Tested: python -m compileall -q maple; python -m pytest -q; python example/test.py -k opt --no-color; static greps for logger/RMS/_traj/write_traj/metadata. Not-tested: ruff check, because ruff is not installed in the active environment.
Constraint: OPT now uses _opt_traj.xyz as the single optimization trajectory output. Rejected: Leave README and pbc_cu unchanged | users would see stale _traj expectations and a machine-specific example path. Confidence: high Scope-risk: narrow Directive: Keep example inputs relative unless a test explicitly requires an absolute path. Tested: python example/test.py -k opt --no-color. Not-tested: none.
LBFGS now uses gradient = -force for the two-loop input and history updates, filters non-positive curvature pairs, and falls back to scaled steepest descent when history produces a non-descent direction. SDCG now stores BB/GDIIS history as matched previous-position and previous-force samples. Constraint: Keep example outputs and BLBFGS out of scope per user request. Rejected: Broad OPT refactor or full-file formatting | would obscure the minimal scientific fix. Confidence: high Scope-risk: narrow Directive: Preserve force-gradient sign conventions when changing OPT algorithms; forces are negative gradients. Tested: python -m pytest -q maple/function/dispatcher/optimization/test_optimizer_gradient_semantics.py; python -m compileall -q maple/function/dispatcher/optimization; git diff --cached --check Not-tested: example/test.py -k opt, BLBFGS, external model-backed examples
Constraint: preserve existing optimizer and scan framework while removing duplicated OPT tail handling. Rejected: keep SCAN hardwired to LBFGS | it contradicted parser-visible method support and blocked real OPT reuse. Confidence: high Scope-risk: moderate Directive: keep future SCAN method support routed through Optimization and CommandControl rather than per-method branches. Tested: python -m pytest -q maple/function/dispatcher/optimization maple/function/dispatcher/scan; python -m pytest -q; git diff --cached --check Not-tested: broad example corpus and GPU-backed model runs
Constraint: pytest files must not enter the main repository; the two test modules have already been pushed to origin/debug/OPT and need to be removed from the remote. Rejected: rewrite history to scrub the introducing commits and force push | debug/OPT was already pushed and may have been fetched by collaborators, so rewriting goes beyond the scope of this change. Confidence: high Scope-risk: narrow Directive: Future OPT pytest files stay in the local worktree; do not commit them under dispatcher/optimization or dispatcher/scan. Tested: git show --stat HEAD reports only the deletion of these two test files; other staging and worktree state remain unchanged (partial commit). Not-tested: remote CI (not configured for this repository).
Constraint: Default behavior must remain GPU="default" (safer than turbo for variable-composition workloads); CPU must coerce to default regardless of user choice (Triton GPU kernels unavailable on CPU). Rejected: new UMAParams dataclass + _init_params | calculator layer has no precedent; would diverge from MACE/AIMNet2/ANI lightweight kwarg style. Rejected: auto-select turbo from jobtype (NEB/TS/freq -> turbo) | jobtype-aware calculator violates separation; users should explicitly opt in for performance trade-offs. Confidence: high Scope-risk: narrow Directive: new UMA-specific options always flow through the model_options dict; never plumb model-private knobs through engine.commandcontrol.params top-level. CPU is a hardware constraint, not a user preference -- keep the warn-and-coerce path. Tested: python -m py_compile on both modified Python files (passes); grep confirms no other `inference=` callers in example/ or maple/. Not-tested: end-to-end UMA load through fairchem (no CUDA / fairchem in active environment); ruff not installed.
Constraint: Must not change params["model"] (dispatcher uses
`elif model == "uma"`); must not introduce fairchem import into
command_control (input parser stays heavy-deps-free); other models
(ANI/MACE/AIMNet2/MACEPol) display behavior must remain unchanged.
Rejected: replace params["model"] with resolved variant ("uma-s-1p2") |
breaks dispatcher branch and would force synchronized updates to
SUPPORTED_MODELS, MODEL_HESSIAN_SUPPORT, UNSUPPORTED_CHARGE_MULT_MODELS.
Rejected: store self.checkpoint_name on UMACalculator + log post-build |
summary() runs at parse time before calculator construction; timing
mismatch would force double-logging.
Confidence: high
Scope-risk: narrow
Directive: command_control owns input-layer display; never import
calculator-layer modules here. UMA-specific constants are intentionally
duplicated (UMA_DEFAULT_SIZE, SUPPORTED_UMA_SIZES) to keep input parsing
free of fairchem -- keep them in sync with _uma_calculator.py.
Tested: python -m py_compile maple/function/read/command_control.py
passes; CommandControl.from_settings driven with 5 synthetic cases --
"#model=uma" -> "model: uma(size=uma-s-1p1)"; "#model=uma(size=uma-s-1p2)"
-> round-trips; "#model=uma(size=uma-s-1p2,task=omol,inference=turbo)"
-> renders all opts; "#model=ani2x" / "#model=maceoff23m" unchanged
(no `()`).
Not-tested: live example/test.py run.
Constraint: README never documented `_traj.xyz` before 9235e34, so the "migrate _traj.xyz to _opt_traj.xyz" hint had no prior mention to back it up; the entire section was effectively standalone and conveyed nothing beyond what the file extensions already imply. Rejected: keep the section but drop only the migration sentence | the rest of the paragraph just re-states which extensions exist, which does not warrant a top-level README section. Confidence: high Scope-risk: narrow Directive: keep README centered on input syntax and the task table; output file names belong in CHANGELOG / per-job docs, not the top-level README. Lesson: do not blindly trust AI. The 9235e34 "OPT outputs" section was AI-drafted and treated a never-mentioned `_traj.xyz` as established prior context, fabricating a migration story that the README itself never supported. AI-authored documentation changes must be cross-checked against the file's actual prior content before merging; otherwise they produce self-referential text that has to be reverted later. Tested: grep "Optimization Outputs\|_opt_traj\|_traj\.xyz" README.md returns no match; git diff vs HEAD shows only the 7-line removal, no other README edits.
Remove stale out2xyz helper copies that are not referenced by MAPLE runtime, tests, or packaging. Constraint: out2xyz.py files were unreferenced and outside the packaged maple module. Rejected: keeping per-example hardcoded conversion helpers | they are stale and not part of CLI/test paths Confidence: high Scope-risk: narrow Directive: Restore only if a documented example workflow depends on these helper scripts. Tested: git diff --check; git grep out2xyz/process_coordinates; python example/test.py --list; pytest Not-tested: full example execution corpus; external manual workflows outside the repository
Add method-specific 360-degree relaxed SCAN examples while keeping per-point optimizer logs focused on the final geometry rather than temporary files that SCAN removes. Constraint: SCAN examples use aimnet2nse, gpu0, relaxed mode, and 5-degree torsion steps over 360 degrees. Rejected: leaving RFO without verbose=0 final-frame coordinates | inconsistent with LBFGS/SDCG scan output. Rejected: printing per-point *_opt.xyz and *_opt_traj.xyz paths in SCAN logs | those temporary optimizer artifacts are removed after scan completion. Confidence: high Scope-risk: moderate Directive: Keep standalone OPT artifact path logging enabled; suppress it only for SCAN-owned internal optimizer calls. Tested: python -m compileall -q maple/function/dispatcher/optimization/algorithm/LBFGS.py maple/function/dispatcher/optimization/algorithm/RFO.py maple/function/dispatcher/optimization/algorithm/SDCG.py maple/function/dispatcher/scan/scan.py Tested: ran five new scan examples; each generated a 73-frame *_scan_final.xyz and no ERROR/WARNING/temp optimizer path lines in .out. Not-tested: full example/test.py corpus.
SCAN uses calculator energies that MAPLE normalizes to Hartree, so the summary label should match the project-wide unit contract without changing numerical values or scan control flow. Constraint: Keep the fix to unit labels only; do not refactor SCAN or regenerate trajectories. Rejected: Converting the stored scan energies numerically | calculator outputs are already Hartree. Confidence: high Scope-risk: narrow Directive: Preserve calculator-level Hartree normalization; dispatcher summaries should label those values as Hartree/Eh, not eV. Tested: python -m compileall -q maple/function/dispatcher/scan/scan.py Tested: git diff --check Tested: rg -n Energy range: .* eV maple/function/dispatcher/scan/scan.py example/scan returned no matches Not-tested: full example/test.py corpus.
Normalize UMA inference at parse time and add a central engine fallback so raised calculation errors are recorded in the requested .out without changing CLI failure semantics. Constraint: Preserve existing runtime behavior while improving error visibility in .out files. Rejected: Broad unknown-parameter validation | deferred to the requested second step because it can change accepted inputs. Confidence: high Scope-risk: narrow Directive: Keep future error reporting centralized and avoid silent config failures without .out evidence. Tested: python -m py_compile maple/function/read/input_reader.py maple/function/read/command_control.py maple/function/calculator/set_calculator.py maple/function/engine.py maple/function/read/test_command_control_errors.py maple/function/calculator/test_set_calculator_errors.py maple/function/test_engine_errors.py; pytest -q maple/function/read/test_command_control_errors.py maple/function/calculator/test_set_calculator_errors.py maple/function/test_engine_errors.py; git diff --check; manual bad UMA inference and bad OPT method CLI probes. Not-tested: Full CUDA UMA turbo model initialization.
Catch unknown OPT, SCAN, MD, UMA model-option, and solvation keys during command parsing so typos fail early and are written to the output file instead of being silently ignored. Constraint: Keep validation limited to known stable option surfaces to avoid breaking broad TS/FREQ/IRC expert parameters. Rejected: Global strict validation for every task | too broad for this step and risks rejecting existing advanced inputs. Confidence: high Scope-risk: moderate Directive: Extend the allowlists only when the corresponding runtime surface is verified; do not silently ignore user-facing option typos. Tested: python -m py_compile maple/function/read/command_control.py maple/function/read/test_command_control_errors.py; pytest -q maple/function/read/test_command_control_errors.py maple/function/calculator/test_set_calculator_errors.py maple/function/test_engine_errors.py; parsed 45 example .inp command headers with CommandControl; git diff --check; manual CLI probe for #opt(method=lbfgs,max_itre=10); legacy #opt(lbfgs) normalization probe. Not-tested: Unknown-key validation for TS/FREQ/IRC intentionally left permissive.
Remove in-package pytest files so validation artifacts live outside the project source tree per the new workflow boundary. Constraint: User requested all project-internal test files be deleted and future tests run outside the project. Rejected: Moving tests to tests/ | user explicitly asked to delete them instead. Confidence: high Scope-risk: narrow Directive: Do not add test_*.py files under this project tree; keep temporary or regression probes outside the repository. Tested: find . -not -path './.git/*' -name 'test_*.py' returned no files; python -m py_compile maple/function/read/command_control.py maple/function/calculator/set_calculator.py maple/function/engine.py maple/function/read/input_reader.py; git diff --check. Not-tested: pytest, because all in-project pytest files were intentionally removed.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
UMA / Calculator
turbotodefault), and refactored the if-statement into a_normalize_device()helper for easier maintenance..outfile when UMA used the default model.defaultinterface is sufficient for most small-moleculecalculations.
OPT
Optmization→Optimization, with anoptmization = optimizationalias kept for backwardcompatibility.
.runacross algorithms (previously LBFGS returned atoms, RFO returned int, SDCGreturned atoms); all now return atoms.
sqrt( Σ(step²) / step_cart.size )instead of the originalsqrt( Σ(step²) / (step_cart.size × 3) ), which made RMS 1.732× too large.write_xyz, convergence checks) shared by LBFGS/RFO/SDCG into a single_common.py;algorithms just import from it.
JobABCinstead of each algorithm rolling its own.verboseparameter andwrite_trajdefault; all algorithms now stream-append to_opt_traj.xyz. (No more *_traj.xyz)1p1.adding
saved_forces = forces.copy().SCAN
lbfgs/rfo/cd/cg/sdcg; fixed the trailingStoken being misidentified as asulfur atom.
Error Handling
.outfile; fixes the silent-failure bug.Cleanup
out2xyz.py.