⚡ Bolt: Optimize JSON serialization in aselmdb.py with orjson (#4)

* ⚡ Bolt: Optimize JSON serialization in aselmdb.py with orjson

Replaced ASE's standard `jsonio.encode` and `decode` with `orjson` in `src/lavello_mlips/aselmdb.py` to significantly improve LMDB database write/read speeds. Added fallback logic to handle custom ASE objects (e.g. `__ndarray__`). Also removed unused imports flagged by ruff.

Co-authored-by: alinelena <3306823+alinelena@users.noreply.github.com>

* ⚡ Bolt: Optimize JSON serialization in aselmdb.py with orjson

Replaced ASE's standard `jsonio.encode` and `decode` with `orjson` in `src/lavello_mlips/aselmdb.py` to significantly improve LMDB database write/read speeds. Added fallback logic to handle custom ASE objects (e.g. `__ndarray__`). Also removed unused imports flagged by ruff.

Co-authored-by: alinelena <3306823+alinelena@users.noreply.github.com>

---------

Co-authored-by: google-labs-jules[bot] <161369871+google-labs-jules[bot]@users.noreply.github.com>
Co-authored-by: alinelena <3306823+alinelena@users.noreply.github.com>

Author: alinelena

.jules/bolt.md

@@ -5,3 +5,6 @@
 ## 2024-05-18 - Replacing `json` with `orjson` for large datasets
 **Learning:** In pipelines handling large datasets via dictionaries containing metadata (e.g. millions of prefixes), `json.dump` and `json.load` can become significant bottlenecks, adding seconds or even minutes to startup and checkpointing phases. `orjson` provides a near drop-in replacement that is 4-10x faster for such operations.
 **Action:** When working with large JSON files, especially in a framework requiring frequent disk checkpoints, replace Python's built-in `json` module with `orjson` wrapping `loads`/`dumps` to preserve API compatibility while gaining massive performance boosts.
+## 2024-03-29 - ASE Custom JSON encoding vs standard JSON
+**Learning:** ASE's custom JSON encoder (`ase.io.jsonio.encode`) will generate dicts with special keys like `__ndarray__` or `__complex__` (e.g. `{"__ndarray__": [[5], "int64", ...]}`). When optimizing JSON deserialization using faster alternatives like `orjson`, it's critical to realize that a normal `json.loads` or `orjson.loads` will deserialize this into a Python dictionary, while ASE's custom `decode` will properly reconstruct the underlying numpy array. Bypassing ASE's decoder without checking for these keys leads to downstream type errors (e.g. `KeyError: '__ndarray__'`).
+**Action:** When replacing or wrapping ASE's jsonio with `orjson`, always fall back to ASE's `decode` if the payload string contains `__ndarray__` or `__complex__` markers, to ensure custom objects are correctly reconstructed.

src/lavello_mlips/aselmdb.py

@@ -2,7 +2,7 @@
 import os
 import zlib
 from pathlib import Path
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Union
 
 import ase.db.core
 import ase.db.row
@@ -11,6 +11,7 @@
 from ase.calculators.singlepoint import SinglePointCalculator
 import lmdb
 import numpy as np
+import orjson
 
 logger = logging.getLogger(__name__)
 
@@ -233,17 +234,34 @@ def get_atoms(self, idx: int) -> Atoms:
 def encode_object(obj: Any, compress=True, json_encode=True) -> bytes:
     """Encode object to compressed JSON."""
     if json_encode:
-        obj = encode(obj)
+        try:
+            # OPT_SERIALIZE_NUMPY handles numpy arrays directly
+            obj_bytes = orjson.dumps(obj, option=orjson.OPT_SERIALIZE_NUMPY)
+        except orjson.JSONEncodeError:
+            # Fallback to standard ASE jsonio if orjson fails (e.g. for unsupported complex objects)
+            obj_bytes = encode(obj).encode("utf-8")
+    else:
+        obj_bytes = obj.encode("utf-8") if isinstance(obj, str) else bytes(obj)
+
     if compress:
-        return zlib.compress(obj.encode("utf-8"))
-    return obj.encode("utf-8")
+        return zlib.compress(obj_bytes)
+    return obj_bytes
 
 def decode_bytestream(bytestream: bytes, decompress=True, json_decode=True) -> Any:
     """Decode compressed JSON bytestream."""
     if decompress:
-        bytestream = zlib.decompress(bytestream).decode("utf-8")
-    else:
-        bytestream = bytestream.decode("utf-8")
+        bytestream = zlib.decompress(bytestream)
+
     if json_decode:
-        return decode(bytestream)
-    return bytestream
+        # ASE's custom JSON encoder uses special keys like __ndarray__, __complex__, etc.
+        # If the payload contains these, we must use ASE's decoder to reconstruct the objects.
+        if b'__ndarray__' in bytestream or b'__complex__' in bytestream:
+            return decode(bytestream.decode("utf-8"))
+
+        try:
+            return orjson.loads(bytestream)
+        except orjson.JSONDecodeError:
+            # Fallback to standard ASE jsonio if orjson fails
+            return decode(bytestream.decode("utf-8"))
+
+    return bytestream.decode("utf-8")

src/lavello_mlips/cli.py

@@ -1,7 +1,6 @@
 import argparse
 import logging
 import os
-from typing import Any
 
 os.environ["ASE_MPI"] = "0"
 from pathlib import Path

src/lavello_mlips/distributions.py

@@ -1,5 +1,5 @@
 import argparse
-from pathlib import Path, PosixPath
+from pathlib import PosixPath
 from dataclasses import dataclass, field
 from typing import Optional
 

src/lavello_mlips/process_omol25.py

@@ -4,17 +4,14 @@
 import os
 import re
 import signal
-import sys
 import time
 from io import BytesIO, StringIO
 from pathlib import Path
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any, Dict, List, Optional, Tuple
 
-import boto3
 import numpy as np
 import pandas as pd
 from ase.io import read, write
-from botocore.config import Config
 from mpi4py import MPI
 from tarfile import open as tar_open
 from tqdm import tqdm

src/lavello_mlips/verify_processed_omol25.py

@@ -13,7 +13,6 @@
 
 import pandas as pd
 import numpy as np
-from typing import Any, Optional, Union, Dict, List
 from ase.io import read
 
 from .utils import setup_logging