nomad-forces-export 0.1.0__tar.gz

nomad_forces_export-0.1.0/PKG-INFO

@@ -0,0 +1,12 @@
+Metadata-Version: 2.4
+Name: nomad-forces-export
+Version: 0.1.0
+Summary: Fetch and package NOMAD atomistic data for MLIP training
+Author-email: Sharat Patil <patilsha@physik.hu-berlin.de>
+Maintainer-email: Sharat Patil <patilsha@physik.hu-berlin.de>
+Requires-Python: >=3.10
+Requires-Dist: ase>=3.23
+Requires-Dist: loguru>=0.7.3
+Requires-Dist: requests>=2.28
+Provides-Extra: dev
+Requires-Dist: pytest>=7.0; extra == "dev"

nomad_forces_export-0.1.0/README.md

@@ -0,0 +1,58 @@
+# nomad-mlip-data
+
+Fetch atomistic structures with energies/forces/stress from
+[NOMAD](https://nomad-lab.eu/prod/v1/gui/search/entries) and package them as ASE-DB or
+extxyz datasets for fine-tuning ASE-calculator-compatible MLIPs.
+
+## Install
+
+```bash
+uv venv .venv
+source .venv/bin/activate
+uv pip install -e ".[dev]"
+```
+
+## Usage (Python)
+
+```python
+from nomad_forces_export import NomadQuery, fetch_dataset
+
+query = NomadQuery(
+    elements=["Si", "O"],
+    exclude_datasets=["Alexandria PBEsol"],
+)
+
+fetch_dataset(
+    query=query,
+    properties={"energy", "forces"},
+    output_format="ase_db",
+    output_path="dataset.db",
+    max_entries=5000,
+)
+```
+
+## Usage (CLI)
+
+```bash
+nomad-mlip-data fetch \
+  --elements Si O \
+  --exclude-dataset "Alexandria PBEsol" \
+  --properties energy forces \
+  --format ase_db extxyz \
+  --output dataset \
+  --max-entries 5000
+```
+
+## Testing
+
+```bash
+pytest tests/ -v
+```
+
+All tests run against recorded NOMAD API fixtures in `tests/fixtures/` — no live
+network access is required.
+
+## Main contributors
+| Name | E-mail     |
+|------|------------|
+| Sharat Patil | [patilsha@physik.hu-berlin.de](mailto:patilsha@physik.hu-berlin.de)

nomad_forces_export-0.1.0/pyproject.toml

@@ -0,0 +1,35 @@
+[project]
+name = "nomad-forces-export"
+version = "0.1.0"
+description = "Fetch and package NOMAD atomistic data for MLIP training"
+requires-python = ">=3.10"
+
+authors = [
+    { name = "Sharat Patil", email = "patilsha@physik.hu-berlin.de" },
+]
+maintainers = [
+    { name = "Sharat Patil", email = "patilsha@physik.hu-berlin.de" },
+]
+
+dependencies = [
+    "ase>=3.23",
+    "loguru>=0.7.3",
+    "requests>=2.28",
+]
+
+[project.optional-dependencies]
+dev = ["pytest>=7.0"]
+
+[project.scripts]
+nomad-forces-export = "nomad_forces_export.cli:main"
+
+[build-system]
+requires = ["setuptools>=68"]
+build-backend = "setuptools.build_meta"
+
+[tool.setuptools]
+package-dir = {"" = "src"}
+
+[tool.setuptools.packages.find]
+where = ["src"]
+include = ["nomad_forces_export*"]

nomad_forces_export-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

nomad_forces_export-0.1.0/src/nomad_forces_export/__init__.py

@@ -0,0 +1,130 @@
+"""nomad_forces_export: fetch and package NOMAD data for MLIP fine-tuning."""
+
+import json
+from pathlib import Path
+from typing import List, Optional, Set, Union
+
+from loguru import logger
+
+from nomad_forces_export.api import NomadClient
+from nomad_forces_export.extract import to_atoms
+from nomad_forces_export.query import NomadQuery
+from nomad_forces_export.write import write_atoms
+
+# logger = logging.getLogger(__name__)
+
+__all__ = [
+    'NomadQuery',
+    'atoms_generator',
+    'create_dataset_from_archives',
+    'fetch_dataset',
+    'write_atoms',
+]
+
+
+def atoms_generator(archives, properties, max_frames: int | None = None):
+    n_frames = 0
+    n_entries_skipped = 0
+    entries_processed = 0
+
+    def _atoms_from_entry(archive_entry: dict):
+        nonlocal n_frames, n_entries_skipped, entries_processed
+        frames_for_entry = 0
+        for atoms in to_atoms(archive_entry, properties=properties):
+            frames_for_entry += 1
+            yield atoms
+            n_frames += 1
+            if max_frames is not None and n_frames >= max_frames:
+                logger.info(
+                    f'Maximum number of frames ({max_frames}) reached; stopping extraction.'
+                )
+                return
+        frames_for_entry = bool(frames_for_entry)
+        n_entries_skipped += not frames_for_entry
+        entries_processed += frames_for_entry
+
+    try:
+        for archive_entry in archives:
+            yield from _atoms_from_entry(archive_entry)
+    except Exception as e:
+        import traceback
+
+        logger.error(f'Traceback:\n{traceback.format_exc()}')
+        logger.error(f'Error occurred while processing entries: {e}')
+    logger.info(
+        f'Extracted {n_frames} frames from {entries_processed} entries; {n_entries_skipped} entries had no usable frames for requested properties. Total entries processed: {n_entries_skipped + entries_processed}.'
+    )
+
+
+def fetch_dataset(
+    query: NomadQuery,
+    properties: set[str],
+    output_format: str | list[str],
+    output_path: str,
+    max_entries: int | None = None,
+    max_frames: int | None = None,
+    batch_size: int = 50,
+    use_nomad_pkg: bool = False,
+) -> None:
+    """Fetch entries matching `query` from NOMAD and write them as a dataset.
+
+    Searches NOMAD for matching entry_ids (respecting `query.exclude_datasets`),
+    fetches full archive data in batches, extracts one `ase.Atoms` per calculation
+    frame with the requested `properties`, and writes the result to `output_path`
+    in the given `output_format` ("ase_db", "extxyz", or a list of both).
+    """
+    if use_nomad_pkg:
+        search = query.search_nomad
+        fetch_archives = query.fetch_archives_nomad
+    else:
+        search = query.search
+        fetch_archives = query.fetch_archives
+    entry_ids = list(
+        search(
+            max_entries=max_entries,
+            properties=properties,
+        )
+    )
+    logger.info(
+        f'found {len(entry_ids)} matching entries after dataset exclusion filtering'
+    )
+    logger.info(f'Fetching {len(entry_ids)} entries in batches of {batch_size}...')
+    archives = fetch_archives(entry_ids=entry_ids, batch_size=batch_size)
+    write_atoms(
+        atoms_generator(archives, properties=properties, max_frames=max_frames),
+        output_path=output_path,
+        output_format=output_format,
+    )
+
+
+def create_dataset_from_archives(
+    archives: list[dict] | str | Path,
+    properties: set[str],
+    output_format: str | list[str],
+    output_path: str,
+    max_frames: int | None = None,
+) -> None:
+    """Create a dataset from a list of NOMAD archive entries.
+
+    Extracts one `ase.Atoms` per calculation frame with the requested `properties`,
+    and writes the result to `output_path` in the given `output_format`
+    ("ase_db", "extxyz", or a list of both).
+    """
+    if isinstance(archives, (str, Path)):
+        entries = []
+        if isinstance(archives, str):
+            archives = Path(archives)
+        if archives.is_dir():
+            for file in Path(archives).glob('*.json'):
+                with open(file) as f:
+                    entries.append(json.load(f))
+        elif archives.is_file():
+            with open(archives) as f:
+                entries = json.load(f)
+    else:
+        entries = archives
+    write_atoms(
+        atoms_generator(entries, properties=properties, max_frames=max_frames),
+        output_path=output_path,
+        output_format=output_format,
+    )

nomad_forces_export-0.1.0/src/nomad_forces_export/api.py

@@ -0,0 +1,55 @@
+"""Low-level HTTP client for the NOMAD API with retry/backoff on transient errors."""
+
+import time
+
+import requests
+
+from nomad_forces_export.config import NOMAD_BASE_URL
+
+
+class NomadClient:
+    """Thin wrapper around `requests` for calling the NOMAD API.
+
+    Retries on 5xx responses and on connection/timeout errors, using a fixed
+    backoff between attempts. Raises `requests.HTTPError` (or the underlying
+    `requests` exception) once retries are exhausted.
+    """
+
+    def __init__(
+        self,
+        base_url: str = NOMAD_BASE_URL,
+        max_retries: int = 3,
+        backoff_seconds: float = 5.0,
+        timeout: float = 60.0,
+    ):
+        self.base_url = base_url
+        self.max_retries = max_retries
+        self.backoff_seconds = backoff_seconds
+        self.timeout = timeout
+
+    def post(self, path: str, payload: dict) -> dict:
+        """POST `payload` as JSON to `base_url + path`, returning the parsed JSON body."""
+        url = f"{self.base_url}{path}"
+        last_exc: Exception | None = None
+        for attempt in range(1, self.max_retries + 1):
+            try:
+                response = requests.post(url, json=payload, timeout=self.timeout)
+                response.raise_for_status()
+                return response.json()
+            except requests.HTTPError as exc:
+                if exc.response is None:
+                    raise
+                elif exc.response.status_code == 429:
+                    # Too Many Requests - wait and retry
+                    time.sleep(self.backoff_seconds)
+                    continue
+                elif exc.response.status_code < 500:
+                    raise
+                last_exc = exc
+                if attempt < self.max_retries:
+                    time.sleep(self.backoff_seconds)
+            except (requests.ConnectionError, requests.Timeout) as exc:
+                last_exc = exc
+                if attempt < self.max_retries:
+                    time.sleep(self.backoff_seconds)
+        raise last_exc

nomad_forces_export-0.1.0/src/nomad_forces_export/cli.py

@@ -0,0 +1,60 @@
+"""Command-line interface for nomad_forces_export: thin wrapper over fetch_dataset."""
+
+import argparse
+import sys
+
+from nomad_forces_export import NomadQuery, fetch_dataset
+
+
+def build_arg_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(prog="nomad-mlip-data")
+    subparsers = parser.add_subparsers(dest="command", required=True)
+
+    fetch_parser = subparsers.add_parser("fetch", help="Fetch a dataset from NOMAD")
+    fetch_parser.add_argument("--elements", nargs="+", default=None)
+    fetch_parser.add_argument("--formula", default=None)
+    fetch_parser.add_argument("--method", default=None)
+    fetch_parser.add_argument("--entry-ids", nargs="+", default=None, dest="entry_ids")
+    fetch_parser.add_argument(
+        "--exclude-dataset", nargs="+", default=[], dest="exclude_dataset"
+    )
+    fetch_parser.add_argument(
+        "--properties",
+        nargs="+",
+        default=["energy", "forces"],
+        choices=["energy", "forces", "stress"],
+    )
+    fetch_parser.add_argument(
+        "--format", nargs="+", default=["ase_db"], choices=["ase_db", "extxyz"]
+    )
+    fetch_parser.add_argument("--output", required=True)
+    fetch_parser.add_argument(
+        "--max-entries", type=int, default=None, dest="max_entries"
+    )
+
+    return parser
+
+
+def main(argv: list[str] | None = None) -> None:
+    parser = build_arg_parser()
+    args = parser.parse_args(argv)
+
+    if args.command == "fetch":
+        query = NomadQuery(
+            elements=args.elements,
+            formula=args.formula,
+            method=args.method,
+            entry_ids=args.entry_ids,
+            exclude_datasets=args.exclude_dataset,
+        )
+        fetch_dataset(
+            query=query,
+            properties=set(args.properties),
+            output_format=args.format,
+            output_path=args.output,
+            max_entries=args.max_entries,
+        )
+
+
+if __name__ == "__main__":
+    main(sys.argv[1:])

nomad_forces_export-0.1.0/src/nomad_forces_export/config.py

@@ -0,0 +1,38 @@
+VALID_FORMATS = {'ase_db', 'extxyz'}
+VALID_PROPERTIES = {'energy', 'forces', 'stress'}
+
+NOMAD_BASE_URL = 'https://nomad-lab.eu/prod/v1/api/v1'
+
+WORKFLOWS = [
+    'SinglePoint',
+    'single_point',
+    'GeometryOptimization',
+    'geometry_optimization',
+]
+
+REQUIRED_METADATA = ['entry_id', 'upload_id']
+REQUIRED_ARCHIVE_DATA = {
+    'metadata': {'entry_id': '*', 'upload_id': '*'},
+    'results': {
+        'method': '*',
+    },
+    'workflow2': {'results': {'is_converged_geometry': '*'}},
+    'workflow': {'geometry_optimization': {'is_converged_geometry': '*'}},
+    'run': {
+        'program': '*',
+        'method': '*',
+        'system': {'atoms': '*', 'is_representative': '*'},
+        'calculation': {
+            'energy': '*',
+            'forces': '*',
+            'stress': '*',
+            'system_ref': '*',
+        },
+    },
+}
+
+BASE_QUERY = {
+    'results.method.method_name:any': ['DFT'],
+    'results.method.workflow_name:any': WORKFLOWS,
+    'quantities:all': ['run.calculation', 'run.system'],
+}

nomad_forces_export-0.1.0/src/nomad_forces_export/extract.py

@@ -0,0 +1,178 @@
+"""Convert a NOMAD archive entry into one ase.Atoms object per calculation frame."""
+
+import logging
+from collections.abc import Iterator
+
+import numpy as np
+from ase import Atoms
+from ase.calculators.singlepoint import SinglePointCalculator
+
+from nomad_forces_export.config import VALID_PROPERTIES
+from nomad_forces_export.units import (
+    joule_to_ev,
+    meter_to_angstrom,
+    newton_to_ev_per_angstrom,
+    pascal_to_ev_per_angstrom3,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def _parse_system_ref(system_ref: str) -> tuple[int, int] | None:
+    # system_ref looks like "/run/0/system/3" (leading slash optional) ->
+    # extract both the run index and system index.
+    parts = system_ref.rstrip('/').split('/')
+    try:
+        run_index = int(parts[parts.index('run') + 1])
+        system_index = int(parts[parts.index('system') + 1])
+    except (ValueError, IndexError):
+        return None
+    return run_index, system_index
+
+
+def _build_atoms_from_system(system: dict) -> Atoms:
+    atoms_data = system['atoms']
+    species = atoms_data['species']
+    positions_m = np.array(atoms_data['positions'])
+    positions_ang = meter_to_angstrom(positions_m)
+    magmoms = system.get('magmoms')
+
+    lattice_vectors_m = atoms_data.get('lattice_vectors')
+    periodic = atoms_data.get('periodic', [False, False, False])
+
+    cell = None
+    if lattice_vectors_m is not None:
+        cell = meter_to_angstrom(np.array(lattice_vectors_m))
+
+    return Atoms(
+        numbers=species,
+        positions=positions_ang,
+        cell=cell,
+        pbc=periodic,
+        magmoms=magmoms,
+    )
+
+
+def _extract_energy_ev(calculation: dict) -> float | None:
+    energy = calculation.get('energy')
+    if not energy:
+        return None
+    if 'total_t0' in energy:
+        return joule_to_ev(energy['total_t0']['value'])
+    if 'total' in energy:
+        return joule_to_ev(energy['total']['value'])
+    if 'free' in energy:
+        return joule_to_ev(energy['free']['value'])
+    return None
+
+
+def _extract_forces_ev_per_ang(calculation: dict) -> np.ndarray | None:
+    forces = calculation.get('forces')
+    if not forces or 'total' not in forces:
+        return None
+    forces_n = np.array(forces['total']['value'])
+    return newton_to_ev_per_angstrom(forces_n)
+
+
+def _extract_stress_ev_per_ang3(calculation: dict) -> np.ndarray | None:
+    stress = calculation.get('stress')
+    if not stress or 'total' not in stress:
+        return None
+    stress_pa = np.array(stress['total']['value'])
+    return pascal_to_ev_per_angstrom3(stress_pa)
+
+
+def to_atoms(archive_entry: dict, properties: set | list[str]) -> Iterator[Atoms]:
+    """Yield one `ase.Atoms` per calculation frame in a NOMAD archive entry.
+
+    `properties` is a subset of {"energy", "forces", "stress"}; frames missing any
+    requested property are skipped (and logged) rather than raising.
+    """
+    unknown = set(properties) - VALID_PROPERTIES
+    if unknown:
+        raise ValueError(f'Unknown properties requested: {sorted(unknown)}')
+
+    entry_id = archive_entry.get('entry_id')
+    upload_id = archive_entry.get('upload_id')
+    runs = archive_entry.get('archive', {}).get('run', [])
+
+    for run_index, run in enumerate(runs):
+        systems = run.get('system', [])
+        calculations = run.get('calculation', [])
+        method = run.get('method', [{}])[0]
+
+        for calc_index, calculation in enumerate(calculations):
+            system_ref = calculation.get('system_ref')
+            if system_ref is not None:
+                parsed = _parse_system_ref(system_ref)
+                if parsed is None:
+                    logger.warning(
+                        f'entry {entry_id}: calculation {calc_index} has an unparseable system_ref {system_ref}, skipping'
+                    )
+                    continue
+                ref_run_index, system_index = parsed
+                if ref_run_index != run_index:
+                    logger.warning(
+                        f'entry {entry_id}: calculation {calc_index} has system_ref pointing to a '
+                        f'different run ({ref_run_index}) than its own run ({run_index}), skipping',
+                    )
+                    continue
+            else:
+                system_index = calc_index
+
+            if system_index < 0 or system_index >= len(systems):
+                logger.warning(
+                    'entry %s: calculation %d references missing system %d, skipping',
+                    entry_id,
+                    calc_index,
+                    system_index,
+                )
+                continue
+
+            if method.get('x_vasp_incar_in'):
+                magmom = method['x_vasp_incar_in'].get('MAGMOM')
+                if magmom:
+                    systems[system_index]['magmoms'] = np.array(magmom)
+            results: dict = {}
+            skip = ''
+
+            if 'energy' in properties:
+                energy_ev = _extract_energy_ev(calculation)
+                if energy_ev is None:
+                    skip = 'energy'
+                else:
+                    results['energy'] = energy_ev
+
+            if not skip and 'forces' in properties:
+                forces = _extract_forces_ev_per_ang(calculation)
+                if forces is None:
+                    skip = 'forces'
+                else:
+                    results['forces'] = forces
+
+            if not skip and 'stress' in properties:
+                stress = _extract_stress_ev_per_ang3(calculation)
+                if stress is None:
+                    skip = 'stress'
+                else:
+                    results['stress'] = stress
+
+            if skip:
+                logger.debug(
+                    f'entry {entry_id}: calculation {calc_index} missing {skip}, skipping'
+                )
+                continue
+
+            atoms = _build_atoms_from_system(systems[system_index])
+            atoms.calc = SinglePointCalculator(atoms, **results)
+            atoms.info['nomad_entry_id'] = entry_id
+            atoms.info['nomad_upload_id'] = upload_id
+            atoms.info['is_representative'] = systems[system_index].get(
+                'is_representative', False
+            )
+            atoms.info['is_converged_geometry'] = run.get('workflow2', {}).get(
+                'results', {}
+            ).get('is_converged_geometry', False) or run.get('workflow', {}).get(
+                'geometry_optimization', {}
+            ).get('is_converged_geometry', False)
+            yield atoms