scripts/query_aqcat.py

#!/usr/bin/env python
from __future__ import annotations

import argparse
import logging
import multiprocessing as mp
from collections import defaultdict
from pathlib import Path
from typing import Dict, List, Optional, Tuple

from ase import Atoms
from ase.db import connect
from ase.io import write
from datasets import Dataset, load_dataset
from tqdm import tqdm

NON_METAL_ELEMENTS = {
    "H",
    "O",
    "N",
    "C",
    "S",
    "P",
    "F",
    "Cl",
    "Br",
    "I",
    "Se",
    "B",
    "Si",
    "Ge",
    "As",
    "Sb",
    "Te",
}


def _extraction_worker(args_tuple: Tuple) -> List[Atoms]:
    """
    A parallel worker that connects to a single DB file, extracts a list of
    rows by index, applies secondary filters, and returns the valid Atoms objects.
    """
    db_path, indices, contains_elements, element_mode, material_type = args_tuple

    extracted_atoms = []
    query_elements = set(contains_elements) if contains_elements else set()

    with connect(db_path) as db:
        for index in indices:
            row = db.get(id=index + 1)
            if not row:
                continue

            atoms = row.toatoms()
            atoms.info = row.data

            symbols = atoms.get_chemical_symbols()
            bulk_elements = {
                symbols[i] for i, tag in enumerate(atoms.get_tags()) if tag != 2
            }

            if contains_elements:
                if element_mode == "any" and query_elements.isdisjoint(bulk_elements):
                    continue
                if element_mode == "all" and not query_elements.issubset(bulk_elements):
                    continue

            if material_type:
                is_nonmetal = any(el in NON_METAL_ELEMENTS for el in bulk_elements)
                if material_type == "metal" and is_nonmetal:
                    continue
                if material_type == "nonmetal" and not is_nonmetal:
                    continue

            extracted_atoms.append(atoms)

    return extracted_atoms


class AQcatQuery:
    """An interface to lazily query the AQcat25 dataset."""

    def __init__(
        self,
        repo_id: str = "SandboxAQ/aqcat25",
        data_root: Path | str = "./aqcat_data",
    ):
        self.repo_id = repo_id
        self.data_root = Path(data_root)
        self.filtered_ids: Optional[List[str]] = None
        self.filtered_split: Optional[str] = None
        self._loaded_splits: Dict[str, Dataset] = {}
        logging.info("Initializing query engine...")

    def filter(
        self,
        split: str,
        adsorbates: Optional[List[str]] = None,
        min_energy: Optional[float] = None,
        max_energy: Optional[float] = None,
        magnetism: Optional[str] = None,
    ) -> AQcatQuery:
        """Filters the dataset based on metadata from the Parquet files."""
        if split not in self._loaded_splits:
            logging.info(
                f"Connecting to Hub and loading metadata for split: '{split}'..."
            )
            try:
                data_file = f"parquet/{split}.parquet"
                dataset_split = load_dataset(
                    self.repo_id,
                    data_files=data_file,
                    split="train",
                    trust_remote_code=True,
                )
                self._loaded_splits[split] = dataset_split
                logging.info(
                    f"Successfully loaded metadata for '{split}' ({len(dataset_split):,} total frames)."
                )
            except Exception as e:
                logging.error(
                    f"Could not load metadata for split '{split}'.\nError: {e}"
                )
                return self

        dataset_split = self._loaded_splits[split]
        logging.info(f"\nFiltering '{split}' split...")
        filtered_data = dataset_split

        if adsorbates:
            filtered_data = filtered_data.filter(
                lambda ex: ex["adsorbate"] in adsorbates
            )
            logging.info(f"  - {len(filtered_data):,} frames passed adsorbate filter")

        if min_energy is not None:
            filtered_data = filtered_data.filter(
                lambda ex: ex["adsorption_energy"] >= min_energy
            )
            logging.info(f"  - {len(filtered_data):,} frames passed min energy filter")

        if max_energy is not None:
            filtered_data = filtered_data.filter(
                lambda ex: ex["adsorption_energy"] <= max_energy
            )
            logging.info(f"  - {len(filtered_data):,} frames passed max energy filter")

        if magnetism:
            if magnetism == "magnetic":
                filtered_data = filtered_data.filter(
                    lambda ex: abs(ex["total_magnetization"]) > 1.0
                )
            elif magnetism == "non-magnetic":
                filtered_data = filtered_data.filter(
                    lambda ex: abs(ex["total_magnetization"]) <= 1.0
                )
            logging.info(f"  - {len(filtered_data):,} frames passed magnetism filter")

        self.filtered_ids = filtered_data["frame_id"]
        self.filtered_split = split
        logging.info(
            f"Found {len(self.filtered_ids):,} frames matching metadata criteria."
        )
        return self

    def extract_and_save(
        self,
        output_file: Path,
        contains_elements: Optional[List[str]] = None,
        element_filter_mode: str = "any",
        material_type: Optional[str] = None,
        num_workers: int = 1,
        chunk_size: int = 5000,
    ):
        """
        Extracts the filtered frames in parallel, applies secondary filters,
        and saves the results to chunked .extxyz files.
        """
        if not self.filtered_ids:
            logging.warning(
                "No frames to extract after metadata filtering. Nothing to do."
            )
            return

        logging.info("\nStarting extraction")

        db_to_indices = defaultdict(list)
        for frame_id in self.filtered_ids:
            db_name, index_str = frame_id.split("::")
            db_to_indices[db_name].append(int(index_str))

        split_dir = self.data_root / self.filtered_split
        if not split_dir.is_dir():
            raise FileNotFoundError(
                f"Data directory '{split_dir}' not found. Please run 'python download_split.py --split {self.filtered_split}'"
            )

        tasks = [
            (
                split_dir / db_name,
                indices,
                contains_elements,
                element_filter_mode,
                material_type,
            )
            for db_name, indices in db_to_indices.items()
        ]

        all_atoms = []
        with mp.Pool(num_workers) as pool:
            results = list(
                tqdm(
                    pool.imap(_extraction_worker, tasks),
                    total=len(tasks),
                    desc="Extracting frames",
                )
            )

        for atoms_chunk in results:
            all_atoms.extend(atoms_chunk)

        logging.info(f"Successfully extracted {len(all_atoms):,} final structures.")

        if all_atoms:
            logging.info(
                f"Writing results to chunked files (max {chunk_size} frames per file)..."
            )

            if len(all_atoms) <= chunk_size:
                write(output_file, all_atoms)
                logging.info(f"Saved {len(all_atoms)} structures to {output_file}")
            else:
                num_chunks = (len(all_atoms) + chunk_size - 1) // chunk_size
                for i in range(num_chunks):
                    start_idx = i * chunk_size
                    end_idx = start_idx + chunk_size
                    chunk = all_atoms[start_idx:end_idx]

                    chunk_path = output_file.with_stem(f"{output_file.stem}_{i:03d}")
                    write(chunk_path, chunk)
                    logging.info(
                        f"  - Wrote chunk {i+1}/{num_chunks} ({len(chunk)} frames) to {chunk_path}"
                    )

            logging.info("Save complete.")


if __name__ == "__main__":
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s | %(levelname)s | %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S",
    )

    parser = argparse.ArgumentParser(
        description="Query the AQcat25 dataset and extract atomic structures."
    )
    parser.add_argument("--split", required=True, help="The data split to query.")
    parser.add_argument(
        "--adsorbates", nargs="+", help="One or more adsorbate strings to match."
    )
    parser.add_argument(
        "--min-energy", type=float, help="Minimum adsorption energy in eV."
    )
    parser.add_argument(
        "--max-energy", type=float, help="Maximum adsorption energy in eV."
    )
    parser.add_argument(
        "--magnetism",
        choices=["magnetic", "non-magnetic"],
        help="Filter by magnetic moment.",
    )
    parser.add_argument(
        "--contains-elements",
        nargs="+",
        help="Filter for slabs containing specified elements.",
    )
    parser.add_argument(
        "--element-filter-mode",
        choices=["any", "all"],
        default="any",
        help="Logic for element filter.",
    )
    parser.add_argument(
        "--material-type",
        choices=["metal", "nonmetal"],
        help="Filter by the bulk slab's material type.",
    )
    parser.add_argument(
        "--data-root",
        type=Path,
        default=Path("./aqcat_data"),
        help="Local directory of downloaded data.",
    )
    parser.add_argument(
        "--output-file",
        type=Path,
        default="filtered_results.extxyz",
        help="Path/prefix for the output .extxyz file(s).",
    )
    parser.add_argument(
        "--num-workers",
        type=int,
        default=max(1, mp.cpu_count() - 2),
        help="Number of parallel workers for data extraction.",
    )
    parser.add_argument(
        "--chunk-size",
        type=int,
        default=5000,
        help="Maximum number of frames per output file.",
    )
    parser.add_argument(
        "--repo-id",
        default="SandboxAQ/aqcat25",
        help="The Hugging Face repository ID.",
    )
    args = parser.parse_args()

    query = AQcatQuery(data_root=args.data_root, repo_id=args.repo_id)

    query.filter(
        split=args.split,
        adsorbates=args.adsorbates,
        min_energy=args.min_energy,
        max_energy=args.max_energy,
        magnetism=args.magnetism,
    )

    query.extract_and_save(
        output_file=args.output_file,
        contains_elements=args.contains_elements,
        element_filter_mode=args.element_filter_mode,
        material_type=args.material_type,
        num_workers=args.num_workers,
        chunk_size=args.chunk_size,
    )