Migrate MLFF geometry optimization files to JSON Lines format for fast partial loading

[janosh]

does not fix #230 (yet). the symmetry analysis rerun for all models is pushed back to a subsequent PR. the current lineup of models was tested with mixture of spglib==2.5.0 and several moyopy versions so a rerun would have been in order anyway for consistency. the code to calculate RMSD metric is fixed in this PR but the existing metrics are not updated. my machine keeps crashing trying to rerun all models. may need to rerun models one by one over night.

this PR makes debugging and iterating on future geometry optimization much faster by allowing to load just ~100 DFT/model-relaxed structures instead of all 257k in WBM on test runs

Unpolished conversion script for ML-relaxed structures from JSON to JSON Lines

"""
Convert geometry optimization JSON files to JSON lines format for better debug performance.

JSON lines format allows for reading specific number of lines without loading the entire file
into memory, which is especially useful for large structure files when only a subset is needed.

Example usage:
    python scripts/convert_to_jsonl.py  # Convert all model geo-opt files
    python scripts/convert_to_jsonl.py --models mace_mp_0 m3gnet  # Convert specific models
    python scripts/convert_to_jsonl.py --force  # Overwrite existing output files
"""

import argparse
import glob
import os
import sys

import pandas as pd
from tqdm import tqdm

from matbench_discovery.enums import Model


def convert_to_jsonl(
    input_file: str,
    output_file: str | None = None,
    suffix: str = ".jsonl.gz",
    force: bool = False,
) -> bool:
    """Convert a JSON file to JSON lines format.

    Args:
        input_file: Path to input JSON file
        output_file: Path to output JSON lines file (default: input_file with .jsonl.gz extension)
        suffix: Suffix to use for output file if output_file is not specified
        force: Whether to overwrite existing output file

    Returns:
        bool: True if the file was processed, False if it was skipped
    """
    if not os.path.exists(input_file):
        print(f"Error: Input file {input_file} does not exist")
        return False

    try:
        print(f"Reading {input_file}...")
        df = pd.read_json(input_file)

        # Get row count for filename
        row_count = len(df)

        # Check if index is already a material ID
        if df.index.name is None or df.index.name != "material_id":
            # If not indexed by material_id, try to find material_id column
            if "material_id" in df.columns:
                df = df.set_index("material_id")
            elif df.index.dtype == "object" and df.index[0].startswith("wbm-"):
                # Index already contains material IDs but not named
                df.index.name = "material_id"

        # Ensure index name is material_id for proper loading later
        if df.index.name is None:
            df.index.name = "material_id"

        if output_file is None:
            # By default, replace extension with the line count and .jsonl.gz
            basename = os.path.splitext(input_file)[0]
            basename = basename.removesuffix(".json")  # Remove .json from .json.gz
            output_file = f"{basename}-n={row_count}{suffix}"

        # Silently skip if file exists and not forcing overwrite
        if os.path.exists(output_file) and not force:
            return False

        print(
            f"Writing {row_count} structures to {output_file} in JSON lines format..."
        )
        df.reset_index().to_json(output_file, orient="records", lines=True)
        print(f"Successfully converted {input_file} to {output_file}")

        # Print info on how to verify the conversion
        print("\nTo verify conversion:")
        print(
            f"  python -c \"import pandas as pd; print(pd.read_json('{output_file}', lines=True).head())\""
        )
        return True

    except Exception as exc:
        print(f"Error converting {input_file}: {exc}")
        return False


def main():
    parser = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    parser.add_argument(
        "--models",
        nargs="*",
        type=str,
        help="Specific models to convert (default: all models with geo_opt_path)",
    )
    parser.add_argument(
        "--suffix",
        default=".jsonl.gz",
        help="Suffix for output file (default: .jsonl.gz)",
    )
    parser.add_argument(
        "--force", action="store_true", help="Overwrite existing output files"
    )
    args = parser.parse_args()

    # Get list of models to process
    models_to_process = []
    if args.models:
        # Convert model names to actual Model enum values
        for model_name in args.models:
            try:
                model = getattr(Model, model_name)
                models_to_process.append(model)
            except AttributeError:
                print(f"Warning: Unknown model '{model_name}', skipping")
    else:
        # Process all models with geo_opt_path
        models_to_process = list(Model)

    # Filter to models with geometry optimization data
    geo_opt_models = []
    for model in models_to_process:
        try:
            if model.geo_opt_path:  # Returns None if model doesn't have geo_opt
                geo_opt_models.append(model)
        except ValueError as exc:
            print(f"Warning: {model.label} - {exc}")
        except Exception as exc:
            print(f"Error with {model.label}: {exc}")

    if not geo_opt_models:
        print("No models with geometry optimization data found.")
        sys.exit(1)

    print(f"Found {len(geo_opt_models)} models with geo_opt paths")

    # Process each model's geo_opt file
    processed_count = 0
    skipped_count = 0
    already_exists_count = 0

    for model in tqdm(geo_opt_models, desc="Converting files"):
        try:
            input_file = model.geo_opt_path
            if input_file and os.path.exists(input_file):
                # Check if output file would exist without conversion attempt
                basename = os.path.splitext(input_file)[0]
                basename = basename.removesuffix(".json")  # Remove .json from .json.gz

                # Check if any file matching this pattern exists
                potential_output_files = glob.glob(f"{basename}-n=*{args.suffix}")
                if potential_output_files and not args.force:
                    already_exists_count += 1
                    continue

                # Process file and track results
                if convert_to_jsonl(input_file, None, args.suffix, args.force):
                    processed_count += 1
                else:
                    already_exists_count += 1
            else:
                print(
                    f"Warning: geo_opt_path for {model.label} doesn't exist: {input_file}"
                )
                skipped_count += 1
        except Exception as exc:
            print(f"Error processing {model.label}: {exc}")
            skipped_count += 1

    print(
        f"Processed {processed_count} file(s), skipped {skipped_count} file(s), {already_exists_count} already existed"
    )


if __name__ == "__main__":
    main()