From da7b661451675f5de82a1932c848c67db5ecc7a8 Mon Sep 17 00:00:00 2001 From: Hao Zhang Date: Thu, 17 Jul 2025 23:24:10 +0800 Subject: [PATCH 1/2] Rename imag to haar to keep consistency with the paper. --- qmb/__main__.py | 2 +- qmb/{imag.py => haar.py} | 61 ++++++++++++++++++++++++++-------------- 2 files changed, 41 insertions(+), 22 deletions(-) rename qmb/{imag.py => haar.py} (91%) diff --git a/qmb/__main__.py b/qmb/__main__.py index fd45fe3..094a028 100644 --- a/qmb/__main__.py +++ b/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] diff --git a/qmb/imag.py b/qmb/haar.py similarity index 91% rename from qmb/imag.py rename to qmb/haar.py index 77eb232..4e517ef 100644 --- a/qmb/imag.py +++ b/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,10 +497,10 @@ 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 = [] @@ -504,18 +508,33 @@ def closure() -> torch.Tensor: 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 From 90f7aebdd04427056d9cb77b61c02486f192790a Mon Sep 17 00:00:00 2001 From: Hao Zhang Date: Tue, 22 Jul 2025 15:09:42 +0800 Subject: [PATCH 2/2] Update subcommands in README.md. --- README.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 31b7dc6..0fa83ed 100644 --- a/README.md +++ b/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`.