diff --git a/pyproject.toml b/pyproject.toml index 58e0ecd..8b5b813 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -16,6 +16,7 @@ dependencies = [ "openfermion>=1.7.1", "tensorboard>=2.20.0", "dacite>=1.9.2", + "pyscf>=2.12.1", ] requires-python = ">=3.12" authors = [{ name = "Hao Zhang", email = "hzhangxyz@outlook.com" }] diff --git a/qmp/plugins/pyscf.py b/qmp/plugins/pyscf.py index fe1726b..0bde07b 100644 --- a/qmp/plugins/pyscf.py +++ b/qmp/plugins/pyscf.py @@ -450,7 +450,7 @@ def kernel( # ------------------------------------------------------------------ # 11. Extract energy and build ci output # ------------------------------------------------------------------ - energy = _extract_energy(data, self.config.action.name) + energy = _extract_energy(data, self.config.action.name, model) # Strip the random engine state — it is stale by the next call. full_checkpoint: dict[str, typing.Any] = {k: v for k, v in data.items() if k != "random"} @@ -469,7 +469,7 @@ def kernel( return energy, ci -def _extract_energy(data: dict[str, typing.Any], action_name: str) -> float: +def _extract_energy(data: dict[str, typing.Any], action_name: str, model: ModelProto | None = None) -> float: """ Extract the final ground-state energy from algorithm checkpoint data. @@ -479,6 +479,8 @@ def _extract_energy(data: dict[str, typing.Any], action_name: str) -> float: Checkpoint data loaded from the temporary directory. action_name : str Name of the algorithm that produced the checkpoint. + model : ModelProto, optional + The model instance for computing energy expectation value if needed. Returns ------- @@ -490,15 +492,36 @@ def _extract_energy(data: dict[str, typing.Any], action_name: str) -> float: haar_data = data.get("haar") if haar_data is None: return 0.0 - haar_global_step: int = haar_data.get("global", 0) - # In haar.py, excited states are stored with key = global_step *before* - # the increment, so the last entry key is (global_step - 1). - last_key = haar_global_step - 1 - excited: dict[int, list[typing.Any]] = haar_data.get("excited", {}) - results = excited.get(last_key) - if results: - # results is list[(energy_tensor, configs, psi)]; index 0 is ground state. - energy_val = results[0][0] + + # Try to get the final energy directly if stored + if "final_energy" in haar_data: + energy_val = haar_data["final_energy"] if hasattr(energy_val, "item"): return float(energy_val.item()) + return float(energy_val) + + # Try to compute energy from the pool (configs, psi) if model is available + pool = haar_data.get("pool") + if pool is not None and model is not None: + configs, psi = pool + if configs is not None and psi is not None: + # Compute expectation value / + h_psi = model.apply_within(configs, psi, configs) + energy = ((psi.conj() @ h_psi) / (psi.conj() @ psi)).real + if hasattr(energy, "item"): + return float(energy.item()) + return float(energy) + + # Fallback: try to get from excited state results (Lanczos energy) + haar_global_step: int = haar_data.get("global", 0) + if haar_global_step > 0: + last_key = haar_global_step - 1 + excited: dict[int, list[typing.Any]] = haar_data.get("excited", {}) + results = excited.get(last_key) + if results: + # results is list[(energy_tensor, configs, psi)]; index 0 is ground state. + energy_val = results[0][0] + if hasattr(energy_val, "item"): + return float(energy_val.item()) + return 0.0 diff --git a/uv.lock b/uv.lock index a3824fb..1304fb6 100644 --- a/uv.lock +++ b/uv.lock @@ -1316,6 +1316,7 @@ dependencies = [ { name = "openfermion" }, { name = "platformdirs" }, { name = "pybind11" }, + { name = "pyscf" }, { name = "scipy" }, { name = "tensorboard" }, { name = "torch" }, @@ -1342,6 +1343,7 @@ requires-dist = [ { name = "openfermion", specifier = ">=1.7.1" }, { name = "platformdirs", specifier = ">=4.5.1" }, { name = "pybind11", specifier = ">=3.0.1" }, + { name = "pyscf", specifier = ">=2.12.1" }, { name = "pyscf", marker = "extra == 'pyscf'", specifier = ">=2.12.1" }, { name = "pytest", marker = "extra == 'dev'", specifier = ">=9.0.2" }, { name = "ruff", marker = "extra == 'dev'", specifier = ">=0.14.10" }, @@ -1555,6 +1557,11 @@ dependencies = [ wheels = [ { url = "https://files.pythonhosted.org/packages/d3/54/a2ba279afcca44bbd320d4e73675b282fcee3d81400ea1b53934efca6462/torch-2.10.0-2-cp312-none-macosx_11_0_arm64.whl", hash = "sha256:13ec4add8c3faaed8d13e0574f5cd4a323c11655546f91fbe6afa77b57423574", size = 79498202, upload-time = "2026-02-10T21:44:52.603Z" }, { url = "https://files.pythonhosted.org/packages/ec/23/2c9fe0c9c27f7f6cb865abcea8a4568f29f00acaeadfc6a37f6801f84cb4/torch-2.10.0-2-cp313-none-macosx_11_0_arm64.whl", hash = "sha256:e521c9f030a3774ed770a9c011751fb47c4d12029a3d6522116e48431f2ff89e", size = 79498254, upload-time = "2026-02-10T21:44:44.095Z" }, + { url = "https://files.pythonhosted.org/packages/b3/7a/abada41517ce0011775f0f4eacc79659bc9bc6c361e6bfe6f7052a6b9363/torch-2.10.0-3-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:98c01b8bb5e3240426dcde1446eed6f40c778091c8544767ef1168fc663a05a6", size = 915622781, upload-time = "2026-03-11T14:17:11.354Z" }, + { url = "https://files.pythonhosted.org/packages/ab/c6/4dfe238342ffdcec5aef1c96c457548762d33c40b45a1ab7033bb26d2ff2/torch-2.10.0-3-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:80b1b5bfe38eb0e9f5ff09f206dcac0a87aadd084230d4a36eea5ec5232c115b", size = 915627275, upload-time = "2026-03-11T14:16:11.325Z" }, + { url = "https://files.pythonhosted.org/packages/d8/f0/72bf18847f58f877a6a8acf60614b14935e2f156d942483af1ffc081aea0/torch-2.10.0-3-cp313-cp313t-manylinux_2_28_x86_64.whl", hash = "sha256:46b3574d93a2a8134b3f5475cfb98e2eb46771794c57015f6ad1fb795ec25e49", size = 915523474, upload-time = "2026-03-11T14:17:44.422Z" }, + { url = "https://files.pythonhosted.org/packages/f4/39/590742415c3030551944edc2ddc273ea1fdfe8ffb2780992e824f1ebee98/torch-2.10.0-3-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:b1d5e2aba4eb7f8e87fbe04f86442887f9167a35f092afe4c237dfcaaef6e328", size = 915632474, upload-time = "2026-03-11T14:15:13.666Z" }, + { url = "https://files.pythonhosted.org/packages/b6/8e/34949484f764dde5b222b7fe3fede43e4a6f0da9d7f8c370bb617d629ee2/torch-2.10.0-3-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:0228d20b06701c05a8f978357f657817a4a63984b0c90745def81c18aedfa591", size = 915523882, upload-time = "2026-03-11T14:14:46.311Z" }, { url = "https://files.pythonhosted.org/packages/cc/af/758e242e9102e9988969b5e621d41f36b8f258bb4a099109b7a4b4b50ea4/torch-2.10.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:5fd4117d89ffd47e3dcc71e71a22efac24828ad781c7e46aaaf56bf7f2796acf", size = 145996088, upload-time = "2026-01-21T16:24:44.171Z" }, { url = "https://files.pythonhosted.org/packages/23/8e/3c74db5e53bff7ed9e34c8123e6a8bfef718b2450c35eefab85bb4a7e270/torch-2.10.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:787124e7db3b379d4f1ed54dd12ae7c741c16a4d29b49c0226a89bea50923ffb", size = 915711952, upload-time = "2026-01-21T16:23:53.503Z" }, { url = "https://files.pythonhosted.org/packages/6e/01/624c4324ca01f66ae4c7cd1b74eb16fb52596dce66dbe51eff95ef9e7a4c/torch-2.10.0-cp312-cp312-win_amd64.whl", hash = "sha256:2c66c61f44c5f903046cc696d088e21062644cbe541c7f1c4eaae88b2ad23547", size = 113757972, upload-time = "2026-01-21T16:24:39.516Z" },