update analyzer, exporter to bidirectional

pacs/__main__.py

@@ -1,11 +1,12 @@
 from typing import Tuple
 
 from ._version import __version__
-from .mdrun.analyzer.a_d import A_D
-from .mdrun.analyzer.association import Association
-from .mdrun.analyzer.dissociation import Dissociation
-from .mdrun.analyzer.ee import EdgeExpansion
-from .mdrun.analyzer.rmsd import RMSD
+# from .mdrun.analyzer.a_d import A_D
+# from .mdrun.analyzer.association import Association
+# from .mdrun.analyzer.dissociation import Dissociation
+# from .mdrun.analyzer.ee import EdgeExpansion
+# from .mdrun.analyzer.rmsd import RMSD
+from .mdrun.analyzer.bd_rmsd import BD_RMSD
 from .mdrun.analyzer.superAnalyzer import SuperAnalyzer
 from .mdrun.analyzer.target import Target
 from .mdrun.analyzer.template import Template
@@ -40,13 +41,14 @@ def prepare_md(
     }.get(settings.simulator)
 
     analyzer: SuperAnalyzer = {
-        "target": Target(),
-        "dissociation": Dissociation(),
-        "association": Association(),
-        "rmsd": RMSD(),
-        "ee": EdgeExpansion(),
-        "a_d": A_D(),
-        "template": Template(),
+        "bd_rmsd": BD_RMSD(),
+        # "target": Target(),
+        # "dissociation": Dissociation(),
+        # "association": Association(),
+        # "rmsd": RMSD(),
+        # "ee": EdgeExpansion(),
+        # "a_d": A_D(),
+        # "template": Template(),
     }.get(settings.type)
 
     exporter: SuperExporter = {

pacs/mdrun/Cycle.py

@@ -130,6 +130,10 @@ def make_dir(path: Path) -> None:
             self.prepare_next_direction(cycle=cycle, direction=direction, replica=replica)
 
     def prepare_next_direction(self, cycle: int, direction: str, replica: int) -> None:
+        def make_dir(path: Path) -> None:
+            if not path.exists():
+                path.mkdir(parents=True)
+
         next_direction_dir = Path(
             f"{self.settings.each_direction(_cycle=cycle + 1, _direction=direction)}"
         )

pacs/mdrun/analyzer/bd_rmsd.py

@@ -0,0 +1,107 @@
+import subprocess
+from typing import List
+
+import numpy as np
+from pacs.mdrun.analyzer.superAnalyzer import SuperAnalyzer
+from pacs.models.settings import MDsettings, Snapshot, ScoresInCycle, ScoresInOnePair
+from pacs.utils.logger import generate_logger
+
+LOGGER = generate_logger(__name__)
+
+
+class BD_RMSD(SuperAnalyzer):
+    def calculate_scores_in_one_pair(
+        self, settings: MDsettings, cycle: int, fore_replica: int, back_replica: int, send_rev
+    ) -> ScoresInOnePair:
+        if settings.analyzer == "mdtraj":
+            ret = self.cal_by_mdtraj(settings, cycle, fore_replica, back_replica)
+        elif settings.analyzer == "gromacs":
+            raise NotImplementedError
+        elif settings.analyzer == "cpptraj":
+            raise NotImplementedError
+        else:
+            raise NotImplementedError
+
+        LOGGER.info(f"cycle {cycle} replica ({fore_replica}, {back_replica}) calculated.")
+        if send_rev is not None:
+            send_rev.send(ret)
+        return ret
+
+    def aggregate(self, settings: MDsettings, scores_in_cycle: ScoresInCycle, direction: str) -> np.ndarray:
+        aggr_func = np.mean
+
+        if direction == "fore":
+            return aggr_func(scores_in_cycle.cv_data, axis=(1, 3))
+        elif direction == "back":
+            return aggr_func(scores_in_cycle.cv_data, axis=(0, 2))
+        else:
+            LOGGER.error(f"invalid direction: {direction}")
+
+    def ranking(self, settings: MDsettings, CVs: List[Snapshot]) -> List[Snapshot]:
+        sorted_cv = sorted(CVs, key=lambda x: x.cv, reverse=False)
+        return sorted_cv
+
+    def is_threshold(self, settings: MDsettings, CVs: List[Snapshot] = None) -> bool:
+        if CVs is None:
+            CVs = self.CVs
+        return CVs[0].cv > settings.threshold
+
+    def cal_by_mdtraj(self, settings: MDsettings, cycle: int, fore_replica: int, back_replica: int) -> ScoresInOnePair:
+        import mdtraj as md
+
+        fore_dir = settings.each_replica(_cycle=cycle, _direction="fore", _replica=fore_replica)
+        back_dir = settings.each_replica(_cycle=cycle, _direction="back", _replica=back_replica)
+        fore_trj = md.load(
+            f"{fore_dir}/prd{settings.trajectory_extension}",
+            top=settings.top_mdtraj_fore,
+        )
+        back_trj = md.load(
+            f"{back_dir}/prd{settings.trajectory_extension}",
+            top=settings.top_mdtraj_back,
+        )
+
+        # slice (exclude the first frame for gromacs)
+        if settings.simulator == "gromacs":
+            n_frames_fore = fore_trj.n_frames
+            n_frames_back = back_trj.n_frames
+            fore_trj = fore_trj.slice(np.arange(1, n_frames_fore))
+            back_trj = back_trj.slice(np.arange(1, n_frames_back))
+        n_frames_fore = fore_trj.n_frames
+        n_frames_back = back_trj.n_frames
+
+        # fit and compute rmsd
+        rmsd = np.zeros((n_frames_fore, n_frames_back))
+        sel1_arr = fore_trj.top.select(settings.selection1)
+        sel2_arr = fore_trj.top.select(settings.selection2)
+        sel3_arr = back_trj.top.select(settings.selection3)
+        sel4_arr = back_trj.top.select(settings.selection4)
+
+        for frame_i_back in range(n_frames_back):
+            back_trj.superpose(
+                fore_trj,
+                frame_i_back,
+                atom_indices=sel1_arr,
+                ref_atom_indices=sel3_arr,
+            )
+            # もうちょい効率化できるかも
+            for frame_i_fore in range(n_frames_fore):
+                tmp = np.sqrt(
+                    3
+                    * np.mean(
+                        np.square(
+                            fore_trj.xyz[frame_i_fore, sel2_arr]
+                            - back_trj.xyz[frame_i_back, sel4_arr]
+                        ),
+                        axis=(0, 1), # xyzのみで平均を取る
+                    )
+                )
+                rmsd[frame_i_fore, frame_i_back] = tmp
+
+        scores_in_one_pair = ScoresInOnePair(
+            cycle=cycle,
+            fore_replica=fore_replica,
+            back_replica=back_replica,
+            cv_data=rmsd,
+        )
+
+        return scores_in_one_pair