Rename subcommand_dict → action_dict and SubcommandProto → ActionProto

qmp/__main__.py

@@ -9,7 +9,7 @@
 import omegaconf
 
 from .utility.context import RuntimeContext, DACITE_CAST
-from .utility.subcommand_dict import subcommand_dict
+from .utility.action_dict import action_dict
 
 
 @hydra.main(version_base=None, config_path=str(pathlib.Path().resolve()), config_name="config")
@@ -30,7 +30,7 @@ def main(runtime_config: omegaconf.DictConfig) -> None:
 
     # 2. Instantiate Algorithm
     run = dacite.from_dict(
-        data_class=subcommand_dict[runtime_config.action.name],  # type: ignore[arg-type]
+        data_class=action_dict[runtime_config.action.name],  # type: ignore[arg-type]
         data=omegaconf.OmegaConf.to_container(runtime_config.action.params, resolve=True),  # type: ignore[arg-type]
         config=dacite.Config(cast=DACITE_CAST),
     )

qmp/algorithms/chop_imag.py

@@ -8,7 +8,7 @@
 import omegaconf
 import torch.utils.tensorboard
 from ..utility.context import RuntimeContext
-from ..utility.subcommand_dict import subcommand_dict
+from ..utility.action_dict import action_dict
 
 
 @dataclasses.dataclass
@@ -98,4 +98,4 @@ def main(
         context.save(data, 0)
 
 
-subcommand_dict["chop_imag"] = ChopImagConfig
+action_dict["chop_imag"] = ChopImagConfig

qmp/algorithms/guide.py

@@ -9,7 +9,7 @@
 import torch
 import torch.utils.tensorboard
 from ..utility.context import RuntimeContext
-from ..utility.subcommand_dict import subcommand_dict
+from ..utility.action_dict import action_dict
 
 
 @dataclasses.dataclass
@@ -188,4 +188,4 @@ def distribution() -> torch.Tensor:
             logging.info("Current optimization cycle completed")
 
 
-subcommand_dict["guide"] = GuideConfig
+action_dict["guide"] = GuideConfig

qmp/algorithms/haar.py

@@ -13,7 +13,7 @@
 import torch.utils.tensorboard
 from ..utility import losses
 from ..utility.context import RuntimeContext
-from ..utility.subcommand_dict import subcommand_dict
+from ..utility.action_dict import action_dict
 from ..utility.model_dict import ModelProto
 from ..utility.optimizer import scale_learning_rate
 
@@ -392,7 +392,7 @@ def main(
         )
 
         if "haar" not in data and "imag" in data:
-            logging.warning("The 'imag' subcommand is deprecated, please use 'haar' instead.")
+            logging.warning("The 'imag' action is deprecated, please use 'haar' instead.")
             data["haar"] = data["imag"]
             del data["imag"]
         if "haar" not in data:
@@ -587,5 +587,5 @@ def closure() -> torch.Tensor:
             logging.info("Current optimization cycle completed")
 
 
-subcommand_dict["haar"] = HaarConfig
-subcommand_dict["imag"] = HaarConfig
+action_dict["haar"] = HaarConfig
+action_dict["imag"] = HaarConfig

qmp/algorithms/haar_with_orbit.py

@@ -13,7 +13,7 @@
 from scipy.optimize import linear_sum_assignment
 from ..utility import losses
 from ..utility.context import RuntimeContext
-from ..utility.subcommand_dict import subcommand_dict
+from ..utility.action_dict import action_dict
 from ..utility.model_dict import ModelProto
 from ..utility.optimizer import scale_learning_rate
 from ..hamiltonian import Hamiltonian
@@ -211,7 +211,7 @@ def main(
         )
 
         if "haar" not in data and "imag" in data:
-            logging.warning("The 'imag' subcommand is deprecated, please use 'haar' instead.")
+            logging.warning("The 'imag' action is deprecated, please use 'haar' instead.")
             data["haar"] = data["imag"]
             del data["imag"]
         if "haar" not in data:
@@ -412,4 +412,4 @@ def closure() -> torch.Tensor:
             logging.info("Checkpoint successfully saved")
 
 
-subcommand_dict["haar_with_orbit"] = HaarWithOrbitConfig
+action_dict["haar_with_orbit"] = HaarWithOrbitConfig

qmp/algorithms/orbit.py

@@ -12,7 +12,7 @@
 import omegaconf
 from ..hamiltonian import Hamiltonian
 from ..utility.context import RuntimeContext
-from ..utility.subcommand_dict import subcommand_dict
+from ..utility.action_dict import action_dict
 
 
 @dataclasses.dataclass
@@ -202,4 +202,4 @@ def calculate_rdm(self, configs: torch.Tensor, psi: torch.Tensor) -> torch.Tenso
         return rdm
 
 
-subcommand_dict["orbit"] = OrbitConfig
+action_dict["orbit"] = OrbitConfig

qmp/algorithms/pert.py

@@ -7,7 +7,7 @@
 import dataclasses
 import omegaconf
 from ..utility.context import RuntimeContext
-from ..utility.subcommand_dict import subcommand_dict
+from ..utility.action_dict import action_dict
 
 
 @dataclasses.dataclass
@@ -64,4 +64,4 @@ def main(
         logging.info("Error is reduced from %.8f to %.8f", energy0 - model.ref_energy, energy2 - model.ref_energy)
 
 
-subcommand_dict["pert"] = PerturbationConfig
+action_dict["pert"] = PerturbationConfig

qmp/algorithms/pretrain.py

@@ -9,7 +9,7 @@
 import torch
 from ..utility import losses
 from ..utility.context import RuntimeContext
-from ..utility.subcommand_dict import subcommand_dict
+from ..utility.action_dict import action_dict
 
 
 @dataclasses.dataclass
@@ -74,4 +74,4 @@ def closure() -> torch.Tensor:
             logging.info("Current optimization cycle completed")
 
 
-subcommand_dict["pretrain"] = PretrainConfig
+action_dict["pretrain"] = PretrainConfig

qmp/algorithms/vmc.py

@@ -9,7 +9,7 @@
 import torch
 import torch.utils.tensorboard
 from ..utility.context import RuntimeContext
-from ..utility.subcommand_dict import subcommand_dict
+from ..utility.action_dict import action_dict
 
 
 @dataclasses.dataclass
@@ -117,4 +117,4 @@ def closure() -> torch.Tensor:
             logging.info("Current optimization cycle completed")
 
 
-subcommand_dict["vmc"] = VmcConfig
+action_dict["vmc"] = VmcConfig

qmp/utility/action_dict.py

@@ -0,0 +1,32 @@
+"""
+This module is used to store a dictionary that maps action names to their corresponding dataclass types.
+
+Other packages or subpackages can register their actions by adding entries to this dictionary, such as
+```
+from qmp.utility.action_dict import action_dict
+action_dict["my_action"] = MyAction
+```
+"""
+
+import typing
+import omegaconf
+from .context import RuntimeContext
+
+
+class ActionProto(typing.Protocol):
+    """
+    This protocol defines a dataclass with a `main` method, which will be called when the action is executed.
+    """
+
+    def main(
+        self,
+        context: RuntimeContext,
+        runtime_config: omegaconf.DictConfig,
+        checkpoint_data: dict[str, typing.Any],
+    ) -> None:
+        """
+        The main method to be called when the action is executed.
+        """
+
+
+action_dict: dict[str, type[ActionProto]] = {}