use DVC Filesystem to open files (#120)

* use DVC Filesystem to open files * use state.fs * wrap in try..except * remove unused import
zincware · Jun 12, 2023 · 03d7c9c · 03d7c9c
1 parent feca534
commit 03d7c9c
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 6 deletions.
diff --git a/ipsuite/base/base.py b/ipsuite/base/base.py
@@ -1,5 +1,6 @@
 import abc
 import collections.abc
+import pathlib
 import typing
 
 import ase
@@ -51,7 +52,12 @@ def get_data(self) -> list[ase.Atoms]:
         if self.data is not None:
             return self.data
         elif self.data_file is not None:
-            return list(ase.io.iread(self.data_file))
+            try:
+                with self.state.fs.open(pathlib.Path(self.data_file).as_posix()) as f:
+                    return list(ase.io.iread(f))
+            except FileNotFoundError:
+                # File can not be opened with DVCFileSystem, try normal open
+                return list(ase.io.iread(self.data_file))
         else:
             raise ValueError("No data given.")
 
@@ -101,7 +107,12 @@ def get_data(self) -> ase.Atoms:
             else:
                 atoms = self.data.copy()
         elif self.data_file is not None:
-            atoms = list(ase.io.iread(self.data_file))[self.data_id]
+            try:
+                with self.state.fs.open(pathlib.Path(self.data_file).as_posix()) as f:
+                    atoms = list(ase.io.iread(f))[self.data_id]
+            except FileNotFoundError:
+                # File can not be opened with DVCFileSystem, try normal open
+                atoms = list(ase.io.iread(self.data_file))[self.data_id]
         else:
             raise ValueError("No data given.")
         return atoms

diff --git a/ipsuite/calculators/apax_jax_md.py b/ipsuite/calculators/apax_jax_md.py
@@ -91,6 +91,6 @@ def run(self):
 
     @functools.cached_property
     def atoms(self) -> typing.List[ase.Atoms]:
-        return list(
-            ase.io.iread(self.sim_dir / "md.traj")
-        )  # filename should be changeable
+        # filename should be changeable
+        with self.state.fs.open((self.sim_dir / "md.traj").as_posix()) as f:
+            return list(ase.io.iread(f))
diff --git a/ipsuite/calculators/ase_geoopt.py b/ipsuite/calculators/ase_geoopt.py
@@ -43,4 +43,5 @@ def run(self):
 
     @property
     def atoms(self):
-        return list(ase.io.iread(self.traj.as_posix()))
+        with self.state.fs.open(self.traj.as_posix()) as f:
+            return list(ase.io.iread(f))