nomad-forces-export 0.1.0__py3-none-any.whl

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/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/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/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/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

nomad_forces_export/query.py

@@ -0,0 +1,269 @@
+"""Build NOMAD search queries and list matching entry_ids, filtering excluded datasets."""
+
+import os
+import time
+from collections.abc import Iterable, Iterator
+from dataclasses import dataclass, field
+
+from loguru import logger
+
+from nomad_forces_export.api import NomadClient
+from nomad_forces_export.config import (
+    BASE_QUERY,
+    REQUIRED_ARCHIVE_DATA,
+    REQUIRED_METADATA,
+)
+
+
+@dataclass
+class NomadQuery:
+    """Describes a NOMAD entries search.
+
+    If `entry_ids` is given, it takes precedence over all other filters (elements,
+    formula, method, raw_query) -- the query becomes a direct id lookup.
+    """
+
+    elements: list[str] | None = None
+    formula: str | None = None
+    method: str | None = None
+    dft_xc_functionals: list[str] | None = None
+    program: str | None = None
+    entry_ids: list[str] | None = None
+    exclude_datasets: list[str] = field(default_factory=list)
+    raw_query: dict | None = None
+    client: NomadClient | None = None
+
+    def to_query_dict(self, properties: set[str] | None = None) -> dict:
+        """Build the NOMAD query dict (the `query` field of an API request body)."""
+        if self.entry_ids:
+            return {'entry_id': {'any': list(self.entry_ids)}}
+
+        query_dict: dict = BASE_QUERY.copy()
+
+        for prop in properties or []:
+            if prop not in {'energy', 'forces', 'stress'}:
+                raise ValueError(f'Unknown property requested: {prop}')
+            if prop == 'energy':
+                query_dict['quantities:all'].append(
+                    'run.calculation.energy.total_t0.value'
+                )
+            else:
+                query_dict['quantities:all'].append(
+                    f'run.calculation.{prop}.total.value'
+                )
+        if self.elements:
+            query_dict['results.material.elements'] = {'all': list(self.elements)}
+        if self.formula:
+            query_dict['results.material.chemical_formula_hill'] = self.formula
+        if self.method:
+            query_dict['results.method.method_name'] = self.method
+        if self.dft_xc_functionals:
+            query_dict['results.method.simulation.dft.xc_functional_names:any'] = list(
+                self.dft_xc_functionals
+            )
+
+        if self.program:
+            query_dict['results.method.simulation.program_name'] = self.program
+        if self.exclude_datasets:
+            query_dict['datasets.dataset_name:none'] = list(self.exclude_datasets)
+            query_dict['external_db:none'] = list(self.exclude_datasets)
+        if self.raw_query:
+            query_dict.update(self.raw_query)
+        return query_dict
+
+    def _is_excluded(self, entry: dict) -> bool:
+        if not self.exclude_datasets:
+            return False
+        dataset_names = [d.get('dataset_name', '') for d in entry.get('datasets', [])]
+        for excluded in self.exclude_datasets:
+            for name in dataset_names:
+                if excluded.lower() in name.lower():
+                    return True
+        return False
+
+    def search(
+        self,
+        max_entries: int | None = None,
+        properties: set[str] | None = None,
+    ) -> Iterator[str]:
+        """Yield entry_ids matching this query, skipping excluded datasets.
+
+        Paginates through `/entries/query` using NOMAD's `page_after_value` cursor
+        until no more pages remain, the cursor stops advancing, or `max_entries`
+        yielded entry_ids is reached.
+        """
+        query_dict = self.to_query_dict(properties=properties)
+        page_after_value = None
+        yielded = 0
+        page_size = min(10000, max_entries) if max_entries is not None else 10000
+        client = self.client or NomadClient()
+        while True:
+            payload = {
+                'query': query_dict,
+                'pagination': {'page_size': page_size},
+                'required': {'include': REQUIRED_METADATA},
+            }
+            if page_after_value:
+                payload['pagination']['page_after_value'] = page_after_value
+
+            response = client.post('/entries/query', payload)
+            pagination = response.get('pagination', {})
+            total_entries = pagination.get('total')
+            logger.info(f'query returned total {total_entries})')
+            next_page_after_value = pagination.get('next_page_after_value')
+
+            # If the cursor returned for this page is identical to the cursor we
+            # just requested with, the server (or a mock) is not making progress
+            # (e.g. returning the same page repeatedly) -- stop before processing
+            # this page's data to avoid re-yielding already-seen entries.
+            if (
+                page_after_value is not None
+                and next_page_after_value == page_after_value
+            ):
+                return
+
+            data = response.get('data', [])
+            for entry in data:
+                if self._is_excluded(entry):
+                    continue
+                yield entry['entry_id']
+                yielded += 1
+                if max_entries is not None and yielded >= max_entries:
+                    return
+
+            # Stop if there's no cursor or no data (no more pages).
+            if not next_page_after_value or not data:
+                return
+
+            page_after_value = next_page_after_value
+
+    def search_nomad(
+        self,
+        max_entries: int | None = None,
+        properties: set[str] | None = None,
+    ) -> Iterator[str]:
+        """Yield entry_ids matching this query, skipping excluded datasets.
+
+        Paginates through `/entries/query` using NOMAD's `page_after_value` cursor
+        until no more pages remain, the cursor stops advancing, or `max_entries`
+        yielded entry_ids is reached.
+        """
+        from nomad.app.v1.models.models import (
+            MetadataRequired,
+        )
+        from nomad.client.archive import ArchiveQuery
+
+        query_dict = self.to_query_dict(properties=properties)
+        page_size = min(10000, max_entries) if max_entries is not None else 10000
+        yielded = 0
+        page_after_value = None
+
+        while True:
+            archive_query = ArchiveQuery(
+                query=query_dict,
+                required=MetadataRequired(include=REQUIRED_METADATA),  # type: ignore
+                page_size=page_size,
+                after=page_after_value,  # type: ignore
+            )
+            num_queries = archive_query.fetch()
+            if num_queries == 0:
+                return
+            for entry in archive_query.entry_list():
+                page_after_value = entry[0]
+                yield entry[0]
+                yielded += 1
+                if max_entries is not None and yielded >= max_entries:
+                    return
+
+    def fetch_archives(
+        self,
+        entry_ids: Iterable[str],
+        batch_size: int = 50,
+        save_dir: str | None = None,
+    ) -> Iterator[dict]:
+        """Fetch full archive data (system + calculation sections) for `entry_ids`.
+
+        Requests are batched (`batch_size` entry_ids per request) to keep payload
+        sizes reasonable. Yields one archive dict (as returned in the response
+        `data` list) per entry.
+        """
+        client = self.client or NomadClient()
+        entry_ids = list(entry_ids)
+        required = REQUIRED_ARCHIVE_DATA
+        batch_size = min(batch_size, 1000)
+        for i in range(0, len(entry_ids), batch_size):
+            batch = entry_ids[i : i + batch_size]
+            payload = {
+                'query': {'entry_id': {'any': batch}},
+                'required': required,
+                'pagination': {'page_size': batch_size},
+            }
+            response = client.post('/entries/archive/query', payload)
+            time.sleep(1)  # avoid overwhelming the server with back-to-back requests
+            logger.info(
+                f'Batch no {i // batch_size + 1}: {len(response.get("data", []))} archive entries.'
+            )
+            for archive_entry in response.get('data', []):
+                if save_dir:
+                    os.makedirs(f'{save_dir}/archives', exist_ok=True)
+                    entry_id = archive_entry.get('entry_id')
+                    if entry_id:
+                        with open(
+                            f'{save_dir}/archives/archive_{entry_id}.json', 'w'
+                        ) as f:
+                            import json
+
+                            json.dump(archive_entry, f, indent=2)
+                yield archive_entry
+
+    def fetch_archives_nomad(
+        self,
+        entry_ids: Iterable[str],
+        batch_size: int = 50,
+        save_dir: str | None = None,
+    ) -> Iterator[dict]:
+        """Fetch full archive data (system + calculation sections) for `entry_ids`.
+
+        Requests are batched (`batch_size` entry_ids per request) to keep payload
+        sizes reasonable. Yields one archive dict (as returned in the response
+        `data` list) per entry.
+        """
+        from nomad.client.archive import ArchiveQuery
+
+        entry_ids = list(entry_ids)
+        required = REQUIRED_ARCHIVE_DATA.copy()
+        # required.pop(
+        #     'workflow2', None
+        # )  # workflow2 is not available in the nomad package
+        for i in range(0, len(entry_ids), batch_size):
+            batch = entry_ids[i : i + batch_size]
+            query = {'entry_id': {'any': batch}}
+            archive_query = ArchiveQuery(
+                query=query,
+                # required=required,
+                page_size=batch_size,
+                batch_size=batch_size,
+            )
+            number_of_entries = archive_query.fetch()
+            logger.info(
+                f'Batch no {i // batch_size + 1}: {number_of_entries} archive entries.'
+            )
+            archive_query._required = required
+            archives = archive_query.download(as_plain_dict=True)
+            for archive_entry in archives:
+                archive = {
+                    'entry_id': archive_entry['metadata']['entry_id'],
+                    'upload_id': archive_entry['metadata']['upload_id'],
+                    'archive': archive_entry,
+                }
+                if save_dir:
+                    os.makedirs(f'{save_dir}/archives', exist_ok=True)
+                    entry_id = archive_entry['metadata']['entry_id']
+                    if entry_id:
+                        with open(
+                            f'{save_dir}/archives/archive_{entry_id}.json', 'w'
+                        ) as f:
+                            import json
+
+                            json.dump(archive, f, indent=2)
+                yield archive

nomad_forces_export/units.py

@@ -0,0 +1,32 @@
+"""Unit conversions from NOMAD's SI-based archive units to ASE's units.
+
+ASE's internal unit system uses electron-volt (eV) for energy, Angstrom (Ang) for
+length, eV/Ang for force, and eV/Ang^3 for stress/pressure. NOMAD archives store
+raw SI values (Joule, meter, Newton, Pascal). These helpers convert one direction:
+NOMAD (SI) -> ASE units.
+"""
+from ase import units as ase_units
+
+
+def joule_to_ev(value: float) -> float:
+    """Convert an energy value in Joules to electron-volts."""
+    # ase_units.J is the number of eV in one Joule.
+    return value * ase_units.J
+
+
+def meter_to_angstrom(value: float) -> float:
+    """Convert a length value in meters to Angstrom."""
+    # ase_units.m is the number of Angstrom in one meter.
+    return value * ase_units.m
+
+
+def newton_to_ev_per_angstrom(value: float) -> float:
+    """Convert a force value in Newtons to eV/Angstrom."""
+    # 1 N = 1 J/m, so N -> eV/Ang is ase_units.J / ase_units.m
+    return value * ase_units.J / ase_units.m
+
+
+def pascal_to_ev_per_angstrom3(value: float) -> float:
+    """Convert a stress/pressure value in Pascals to eV/Angstrom^3."""
+    # 1 Pa = 1 J/m^3, so Pa -> eV/Ang^3 is ase_units.J / ase_units.m**3
+    return value * ase_units.J / (ase_units.m ** 3)

nomad_forces_export/write.py

@@ -0,0 +1,59 @@
+"""Write ase.Atoms iterables to ASE-DB and/or extxyz files on disk."""
+
+import os
+from collections.abc import Iterable
+
+from ase import Atoms
+from ase.db import connect
+from ase.io import write as ase_write
+
+from nomad_forces_export.config import VALID_FORMATS
+
+
+def _path_for_format(output_path: str, output_format: str) -> str:
+    root, ext = os.path.splitext(output_path)
+    if ext not in ('.db', '.xyz', '.extxyz'):
+        root = output_path
+    if output_format == 'ase_db':
+        return output_path if ext == '.db' else f'{root}.db'
+    if output_format == 'extxyz':
+        return output_path if ext in ('.xyz', '.extxyz') else f'{root}.xyz'
+    raise ValueError(f'Unknown output format: {output_format}')
+
+
+def write_atoms(
+    atoms_iter: Iterable[Atoms],
+    output_path: str,
+    output_format: str | list[str],
+) -> None:
+    """Write `atoms_iter` to disk in one or more formats.
+
+    `output_format` is `"ase_db"`, `"extxyz"`, or a list containing both. When a
+    list is given, `output_path`'s extension is replaced per-format (e.g.
+    `"dataset"` becomes `dataset.db` and `dataset.extxyz`).
+    """
+    formats = [output_format] if isinstance(output_format, str) else list(output_format)
+    unknown = set(formats) - VALID_FORMATS
+    if unknown:
+        raise ValueError(f'Unknown output format(s): {sorted(unknown)}')
+
+    atoms_list = list(atoms_iter)
+
+    if 'ase_db' in formats:
+        db_path = _path_for_format(output_path, 'ase_db')
+        if os.path.exists(db_path):
+            os.remove(db_path)
+        with connect(db_path) as db:
+            for atoms in atoms_list:
+                key_value_pairs = {
+                    'nomad_entry_id': atoms.info.get('nomad_entry_id', ''),
+                    'nomad_upload_id': atoms.info.get('nomad_upload_id', ''),
+                }
+                db.write(atoms, key_value_pairs=key_value_pairs, data=atoms.info)
+
+    if 'extxyz' in formats:
+        xyz_path = _path_for_format(output_path, 'extxyz')
+        if os.path.exists(xyz_path):
+            os.remove(xyz_path)
+        for atoms in atoms_list:
+            ase_write(xyz_path, atoms, format='extxyz', append=True)

nomad_forces_export-0.1.0.dist-info/METADATA

@@ -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.dist-info/RECORD

@@ -0,0 +1,13 @@
+nomad_forces_export/__init__.py,sha256=gchMcvlMFEeU5M2ZtfeUU8IrMAuQFo-d41ZgcNJtEPE,4446
+nomad_forces_export/api.py,sha256=edpHofny42aqmw9P0bGbgDauGctipzLQulh9S8xhQrg,2009
+nomad_forces_export/cli.py,sha256=wsNFm-wDFrRKdV2fqdXA4rVKRQ0HUsLwT5v1tA9qsMs,1943
+nomad_forces_export/config.py,sha256=-qfxS0gSAhAv7JsdLnLvyw3BzEdqeyQ5fLIHdm-tKwo,1016
+nomad_forces_export/extract.py,sha256=EAIJRIjq8EWxlUJYYivoIK6DK6QKIYV78NIXamsW8eA,6458
+nomad_forces_export/query.py,sha256=_C99aDnYT4rWtOw61leFRSqA1Vp0USwozfdVqpn2oe8,10630
+nomad_forces_export/units.py,sha256=c4u5Wr6e4UdAL_pr0QRFTen6PpI9nI4ZrianjrpiooM,1239
+nomad_forces_export/write.py,sha256=yE_esOVeeKAXMa72NwR9Ba6zmUDEgk6sU3_0RGEed0E,2154
+nomad_forces_export-0.1.0.dist-info/METADATA,sha256=B32W0K3qekLTlL0326A-czCAxzfZhndqZe54K6YbePc,420
+nomad_forces_export-0.1.0.dist-info/WHEEL,sha256=YVMoNqKzERt-wjUZwJ33xBGAwnFl-4cqbYkTtWa4itE,91
+nomad_forces_export-0.1.0.dist-info/entry_points.txt,sha256=NLgBLMIQ-co49PXgcpLawDrvQZHcO2K6WLY9141gPBQ,69
+nomad_forces_export-0.1.0.dist-info/top_level.txt,sha256=eIvx4I2uSmOJ-w-EcHzTIbZ18D1CVaWrUI1nBkrqMmU,20
+nomad_forces_export-0.1.0.dist-info/RECORD,,

nomad_forces_export-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (84.0.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

nomad_forces_export-0.1.0.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+nomad-forces-export = nomad_forces_export.cli:main

nomad_forces_export-0.1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+nomad_forces_export