Rename imag to haar to keep consistency with the paper.

PR: USTC-KnowledgeComputingLab/qmb#64
Signed-off-by: Hao Zhang <[email protected]>

Author: hzhangxyz

README.md

@@ -67,7 +67,7 @@ qmb --help
 
 This command provides a collection of subcommands, such as `imag`.
 To access detailed help for a specific subcommand, users can append `--help` to the command.
-For example, use `qmb imag --help` to view the help information for the `imag` subcommand.
+For example, use `qmb haar --help` to view the help information for the `imag` subcommand.
 
 Typically, `qmb` requires a specific descriptor for a particular physical or chemical model to execute.
 We have collected a set of such models [here][models-url].
@@ -77,7 +77,7 @@ Alternatively, users can specify a custom path by setting the `$QMB_MODEL_PATH`
 
 After cloning or downloading the dataset, users can calculate the ground state of the $N_2$ system by running the command:
 ```
-qmb imag openfermion mlp -PN2
+qmb haar openfermion mlp -PN2
 ```
 This command utilizes the `imag` subcommand with the descriptor in OpenFermion format and the [mlp network][naqs-url],
 It specifies the $N_2$ model via the `-PN2` flag since the $N_2$ model is loaded from the file `N2.hdf5` in the folder `models`.

qmb/__main__.py

@@ -10,7 +10,7 @@
 from . import hubbard as _  # type: ignore[no-redef]
 from . import ising as _  # type: ignore[no-redef]
 from . import vmc as _  # type: ignore[no-redef]
-from . import imag as _  # type: ignore[no-redef]
+from . import haar as _  # type: ignore[no-redef]
 from . import rldiag as _  # type: ignore[no-redef]
 from . import precompile as _  # type: ignore[no-redef]
 from . import list_loss as _  # type: ignore[no-redef]

qmb/imag.py renamed to qmb/haar.py

@@ -247,7 +247,7 @@ def _merge_pool_from_neural_network_and_pool_from_last_iteration(
 
 
 @dataclasses.dataclass
-class ImaginaryConfig:
+class HaarConfig:
     """
     The two-step optimization process for solving quantum many-body problems based on imaginary time.
     """
@@ -354,11 +354,15 @@ def main(self, *, model_param: typing.Any = None, network_param: typing.Any = No
             state_dict=data.get("optimizer"),
         )
 
-        if "imag" not in data:
-            data["imag"] = {"global": 0, "local": 0, "lanczos": 0, "pool": None}
+        if "haar" not in data and "imag" in data:
+            logging.warning("The 'imag' subcommand is deprecated, please use 'haar' instead.")
+            data["haar"] = data["imag"]
+            del data["imag"]
+        if "haar" not in data:
+            data["haar"] = {"global": 0, "local": 0, "lanczos": 0, "pool": None}
         else:
-            pool_configs, pool_psi = data["imag"]["pool"]
-            data["imag"]["pool"] = (pool_configs.to(device=self.common.device), pool_psi.to(device=self.common.device))
+            pool_configs, pool_psi = data["haar"]["pool"]
+            data["haar"]["pool"] = (pool_configs.to(device=self.common.device), pool_psi.to(device=self.common.device))
 
         writer = torch.utils.tensorboard.SummaryWriter(log_dir=self.common.folder())  # type: ignore[no-untyped-call]
 
@@ -368,7 +372,7 @@ def main(self, *, model_param: typing.Any = None, network_param: typing.Any = No
             logging.info("Sampling configurations from neural network")
             configs_from_neural_network, psi_from_neural_network, _, _ = network.generate_unique(self.sampling_count_from_neural_network, self.local_batch_count_generation)
             logging.info("Sampling configurations from last iteration")
-            configs_from_last_iteration, psi_from_last_iteration = _sampling_from_last_iteration(data["imag"]["pool"], self.sampling_count_from_last_iteration)
+            configs_from_last_iteration, psi_from_last_iteration = _sampling_from_last_iteration(data["haar"]["pool"], self.sampling_count_from_last_iteration)
             logging.info("Merging configurations from neural network and last iteration")
             configs, original_psi = _merge_pool_from_neural_network_and_pool_from_last_iteration(
                 configs_from_neural_network,
@@ -392,9 +396,9 @@ def main(self, *, model_param: typing.Any = None, network_param: typing.Any = No
                     first_extend=self.krylov_extend_first,
             ).run():
                 logging.info("The current energy is %.10f where the sampling count is %d", target_energy.item(), len(configs))
-                writer.add_scalar("imag/lanczos/energy", target_energy, data["imag"]["lanczos"])  # type: ignore[no-untyped-call]
-                writer.add_scalar("imag/lanczos/error", target_energy - model.ref_energy, data["imag"]["lanczos"])  # type: ignore[no-untyped-call]
-                data["imag"]["lanczos"] += 1
+                writer.add_scalar("haar/lanczos/energy", target_energy, data["haar"]["lanczos"])  # type: ignore[no-untyped-call]
+                writer.add_scalar("haar/lanczos/error", target_energy - model.ref_energy, data["haar"]["lanczos"])  # type: ignore[no-untyped-call]
+                data["haar"]["lanczos"] += 1
             max_index = original_psi.abs().argmax()
             target_psi = original_psi / original_psi[max_index]
             logging.info("Local optimization target calculated, the target energy is %.10f, the sampling count is %d", target_energy.item(), len(configs))
@@ -449,12 +453,12 @@ def closure() -> torch.Tensor:
                 logging.info("Starting local optimization process")
                 success = True
                 last_loss: float = 0.0
-                local_step: int = data["imag"]["local"]
+                local_step: int = data["haar"]["local"]
                 scale_learning_rate(optimizer, 1 / (1 << try_index))
                 for i in range(self.local_step):
                     loss = optimizer.step(closure)  # type: ignore[assignment,arg-type]
                     logging.info("Local optimization in progress, step %d, current loss: %.10f", i, loss.item())
-                    writer.add_scalar(f"imag/loss/{self.loss_name}", loss, local_step)  # type: ignore[no-untyped-call]
+                    writer.add_scalar(f"haar/loss/{self.loss_name}", loss, local_step)  # type: ignore[no-untyped-call]
                     local_step += 1
                     if torch.isnan(loss) or torch.isinf(loss):
                         logging.warning("Loss is NaN, restoring the previous state and exiting the optimization loop")
@@ -474,7 +478,7 @@ def closure() -> torch.Tensor:
                         success = False
                 if success:
                     logging.info("Local optimization process completed")
-                    data["imag"]["local"] = local_step
+                    data["haar"]["local"] = local_step
                     break
                 network.load_state_dict(state_backup)
                 optimizer.load_state_dict(optimizer_backup)
@@ -493,29 +497,44 @@ def closure() -> torch.Tensor:
                 model.ref_energy,
                 final_energy.item() - model.ref_energy,
             )
-            writer.add_scalar("imag/energy/state", final_energy, data["imag"]["global"])  # type: ignore[no-untyped-call]
-            writer.add_scalar("imag/energy/target", target_energy, data["imag"]["global"])  # type: ignore[no-untyped-call]
-            writer.add_scalar("imag/error/state", final_energy - model.ref_energy, data["imag"]["global"])  # type: ignore[no-untyped-call]
-            writer.add_scalar("imag/error/target", target_energy - model.ref_energy, data["imag"]["global"])  # type: ignore[no-untyped-call]
+            writer.add_scalar("haar/energy/state", final_energy, data["haar"]["global"])  # type: ignore[no-untyped-call]
+            writer.add_scalar("haar/energy/target", target_energy, data["haar"]["global"])  # type: ignore[no-untyped-call]
+            writer.add_scalar("haar/error/state", final_energy - model.ref_energy, data["haar"]["global"])  # type: ignore[no-untyped-call]
+            writer.add_scalar("haar/error/target", target_energy - model.ref_energy, data["haar"]["global"])  # type: ignore[no-untyped-call]
             logging.info("Displaying the largest amplitudes")
             indices = target_psi.abs().argsort(descending=True)
             text = []
             for index in indices[:self.logging_psi]:
                 this_config = model.show_config(configs[index])
                 logging.info("Configuration: %s, Target amplitude: %s, Final amplitude: %s", this_config, f"{target_psi[index].item():.8f}", f"{psi[index].item():.8f}")
                 text.append(f"Configuration: {this_config}, Target amplitude: {target_psi[index].item():.8f}, Final amplitude: {psi[index].item():.8f}")
-            writer.add_text("config", "\n".join(text), data["imag"]["global"])  # type: ignore[no-untyped-call]
+            writer.add_text("config", "\n".join(text), data["haar"]["global"])  # type: ignore[no-untyped-call]
             writer.flush()  # type: ignore[no-untyped-call]
 
             logging.info("Saving model checkpoint")
-            data["imag"]["pool"] = (configs, original_psi)
-            data["imag"]["global"] += 1
+            data["haar"]["pool"] = (configs, original_psi)
+            data["haar"]["global"] += 1
             data["network"] = network.state_dict()
             data["optimizer"] = optimizer.state_dict()
-            self.common.save(data, data["imag"]["global"])
+            self.common.save(data, data["haar"]["global"])
             logging.info("Checkpoint successfully saved")
 
             logging.info("Current optimization cycle completed")
 
 
-subcommand_dict["imag"] = ImaginaryConfig
+subcommand_dict["haar"] = HaarConfig
+
+
+class ImagConfig(HaarConfig):
+    """
+    Deprecated, use "haar" instead.
+    """
+
+    # pylint: disable=too-few-public-methods
+
+    def __post_init__(self) -> None:
+        logging.warning("The 'imag' subcommand is deprecated, please use 'haar' instead.")
+        super().__post_init__()
+
+
+subcommand_dict["imag"] = ImagConfig