refactorize ExplorationTaskGroup. provide abs method make_task. provi…

…de ExplorationTaskGroupData as a data holder
wanghan-iapcm · Feb 2, 2024 · cc8a7a2 · cc8a7a2
1 parent ebe7a17
commit cc8a7a2
Show file tree

Hide file tree

Showing 10 changed files with 128 additions and 86 deletions.
diff --git a/dpgen2/exploration/task/__init__.py b/dpgen2/exploration/task/__init__.py
@@ -23,5 +23,8 @@
 )
 from .task import (
     ExplorationTask,
+)
+from .task_group import (
     ExplorationTaskGroup,
+    ExplorationTaskGroupData,
 )
diff --git a/dpgen2/exploration/task/conf_sampling_task_group.py b/dpgen2/exploration/task/conf_sampling_task_group.py
@@ -13,6 +13,8 @@
 
 from .task import (
     ExplorationTask,
+)
+from .task_group import (
     ExplorationTaskGroup,
 )
 

diff --git a/dpgen2/exploration/task/customized_lmp_template_task_group.py b/dpgen2/exploration/task/customized_lmp_template_task_group.py
@@ -37,7 +37,6 @@
 )
 from .task import (
     ExplorationTask,
-    ExplorationTaskGroup,
 )
 
 
@@ -147,7 +146,7 @@ def set_lmp(
 
     def make_task(
         self,
-    ) -> ExplorationTaskGroup:
+    ) -> "CustomizedLmpTemplateTaskGroup":
         if not self.conf_set:
             raise RuntimeError("confs are not set")
         if not self.lmp_set:
@@ -166,7 +165,7 @@ def make_task(
     def _make_customized_task_group(
         self,
         conf,
-    ) -> ExplorationTaskGroup:
+    ) -> "CustomizedLmpTemplateTaskGroup":
         with tempfile.TemporaryDirectory() as tmpdir:
             with set_directory(Path(tmpdir)):
                 Path(self.input_lmp_conf_name).write_text(conf)

diff --git a/dpgen2/exploration/task/lmp_template_task_group.py b/dpgen2/exploration/task/lmp_template_task_group.py
@@ -25,7 +25,6 @@
 )
 from .task import (
     ExplorationTask,
-    ExplorationTaskGroup,
 )
 
 
@@ -60,7 +59,7 @@ def set_lmp(
 
     def make_task(
         self,
-    ) -> ExplorationTaskGroup:
+    ) -> "LmpTemplateTaskGroup":
         if not self.conf_set:
             raise RuntimeError("confs are not set")
         if not self.lmp_set:

diff --git a/dpgen2/exploration/task/npt_task_group.py b/dpgen2/exploration/task/npt_task_group.py
@@ -19,7 +19,6 @@
 )
 from .task import (
     ExplorationTask,
-    ExplorationTaskGroup,
 )
 
 
@@ -76,7 +75,7 @@ def set_md(
 
     def make_task(
         self,
-    ) -> ExplorationTaskGroup:
+    ) -> "NPTTaskGroup":
         """
         Make the LAMMPS task group.
 

diff --git a/dpgen2/exploration/task/stage.py b/dpgen2/exploration/task/stage.py
@@ -14,7 +14,10 @@
 
 from .task import (
     ExplorationTask,
+)
+from .task_group import (
     ExplorationTaskGroup,
+    ExplorationTaskGroupData,
 )
 
 
@@ -65,7 +68,7 @@ def make_task(
 
         """
 
-        lmp_task_grp = ExplorationTaskGroup()
+        lmp_task_grp = ExplorationTaskGroupData()
         for ii in self.explor_groups:
             # lmp_task_grp.add_group(ii.make_task())
             lmp_task_grp += ii.make_task()

diff --git a/dpgen2/exploration/task/task.py b/dpgen2/exploration/task/task.py
@@ -1,8 +1,4 @@
 import os
-from abc import (
-    ABC,
-    abstractmethod,
-)
 from collections.abc import (
     Sequence,
 )
@@ -60,51 +56,6 @@ def files(self) -> Dict:
         return self._files
 
 
-class ExplorationTaskGroup(Sequence):
-    """A group of exploration tasks. Implemented as a `list` of `ExplorationTask`."""
-
-    def __init__(self):
-        super().__init__()
-        self.clear()
-
-    def __getitem__(self, ii: int) -> ExplorationTask:
-        """Get the `ii`th task"""
-        return self.task_list[ii]
-
-    def __len__(self) -> int:
-        """Get the number of tasks in the group"""
-        return len(self.task_list)
-
-    def clear(self) -> None:
-        self._task_list = []
-
-    @property
-    def task_list(self) -> List[ExplorationTask]:
-        """Get the `list` of `ExplorationTask`"""
-        return self._task_list
-
-    def add_task(self, task: ExplorationTask):
-        """Add one task to the group."""
-        self.task_list.append(task)
-        return self
-
-    def add_group(
-        self,
-        group: "ExplorationTaskGroup",
-    ):
-        """Add another group to the group."""
-        # see https://www.python.org/dev/peps/pep-0484/#forward-references for forward references
-        self._task_list = self._task_list + group._task_list
-        return self
-
-    def __add__(
-        self,
-        group: "ExplorationTaskGroup",
-    ):
-        """Add another group to the group."""
-        return self.add_group(group)
-
-
 class FooTask(ExplorationTask):
     def __init__(
         self,
@@ -118,30 +69,3 @@ def __init__(
             conf_name: conf_cont,
             inpu_name: inpu_cont,
         }
-
-
-class FooTaskGroup(ExplorationTaskGroup):
-    def __init__(self, numb_task):
-        super().__init__()
-        # TODO: confirm the following is correct
-        self.tlist = ExplorationTaskGroup()
-        for ii in range(numb_task):
-            self.tlist.add_task(
-                FooTask(
-                    f"conf.{ii}",
-                    f"this is conf.{ii}",
-                    f"input.{ii}",
-                    f"this is input.{ii}",
-                )
-            )
-
-    @property
-    def task_list(self):
-        return self.tlist
-
-
-if __name__ == "__main__":
-    grp = FooTaskGroup(3)
-    for ii in grp:
-        fcs = ii.files()
-        print(fcs)
diff --git a/dpgen2/exploration/task/task_group.py b/dpgen2/exploration/task/task_group.py
@@ -0,0 +1,104 @@
+from abc import (
+    ABC,
+    abstractmethod,
+)
+from collections.abc import (
+    Sequence,
+)
+from typing import (
+    Dict,
+    List,
+    Tuple,
+)
+
+from .task import (
+    ExplorationTask,
+)
+
+
+class ExplorationTaskGroup(Sequence):
+    """A group of exploration tasks. Implemented as a `list` of `ExplorationTask`."""
+
+    def __init__(self):
+        super().__init__()
+        self.clear()
+
+    def __getitem__(self, ii: int) -> ExplorationTask:
+        """Get the `ii`th task"""
+        return self.task_list[ii]
+
+    def __len__(self) -> int:
+        """Get the number of tasks in the group"""
+        return len(self.task_list)
+
+    def clear(self) -> None:
+        self._task_list = []
+
+    @property
+    def task_list(self) -> List[ExplorationTask]:
+        """Get the `list` of `ExplorationTask`"""
+        return self._task_list
+
+    def add_task(self, task: ExplorationTask):
+        """Add one task to the group."""
+        self.task_list.append(task)
+        return self
+
+    def add_group(
+        self,
+        group: "ExplorationTaskGroup",
+    ):
+        """Add another group to the group."""
+        # see https://www.python.org/dev/peps/pep-0484/#forward-references for forward references
+        self._task_list = self._task_list + group._task_list
+        return self
+
+    def __add__(
+        self,
+        group: "ExplorationTaskGroup",
+    ):
+        """Add another group to the group."""
+        return self.add_group(group)
+
+    @abstractmethod
+    def make_task(self) -> "ExplorationTaskGroup":
+        """Make the task group."""
+        pass
+
+
+class ExplorationTaskGroupData(ExplorationTaskGroup):
+    """Data-only exploration task group."""
+
+    def __init__(self):
+        super().__init__()
+
+    def make_task(self):
+        """Make the task group."""
+        raise NotImplementedError("This class is not supposed to supply make_task")
+
+
+class FooTaskGroup(ExplorationTaskGroup):
+    def __init__(self, numb_task):
+        super().__init__()
+        # TODO: confirm the following is correct
+        self.tlist = ExplorationTaskGroup()
+        for ii in range(numb_task):
+            self.tlist.add_task(
+                FooTask(
+                    f"conf.{ii}",
+                    f"this is conf.{ii}",
+                    f"input.{ii}",
+                    f"this is input.{ii}",
+                )
+            )
+
+    @property
+    def task_list(self):
+        return self.tlist
+
+
+if __name__ == "__main__":
+    grp = FooTaskGroup(3)
+    for ii in grp:
+        fcs = ii.files()
+        print(fcs)
diff --git a/tests/mocked_ops.py b/tests/mocked_ops.py
@@ -789,6 +789,9 @@ def __init__(self):
             )
             self.add_task(tt)
 
+    def make_task(self):
+        raise NotImplementedError
+
 
 class MockedExplorationTaskGroup1(ExplorationTaskGroup):
     def __init__(self):
@@ -801,6 +804,9 @@ def __init__(self):
             )
             self.add_task(tt)
 
+    def make_task(self):
+        raise NotImplementedError
+
 
 class MockedExplorationTaskGroup2(ExplorationTaskGroup):
     def __init__(self):
@@ -813,6 +819,9 @@ def __init__(self):
             )
             self.add_task(tt)
 
+    def make_task(self):
+        raise NotImplementedError
+
 
 class MockedStage(ExplorationStage):
     def make_task(self):

diff --git a/tests/test_prep_run_lmp.py b/tests/test_prep_run_lmp.py
@@ -70,7 +70,7 @@
 )
 from dpgen2.exploration.task import (
     ExplorationTask,
-    ExplorationTaskGroup,
+    ExplorationTaskGroupData,
 )
 from dpgen2.op.prep_lmp import (
     PrepLmp,
@@ -90,7 +90,7 @@
 
 
 def make_task_group_list(ngrp, ntask_per_grp):
-    tgrp = ExplorationTaskGroup()
+    tgrp = ExplorationTaskGroupData()
     for ii in range(ngrp):
         for jj in range(ntask_per_grp):
             tt = ExplorationTask()