From 8f60ddf7aa9ffb787d2ce0ff56e4b5be83e63aa3 Mon Sep 17 00:00:00 2001
From: aychachouchene <aycha.chouchene@etudiant-enit.utm.tn>
Date: Wed, 8 Jul 2026 11:51:14 +0200
Subject: [PATCH] Address review comments for AseDataLoader feat(data): Add
 AseDataLoader utility with lazy loading

---
 src/tensorial/gcnn/data/__init__.py |  7 ++-
 src/tensorial/gcnn/data/_ase.py     | 88 +++++++++++++++++++++++++++++
 2 files changed, 93 insertions(+), 2 deletions(-)
 create mode 100644 src/tensorial/gcnn/data/_ase.py

diff --git a/src/tensorial/gcnn/data/__init__.py b/src/tensorial/gcnn/data/__init__.py
index cddf8eb..b4f2ccf 100644
--- a/src/tensorial/gcnn/data/__init__.py
+++ b/src/tensorial/gcnn/data/__init__.py
@@ -1,7 +1,10 @@
-from . import _batching, _common, _dataloader, _datamodule
+from . import _ase, _batching, _common, _dataloader, _datamodule
+from ._ase import *
 from ._batching import *
 from ._common import *
 from ._dataloader import *
 from ._datamodule import *
 
-__all__ = _batching.__all__ + _common.__all__ + _datamodule.__all__ + _dataloader.__all__
+__all__ = (
+    _ase.__all__ + _batching.__all__ + _common.__all__ + _datamodule.__all__ + _dataloader.__all__
+)
diff --git a/src/tensorial/gcnn/data/_ase.py b/src/tensorial/gcnn/data/_ase.py
new file mode 100644
index 0000000..93a5b99
--- /dev/null
+++ b/src/tensorial/gcnn/data/_ase.py
@@ -0,0 +1,88 @@
+"""Module for loading ase.Atoms objects as graphs"""
+
+from collections.abc import Sequence
+import importlib.util
+import sys
+from typing import Any, Final
+
+import jraph
+
+from .. import atomic
+
+__all__ = ("AseDataLoader",)
+
+
+def lazy_import(name: str):
+    """Lazily import a module using the standard library."""
+    spec = importlib.util.find_spec(name)
+    if spec is None:
+        # Optimization: if it's completely missing, we can fail early
+        # or return a dummy object.
+        pass
+    loader = importlib.util.LazyLoader(spec.loader)
+    module = importlib.util.module_from_spec(spec)
+    spec.loader = loader
+    sys.modules[name] = module
+    return module
+
+
+# Lazy top-level imports: real load is deferred until first attribute access,
+# but the dependency is still visible here at the top of the module.
+ase = lazy_import("ase")
+ase_io = lazy_import("ase.io")
+
+
+class AseDataLoader(Sequence[jraph.GraphsTuple]):
+    def __init__(
+        self,
+        path: str | Sequence[str],
+        limit: int | None = None,
+        read_kwargs: dict[str, Any] | None = None,
+        as_graphs: dict[str, Any] | None = None,
+    ):
+        # Params
+        self._filepath: Final[tuple[str]] = (path,) if isinstance(path, str) else tuple(path)
+        self._limit: Final[int | None] = limit
+        self._to_graphs: Final[dict[str, Any]] = as_graphs
+        self._read_kwargs: Final[dict[str, Any]] = self._init_kwargs(limit, read_kwargs)
+
+        try:
+            loaded: list["ase.Atoms"] = []
+            for entry in self._filepath:
+                loaded.extend(ase_io.read(entry, **self._read_kwargs))
+            self._data: list["ase.Atoms" | jraph.GraphsTuple] = (
+                [loaded] if isinstance(loaded, ase.Atoms) else loaded
+            )
+        except FileNotFoundError:
+            raise ValueError(
+                f"Could not load ASE structures, the passed path does not exist: {path}"
+            ) from None
+
+        if self._to_graphs and len(self) > 0:
+            # Check that the parameters they passed are OK by requesting the first structure
+            # to be converted
+            self[0]  # noqa, pylint: disable=pointless-statement
+
+    def __len__(self) -> int:
+        return len(self._data)
+
+    def __getitem__(self, item: int) -> Any:
+        entry = self._data[item]
+        if self._to_graphs and not isinstance(entry, jraph.GraphsTuple):
+            # Lazily convert the first time
+            entry = atomic.graph_from_ase(entry, **self._to_graphs)
+            self._data[item] = entry
+        return entry
+
+    @staticmethod
+    def _init_kwargs(
+        limit: int | None, read_kwargs: dict[str, Any] | None = None
+    ) -> dict[str, Any]:
+        if read_kwargs is None:
+            read_kwargs = {"index": f":{limit}" if limit is not None else ":"}
+        else:
+            if limit is not None:
+                read_kwargs["index"] = f":{limit}"
+            elif "index" not in read_kwargs:
+                read_kwargs["index"] = ":"
+        return read_kwargs