Skip to content

Add run_stats calculations and corresponding tests for FHI-aims - #1527

Open
TomatoTurtlez wants to merge 2 commits into
materialsproject:mainfrom
TomatoTurtlez:main
Open

Add run_stats calculations and corresponding tests for FHI-aims#1527
TomatoTurtlez wants to merge 2 commits into
materialsproject:mainfrom
TomatoTurtlez:main

Conversation

@TomatoTurtlez

Copy link
Copy Markdown

Summary

  • Add _parse_run_stats helper to atomate2/aims/schemas/calculation.py that parses timing, peak memory, and core-count stats directly from aims.out (CPU/elapsed time, min/max/avg tracked memory, core count), analogous to VASP's Outcar.run_stats.
  • Add run_stats: dict[str, Any] | None field to CalculationOutput (schemas/calculation.py), populated in Calculation.from_aims_files via _parse_run_stats.
  • Add run_stats: dict[str, Any] | None field to AimsTaskDoc (schemas/task.py), populated in from_directory from calcs_reversed[-1].output.run_stats — mirrors VASP's two-level TaskDoc/CalculationOutput pattern, no re-parsing.
  • Fix _parse_bandstructure (schemas/calculation.py) raising AttributeError: 'AimsOutput' object has no attribute 'band_structure' whenever parse_bandstructure was truthy. Now guards with hasattr and warns + returns None if band structure parsing isn't yet supported by the installed pymatgen.io.aims.outputs.AimsOutput.

Additional dependencies introduced (if any)

None. Uses re and pathlib.Path, both already in the standard library and already imported elsewhere in this codebase.

TODO (if any)

  • run_stats is parsed directly from aims.out text in atomate2 as a stopgap. This should eventually be upstreamed so pyfhiaims.outputs.stdout.AimsStdout and pymatgen.io.aims.outputs.AimsOutput expose run_stats natively (matching how VASP's Outcar.run_stats is parsed in pymatgen, not atomate2). Once that lands, _parse_run_stats here should be removed and Calculation.from_aims_files updated to read aims_output.run_stats directly.
  • Band structure parsing (AimsOutput.band_structure) is not implemented upstream in pymatgen. This PR only prevents the crash when it's requested; actual band structure support is out of scope here and tracked as a separate follow-up.

Checklist

  • [ X] Code is in the standard Python style. The easiest way to handle this is to run the following in the correct sequence on your local machine. Start with running ruff and ruff format on your new code. This will automatically reformat your code to PEP8 conventions and fix many linting issues.
  • [ X] Doc strings have been added in the [Numpy docstring format] (https://sphinxcontrib-napoleon.readthedocs.io/en/latest/example_numpy.html). Run ruff on your code.
  • [ X] Type annotations are highly encouraged. Run [mypy] (http://mypy-lang.org) to type check your code.
  • [ X] Tests have been added for any new functionality or bug fixes.
  • [ X] All linting and tests pass.

@JaGeo

JaGeo commented Aug 17, 2026

Copy link
Copy Markdown
Member

@tpurcell90 would you mind taking a look or refer to a better suited reviewer? If not, I am happy to take over but I am not an aims expert, of course!

@tpurcell90

Copy link
Copy Markdown
Contributor

Ya I'll take a look once my laptop can log in. There send to be technical issues with the servers.

@JaGeo

JaGeo commented Aug 17, 2026

Copy link
Copy Markdown
Member

@tpurcell90 sure!

return output_file_paths


def _parse_run_stats(aims_output_file: Path | str) -> dict[str, Any]:

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why is this function located here and not the pymatgen-io-aims plugin? I think it would be better for parsing of output flies to be done outside of atomate2 as much as possible.

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hi @tpurcell90! Agreed. The current _parse_run_stats was a stopgap so I could benchmark a running job before pyfhiaims/pymatgen upstream support existed (noted in the TODO section of this PR).

I've already extended pyfhiaims.parsers (FinalMemoriesParser for peak memory, feeding a new Aims.Stdout.run_stats property) to replicate this behavior at the source, and plan to open that PR next.

Question: Would it make sense for atomate2 to depend on pyfhiaims.stdout directly, rather than going through pymatgen.io.aims.outputs.AimsOutput? That would let this land after just one upstream PR (pyfhiaims) instead of two (pyfhiaims then pymatgen).

I know the rest of the AIMS code in atomate2 currently goes through pymatgen exclusively, so I wanted to check whether that's a hard convention here or if a direct pyfhiaims dependency would be acceptable.

re.DOTALL,
)

return {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This will go for an entire calculation file or individual steps? This looks like it is for an entire calculation, but it would be good if this is done for each individual calculation

@TomatoTurtlez TomatoTurtlez Aug 18, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good catch. FHI-aims does print per-SCF-iteration timing/memory (End self-consistency iteration #N) as well as the SCF-initialization phase, so per-step stats are parseable.

Currently, run_stats only captures the final summary once per calculation.

I can add an opt-in flag (e.g., parse_run_stats="all" | "final" | None) to additionally capture per-iteration stats as a list, keeping the default behavior as calculation-level only so it doesn't add disk overhead for people who don't need it.

Let me know if that granularity is worth adding now or if the calculation-level summary is sufficient for this PR, and I can scope it as a fast follow-up either way.

"tags": tags,
"completed": calcs_reversed[-1].completed,
"completed_at": calcs_reversed[-1].completed_at,
"run_stats": calcs_reversed[-1].output.run_stats,

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I would make this an optional tag since in high-throughput this can add extra disk space usage and not everyone will need to have this information.

@tpurcell90

Copy link
Copy Markdown
Contributor

Thanks for this pull request! I made some comments, but once resolved it should be good to go

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants