diff --git a/qmp/__main__.py b/qmp/__main__.py index 29d7913..ee64a78 100644 --- a/qmp/__main__.py +++ b/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), ) diff --git a/qmp/algorithms/chop_imag.py b/qmp/algorithms/chop_imag.py index 6f807bb..5059e56 100644 --- a/qmp/algorithms/chop_imag.py +++ b/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 diff --git a/qmp/algorithms/guide.py b/qmp/algorithms/guide.py index d8762bb..8610a00 100644 --- a/qmp/algorithms/guide.py +++ b/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 diff --git a/qmp/algorithms/haar.py b/qmp/algorithms/haar.py index fab0a63..31e7833 100644 --- a/qmp/algorithms/haar.py +++ b/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 diff --git a/qmp/algorithms/haar_with_orbit.py b/qmp/algorithms/haar_with_orbit.py index c0c8611..fb68d76 100644 --- a/qmp/algorithms/haar_with_orbit.py +++ b/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 diff --git a/qmp/algorithms/orbit.py b/qmp/algorithms/orbit.py index 61ddcf7..dddeb12 100644 --- a/qmp/algorithms/orbit.py +++ b/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 diff --git a/qmp/algorithms/pert.py b/qmp/algorithms/pert.py index e93e0e1..a48c04a 100644 --- a/qmp/algorithms/pert.py +++ b/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 diff --git a/qmp/algorithms/pretrain.py b/qmp/algorithms/pretrain.py index 864ec30..8e099e0 100644 --- a/qmp/algorithms/pretrain.py +++ b/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 diff --git a/qmp/algorithms/vmc.py b/qmp/algorithms/vmc.py index 1c19c83..89b5eba 100644 --- a/qmp/algorithms/vmc.py +++ b/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 diff --git a/qmp/utility/action_dict.py b/qmp/utility/action_dict.py new file mode 100644 index 0000000..86e5146 --- /dev/null +++ b/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]] = {} diff --git a/qmp/utility/subcommand_dict.py b/qmp/utility/subcommand_dict.py deleted file mode 100644 index 7a7c72d..0000000 --- a/qmp/utility/subcommand_dict.py +++ /dev/null @@ -1,32 +0,0 @@ -""" -This module is used to store a dictionary that maps subcommand names to their corresponding dataclass types. - -Other packages or subpackages can register their subcommands by adding entries to this dictionary, such as -``` -from qmp.utility.subcommand_dict import subcommand_dict -subcommand_dict["my_subcommand"] = MySubcommand -``` -""" - -import typing -import omegaconf -from .context import RuntimeContext - - -class SubcommandProto(typing.Protocol): - """ - This protocol defines a dataclass with a `main` method, which will be called when the subcommand 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 subcommand is executed. - """ - - -subcommand_dict: dict[str, type[SubcommandProto]] = {}