From e7c19b9b881220fe961eba57ac2b7b40feee8fbe Mon Sep 17 00:00:00 2001 From: kousuke-nakano <37653569+kousuke-nakano@users.noreply.github.com> Date: Sat, 10 Jan 2026 23:22:49 +0900 Subject: [PATCH 1/6] Implemented write 1e and 2e integrals --- pyscf/tools/test/test_trexio.py | 311 +++++++++++++++++++++++ pyscf/tools/trexio.py | 422 ++++++++++++++++++++++++++++++-- 2 files changed, 711 insertions(+), 22 deletions(-) diff --git a/pyscf/tools/test/test_trexio.py b/pyscf/tools/test/test_trexio.py index a408db616..dd61bedaf 100644 --- a/pyscf/tools/test/test_trexio.py +++ b/pyscf/tools/test/test_trexio.py @@ -5,9 +5,14 @@ import numpy as np import tempfile import pytest +import trexio as trexio_lib +from pyscf import ao2mo +from pyscf import pbc DIFF_TOL = 1e-10 +_write_2e_int_eri = trexio._write_2e_int_eri + ################################################################# # reading/writing `mol` from/to trexio file ################################################################# @@ -545,3 +550,309 @@ def test_mol_rhf_ccecp_ccpvqz(cart): mf1.run() e1 = mf1.e_tot assert abs(e0 - e1).max() < DIFF_TOL + +################################################################# +# writing `1e_int` and `2e_int` to trexio file +################################################################# + +def _trexio_pack_eri(eri, basis): + basis = basis.upper() + if basis not in ('AO', 'MO'): + raise ValueError("basis must be 'AO' or 'MO'") + + with tempfile.TemporaryDirectory() as tmpdir: + filename = os.path.join(tmpdir, 'pack.h5') + _write_2e_int_eri(eri, filename, backend='h5', basis=basis) + with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: + if basis == 'AO': + size = trexio_lib.read_ao_2e_int_eri_size(tf) + idx, val, n_read, _ = trexio_lib.read_ao_2e_int_eri(tf, 0, size) + else: + size = trexio_lib.read_mo_2e_int_eri_size(tf) + idx, val, n_read, _ = trexio_lib.read_mo_2e_int_eri(tf, 0, size) + assert n_read == size + return np.asarray(idx, dtype=np.int32).ravel(), np.asarray(val) + + +def _hermitize(mat): + mat = np.asarray(mat) + return 0.5 * (mat + mat.T.conj()) + + +@pytest.mark.parametrize("cart", [False, True], ids=["cart=false", "cart=true"]) +def test_write_molecule_integrals_to_trexio_rks(cart): + with tempfile.TemporaryDirectory() as d: + filename = os.path.join(d, 'mol_integrals.h5') + + mol0 = pyscf.M(atom='H 0 0 0; F 0 0 1', basis='6-31g*', cart=cart) + mf0 = mol0.RHF().run() + + overlap = _hermitize(mf0.get_ovlp()) + kinetic = _hermitize(mol0.intor('int1e_kin')) + potential = _hermitize(mol0.intor('int1e_nuc')) + if mol0._ecp: + from pyscf.gto import ecp + potential += _hermitize(ecp.ecp_int(mol0)) + core = kinetic + potential + + trexio.to_trexio(mf0, filename) + + trexio.write_scf_1e_int_eri(mf0, filename, basis='AO') + with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: + np.testing.assert_allclose(trexio_lib.read_ao_1e_int_overlap(tf), overlap, atol=DIFF_TOL) + np.testing.assert_allclose(trexio_lib.read_ao_1e_int_kinetic(tf), kinetic, atol=DIFF_TOL) + np.testing.assert_allclose(trexio_lib.read_ao_1e_int_potential_n_e(tf), potential, atol=DIFF_TOL) + np.testing.assert_allclose(trexio_lib.read_ao_1e_int_core_hamiltonian(tf), core, atol=DIFF_TOL) + + coeff = mf0.mo_coeff + mo_overlap = _hermitize(coeff.conj().T @ overlap @ coeff) + mo_kinetic = _hermitize(coeff.conj().T @ kinetic @ coeff) + mo_potential = _hermitize(coeff.conj().T @ potential @ coeff) + mo_core = _hermitize(coeff.conj().T @ core @ coeff) + + trexio.write_scf_1e_int_eri(mf0, filename, basis='MO') + with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: + np.testing.assert_allclose(trexio_lib.read_mo_1e_int_overlap(tf), mo_overlap, atol=DIFF_TOL) + np.testing.assert_allclose(trexio_lib.read_mo_1e_int_kinetic(tf), mo_kinetic, atol=DIFF_TOL) + np.testing.assert_allclose(trexio_lib.read_mo_1e_int_potential_n_e(tf), mo_potential, atol=DIFF_TOL) + np.testing.assert_allclose(trexio_lib.read_mo_1e_int_core_hamiltonian(tf), mo_core, atol=DIFF_TOL) + + ao_eri = mol0.intor('int2e', aosym='s8') + ao_idx_exp, ao_val_exp = _trexio_pack_eri(ao_eri, 'AO') + trexio.write_scf_2e_int_eri(mf0, filename, basis='AO', aosym='s8') + with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: + assert trexio_lib.has_ao_2e_int_eri(tf) + size = trexio_lib.read_ao_2e_int_eri_size(tf) + idx, val, n_read, _ = trexio_lib.read_ao_2e_int_eri(tf, 0, size) + assert n_read == size + np.testing.assert_array_equal(np.asarray(idx, dtype=np.int32).ravel(), ao_idx_exp) + np.testing.assert_allclose(np.asarray(val), ao_val_exp, atol=DIFF_TOL) + + mo_eri = ao2mo.kernel(mol0, coeff, compact=True) + mo_idx_exp, mo_val_exp = _trexio_pack_eri(mo_eri, 'MO') + trexio.write_scf_2e_int_eri(mf0, filename, basis='MO', aosym='s8') + with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: + assert trexio_lib.has_mo_2e_int_eri(tf) + size = trexio_lib.read_mo_2e_int_eri_size(tf) + idx, val, n_read, _ = trexio_lib.read_mo_2e_int_eri(tf, 0, size) + assert n_read == size + np.testing.assert_array_equal(np.asarray(idx, dtype=np.int32).ravel(), mo_idx_exp) + np.testing.assert_allclose(np.asarray(val), mo_val_exp, atol=DIFF_TOL) + + +@pytest.mark.parametrize("cart", [False, True], ids=["cart=false", "cart=true"]) +def test_write_molecule_integrals_to_trexio_uks(cart): + with tempfile.TemporaryDirectory() as d: + filename = os.path.join(d, 'mol_uks_integrals.h5') + + mol0 = pyscf.M(atom='O 0 0 0', basis='6-31g*', spin=2, cart=cart) + mf0 = mol0.UKS() + mf0.xc = 'LDA' + mf0.kernel() + assert mf0.converged + + overlap = _hermitize(mf0.get_ovlp()) + kinetic = _hermitize(mol0.intor('int1e_kin')) + potential = _hermitize(mol0.intor('int1e_nuc')) + if mol0._ecp: + from pyscf.gto import ecp + potential += _hermitize(ecp.ecp_int(mol0)) + core = kinetic + potential + + trexio.to_trexio(mf0, filename) + + trexio.write_scf_1e_int_eri(mf0, filename, basis='AO') + with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: + np.testing.assert_allclose(trexio_lib.read_ao_1e_int_overlap(tf), overlap, atol=DIFF_TOL) + np.testing.assert_allclose(trexio_lib.read_ao_1e_int_kinetic(tf), kinetic, atol=DIFF_TOL) + np.testing.assert_allclose( + trexio_lib.read_ao_1e_int_potential_n_e(tf), potential, atol=DIFF_TOL + ) + np.testing.assert_allclose( + trexio_lib.read_ao_1e_int_core_hamiltonian(tf), core, atol=DIFF_TOL + ) + + coeff_alpha, coeff_beta = mf0.mo_coeff + coeff = np.concatenate([coeff_alpha, coeff_beta], axis=1) + mo_overlap = _hermitize(coeff.conj().T @ overlap @ coeff) + mo_kinetic = _hermitize(coeff.conj().T @ kinetic @ coeff) + mo_potential = _hermitize(coeff.conj().T @ potential @ coeff) + mo_core = _hermitize(coeff.conj().T @ core @ coeff) + + trexio.write_scf_1e_int_eri(mf0, filename, basis='MO') + with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: + np.testing.assert_allclose(trexio_lib.read_mo_1e_int_overlap(tf), mo_overlap, atol=DIFF_TOL) + np.testing.assert_allclose(trexio_lib.read_mo_1e_int_kinetic(tf), mo_kinetic, atol=DIFF_TOL) + np.testing.assert_allclose( + trexio_lib.read_mo_1e_int_potential_n_e(tf), mo_potential, atol=DIFF_TOL + ) + np.testing.assert_allclose( + trexio_lib.read_mo_1e_int_core_hamiltonian(tf), mo_core, atol=DIFF_TOL + ) + + ao_eri = mol0.intor('int2e', aosym='s8') + ao_idx_exp, ao_val_exp = _trexio_pack_eri(ao_eri, 'AO') + trexio.write_scf_2e_int_eri(mf0, filename, basis='AO', aosym='s8') + with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: + assert trexio_lib.has_ao_2e_int_eri(tf) + size = trexio_lib.read_ao_2e_int_eri_size(tf) + idx, val, n_read, _ = trexio_lib.read_ao_2e_int_eri(tf, 0, size) + assert n_read == size + np.testing.assert_array_equal(np.asarray(idx, dtype=np.int32).ravel(), ao_idx_exp) + np.testing.assert_allclose(np.asarray(val), ao_val_exp, atol=DIFF_TOL) + + mo_eri = ao2mo.kernel(mol0, coeff, compact=True) + mo_idx_exp, mo_val_exp = _trexio_pack_eri(mo_eri, 'MO') + trexio.write_scf_2e_int_eri(mf0, filename, basis='MO', aosym='s8') + with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: + assert trexio_lib.has_mo_2e_int_eri(tf) + size = trexio_lib.read_mo_2e_int_eri_size(tf) + idx, val, n_read, _ = trexio_lib.read_mo_2e_int_eri(tf, 0, size) + assert n_read == size + np.testing.assert_array_equal(np.asarray(idx, dtype=np.int32).ravel(), mo_idx_exp) + np.testing.assert_allclose(np.asarray(val), mo_val_exp, atol=DIFF_TOL) + + +@pytest.mark.parametrize("cart", [False, True], ids=["cart=false", "cart=true"]) +def test_write_cell_gamma_integrals_to_trexio_rks(cart): + with tempfile.TemporaryDirectory() as d: + filename = os.path.join(d, 'cell_integrals.h5') + + cell0 = pbc.gto.Cell() + cell0.cart = cart + cell0.build(atom='H 0 0 0; H 0 0 1', basis='6-31g*', a=np.diag([3.0, 3.0, 5.0])) + mf0 = pbc.scf.RHF(cell0).density_fit() + mf0.kernel() + assert mf0.converged + + overlap = _hermitize(mf0.get_ovlp()) + kinetic = _hermitize(cell0.pbc_intor('int1e_kin', 1, 1)) + df_builder = mf0.with_df.build() if mf0.with_df is not None else pbc.df.MDF(cell0).build() + potential = _hermitize(df_builder.get_nuc()) + if len(getattr(cell0, '_ecpbas', [])) > 0: + from pyscf.pbc.gto import ecp + potential += _hermitize(ecp.ecp_int(cell0)) + core = kinetic + potential + + trexio.to_trexio(mf0, filename) + + trexio.write_scf_1e_int_eri(mf0, filename, basis='AO') + with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: + np.testing.assert_allclose(trexio_lib.read_ao_1e_int_overlap(tf), overlap, atol=DIFF_TOL) + np.testing.assert_allclose(trexio_lib.read_ao_1e_int_kinetic(tf), kinetic, atol=DIFF_TOL) + np.testing.assert_allclose(trexio_lib.read_ao_1e_int_potential_n_e(tf), potential, atol=DIFF_TOL) + np.testing.assert_allclose(trexio_lib.read_ao_1e_int_core_hamiltonian(tf), core, atol=DIFF_TOL) + + coeff = mf0.mo_coeff + if coeff.ndim == 3 and coeff.shape[0] == 1: + coeff = coeff[0] + mo_overlap = _hermitize(coeff.conj().T @ overlap @ coeff) + mo_kinetic = _hermitize(coeff.conj().T @ kinetic @ coeff) + mo_potential = _hermitize(coeff.conj().T @ potential @ coeff) + mo_core = _hermitize(coeff.conj().T @ core @ coeff) + + trexio.write_scf_1e_int_eri(mf0, filename, basis='MO') + with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: + np.testing.assert_allclose(trexio_lib.read_mo_1e_int_overlap(tf), mo_overlap, atol=DIFF_TOL) + np.testing.assert_allclose(trexio_lib.read_mo_1e_int_kinetic(tf), mo_kinetic, atol=DIFF_TOL) + np.testing.assert_allclose(trexio_lib.read_mo_1e_int_potential_n_e(tf), mo_potential, atol=DIFF_TOL) + np.testing.assert_allclose(trexio_lib.read_mo_1e_int_core_hamiltonian(tf), mo_core, atol=DIFF_TOL) + + df_obj = pbc.df.MDF(cell0).build() + eri_ao = pbc.df.df_ao2mo.get_eri(df_obj, compact=False) + nao = cell0.nao_nr() + eri_ao = eri_ao.reshape(nao, nao, nao, nao) + ao_idx_exp, ao_val_exp = _trexio_pack_eri(eri_ao, 'AO') + trexio.write_scf_2e_int_eri(mf0, filename, basis='AO') + with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: + assert trexio_lib.has_ao_2e_int_eri(tf) + size = trexio_lib.read_ao_2e_int_eri_size(tf) + idx, val, n_read, _ = trexio_lib.read_ao_2e_int_eri(tf, 0, size) + assert n_read == size + np.testing.assert_array_equal(np.asarray(idx, dtype=np.int32).ravel(), ao_idx_exp) + np.testing.assert_allclose(np.asarray(val), ao_val_exp, atol=DIFF_TOL) + + df_obj_mo = pbc.df.MDF(cell0).build() + mo_eri = df_obj_mo.get_mo_eri((coeff, coeff, coeff, coeff)) + mo_idx_exp, mo_val_exp = _trexio_pack_eri(np.real_if_close (mo_eri), 'MO') + trexio.write_scf_2e_int_eri(mf0, filename, basis='MO') + with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: + assert trexio_lib.has_mo_2e_int_eri(tf) + size = trexio_lib.read_mo_2e_int_eri_size(tf) + idx, val, n_read, _ = trexio_lib.read_mo_2e_int_eri(tf, 0, size) + assert n_read == size + np.testing.assert_array_equal(np.asarray(idx, dtype=np.int32).ravel(), mo_idx_exp) + np.testing.assert_allclose(np.asarray(val), mo_val_exp, atol=DIFF_TOL) + + +@pytest.mark.parametrize("cart", [False, True], ids=["cart=false", "cart=true"]) +def test_write_cell_gamma_integrals_to_trexio_uks(cart): + with tempfile.TemporaryDirectory() as d: + filename = os.path.join(d, 'cell_uks_integrals.h5') + + cell0 = pbc.gto.Cell() + cell0.spin = 2 + cell0.cart = cart + cell0.build(atom='H 0 0 0; H 0 0 1', basis='6-31g*', a=np.diag([3.0, 3.0, 5.0])) + mf0 = pbc.scf.UKS(cell0).density_fit() + mf0.xc = 'LDA' + mf0.kernel() + assert mf0.converged + + overlap = _hermitize(mf0.get_ovlp()) + kinetic = _hermitize(cell0.pbc_intor('int1e_kin', 1, 1)) + df_builder = mf0.with_df.build() if mf0.with_df is not None else pbc.df.MDF(cell0).build() + potential = _hermitize(df_builder.get_nuc()) + if len(getattr(cell0, '_ecpbas', [])) > 0: + from pyscf.pbc.gto import ecp + potential += _hermitize(ecp.ecp_int(cell0)) + core = kinetic + potential + + trexio.to_trexio(mf0, filename) + + trexio.write_scf_1e_int_eri(mf0, filename, basis='AO') + with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: + np.testing.assert_allclose(trexio_lib.read_ao_1e_int_overlap(tf), overlap, atol=DIFF_TOL) + np.testing.assert_allclose(trexio_lib.read_ao_1e_int_kinetic(tf), kinetic, atol=DIFF_TOL) + np.testing.assert_allclose(trexio_lib.read_ao_1e_int_potential_n_e(tf), potential, atol=DIFF_TOL) + np.testing.assert_allclose(trexio_lib.read_ao_1e_int_core_hamiltonian(tf), core, atol=DIFF_TOL) + + coeff_alpha, coeff_beta = mf0.mo_coeff + coeff = np.concatenate([coeff_alpha, coeff_beta], axis=1) + mo_overlap = _hermitize(coeff.conj().T @ overlap @ coeff) + mo_kinetic = _hermitize(coeff.conj().T @ kinetic @ coeff) + mo_potential = _hermitize(coeff.conj().T @ potential @ coeff) + mo_core = _hermitize(coeff.conj().T @ core @ coeff) + + trexio.write_scf_1e_int_eri(mf0, filename, basis='MO') + with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: + np.testing.assert_allclose(trexio_lib.read_mo_1e_int_overlap(tf), mo_overlap, atol=DIFF_TOL) + np.testing.assert_allclose(trexio_lib.read_mo_1e_int_kinetic(tf), mo_kinetic, atol=DIFF_TOL) + np.testing.assert_allclose(trexio_lib.read_mo_1e_int_potential_n_e(tf), mo_potential, atol=DIFF_TOL) + np.testing.assert_allclose(trexio_lib.read_mo_1e_int_core_hamiltonian(tf), mo_core, atol=DIFF_TOL) + + df_obj = pbc.df.MDF(cell0).build() + eri_ao = pbc.df.df_ao2mo.get_eri(df_obj, compact=False) + nao = cell0.nao_nr() + eri_ao = eri_ao.reshape(nao, nao, nao, nao) + ao_idx_exp, ao_val_exp = _trexio_pack_eri(eri_ao, 'AO') + trexio.write_scf_2e_int_eri(mf0, filename, basis='AO') + with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: + assert trexio_lib.has_ao_2e_int_eri(tf) + size = trexio_lib.read_ao_2e_int_eri_size(tf) + idx, val, n_read, _ = trexio_lib.read_ao_2e_int_eri(tf, 0, size) + assert n_read == size + np.testing.assert_array_equal(np.asarray(idx, dtype=np.int32).ravel(), ao_idx_exp) + np.testing.assert_allclose(np.asarray(val), ao_val_exp, atol=DIFF_TOL) + + df_obj_mo = pbc.df.MDF(cell0).build() + mo_eri = df_obj_mo.get_mo_eri((coeff, coeff, coeff, coeff)) + mo_idx_exp, mo_val_exp = _trexio_pack_eri(np.real_if_close(mo_eri), 'MO') + trexio.write_scf_2e_int_eri(mf0, filename, basis='MO') + with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: + assert trexio_lib.has_mo_2e_int_eri(tf) + size = trexio_lib.read_mo_2e_int_eri_size(tf) + idx, val, n_read, _ = trexio_lib.read_mo_2e_int_eri(tf, 0, size) + assert n_read == size + np.testing.assert_array_equal(np.asarray(idx, dtype=np.int32).ravel(), mo_idx_exp) + np.testing.assert_allclose(np.asarray(val), mo_val_exp, atol=DIFF_TOL) diff --git a/pyscf/tools/trexio.py b/pyscf/tools/trexio.py index 4444d7d91..cdf8d81cd 100644 --- a/pyscf/tools/trexio.py +++ b/pyscf/tools/trexio.py @@ -14,7 +14,6 @@ import re import math import numpy as np -import scipy.linalg from collections import defaultdict import pyscf @@ -26,6 +25,8 @@ from pyscf import mcscf from pyscf import fci from pyscf.pbc import gto as pbcgto +from pyscf import ao2mo +from pyscf.pbc import df as pbcdf import trexio @@ -777,13 +778,171 @@ def scf_from_trexio(filename): else: raise ValueError(f'Unknown spin multiplicity {uniq}') -def write_ao_2e_int_eri(eri, filename, backend='h5'): - raise NotImplementedError - -def read_ao_2e_int_eri(filename): - raise NotImplementedError - -def write_mo_2e_int_eri(eri, filename, backend='h5'): +def write_scf_2e_int_eri( + mf, filename, backend='h5', basis='mo', mo_compact=True, aosym='s8', + df_engine='MDF', +): + """ + Write two-electron integrals (ERI) in AO or MO basis. + + Rules + ----- + - Backend is real-only: complex ERIs are not supported. + - PBC: Only single-k Gamma is supported (non-Gamma -> NotImplementedError). + - Molecular RHF/RKS: AO and MO supported. + - Molecular/PBC UKS/UHF: + * AO: spin-independent (single tensor). + * MO: build a combined MO space by column-wise concatenation [alpha | beta] + and compute full ERIs including cross-spin terms. + + Notes + ----- + - Immediately before writing, arrays are forced to be C-contiguous *without* + changing dtype: + np.ascontiguousarray(arr) + """ + + # ----- internal constants ----- + TOL_IMAG = 1e-12 # drop imag parts <= this; otherwise raise (real-only backend) + + basis = basis.upper() + if aosym not in ('s8', 's4', 's1'): + raise ValueError("aosym must be one of {'s8','s4','s1'}") + is_pbc = hasattr(mf, 'cell') + + # ---------- helpers ---------- + def _is_gamma_single_k(mf_obj) -> bool: + """True if this PBC calculation is single-k at Gamma.""" + if not hasattr(mf_obj, 'cell'): + return False + if hasattr(mf_obj, 'kpt'): + return np.allclose(np.asarray(mf_obj.kpt), 0.0) + if hasattr(mf_obj, 'kpts'): + kpts = np.asarray(mf_obj.kpts) + return (kpts.shape[0] == 1) and np.allclose(kpts[0], 0.0) + return True # treat no-kpt attribute as Gamma + + def _ensure_real(x): + """Ensure array is real; drop tiny imaginary part else raise (dtype preserved).""" + if np.iscomplexobj(x): + if np.all(np.abs(np.imag(x)) <= TOL_IMAG): + # np.real keeps precision: complex128->float64, complex64->float32 + x = np.real(x) + else: + raise NotImplementedError( + "Complex ERI encountered but the backend is real-only. " + "Use Gamma-point (k=0) or a complex-capable backend." + ) + return x + + def _df_obj(): + """Construct a DF engine for PBC Γ-point.""" + if df_engine.upper() == 'MDF': + return pbcdf.MDF(mf.cell).build() + if df_engine.upper() == 'GDF': + return pbcdf.GDF(mf.cell).build() + raise ValueError("df_engine must be 'MDF' or 'GDF'.") + + def _get_uks_coeff_pair(mf_obj): + """ + Extract (Ca, Cb) as 2D arrays (nao, nalpha), (nao, nbeta), squeezing possible k-dims. + Supports: + - list/tuple: (Ca, Cb) + - ndarray shapes: (2, nao, nmo) or (2, 1, nao, nmo) + """ + C = mf_obj.mo_coeff + if isinstance(C, (list, tuple)) and len(C) == 2: + Ca, Cb = C + elif isinstance(C, np.ndarray) and C.ndim >= 3 and C.shape[0] == 2: + if C.ndim == 3: + Ca, Cb = C[0], C[1] + elif C.ndim == 4: + if C.shape[1] != 1: + raise NotImplementedError("Only single-k UKS/UHF is supported.") + Ca, Cb = C[0, 0], C[1, 0] + else: + raise ValueError(f"Unexpected mo_coeff shape: {C.shape}") + else: + raise TypeError("Not a UKS/UHF object or unsupported mo_coeff layout.") + if Ca.ndim != 2 or Cb.ndim != 2: + raise ValueError(f"Unexpected UKS/UHF mo_coeff shapes: Ca {Ca.shape}, Cb {Cb.shape}") + return Ca, Cb + + # --------------------- + # MO-basis ERI writing + # --------------------- + if basis == 'MO': + # Molecular + if not is_pbc: + if (isinstance(mf.mo_coeff, (list, tuple)) or + (isinstance(mf.mo_coeff, np.ndarray) and mf.mo_coeff.ndim >= 3 and mf.mo_coeff.shape[0] == 2)): + # UKS/UHF -> concatenate [alpha | beta], include cross-spin terms + Ca, Cb = _get_uks_coeff_pair(mf) + C = np.concatenate([Ca, Cb], axis=1) # (nao, nalpha+nbeta) + if getattr(mf, '_eri', None) is not None: + eri_mo = ao2mo.incore.full(mf._eri, C, compact=mo_compact) + else: + eri_mo = ao2mo.kernel(mf.mol, C, compact=mo_compact) + eri_mo = _ensure_real(eri_mo) + _write_2e_int_eri(np.ascontiguousarray(eri_mo), filename, backend, 'MO') + else: # RHF/RKS + C = mf.mo_coeff + if getattr(mf, '_eri', None) is not None: + eri_mo = ao2mo.incore.full(mf._eri, C, compact=mo_compact) + else: + eri_mo = ao2mo.kernel(mf.mol, C, compact=mo_compact) + eri_mo = _ensure_real(eri_mo) + _write_2e_int_eri(np.ascontiguousarray(eri_mo), filename, backend, 'MO') + return + + # PBC (Gamma only) + if not _is_gamma_single_k(mf): + raise NotImplementedError("PBC MO-ERI: non-Gamma k-points are not supported (real-only backend).") + dfobj = _df_obj() + + if (isinstance(mf.mo_coeff, (list, tuple)) or + (isinstance(mf.mo_coeff, np.ndarray) and mf.mo_coeff.ndim >= 3 and mf.mo_coeff.shape[0] == 2)): + # UKS/UHF @ Gamma: combined MO matrix [Ca | Cb] + Ca, Cb = _get_uks_coeff_pair(mf) + C = np.concatenate([Ca, Cb], axis=1) + eri_mo = dfobj.get_mo_eri((C, C, C, C)) + eri_mo = _ensure_real(eri_mo) + _write_2e_int_eri(np.ascontiguousarray(eri_mo), filename, backend, 'MO') + else: # RHF/RKS @ Gamma + C = mf.mo_coeff + if C.ndim == 3 and C.shape[0] == 1: # normalize (1,nao,nmo) -> (nao,nmo) + C = C[0] + eri_mo = dfobj.get_mo_eri((C, C, C, C)) + eri_mo = _ensure_real(eri_mo) + _write_2e_int_eri(np.ascontiguousarray(eri_mo), filename, backend, 'MO') + return + + # --------------------- + # AO-basis ERI writing (spin-independent even for UKS/UHF) + # --------------------- + if is_pbc: + # PBC AO: Gamma only via DF (real-only) + if not _is_gamma_single_k(mf): + raise NotImplementedError("PBC AO-ERI: non-Gamma k-points are not supported (real-only backend).") + dfobj = _df_obj() + eri2 = pbcdf.df_ao2mo.get_eri(dfobj, compact=False) # (nao^2, nao^2) + nao = mf.cell.nao_nr() + if eri2.shape != (nao * nao, nao * nao): + raise RuntimeError(f"Unexpected ERI shape {eri2.shape}; expected ({nao*nao}, {nao*nao}) at Gamma.") + eri_ao = eri2.reshape(nao, nao, nao, nao) + eri_ao = _ensure_real(eri_ao) + _write_2e_int_eri(np.ascontiguousarray(eri_ao), filename, backend, 'AO') + else: + # Molecular AO + eri_ao = getattr(mf, '_eri', None) + if eri_ao is None: + # 'int2e' follows mol.cart automatically (spherical vs Cartesian) + eri_ao = mf.mol.intor('int2e', aosym=aosym) + eri_ao = _ensure_real(eri_ao) + _write_2e_int_eri(np.ascontiguousarray(eri_ao), filename, backend, 'AO') + +def _write_2e_int_eri(eri, filename, backend='h5', basis='MO'): + assert basis.upper() in ['MO','AO'] num_integrals = eri.size if eri.ndim == 4: n = eri.shape[0] @@ -811,23 +970,242 @@ def write_mo_2e_int_eri(eri, filename, backend='h5'): idx=idx.reshape((num_integrals,4)) for i in range(num_integrals): idx[i,1],idx[i,2]=idx[i,2],idx[i,1] - idx=idx.flatten() - with trexio.File(filename, 'w', back_end=_mode(backend)) as tf: - trexio.write_mo_2e_int_eri(tf, 0, num_integrals, idx, eri.ravel()) + # write ERI + with trexio.File(filename, 'u', back_end=_mode(backend)) as tf: + if basis.upper() == 'AO': + if not trexio.has_ao_num(tf): + trexio.write_ao_num(tf, n) + else: + if not trexio.has_mo_num(tf): + trexio.write_mo_num(tf, n) + if basis.upper() == 'MO': + trexio.write_mo_2e_int_eri(tf, 0, num_integrals, idx, eri.ravel()) + else: + trexio.write_ao_2e_int_eri(tf, 0, num_integrals, idx, eri.ravel()) + +def write_scf_1e_int_eri( + mf, filename, backend='h5', basis='AO', df_engine='MDF', +): + """ + Write one-electron integrals (overlap, kinetic, nuclear-electron potential, + and core Hamiltonian) in AO or MO basis. + + Rules + ----- + - Backend is real-only: negligible imaginary parts are discarded, otherwise + a NotImplementedError is raised. + - PBC: only single-k Gamma calculations are supported. + - MO basis: UHF/UKS uses a concatenated MO matrix [alpha | beta]. + """ + + TOL_IMAG = 1e-12 + basis = basis.upper() + if basis not in ('AO', 'MO'): + raise ValueError("basis must be either 'AO' or 'MO'") + + def _ensure_real(x): + if np.iscomplexobj(x): + if np.all(np.abs(np.imag(x)) <= TOL_IMAG): + x = np.real(x) + else: + raise NotImplementedError( + "Complex one-electron integrals encountered but the backend is real-only." + ) + return x + + def _is_gamma_single_k(mf_obj) -> bool: + if not hasattr(mf_obj, 'cell'): + return False + if hasattr(mf_obj, 'kpt'): + return np.allclose(np.asarray(mf_obj.kpt), 0.0) + if hasattr(mf_obj, 'kpts'): + kpts = np.asarray(mf_obj.kpts) + return (kpts.shape[0] == 1) and np.allclose(kpts[0], 0.0) + return True + + def _as_matrix(mat, label): + if isinstance(mat, (tuple, list)): + if len(mat) == 0: + raise ValueError(f"Empty data for {label}") + if len(mat) > 1: + raise NotImplementedError( + f"{label}: multiple blocks are not supported in this helper" + ) + mat = mat[0] + mat = np.asarray(mat) + if mat.ndim == 3: + if mat.shape[0] != 1: + raise NotImplementedError( + f"{label}: Gamma-only support; received shape {mat.shape}" + ) + mat = mat[0] + if mat.ndim != 2: + raise ValueError(f"{label} must be a 2D matrix, got shape {mat.shape}") + return mat + + def _hermitize(mat): + return 0.5 * (mat + mat.T.conj()) + + is_pbc = hasattr(mf, 'cell') + + if is_pbc: + if not _is_gamma_single_k(mf): + raise NotImplementedError( + "PBC one-electron integrals are implemented for Gamma-point only." + ) + + cell = mf.cell + overlap = _as_matrix(mf.get_ovlp(), 'AO overlap') + kinetic = _as_matrix(cell.pbc_intor('int1e_kin', 1, 1), 'AO kinetic') + + df_builder = getattr(mf, 'with_df', None) + if df_builder is None: + if df_engine.upper() == 'MDF': + df_builder = pbcdf.MDF(cell).build() + elif df_engine.upper() == 'GDF': + df_builder = pbcdf.GDF(cell).build() + else: + raise ValueError("df_engine must be 'MDF' or 'GDF'") + else: + df_builder = df_builder.build() -def read_mo_2e_int_eri(filename): - with trexio.File(filename, 'r', back_end=trexio.TREXIO_AUTO) as tf: - nmo = trexio.read_mo_num(tf) - nao_pair = nmo * (nmo+1) // 2 - eri_size = nao_pair * (nao_pair+1) // 2 - idx, data, n_read, eof_flag = trexio.read_mo_2e_int_eri(tf, 0, eri_size) - eri = np.zeros(eri_size) - x = idx[:,0]*(idx[:,0]+1)//2 + idx[:,2] - y = idx[:,1]*(idx[:,1]+1)//2 + idx[:,3] - eri[x*(x+1)//2+y] = data - return eri + if cell.pseudo: + potential = _as_matrix(df_builder.get_pp(), 'AO potential') + else: + potential = _as_matrix(df_builder.get_nuc(), 'AO potential') + + if len(getattr(cell, '_ecpbas', [])) > 0: + from pyscf.pbc.gto import ecp + potential += _as_matrix(ecp.ecp_int(cell), 'AO ECP potential') + + core = kinetic + potential + else: + mol = mf.mol + overlap = _as_matrix(mf.get_ovlp(), 'AO overlap') + kinetic = _as_matrix(mol.intor('int1e_kin'), 'AO kinetic') + potential = _as_matrix(mol.intor('int1e_nuc'), 'AO potential') + if mol._ecp: + from pyscf.gto import ecp + potential += _as_matrix(ecp.ecp_int(mol), 'AO ECP potential') + core = kinetic + potential + + overlap = np.ascontiguousarray(_ensure_real(_hermitize(overlap))) + kinetic = np.ascontiguousarray(_ensure_real(_hermitize(kinetic))) + potential = np.ascontiguousarray(_ensure_real(_hermitize(potential))) + core = np.ascontiguousarray(_ensure_real(_hermitize(core))) + + if basis == 'AO': + _write_1e_int_eri(overlap, kinetic, potential, core, filename, backend, 'AO') + return + + def _get_uks_coeff_pair(mf_obj): + coeff = mf_obj.mo_coeff + if isinstance(coeff, (list, tuple)) and len(coeff) == 2: + Ca, Cb = coeff + elif isinstance(coeff, np.ndarray) and coeff.ndim >= 3 and coeff.shape[0] == 2: + if coeff.ndim == 3: + Ca, Cb = coeff[0], coeff[1] + elif coeff.ndim == 4: + if coeff.shape[1] != 1: + raise NotImplementedError( + "Only single-k UKS/UHF is supported for MO one-electron integrals." + ) + Ca, Cb = coeff[0, 0], coeff[1, 0] + else: + raise ValueError(f"Unexpected mo_coeff shape: {coeff.shape}") + else: + raise TypeError("Unsupported mo_coeff layout for UKS/UHF object") + if Ca.ndim != 2 or Cb.ndim != 2: + raise ValueError( + f"Unexpected UKS/UHF mo_coeff shapes: Ca {Ca.shape}, Cb {Cb.shape}" + ) + return Ca, Cb + + def _get_rhf_coeff(mf_obj): + coeff = mf_obj.mo_coeff + if isinstance(coeff, np.ndarray): + if coeff.ndim == 2: + return coeff + if coeff.ndim == 3 and coeff.shape[0] == 1: + return coeff[0] + if isinstance(coeff, (list, tuple)) and len(coeff) == 1: + arr = np.asarray(coeff[0]) + if arr.ndim == 2: + return arr + raise TypeError( + "Unsupported mo_coeff layout for RHF/RKS object in MO one-electron integrals" + ) + + if ( + isinstance(mf.mo_coeff, (list, tuple)) and len(mf.mo_coeff) == 2 + ) or ( + isinstance(mf.mo_coeff, np.ndarray) and mf.mo_coeff.ndim >= 3 and mf.mo_coeff.shape[0] == 2 + ): + Ca, Cb = _get_uks_coeff_pair(mf) + if is_pbc: + if Ca.ndim == 3 and Ca.shape[0] == 1: + Ca = Ca[0] + if Cb.ndim == 3 and Cb.shape[0] == 1: + Cb = Cb[0] + if Ca.ndim != 2 or Cb.ndim != 2: + raise NotImplementedError( + "Only Gamma-point data are supported for MO one-electron integrals" + ) + C = np.concatenate([Ca, Cb], axis=1) + else: + C = _get_rhf_coeff(mf) + + if is_pbc and C.ndim == 3: + if C.shape[0] != 1: + raise NotImplementedError( + "MO one-electron integrals currently support single-k Gamma calculations only." + ) + C = C[0] + + if C.ndim != 2: + raise ValueError(f"MO coefficient matrix must be 2D, got shape {C.shape}") + + def _ao_to_mo(mat, coeff): + return _hermitize(coeff.conj().T @ mat @ coeff) + + mo_overlap = np.ascontiguousarray(_ensure_real(_ao_to_mo(overlap, C))) + mo_kinetic = np.ascontiguousarray(_ensure_real(_ao_to_mo(kinetic, C))) + mo_potential = np.ascontiguousarray(_ensure_real(_ao_to_mo(potential, C))) + mo_core = np.ascontiguousarray(_ensure_real(_ao_to_mo(core, C))) + + _write_1e_int_eri( + mo_overlap, mo_kinetic, mo_potential, mo_core, filename, backend, 'MO' + ) + +def _write_1e_int_eri(overlap, kinetic, potential, core, filename, backend='h5', basis='AO'): + basis = basis.upper() + if basis not in ('AO', 'MO'): + raise ValueError("basis must be either 'AO' or 'MO'") + + overlap = np.ascontiguousarray(overlap) + kinetic = np.ascontiguousarray(kinetic) + potential = np.ascontiguousarray(potential) + core = np.ascontiguousarray(core) + + with trexio.File(filename, 'u', back_end=_mode(backend)) as tf: + if basis == 'AO': + ao_dim = overlap.shape[0] + if not trexio.has_ao_num(tf): + trexio.write_ao_num(tf, ao_dim) + trexio.write_ao_1e_int_overlap(tf, overlap) + trexio.write_ao_1e_int_kinetic(tf, kinetic) + trexio.write_ao_1e_int_potential_n_e(tf, potential) + trexio.write_ao_1e_int_core_hamiltonian(tf, core) + else: + mo_dim = overlap.shape[0] + if not trexio.has_mo_num(tf): + trexio.write_mo_num(tf, mo_dim) + trexio.write_mo_1e_int_overlap(tf, overlap) + trexio.write_mo_1e_int_kinetic(tf, kinetic) + trexio.write_mo_1e_int_potential_n_e(tf, potential) + trexio.write_mo_1e_int_core_hamiltonian(tf, core) def _order_ao_index(mol): if mol.cart: From 0b47d33e0e39432ba9940aeba64214319188ddd3 Mon Sep 17 00:00:00 2001 From: kousuke-nakano <37653569+kousuke-nakano@users.noreply.github.com> Date: Sun, 11 Jan 2026 00:06:08 +0900 Subject: [PATCH 2/6] Fix TREXIO gamma-point tests by squeezing singleton k-axis --- pyscf/tools/test/test_trexio.py | 30 ++++++++++++++++++------------ 1 file changed, 18 insertions(+), 12 deletions(-) diff --git a/pyscf/tools/test/test_trexio.py b/pyscf/tools/test/test_trexio.py index dd61bedaf..5989563f2 100644 --- a/pyscf/tools/test/test_trexio.py +++ b/pyscf/tools/test/test_trexio.py @@ -579,6 +579,11 @@ def _hermitize(mat): return 0.5 * (mat + mat.T.conj()) +def _squeeze_k1(mat): + mat = np.asarray(mat) + return mat[0] if mat.ndim == 3 and mat.shape[0] == 1 else mat + + @pytest.mark.parametrize("cart", [False, True], ids=["cart=false", "cart=true"]) def test_write_molecule_integrals_to_trexio_rks(cart): with tempfile.TemporaryDirectory() as d: @@ -721,17 +726,17 @@ def test_write_cell_gamma_integrals_to_trexio_rks(cart): cell0 = pbc.gto.Cell() cell0.cart = cart cell0.build(atom='H 0 0 0; H 0 0 1', basis='6-31g*', a=np.diag([3.0, 3.0, 5.0])) - mf0 = pbc.scf.RHF(cell0).density_fit() + mf0 = pbc.scf.RHF(cell0, kpt=np.zeros(3)).density_fit() mf0.kernel() assert mf0.converged - overlap = _hermitize(mf0.get_ovlp()) - kinetic = _hermitize(cell0.pbc_intor('int1e_kin', 1, 1)) + overlap = _squeeze_k1(_hermitize(mf0.get_ovlp())) + kinetic = _squeeze_k1(_hermitize(cell0.pbc_intor('int1e_kin', 1, 1))) df_builder = mf0.with_df.build() if mf0.with_df is not None else pbc.df.MDF(cell0).build() - potential = _hermitize(df_builder.get_nuc()) + potential = _squeeze_k1(_hermitize(df_builder.get_nuc())) if len(getattr(cell0, '_ecpbas', [])) > 0: from pyscf.pbc.gto import ecp - potential += _hermitize(ecp.ecp_int(cell0)) + potential += _squeeze_k1(_hermitize(ecp.ecp_int(cell0))) core = kinetic + potential trexio.to_trexio(mf0, filename) @@ -744,8 +749,7 @@ def test_write_cell_gamma_integrals_to_trexio_rks(cart): np.testing.assert_allclose(trexio_lib.read_ao_1e_int_core_hamiltonian(tf), core, atol=DIFF_TOL) coeff = mf0.mo_coeff - if coeff.ndim == 3 and coeff.shape[0] == 1: - coeff = coeff[0] + coeff = _squeeze_k1(coeff) if getattr(coeff, 'ndim', 0) == 3 else coeff mo_overlap = _hermitize(coeff.conj().T @ overlap @ coeff) mo_kinetic = _hermitize(coeff.conj().T @ kinetic @ coeff) mo_potential = _hermitize(coeff.conj().T @ potential @ coeff) @@ -794,18 +798,18 @@ def test_write_cell_gamma_integrals_to_trexio_uks(cart): cell0.spin = 2 cell0.cart = cart cell0.build(atom='H 0 0 0; H 0 0 1', basis='6-31g*', a=np.diag([3.0, 3.0, 5.0])) - mf0 = pbc.scf.UKS(cell0).density_fit() + mf0 = pbc.scf.UKS(cell0, kpt=np.zeros(3)).density_fit() mf0.xc = 'LDA' mf0.kernel() assert mf0.converged - overlap = _hermitize(mf0.get_ovlp()) - kinetic = _hermitize(cell0.pbc_intor('int1e_kin', 1, 1)) + overlap = _squeeze_k1(_hermitize(mf0.get_ovlp())) + kinetic = _squeeze_k1(_hermitize(cell0.pbc_intor('int1e_kin', 1, 1))) df_builder = mf0.with_df.build() if mf0.with_df is not None else pbc.df.MDF(cell0).build() - potential = _hermitize(df_builder.get_nuc()) + potential = _squeeze_k1(_hermitize(df_builder.get_nuc())) if len(getattr(cell0, '_ecpbas', [])) > 0: from pyscf.pbc.gto import ecp - potential += _hermitize(ecp.ecp_int(cell0)) + potential += _squeeze_k1(_hermitize(ecp.ecp_int(cell0))) core = kinetic + potential trexio.to_trexio(mf0, filename) @@ -818,6 +822,8 @@ def test_write_cell_gamma_integrals_to_trexio_uks(cart): np.testing.assert_allclose(trexio_lib.read_ao_1e_int_core_hamiltonian(tf), core, atol=DIFF_TOL) coeff_alpha, coeff_beta = mf0.mo_coeff + coeff_alpha = _squeeze_k1(coeff_alpha) if getattr(coeff_alpha, 'ndim', 0) == 3 else coeff_alpha + coeff_beta = _squeeze_k1(coeff_beta) if getattr(coeff_beta, 'ndim', 0) == 3 else coeff_beta coeff = np.concatenate([coeff_alpha, coeff_beta], axis=1) mo_overlap = _hermitize(coeff.conj().T @ overlap @ coeff) mo_kinetic = _hermitize(coeff.conj().T @ kinetic @ coeff) From ed8ca7e48235a792584072644df1909eb3dc1c58 Mon Sep 17 00:00:00 2001 From: kousuke-nakano <37653569+kousuke-nakano@users.noreply.github.com> Date: Sun, 11 Jan 2026 00:39:04 +0900 Subject: [PATCH 3/6] Fix TREXIO gamma-point tests by squeezing singleton k-axis again. --- pyscf/tools/test/test_trexio.py | 18 ++++++++++++++---- 1 file changed, 14 insertions(+), 4 deletions(-) diff --git a/pyscf/tools/test/test_trexio.py b/pyscf/tools/test/test_trexio.py index 5989563f2..f261e272a 100644 --- a/pyscf/tools/test/test_trexio.py +++ b/pyscf/tools/test/test_trexio.py @@ -726,13 +726,18 @@ def test_write_cell_gamma_integrals_to_trexio_rks(cart): cell0 = pbc.gto.Cell() cell0.cart = cart cell0.build(atom='H 0 0 0; H 0 0 1', basis='6-31g*', a=np.diag([3.0, 3.0, 5.0])) - mf0 = pbc.scf.RHF(cell0, kpt=np.zeros(3)).density_fit() + gamma_kpt = np.zeros(3) + mf0 = pbc.scf.RHF(cell0, kpt=gamma_kpt).density_fit() mf0.kernel() assert mf0.converged overlap = _squeeze_k1(_hermitize(mf0.get_ovlp())) kinetic = _squeeze_k1(_hermitize(cell0.pbc_intor('int1e_kin', 1, 1))) - df_builder = mf0.with_df.build() if mf0.with_df is not None else pbc.df.MDF(cell0).build() + df_builder = ( + mf0.with_df.build() + if mf0.with_df is not None + else pbc.df.MDF(cell0, kpts=[gamma_kpt]).build() + ) potential = _squeeze_k1(_hermitize(df_builder.get_nuc())) if len(getattr(cell0, '_ecpbas', [])) > 0: from pyscf.pbc.gto import ecp @@ -798,14 +803,19 @@ def test_write_cell_gamma_integrals_to_trexio_uks(cart): cell0.spin = 2 cell0.cart = cart cell0.build(atom='H 0 0 0; H 0 0 1', basis='6-31g*', a=np.diag([3.0, 3.0, 5.0])) - mf0 = pbc.scf.UKS(cell0, kpt=np.zeros(3)).density_fit() + gamma_kpt = np.zeros(3) + mf0 = pbc.scf.UKS(cell0, kpt=gamma_kpt).density_fit() mf0.xc = 'LDA' mf0.kernel() assert mf0.converged overlap = _squeeze_k1(_hermitize(mf0.get_ovlp())) kinetic = _squeeze_k1(_hermitize(cell0.pbc_intor('int1e_kin', 1, 1))) - df_builder = mf0.with_df.build() if mf0.with_df is not None else pbc.df.MDF(cell0).build() + df_builder = ( + mf0.with_df.build() + if mf0.with_df is not None + else pbc.df.MDF(cell0, kpts=[gamma_kpt]).build() + ) potential = _squeeze_k1(_hermitize(df_builder.get_nuc())) if len(getattr(cell0, '_ecpbas', [])) > 0: from pyscf.pbc.gto import ecp From 11dcdd1963b700ece78d58b1f86746ed4ac3fcd2 Mon Sep 17 00:00:00 2001 From: kousuke-nakano <37653569+kousuke-nakano@users.noreply.github.com> Date: Sun, 11 Jan 2026 01:04:35 +0900 Subject: [PATCH 4/6] Fix TREXIO gamma-point tests by squeezing singleton k-axis again. --- pyscf/tools/test/test_trexio.py | 26 ++++++++++++++++---------- 1 file changed, 16 insertions(+), 10 deletions(-) diff --git a/pyscf/tools/test/test_trexio.py b/pyscf/tools/test/test_trexio.py index f261e272a..21b4bc8c3 100644 --- a/pyscf/tools/test/test_trexio.py +++ b/pyscf/tools/test/test_trexio.py @@ -584,6 +584,12 @@ def _squeeze_k1(mat): return mat[0] if mat.ndim == 3 and mat.shape[0] == 1 else mat +def _take_gamma(mat): + """Pick the gamma (first) k-block if present.""" + mat = np.asarray(mat) + return mat[0] if mat.ndim >= 3 else mat + + @pytest.mark.parametrize("cart", [False, True], ids=["cart=false", "cart=true"]) def test_write_molecule_integrals_to_trexio_rks(cart): with tempfile.TemporaryDirectory() as d: @@ -731,17 +737,17 @@ def test_write_cell_gamma_integrals_to_trexio_rks(cart): mf0.kernel() assert mf0.converged - overlap = _squeeze_k1(_hermitize(mf0.get_ovlp())) - kinetic = _squeeze_k1(_hermitize(cell0.pbc_intor('int1e_kin', 1, 1))) + overlap = _hermitize(_take_gamma(mf0.get_ovlp())) + kinetic = _hermitize(_take_gamma(cell0.pbc_intor('int1e_kin', 1, 1))) df_builder = ( mf0.with_df.build() if mf0.with_df is not None else pbc.df.MDF(cell0, kpts=[gamma_kpt]).build() ) - potential = _squeeze_k1(_hermitize(df_builder.get_nuc())) + potential = _hermitize(_take_gamma(df_builder.get_nuc())) if len(getattr(cell0, '_ecpbas', [])) > 0: from pyscf.pbc.gto import ecp - potential += _squeeze_k1(_hermitize(ecp.ecp_int(cell0))) + potential += _hermitize(_take_gamma(ecp.ecp_int(cell0))) core = kinetic + potential trexio.to_trexio(mf0, filename) @@ -809,17 +815,17 @@ def test_write_cell_gamma_integrals_to_trexio_uks(cart): mf0.kernel() assert mf0.converged - overlap = _squeeze_k1(_hermitize(mf0.get_ovlp())) - kinetic = _squeeze_k1(_hermitize(cell0.pbc_intor('int1e_kin', 1, 1))) + overlap = _hermitize(_take_gamma(mf0.get_ovlp())) + kinetic = _hermitize(_take_gamma(cell0.pbc_intor('int1e_kin', 1, 1))) df_builder = ( mf0.with_df.build() if mf0.with_df is not None else pbc.df.MDF(cell0, kpts=[gamma_kpt]).build() ) - potential = _squeeze_k1(_hermitize(df_builder.get_nuc())) + potential = _hermitize(_take_gamma(df_builder.get_nuc())) if len(getattr(cell0, '_ecpbas', [])) > 0: from pyscf.pbc.gto import ecp - potential += _squeeze_k1(_hermitize(ecp.ecp_int(cell0))) + potential += _hermitize(_take_gamma(ecp.ecp_int(cell0))) core = kinetic + potential trexio.to_trexio(mf0, filename) @@ -832,8 +838,8 @@ def test_write_cell_gamma_integrals_to_trexio_uks(cart): np.testing.assert_allclose(trexio_lib.read_ao_1e_int_core_hamiltonian(tf), core, atol=DIFF_TOL) coeff_alpha, coeff_beta = mf0.mo_coeff - coeff_alpha = _squeeze_k1(coeff_alpha) if getattr(coeff_alpha, 'ndim', 0) == 3 else coeff_alpha - coeff_beta = _squeeze_k1(coeff_beta) if getattr(coeff_beta, 'ndim', 0) == 3 else coeff_beta + coeff_alpha = _take_gamma(coeff_alpha) + coeff_beta = _take_gamma(coeff_beta) coeff = np.concatenate([coeff_alpha, coeff_beta], axis=1) mo_overlap = _hermitize(coeff.conj().T @ overlap @ coeff) mo_kinetic = _hermitize(coeff.conj().T @ kinetic @ coeff) From d738d479aca6839ef0578f32737626a66270fa7d Mon Sep 17 00:00:00 2001 From: kousuke-nakano <37653569+kousuke-nakano@users.noreply.github.com> Date: Thu, 29 Jan 2026 01:15:19 +0900 Subject: [PATCH 5/6] Completed 1e-ERI, 2e-ERI, 1RDM, and 2RDM implementations. --- pyscf/tools/test/test_trexio.py | 1850 ++++++++++++++++++++++++++++++- pyscf/tools/trexio.py | 567 ++++++++-- 2 files changed, 2306 insertions(+), 111 deletions(-) diff --git a/pyscf/tools/test/test_trexio.py b/pyscf/tools/test/test_trexio.py index 21b4bc8c3..01bb47d47 100644 --- a/pyscf/tools/test/test_trexio.py +++ b/pyscf/tools/test/test_trexio.py @@ -2,12 +2,14 @@ from pyscf import df from pyscf.tools import trexio import os +import itertools import numpy as np import tempfile import pytest import trexio as trexio_lib from pyscf import ao2mo from pyscf import pbc +from pyscf.pbc import df as pbcdf DIFF_TOL = 1e-10 @@ -555,14 +557,15 @@ def test_mol_rhf_ccecp_ccpvqz(cart): # writing `1e_int` and `2e_int` to trexio file ################################################################# -def _trexio_pack_eri(eri, basis): +def _trexio_pack_eri(eri, basis, sym='s1'): basis = basis.upper() + sym = sym.lower() if basis not in ('AO', 'MO'): raise ValueError("basis must be 'AO' or 'MO'") with tempfile.TemporaryDirectory() as tmpdir: filename = os.path.join(tmpdir, 'pack.h5') - _write_2e_int_eri(eri, filename, backend='h5', basis=basis) + _write_2e_int_eri(eri, filename, backend='h5', basis=basis, sym=sym) with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: if basis == 'AO': size = trexio_lib.read_ao_2e_int_eri_size(tf) @@ -574,6 +577,23 @@ def _trexio_pack_eri(eri, basis): return np.asarray(idx, dtype=np.int32).ravel(), np.asarray(val) +def _expand_trexio_eri(idx, val, n): + idx = np.asarray(idx, dtype=np.int64).reshape(-1, 4) + val = np.asarray(val) + w_chem = np.zeros((n, n, n, n)) + for (p, r, q, s), v in zip(idx, val): + i, j, k, l = p, q, r, s + w_chem[i, j, k, l] = v + w_chem[j, i, k, l] = v + w_chem[i, j, l, k] = v + w_chem[j, i, l, k] = v + w_chem[k, l, i, j] = v + w_chem[l, k, i, j] = v + w_chem[k, l, j, i] = v + w_chem[l, k, j, i] = v + return w_chem.transpose(0, 2, 1, 3) + + def _hermitize(mat): mat = np.asarray(mat) return 0.5 * (mat + mat.T.conj()) @@ -591,7 +611,7 @@ def _take_gamma(mat): @pytest.mark.parametrize("cart", [False, True], ids=["cart=false", "cart=true"]) -def test_write_molecule_integrals_to_trexio_rks(cart): +def test_write_molecule_integrals_sym_s1_to_trexio_rhf_ae(cart): with tempfile.TemporaryDirectory() as d: filename = os.path.join(d, 'mol_integrals.h5') @@ -602,10 +622,10 @@ def test_write_molecule_integrals_to_trexio_rks(cart): kinetic = _hermitize(mol0.intor('int1e_kin')) potential = _hermitize(mol0.intor('int1e_nuc')) if mol0._ecp: - from pyscf.gto import ecp - potential += _hermitize(ecp.ecp_int(mol0)) + potential += _hermitize(mol0.intor('ECPscalar')) core = kinetic + potential + trexio.to_trexio(mf0, filename) trexio.write_scf_1e_int_eri(mf0, filename, basis='AO') @@ -628,9 +648,9 @@ def test_write_molecule_integrals_to_trexio_rks(cart): np.testing.assert_allclose(trexio_lib.read_mo_1e_int_potential_n_e(tf), mo_potential, atol=DIFF_TOL) np.testing.assert_allclose(trexio_lib.read_mo_1e_int_core_hamiltonian(tf), mo_core, atol=DIFF_TOL) - ao_eri = mol0.intor('int2e', aosym='s8') + ao_eri = mol0.intor('int2e', aosym='s1') ao_idx_exp, ao_val_exp = _trexio_pack_eri(ao_eri, 'AO') - trexio.write_scf_2e_int_eri(mf0, filename, basis='AO', aosym='s8') + trexio.write_scf_2e_int_eri(mf0, filename, basis='AO') with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: assert trexio_lib.has_ao_2e_int_eri(tf) size = trexio_lib.read_ao_2e_int_eri_size(tf) @@ -639,9 +659,11 @@ def test_write_molecule_integrals_to_trexio_rks(cart): np.testing.assert_array_equal(np.asarray(idx, dtype=np.int32).ravel(), ao_idx_exp) np.testing.assert_allclose(np.asarray(val), ao_val_exp, atol=DIFF_TOL) - mo_eri = ao2mo.kernel(mol0, coeff, compact=True) + mo_eri = ao2mo.kernel(mol0, coeff, compact=False) + nmo = coeff.shape[1] + mo_eri = mo_eri.reshape(nmo, nmo, nmo, nmo) mo_idx_exp, mo_val_exp = _trexio_pack_eri(mo_eri, 'MO') - trexio.write_scf_2e_int_eri(mf0, filename, basis='MO', aosym='s8') + trexio.write_scf_2e_int_eri(mf0, filename, basis='MO') with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: assert trexio_lib.has_mo_2e_int_eri(tf) size = trexio_lib.read_mo_2e_int_eri_size(tf) @@ -652,22 +674,19 @@ def test_write_molecule_integrals_to_trexio_rks(cart): @pytest.mark.parametrize("cart", [False, True], ids=["cart=false", "cart=true"]) -def test_write_molecule_integrals_to_trexio_uks(cart): +def test_write_molecule_integrals_sym_s1_to_trexio_uhf_ae(cart): with tempfile.TemporaryDirectory() as d: - filename = os.path.join(d, 'mol_uks_integrals.h5') + filename = os.path.join(d, 'mol_uhf_integrals.h5') mol0 = pyscf.M(atom='O 0 0 0', basis='6-31g*', spin=2, cart=cart) - mf0 = mol0.UKS() - mf0.xc = 'LDA' - mf0.kernel() + mf0 = mol0.UHF().run() assert mf0.converged overlap = _hermitize(mf0.get_ovlp()) kinetic = _hermitize(mol0.intor('int1e_kin')) potential = _hermitize(mol0.intor('int1e_nuc')) if mol0._ecp: - from pyscf.gto import ecp - potential += _hermitize(ecp.ecp_int(mol0)) + potential += _hermitize(mol0.intor('ECPscalar')) core = kinetic + potential trexio.to_trexio(mf0, filename) @@ -701,9 +720,9 @@ def test_write_molecule_integrals_to_trexio_uks(cart): trexio_lib.read_mo_1e_int_core_hamiltonian(tf), mo_core, atol=DIFF_TOL ) - ao_eri = mol0.intor('int2e', aosym='s8') + ao_eri = mol0.intor('int2e', aosym='s1') ao_idx_exp, ao_val_exp = _trexio_pack_eri(ao_eri, 'AO') - trexio.write_scf_2e_int_eri(mf0, filename, basis='AO', aosym='s8') + trexio.write_scf_2e_int_eri(mf0, filename, basis='AO') with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: assert trexio_lib.has_ao_2e_int_eri(tf) size = trexio_lib.read_ao_2e_int_eri_size(tf) @@ -712,9 +731,45 @@ def test_write_molecule_integrals_to_trexio_uks(cart): np.testing.assert_array_equal(np.asarray(idx, dtype=np.int32).ravel(), ao_idx_exp) np.testing.assert_allclose(np.asarray(val), ao_val_exp, atol=DIFF_TOL) - mo_eri = ao2mo.kernel(mol0, coeff, compact=True) + mo_eri = ao2mo.kernel(mol0, coeff, compact=False) + nmo = coeff.shape[1] + mo_eri = mo_eri.reshape(nmo, nmo, nmo, nmo) mo_idx_exp, mo_val_exp = _trexio_pack_eri(mo_eri, 'MO') - trexio.write_scf_2e_int_eri(mf0, filename, basis='MO', aosym='s8') + trexio.write_scf_2e_int_eri(mf0, filename, basis='MO') + with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: + assert trexio_lib.has_mo_2e_int_eri(tf) + size = trexio_lib.read_mo_2e_int_eri_size(tf) + idx, val, n_read, _ = trexio_lib.read_mo_2e_int_eri(tf, 0, size) + assert n_read == size + np.testing.assert_array_equal(np.asarray(idx, dtype=np.int32).ravel(), mo_idx_exp) + np.testing.assert_allclose(np.asarray(val), mo_val_exp, atol=DIFF_TOL) + + +@pytest.mark.parametrize("cart", [False, True], ids=["cart=false", "cart=true"]) +def test_write_molecule_integrals_sym_s4_to_trexio_rhf_ae(cart): + with tempfile.TemporaryDirectory() as d: + filename = os.path.join(d, 'mol_integrals_s4.h5') + + mol0 = pyscf.M(atom='H 0 0 0; F 0 0 1', basis='6-31g*', cart=cart) + mf0 = mol0.RHF().run() + + trexio.to_trexio(mf0, filename) + + ao_eri = mol0.intor('int2e', aosym='s4') + ao_idx_exp, ao_val_exp = _trexio_pack_eri(ao_eri, 'AO', sym='s4') + trexio.write_scf_2e_int_eri(mf0, filename, basis='AO', sym='s4') + with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: + assert trexio_lib.has_ao_2e_int_eri(tf) + size = trexio_lib.read_ao_2e_int_eri_size(tf) + idx, val, n_read, _ = trexio_lib.read_ao_2e_int_eri(tf, 0, size) + assert n_read == size + np.testing.assert_array_equal(np.asarray(idx, dtype=np.int32).ravel(), ao_idx_exp) + np.testing.assert_allclose(np.asarray(val), ao_val_exp, atol=DIFF_TOL) + + coeff = mf0.mo_coeff + mo_eri = ao2mo.kernel(mol0, coeff, compact=True) + mo_idx_exp, mo_val_exp = _trexio_pack_eri(mo_eri, 'MO', sym='s4') + trexio.write_scf_2e_int_eri(mf0, filename, basis='MO', sym='s4') with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: assert trexio_lib.has_mo_2e_int_eri(tf) size = trexio_lib.read_mo_2e_int_eri_size(tf) @@ -725,7 +780,88 @@ def test_write_molecule_integrals_to_trexio_uks(cart): @pytest.mark.parametrize("cart", [False, True], ids=["cart=false", "cart=true"]) -def test_write_cell_gamma_integrals_to_trexio_rks(cart): +def test_write_molecule_integrals_sym_s4_to_trexio_uhf_ae(cart): + with tempfile.TemporaryDirectory() as d: + filename = os.path.join(d, 'mol_uhf_integrals_s4.h5') + + mol0 = pyscf.M(atom='O 0 0 0', basis='6-31g*', spin=2, cart=cart) + mf0 = mol0.UHF().run() + assert mf0.converged + + trexio.to_trexio(mf0, filename) + + ao_eri = mol0.intor('int2e', aosym='s4') + ao_idx_exp, ao_val_exp = _trexio_pack_eri(ao_eri, 'AO', sym='s4') + trexio.write_scf_2e_int_eri(mf0, filename, basis='AO', sym='s4') + with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: + assert trexio_lib.has_ao_2e_int_eri(tf) + size = trexio_lib.read_ao_2e_int_eri_size(tf) + idx, val, n_read, _ = trexio_lib.read_ao_2e_int_eri(tf, 0, size) + assert n_read == size + np.testing.assert_array_equal(np.asarray(idx, dtype=np.int32).ravel(), ao_idx_exp) + np.testing.assert_allclose(np.asarray(val), ao_val_exp, atol=DIFF_TOL) + + coeff_alpha, coeff_beta = mf0.mo_coeff + coeff = np.concatenate([coeff_alpha, coeff_beta], axis=1) + mo_eri = ao2mo.kernel(mol0, coeff, compact=True) + mo_idx_exp, mo_val_exp = _trexio_pack_eri(mo_eri, 'MO', sym='s4') + trexio.write_scf_2e_int_eri(mf0, filename, basis='MO', sym='s4') + with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: + assert trexio_lib.has_mo_2e_int_eri(tf) + size = trexio_lib.read_mo_2e_int_eri_size(tf) + idx, val, n_read, _ = trexio_lib.read_mo_2e_int_eri(tf, 0, size) + assert n_read == size + np.testing.assert_array_equal(np.asarray(idx, dtype=np.int32).ravel(), mo_idx_exp) + np.testing.assert_allclose(np.asarray(val), mo_val_exp, atol=DIFF_TOL) + + +@pytest.mark.parametrize("cart", [False, True], ids=["cart=false", "cart=true"]) +def test_write_molecule_integrals_sym_s8_to_trexio_rhf_ae(cart): + with tempfile.TemporaryDirectory() as d: + filename = os.path.join(d, 'mol_integrals_s8.h5') + + mol0 = pyscf.M(atom='H 0 0 0; F 0 0 1', basis='6-31g*', cart=cart) + mf0 = mol0.RHF().run() + + trexio.to_trexio(mf0, filename) + + ao_eri = mol0.intor('int2e', aosym='s8') + ao_idx_exp, ao_val_exp = _trexio_pack_eri(ao_eri, 'AO', sym='s8') + trexio.write_scf_2e_int_eri(mf0, filename, basis='AO', sym='s8') + with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: + assert trexio_lib.has_ao_2e_int_eri(tf) + size = trexio_lib.read_ao_2e_int_eri_size(tf) + idx, val, n_read, _ = trexio_lib.read_ao_2e_int_eri(tf, 0, size) + assert n_read == size + np.testing.assert_array_equal(np.asarray(idx, dtype=np.int32).ravel(), ao_idx_exp) + np.testing.assert_allclose(np.asarray(val), ao_val_exp, atol=DIFF_TOL) + + +@pytest.mark.parametrize("cart", [False, True], ids=["cart=false", "cart=true"]) +def test_write_molecule_integrals_sym_s8_to_trexio_uhf_ae(cart): + with tempfile.TemporaryDirectory() as d: + filename = os.path.join(d, 'mol_uhf_integrals_s8.h5') + + mol0 = pyscf.M(atom='O 0 0 0', basis='6-31g*', spin=2, cart=cart) + mf0 = mol0.UHF().run() + assert mf0.converged + + trexio.to_trexio(mf0, filename) + + ao_eri = mol0.intor('int2e', aosym='s8') + ao_idx_exp, ao_val_exp = _trexio_pack_eri(ao_eri, 'AO', sym='s8') + trexio.write_scf_2e_int_eri(mf0, filename, basis='AO', sym='s8') + with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: + assert trexio_lib.has_ao_2e_int_eri(tf) + size = trexio_lib.read_ao_2e_int_eri_size(tf) + idx, val, n_read, _ = trexio_lib.read_ao_2e_int_eri(tf, 0, size) + assert n_read == size + np.testing.assert_array_equal(np.asarray(idx, dtype=np.int32).ravel(), ao_idx_exp) + np.testing.assert_allclose(np.asarray(val), ao_val_exp, atol=DIFF_TOL) + + +@pytest.mark.parametrize("cart", [False, True], ids=["cart=false", "cart=true"]) +def test_write_cell_gamma_integrals_sym_s1_to_trexio_rhf_ae(cart): with tempfile.TemporaryDirectory() as d: filename = os.path.join(d, 'cell_integrals.h5') @@ -789,7 +925,14 @@ def test_write_cell_gamma_integrals_to_trexio_rks(cart): df_obj_mo = pbc.df.MDF(cell0).build() mo_eri = df_obj_mo.get_mo_eri((coeff, coeff, coeff, coeff)) - mo_idx_exp, mo_val_exp = _trexio_pack_eri(np.real_if_close (mo_eri), 'MO') + mo_eri = np.real_if_close(mo_eri) + nmo = coeff.shape[1] + if mo_eri.ndim == 2: + # Pair-pair (s4) -> expand to full tensor + mo_eri = ao2mo.restore(1, ao2mo.restore(4, mo_eri, nmo), nmo) + elif mo_eri.ndim < 4: + mo_eri = ao2mo.restore(1, mo_eri, nmo) + mo_idx_exp, mo_val_exp = _trexio_pack_eri(mo_eri, 'MO') trexio.write_scf_2e_int_eri(mf0, filename, basis='MO') with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: assert trexio_lib.has_mo_2e_int_eri(tf) @@ -801,17 +944,16 @@ def test_write_cell_gamma_integrals_to_trexio_rks(cart): @pytest.mark.parametrize("cart", [False, True], ids=["cart=false", "cart=true"]) -def test_write_cell_gamma_integrals_to_trexio_uks(cart): +def test_write_cell_gamma_integrals_sym_s1_to_trexio_uhf_ae(cart): with tempfile.TemporaryDirectory() as d: - filename = os.path.join(d, 'cell_uks_integrals.h5') + filename = os.path.join(d, 'cell_uhf_integrals.h5') cell0 = pbc.gto.Cell() cell0.spin = 2 cell0.cart = cart cell0.build(atom='H 0 0 0; H 0 0 1', basis='6-31g*', a=np.diag([3.0, 3.0, 5.0])) gamma_kpt = np.zeros(3) - mf0 = pbc.scf.UKS(cell0, kpt=gamma_kpt).density_fit() - mf0.xc = 'LDA' + mf0 = pbc.scf.UHF(cell0, kpt=gamma_kpt).density_fit() mf0.kernel() assert mf0.converged @@ -869,7 +1011,13 @@ def test_write_cell_gamma_integrals_to_trexio_uks(cart): df_obj_mo = pbc.df.MDF(cell0).build() mo_eri = df_obj_mo.get_mo_eri((coeff, coeff, coeff, coeff)) - mo_idx_exp, mo_val_exp = _trexio_pack_eri(np.real_if_close(mo_eri), 'MO') + mo_eri = np.real_if_close(mo_eri) + nmo = coeff.shape[1] + if mo_eri.ndim == 2: + mo_eri = ao2mo.restore(1, ao2mo.restore(4, mo_eri, nmo), nmo) + elif mo_eri.ndim < 4: + mo_eri = ao2mo.restore(1, mo_eri, nmo) + mo_idx_exp, mo_val_exp = _trexio_pack_eri(mo_eri, 'MO') trexio.write_scf_2e_int_eri(mf0, filename, basis='MO') with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: assert trexio_lib.has_mo_2e_int_eri(tf) @@ -878,3 +1026,1651 @@ def test_write_cell_gamma_integrals_to_trexio_uks(cart): assert n_read == size np.testing.assert_array_equal(np.asarray(idx, dtype=np.int32).ravel(), mo_idx_exp) np.testing.assert_allclose(np.asarray(val), mo_val_exp, atol=DIFF_TOL) + + +@pytest.mark.parametrize("cart", [False, True], ids=["cart=false", "cart=true"]) +def test_energy_molecule_integrals_sym_s1_in_trexio_rhf_ae(cart): + with tempfile.TemporaryDirectory() as d: + filename = os.path.join(d, 'mol_integrals_energy.h5') + + mol0 = pyscf.M(atom='H 0 0 0; F 0 0 1', basis='6-31g*', cart=cart) + mf0 = mol0.RHF().run() + + overlap = _hermitize(mf0.get_ovlp()) + kinetic = _hermitize(mol0.intor('int1e_kin')) + potential = _hermitize(mol0.intor('int1e_nuc')) + if mol0._ecp: + potential += _hermitize(mol0.intor('ECPscalar')) + core = kinetic + potential + + trexio.to_trexio(mf0, filename) + + trexio.write_scf_1e_int_eri(mf0, filename, basis='AO') + with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: + np.testing.assert_allclose(trexio_lib.read_ao_1e_int_overlap(tf), overlap, atol=DIFF_TOL) + np.testing.assert_allclose(trexio_lib.read_ao_1e_int_kinetic(tf), kinetic, atol=DIFF_TOL) + np.testing.assert_allclose(trexio_lib.read_ao_1e_int_potential_n_e(tf), potential, atol=DIFF_TOL) + np.testing.assert_allclose(trexio_lib.read_ao_1e_int_core_hamiltonian(tf), core, atol=DIFF_TOL) + + trexio.write_scf_2e_int_eri(mf0, filename, basis='AO') + with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: + assert trexio_lib.has_ao_2e_int_eri(tf) + size = trexio_lib.read_ao_2e_int_eri_size(tf) + idx, val, n_read, _ = trexio_lib.read_ao_2e_int_eri(tf, 0, size) + assert n_read == size + + trexio.write_scf_1e_int_eri(mf0, filename, basis='MO') + trexio.write_scf_2e_int_eri(mf0, filename, basis='MO') + trexio.write_scf_rdm_1e(mf0, filename) + trexio.write_scf_rdm_2e(mf0, filename) + + BUFSIZE = 100000 + with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: + e_nn = trexio_lib.read_nucleus_repulsion(tf) + + # AO energy reconstruction + nao = trexio_lib.read_ao_num(tf) + core_ao = trexio_lib.read_ao_1e_int_core_hamiltonian(tf) + core_ao = np.asarray(core_ao) + if core_ao.ndim == 1: + core_ao = core_ao.reshape(nao, nao) + + dm_ao = mf0.make_rdm1() + + ao_eri_size = trexio_lib.read_ao_2e_int_eri_size(tf) + assert ao_eri_size == nao ** 4 + offset = 0 + W_ao_phys = np.zeros((nao, nao, nao, nao)) # raw order as stored + while offset < ao_eri_size: + bufsize = min(BUFSIZE, ao_eri_size - offset) + buf_idx, buf_val, icount, feof = trexio_lib.read_ao_2e_int_eri(tf, offset, bufsize) + offset += icount + if icount: + p, r, q, s = np.asarray(buf_idx, dtype=np.int64).T # stored as (p,r,q,s) + W_ao_phys[p, r, q, s] = buf_val + if feof: + break + + W_ao_chem = W_ao_phys.transpose(0, 2, 1, 3) + + J = np.einsum('pqrs,rs->pq', W_ao_chem, dm_ao) + K = np.einsum('prqs,rs->pq', W_ao_chem, dm_ao) + vhf = J - 0.5 * K + e_ao = e_nn + e_ao += np.einsum('pq,pq->', dm_ao, core_ao) + e_ao += 0.5 * np.einsum('pq,pq->', dm_ao, vhf) + + # MO energy reconstruction + n = trexio_lib.read_mo_num(tf) + core = trexio_lib.read_mo_1e_int_core_hamiltonian(tf) + core = np.asarray(core) + if core.ndim == 1: + core = core.reshape(n, n) + + eri_size = trexio_lib.read_mo_2e_int_eri_size(tf) + assert eri_size == n ** 4 + offset = 0 + W_mo_phys = np.zeros((n, n, n, n)) + while offset < eri_size: + bufsize = min(BUFSIZE, eri_size - offset) + buf_idx, buf_val, icount, feof = trexio_lib.read_mo_2e_int_eri(tf, offset, bufsize) + offset += icount + if icount: + i, j, k, l = np.asarray(buf_idx, dtype=np.int64).T + W_mo_phys[i, j, k, l] = buf_val + if feof: + break + dens = trexio_lib.read_rdm_1e(tf, doReshape=True) + dens = np.asarray(dens) + if dens.ndim == 1: + dens = dens.reshape(n, n) + + offset = 0 + G_mo_phys = np.zeros((n, n, n, n)) + while True: + buf_idx, buf_val, icount, feof = trexio_lib.read_rdm_2e(tf, offset, BUFSIZE) + offset += icount + if icount: + i, j, k, l = np.asarray(buf_idx, dtype=np.int64).T + G_mo_phys[i, j, k, l] = buf_val + if feof: + break + + e_mo = e_nn + e_mo += np.einsum('pq,pq->', dens, core) + e_mo += 0.5 * np.tensordot(G_mo_phys, W_mo_phys, axes=4) + + assert abs(e_ao - mf0.e_tot) < 1e-8 + assert abs(e_mo - mf0.e_tot) < 1e-8 + + +@pytest.mark.parametrize("cart", [False, True], ids=["cart=false", "cart=true"]) +def test_energy_crystal_integrals_sym_s1_in_trexio_rhf_ae(cart): + with tempfile.TemporaryDirectory() as d: + filename = os.path.join(d, 'pbc_gamma_rdm_energy.h5') + + cell = pyscf.pbc.gto.Cell() + cell.cart = cart + cell.unit = 'Bohr' + cell.build( + atom='H 0 0 0; H 0 0 1.4', + basis='sto-3g', + a=np.diag([3.0, 3.0, 5.0]), + ) + + mf0 = pyscf.pbc.scf.RHF(cell) + mf0.with_df = pbcdf.MDF(cell).build() + mf0.run() + assert mf0.converged + + trexio.to_trexio(mf0, filename) + trexio.write_scf_1e_int_eri(mf0, filename, basis='AO') + trexio.write_scf_2e_int_eri(mf0, filename, basis='AO') + trexio.write_scf_1e_int_eri(mf0, filename, basis='MO') + trexio.write_scf_2e_int_eri(mf0, filename, basis='MO') + trexio.write_scf_rdm_1e(mf0, filename) + trexio.write_scf_rdm_2e(mf0, filename) + + BUFSIZE = 100000 + with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: + e_nn = trexio_lib.read_nucleus_repulsion(tf) + if abs(e_nn) < 1e-12: + e_nn_mol = cell.energy_nuc() + if abs(e_nn_mol) > 1e-12: + e_nn = e_nn_mol + madelung = 0.0 + try: + madelung = trexio_lib.read_pbc_madelung(tf) + madelung = float(np.asarray(madelung).ravel()[0]) + except trexio_lib.Error: + madelung = 0.0 + + n = trexio_lib.read_mo_num(tf) + core = trexio_lib.read_mo_1e_int_core_hamiltonian(tf) + core = np.asarray(core) + if core.ndim == 1: + core = core.reshape(n, n) + + eri_size = trexio_lib.read_mo_2e_int_eri_size(tf) + assert eri_size == n ** 4 + offset = 0 + W_mo_phys = np.zeros((n, n, n, n)) + while offset < eri_size: + bufsize = min(BUFSIZE, eri_size - offset) + buf_idx, buf_val, icount, feof = trexio_lib.read_mo_2e_int_eri(tf, offset, bufsize) + offset += icount + if icount: + i, j, k, l = np.asarray(buf_idx, dtype=np.int64).T + W_mo_phys[i, j, k, l] = buf_val + if feof: + break + + dens = trexio_lib.read_rdm_1e(tf, doReshape=True) + dens = np.asarray(dens) + if dens.ndim == 1: + dens = dens.reshape(n, n) + + offset = 0 + G_mo_phys = np.zeros((n, n, n, n)) + while True: + buf_idx, buf_val, icount, feof = trexio_lib.read_rdm_2e(tf, offset, BUFSIZE) + offset += icount + if icount: + i, j, k, l = np.asarray(buf_idx, dtype=np.int64).T + G_mo_phys[i, j, k, l] = buf_val + if feof: + break + + e_mo = e_nn - madelung + e_mo += np.einsum('pq,pq->', dens, core) + e_mo += 0.5 * np.tensordot(G_mo_phys, W_mo_phys, axes=4) + + assert abs(e_mo - mf0.e_tot) < 1e-6 + + +@pytest.mark.parametrize("cart", [False, True], ids=["cart=false", "cart=true"]) +def test_energy_crystal_integrals_sym_s1_in_trexio_uhf_ae(cart): + with tempfile.TemporaryDirectory() as d: + filename = os.path.join(d, 'pbc_gamma_rdm_energy_uhf.h5') + + cell = pyscf.pbc.gto.Cell() + cell.cart = cart + cell.unit = 'Bohr' + cell.spin = 2 + cell.build( + atom='H 0 0 0; H 0 0 1.4', + basis='sto-3g', + a=np.diag([3.0, 3.0, 5.0]), + ) + + mf0 = pyscf.pbc.scf.UHF(cell) + mf0.with_df = pbcdf.MDF(cell).build() + mf0.run() + assert mf0.converged + + trexio.to_trexio(mf0, filename) + trexio.write_scf_1e_int_eri(mf0, filename, basis='AO') + trexio.write_scf_2e_int_eri(mf0, filename, basis='AO') + trexio.write_scf_1e_int_eri(mf0, filename, basis='MO') + trexio.write_scf_2e_int_eri(mf0, filename, basis='MO') + trexio.write_scf_rdm_1e(mf0, filename) + trexio.write_scf_rdm_2e(mf0, filename) + + BUFSIZE = 100000 + with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: + e_nn = trexio_lib.read_nucleus_repulsion(tf) + if abs(e_nn) < 1e-12: + e_nn_mol = cell.energy_nuc() + if abs(e_nn_mol) > 1e-12: + e_nn = e_nn_mol + madelung = 0.0 + try: + madelung = trexio_lib.read_pbc_madelung(tf) + madelung = float(np.asarray(madelung).ravel()[0]) + except trexio_lib.Error: + madelung = 0.0 + + n = trexio_lib.read_mo_num(tf) + core = trexio_lib.read_mo_1e_int_core_hamiltonian(tf) + core = np.asarray(core) + if core.ndim == 1: + core = core.reshape(n, n) + + eri_size = trexio_lib.read_mo_2e_int_eri_size(tf) + assert eri_size == n ** 4 + offset = 0 + W_mo_phys = np.zeros((n, n, n, n)) + while offset < eri_size: + bufsize = min(BUFSIZE, eri_size - offset) + buf_idx, buf_val, icount, feof = trexio_lib.read_mo_2e_int_eri(tf, offset, bufsize) + offset += icount + if icount: + i, j, k, l = np.asarray(buf_idx, dtype=np.int64).T + W_mo_phys[i, j, k, l] = buf_val + if feof: + break + + dens = trexio_lib.read_rdm_1e(tf, doReshape=True) + dens = np.asarray(dens) + if dens.ndim == 1: + dens = dens.reshape(n, n) + + spin_rdm = ( + trexio_lib.has_rdm_1e_up(tf) + and trexio_lib.has_rdm_2e_upup(tf) + and trexio_lib.has_rdm_2e_dndn(tf) + and trexio_lib.has_rdm_2e_updn(tf) + ) + + assert spin_rdm + + rdm2_size = trexio_lib.read_rdm_2e_upup_size(tf) + n_up = int(round(rdm2_size ** 0.25)) + assert rdm2_size == n_up ** 4 + + rdm2_size = trexio_lib.read_rdm_2e_dndn_size(tf) + n_dn = int(round(rdm2_size ** 0.25)) + assert rdm2_size == n_dn ** 4 + + rdm2_size = trexio_lib.read_rdm_2e_updn_size(tf) + assert rdm2_size == (n_up * n_dn) ** 2 + + offset = 0 + G_uu = np.zeros((n_up, n_up, n_up, n_up)) + while True: + buf_idx, buf_val, icount, feof = trexio_lib.read_rdm_2e_upup(tf, offset, BUFSIZE) + offset += icount + if icount: + i, j, k, l = np.asarray(buf_idx, dtype=np.int64).T + G_uu[i, j, k, l] = buf_val + if feof: + break + + offset = 0 + G_dd = np.zeros((n_dn, n_dn, n_dn, n_dn)) + while True: + buf_idx, buf_val, icount, feof = trexio_lib.read_rdm_2e_dndn(tf, offset, BUFSIZE) + offset += icount + if icount: + i, j, k, l = np.asarray(buf_idx, dtype=np.int64).T + G_dd[i, j, k, l] = buf_val + if feof: + break + + offset = 0 + G_ud = np.zeros((n_up, n_dn, n_up, n_dn)) + while True: + buf_idx, buf_val, icount, feof = trexio_lib.read_rdm_2e_updn(tf, offset, BUFSIZE) + offset += icount + if icount: + i, j, k, l = np.asarray(buf_idx, dtype=np.int64).T + G_ud[i, j, k, l] = buf_val + if feof: + break + + e_mo = e_nn - madelung + off = n_up + + W_uu = W_mo_phys[:n_up, :n_up, :n_up, :n_up] + W_dd = W_mo_phys[off:off + n_dn, off:off + n_dn, off:off + n_dn, off:off + n_dn] + W_ud = W_mo_phys[:n_up, off:off + n_dn, :n_up, off:off + n_dn] + + e_mo += np.einsum('pq,pq->', dens, core) + e_mo += 0.5 * ( + np.tensordot(G_uu, W_uu, axes=4) + + np.tensordot(G_dd, W_dd, axes=4) + + 2.0 * np.tensordot(G_ud, W_ud, axes=4) + ) + + assert abs(e_mo - mf0.e_tot) < 1e-6 + + +@pytest.mark.parametrize("cart", [False, True], ids=["cart=false", "cart=true"]) +def test_energy_crystal_integrals_sym_s1_in_trexio_rhf_ccecp(cart): + with tempfile.TemporaryDirectory() as d: + filename = os.path.join(d, 'pbc_gamma_rdm_energy_ccecp.h5') + + cell = pyscf.pbc.gto.Cell() + cell.cart = cart + cell.unit = 'Bohr' + cell.build( + atom='H 0 0 0; H 0 0 2.6', + basis='ccecp-ccpvdz', + ecp='ccecp', + a=np.diag([6.0, 6.0, 6.0]), + ) + cell.exp_to_discard=0.2 + + mf0 = pyscf.pbc.scf.RHF(cell) + mf0.with_df = pbcdf.MDF(cell).build() + mf0.run() + assert mf0.converged + + trexio.to_trexio(mf0, filename) + trexio.write_scf_1e_int_eri(mf0, filename, basis='AO') + trexio.write_scf_2e_int_eri(mf0, filename, basis='AO') + trexio.write_scf_1e_int_eri(mf0, filename, basis='MO') + trexio.write_scf_2e_int_eri(mf0, filename, basis='MO') + trexio.write_scf_rdm_1e(mf0, filename) + trexio.write_scf_rdm_2e(mf0, filename) + + BUFSIZE = 100000 + with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: + e_nn = trexio_lib.read_nucleus_repulsion(tf) + if abs(e_nn) < 1e-12: + e_nn_mol = cell.energy_nuc() + if abs(e_nn_mol) > 1e-12: + e_nn = e_nn_mol + madelung = 0.0 + try: + madelung = trexio_lib.read_pbc_madelung(tf) + madelung = float(np.asarray(madelung).ravel()[0]) + except trexio_lib.Error: + madelung = 0.0 + + n = trexio_lib.read_mo_num(tf) + core = trexio_lib.read_mo_1e_int_core_hamiltonian(tf) + core = np.asarray(core) + if core.ndim == 1: + core = core.reshape(n, n) + + eri_size = trexio_lib.read_mo_2e_int_eri_size(tf) + assert eri_size == n ** 4 + offset = 0 + W_mo_phys = np.zeros((n, n, n, n)) + while offset < eri_size: + bufsize = min(BUFSIZE, eri_size - offset) + buf_idx, buf_val, icount, feof = trexio_lib.read_mo_2e_int_eri(tf, offset, bufsize) + offset += icount + if icount: + i, j, k, l = np.asarray(buf_idx, dtype=np.int64).T + W_mo_phys[i, j, k, l] = buf_val + if feof: + break + + dens = trexio_lib.read_rdm_1e(tf, doReshape=True) + dens = np.asarray(dens) + if dens.ndim == 1: + dens = dens.reshape(n, n) + + offset = 0 + G_mo_phys = np.zeros((n, n, n, n)) + while True: + buf_idx, buf_val, icount, feof = trexio_lib.read_rdm_2e(tf, offset, BUFSIZE) + offset += icount + if icount: + i, j, k, l = np.asarray(buf_idx, dtype=np.int64).T + G_mo_phys[i, j, k, l] = buf_val + if feof: + break + + e_mo = e_nn - madelung + e_mo += np.einsum('pq,pq->', dens, core) + e_mo += 0.5 * np.tensordot(G_mo_phys, W_mo_phys, axes=4) + + assert abs(e_mo - mf0.e_tot) < 1e-6 + + +@pytest.mark.parametrize("cart", [False, True], ids=["cart=false", "cart=true"]) +def test_energy_crystal_integrals_sym_s1_in_trexio_uhf_ccecp(cart): + with tempfile.TemporaryDirectory() as d: + filename = os.path.join(d, 'pbc_gamma_rdm_energy_uhf_ccecp.h5') + + cell = pyscf.pbc.gto.Cell() + cell.cart = cart + cell.unit = 'Bohr' + cell.spin = 2 + cell.build( + atom='H 0 0 0; H 0 0 2.6', + basis='ccecp-ccpvdz', + ecp='ccecp', + a=np.diag([6.0, 6.0, 6.0]), + ) + cell.exp_to_discard=0.2 + + mf0 = pyscf.pbc.scf.UHF(cell) + mf0.with_df = pbcdf.MDF(cell).build() + mf0.run() + assert mf0.converged + + trexio.to_trexio(mf0, filename) + trexio.write_scf_1e_int_eri(mf0, filename, basis='AO') + trexio.write_scf_2e_int_eri(mf0, filename, basis='AO') + trexio.write_scf_1e_int_eri(mf0, filename, basis='MO') + trexio.write_scf_2e_int_eri(mf0, filename, basis='MO') + trexio.write_scf_rdm_1e(mf0, filename) + trexio.write_scf_rdm_2e(mf0, filename) + + BUFSIZE = 100000 + with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: + e_nn = trexio_lib.read_nucleus_repulsion(tf) + if abs(e_nn) < 1e-12: + e_nn_mol = cell.energy_nuc() + if abs(e_nn_mol) > 1e-12: + e_nn = e_nn_mol + madelung = 0.0 + try: + madelung = trexio_lib.read_pbc_madelung(tf) + madelung = float(np.asarray(madelung).ravel()[0]) + except trexio_lib.Error: + madelung = 0.0 + + n = trexio_lib.read_mo_num(tf) + core = trexio_lib.read_mo_1e_int_core_hamiltonian(tf) + core = np.asarray(core) + if core.ndim == 1: + core = core.reshape(n, n) + + eri_size = trexio_lib.read_mo_2e_int_eri_size(tf) + assert eri_size == n ** 4 + offset = 0 + W_mo_phys = np.zeros((n, n, n, n)) + while offset < eri_size: + bufsize = min(BUFSIZE, eri_size - offset) + buf_idx, buf_val, icount, feof = trexio_lib.read_mo_2e_int_eri(tf, offset, bufsize) + offset += icount + if icount: + i, j, k, l = np.asarray(buf_idx, dtype=np.int64).T + W_mo_phys[i, j, k, l] = buf_val + if feof: + break + + dens = trexio_lib.read_rdm_1e(tf, doReshape=True) + dens = np.asarray(dens) + if dens.ndim == 1: + dens = dens.reshape(n, n) + + spin_rdm = ( + trexio_lib.has_rdm_1e_up(tf) + and trexio_lib.has_rdm_2e_upup(tf) + and trexio_lib.has_rdm_2e_dndn(tf) + and trexio_lib.has_rdm_2e_updn(tf) + ) + + assert spin_rdm + + rdm2_size = trexio_lib.read_rdm_2e_upup_size(tf) + n_up = int(round(rdm2_size ** 0.25)) + assert rdm2_size == n_up ** 4 + + rdm2_size = trexio_lib.read_rdm_2e_dndn_size(tf) + n_dn = int(round(rdm2_size ** 0.25)) + assert rdm2_size == n_dn ** 4 + + rdm2_size = trexio_lib.read_rdm_2e_updn_size(tf) + assert rdm2_size == (n_up * n_dn) ** 2 + + offset = 0 + G_uu = np.zeros((n_up, n_up, n_up, n_up)) + while True: + buf_idx, buf_val, icount, feof = trexio_lib.read_rdm_2e_upup(tf, offset, BUFSIZE) + offset += icount + if icount: + i, j, k, l = np.asarray(buf_idx, dtype=np.int64).T + G_uu[i, j, k, l] = buf_val + if feof: + break + + offset = 0 + G_dd = np.zeros((n_dn, n_dn, n_dn, n_dn)) + while True: + buf_idx, buf_val, icount, feof = trexio_lib.read_rdm_2e_dndn(tf, offset, BUFSIZE) + offset += icount + if icount: + i, j, k, l = np.asarray(buf_idx, dtype=np.int64).T + G_dd[i, j, k, l] = buf_val + if feof: + break + + offset = 0 + G_ud = np.zeros((n_up, n_dn, n_up, n_dn)) + while True: + buf_idx, buf_val, icount, feof = trexio_lib.read_rdm_2e_updn(tf, offset, BUFSIZE) + offset += icount + if icount: + i, j, k, l = np.asarray(buf_idx, dtype=np.int64).T + G_ud[i, j, k, l] = buf_val + if feof: + break + + e_mo = e_nn - madelung + off = n_up + + W_uu = W_mo_phys[:n_up, :n_up, :n_up, :n_up] + W_dd = W_mo_phys[off:off + n_dn, off:off + n_dn, off:off + n_dn, off:off + n_dn] + W_ud = W_mo_phys[:n_up, off:off + n_dn, :n_up, off:off + n_dn] + + e_mo += np.einsum('pq,pq->', dens, core) + e_mo += 0.5 * ( + np.tensordot(G_uu, W_uu, axes=4) + + np.tensordot(G_dd, W_dd, axes=4) + + 2.0 * np.tensordot(G_ud, W_ud, axes=4) + ) + + assert abs(e_mo - mf0.e_tot) < 1e-6 + + +@pytest.mark.parametrize("cart", [False, True], ids=["cart=false", "cart=true"]) +def test_energy_crystal_integrals_sym_s4_in_trexio_rhf_ae(cart): + with tempfile.TemporaryDirectory() as d: + filename = os.path.join(d, 'pbc_gamma_rdm_energy_s4.h5') + + cell = pyscf.pbc.gto.Cell() + cell.cart = cart + cell.unit = 'Bohr' + cell.build( + atom='H 0 0 0; H 0 0 1.4', + basis='sto-3g', + a=np.diag([3.0, 3.0, 5.0]), + ) + + mf0 = pyscf.pbc.scf.RHF(cell) + mf0.with_df = pbcdf.MDF(cell).build() + mf0.run() + assert mf0.converged + + trexio.to_trexio(mf0, filename) + trexio.write_scf_1e_int_eri(mf0, filename, basis='AO') + trexio.write_scf_2e_int_eri(mf0, filename, basis='AO', sym='s4') + trexio.write_scf_1e_int_eri(mf0, filename, basis='MO') + trexio.write_scf_2e_int_eri(mf0, filename, basis='MO', sym='s4') + trexio.write_scf_rdm_1e(mf0, filename) + trexio.write_scf_rdm_2e(mf0, filename) + + BUFSIZE = 100000 + with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: + e_nn = trexio_lib.read_nucleus_repulsion(tf) + if abs(e_nn) < 1e-12: + e_nn_mol = cell.energy_nuc() + if abs(e_nn_mol) > 1e-12: + e_nn = e_nn_mol + madelung = 0.0 + try: + madelung = trexio_lib.read_pbc_madelung(tf) + madelung = float(np.asarray(madelung).ravel()[0]) + except trexio_lib.Error: + madelung = 0.0 + + n = trexio_lib.read_mo_num(tf) + core = trexio_lib.read_mo_1e_int_core_hamiltonian(tf) + core = np.asarray(core) + if core.ndim == 1: + core = core.reshape(n, n) + + mo_eri_size = trexio_lib.read_mo_2e_int_eri_size(tf) + idx, val, n_read, _ = trexio_lib.read_mo_2e_int_eri(tf, 0, mo_eri_size) + assert n_read == mo_eri_size + W_mo_phys = _expand_trexio_eri(idx, val, n) + + dens = trexio_lib.read_rdm_1e(tf, doReshape=True) + dens = np.asarray(dens) + if dens.ndim == 1: + dens = dens.reshape(n, n) + + offset = 0 + G_mo_phys = np.zeros((n, n, n, n)) + while True: + buf_idx, buf_val, icount, feof = trexio_lib.read_rdm_2e(tf, offset, BUFSIZE) + offset += icount + if icount: + i, j, k, l = np.asarray(buf_idx, dtype=np.int64).T + G_mo_phys[i, j, k, l] = buf_val + if feof: + break + + e_mo = e_nn - madelung + e_mo += np.einsum('pq,pq->', dens, core) + e_mo += 0.5 * np.tensordot(G_mo_phys, W_mo_phys, axes=4) + + assert abs(e_mo - mf0.e_tot) < 1e-6 + + +@pytest.mark.parametrize("cart", [False, True], ids=["cart=false", "cart=true"]) +def test_energy_crystal_integrals_sym_s4_in_trexio_uhf_ae(cart): + with tempfile.TemporaryDirectory() as d: + filename = os.path.join(d, 'pbc_gamma_rdm_energy_uhf_s4.h5') + + cell = pyscf.pbc.gto.Cell() + cell.cart = cart + cell.unit = 'Bohr' + cell.spin = 2 + cell.build( + atom='H 0 0 0; H 0 0 1.4', + basis='sto-3g', + a=np.diag([3.0, 3.0, 5.0]), + ) + + mf0 = pyscf.pbc.scf.UHF(cell) + mf0.with_df = pbcdf.MDF(cell).build() + mf0.run() + assert mf0.converged + + trexio.to_trexio(mf0, filename) + trexio.write_scf_1e_int_eri(mf0, filename, basis='AO') + trexio.write_scf_2e_int_eri(mf0, filename, basis='AO', sym='s4') + trexio.write_scf_1e_int_eri(mf0, filename, basis='MO') + trexio.write_scf_2e_int_eri(mf0, filename, basis='MO', sym='s4') + trexio.write_scf_rdm_1e(mf0, filename) + trexio.write_scf_rdm_2e(mf0, filename) + + BUFSIZE = 100000 + with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: + e_nn = trexio_lib.read_nucleus_repulsion(tf) + if abs(e_nn) < 1e-12: + e_nn_mol = cell.energy_nuc() + if abs(e_nn_mol) > 1e-12: + e_nn = e_nn_mol + madelung = 0.0 + try: + madelung = trexio_lib.read_pbc_madelung(tf) + madelung = float(np.asarray(madelung).ravel()[0]) + except trexio_lib.Error: + madelung = 0.0 + + n = trexio_lib.read_mo_num(tf) + core = trexio_lib.read_mo_1e_int_core_hamiltonian(tf) + core = np.asarray(core) + if core.ndim == 1: + core = core.reshape(n, n) + + mo_eri_size = trexio_lib.read_mo_2e_int_eri_size(tf) + idx, val, n_read, _ = trexio_lib.read_mo_2e_int_eri(tf, 0, mo_eri_size) + assert n_read == mo_eri_size + W_mo_phys = _expand_trexio_eri(idx, val, n) + + dens = trexio_lib.read_rdm_1e(tf, doReshape=True) + dens = np.asarray(dens) + if dens.ndim == 1: + dens = dens.reshape(n, n) + + spin_rdm = ( + trexio_lib.has_rdm_1e_up(tf) + and trexio_lib.has_rdm_2e_upup(tf) + and trexio_lib.has_rdm_2e_dndn(tf) + and trexio_lib.has_rdm_2e_updn(tf) + ) + + assert spin_rdm + + rdm2_size = trexio_lib.read_rdm_2e_upup_size(tf) + n_up = int(round(rdm2_size ** 0.25)) + assert rdm2_size == n_up ** 4 + + rdm2_size = trexio_lib.read_rdm_2e_dndn_size(tf) + n_dn = int(round(rdm2_size ** 0.25)) + assert rdm2_size == n_dn ** 4 + + rdm2_size = trexio_lib.read_rdm_2e_updn_size(tf) + assert rdm2_size == (n_up * n_dn) ** 2 + + offset = 0 + G_uu = np.zeros((n_up, n_up, n_up, n_up)) + while True: + buf_idx, buf_val, icount, feof = trexio_lib.read_rdm_2e_upup(tf, offset, BUFSIZE) + offset += icount + if icount: + i, j, k, l = np.asarray(buf_idx, dtype=np.int64).T + G_uu[i, j, k, l] = buf_val + if feof: + break + + offset = 0 + G_dd = np.zeros((n_dn, n_dn, n_dn, n_dn)) + while True: + buf_idx, buf_val, icount, feof = trexio_lib.read_rdm_2e_dndn(tf, offset, BUFSIZE) + offset += icount + if icount: + i, j, k, l = np.asarray(buf_idx, dtype=np.int64).T + G_dd[i, j, k, l] = buf_val + if feof: + break + + offset = 0 + G_ud = np.zeros((n_up, n_dn, n_up, n_dn)) + while True: + buf_idx, buf_val, icount, feof = trexio_lib.read_rdm_2e_updn(tf, offset, BUFSIZE) + offset += icount + if icount: + i, j, k, l = np.asarray(buf_idx, dtype=np.int64).T + G_ud[i, j, k, l] = buf_val + if feof: + break + + e_mo = e_nn - madelung + off = n_up + + W_uu = W_mo_phys[:n_up, :n_up, :n_up, :n_up] + W_dd = W_mo_phys[off:off + n_dn, off:off + n_dn, off:off + n_dn, off:off + n_dn] + W_ud = W_mo_phys[:n_up, off:off + n_dn, :n_up, off:off + n_dn] + + e_mo += np.einsum('pq,pq->', dens, core) + e_mo += 0.5 * ( + np.tensordot(G_uu, W_uu, axes=4) + + np.tensordot(G_dd, W_dd, axes=4) + + 2.0 * np.tensordot(G_ud, W_ud, axes=4) + ) + + assert abs(e_mo - mf0.e_tot) < 1e-6 + + +@pytest.mark.parametrize("cart", [False, True], ids=["cart=false", "cart=true"]) +def test_energy_molecule_integrals_sym_s1_in_trexio_uhf_ae(cart): + with tempfile.TemporaryDirectory() as d: + filename = os.path.join(d, 'mol_integrals_energy_uks.h5') + + mol0 = pyscf.M(atom='O 0 0 0', basis='6-31g*', spin=2, cart=cart) + mf0 = mol0.UHF().run() + assert mf0.converged + + overlap = _hermitize(mf0.get_ovlp()) + kinetic = _hermitize(mol0.intor('int1e_kin')) + potential = _hermitize(mol0.intor('int1e_nuc')) + if mol0._ecp: + potential += _hermitize(mol0.intor('ECPscalar')) + core = kinetic + potential + + trexio.to_trexio(mf0, filename) + + trexio.write_scf_1e_int_eri(mf0, filename, basis='AO') + with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: + np.testing.assert_allclose(trexio_lib.read_ao_1e_int_overlap(tf), overlap, atol=DIFF_TOL) + np.testing.assert_allclose(trexio_lib.read_ao_1e_int_kinetic(tf), kinetic, atol=DIFF_TOL) + np.testing.assert_allclose(trexio_lib.read_ao_1e_int_potential_n_e(tf), potential, atol=DIFF_TOL) + np.testing.assert_allclose(trexio_lib.read_ao_1e_int_core_hamiltonian(tf), core, atol=DIFF_TOL) + + trexio.write_scf_2e_int_eri(mf0, filename, basis='AO') + with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: + assert trexio_lib.has_ao_2e_int_eri(tf) + size = trexio_lib.read_ao_2e_int_eri_size(tf) + nao = trexio_lib.read_ao_num(tf) + assert size == nao ** 4 + + trexio.write_scf_1e_int_eri(mf0, filename, basis='MO') + trexio.write_scf_2e_int_eri(mf0, filename, basis='MO') + trexio.write_scf_rdm_1e(mf0, filename) + trexio.write_scf_rdm_2e(mf0, filename) + + BUFSIZE = 100000 + with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: + e_nn = trexio_lib.read_nucleus_repulsion(tf) + + # AO energy reconstruction (UKS) + nao = trexio_lib.read_ao_num(tf) + core_ao = trexio_lib.read_ao_1e_int_core_hamiltonian(tf) + core_ao = np.asarray(core_ao) + if core_ao.ndim == 1: + core_ao = core_ao.reshape(nao, nao) + + dm_ao = mf0.make_rdm1() + dm_a = dm_b = None + if isinstance(dm_ao, (tuple, list)) and len(dm_ao) == 2: + dm_a, dm_b = dm_ao + dm_tot = dm_a + dm_b + elif isinstance(dm_ao, np.ndarray) and dm_ao.ndim == 3 and dm_ao.shape[0] == 2: + dm_a, dm_b = dm_ao[0], dm_ao[1] + dm_tot = dm_a + dm_b + else: + dm_tot = dm_ao + + ao_eri_size = trexio_lib.read_ao_2e_int_eri_size(tf) + assert ao_eri_size == nao ** 4 + offset = 0 + W_ao_phys = np.zeros((nao, nao, nao, nao)) # raw order as stored + while offset < ao_eri_size: + bufsize = min(BUFSIZE, ao_eri_size - offset) + buf_idx, buf_val, icount, feof = trexio_lib.read_ao_2e_int_eri(tf, offset, bufsize) + offset += icount + if icount: + p, r, q, s = np.asarray(buf_idx, dtype=np.int64).T # stored as (p,r,q,s) + W_ao_phys[p, r, q, s] = buf_val + if feof: + break + + W_ao_chem = W_ao_phys.transpose(0, 2, 1, 3) + + J = np.einsum('pqrs,rs->pq', W_ao_chem, dm_tot) + e_ao = e_nn + if dm_a is not None and dm_b is not None: + K_a = np.einsum('prqs,rs->pq', W_ao_chem, dm_a) + K_b = np.einsum('prqs,rs->pq', W_ao_chem, dm_b) + + e_ao += np.einsum('pq,pq->', dm_a, core_ao) + e_ao += np.einsum('pq,pq->', dm_b, core_ao) + e_ao += 0.5 * np.einsum('pq,pq->', dm_tot, J) + e_ao -= 0.5 * ( + np.einsum('pq,pq->', dm_a, K_a) + + np.einsum('pq,pq->', dm_b, K_b) + ) + else: + K = np.einsum('prqs,rs->pq', W_ao_chem, dm_tot) + + e_ao += np.einsum('pq,pq->', dm_tot, core_ao) + e_ao += 0.5 * np.einsum('pq,pq->', dm_tot, J) + e_ao -= 0.25 * np.einsum('pq,pq->', dm_tot, K) + + # MO energy reconstruction + n = trexio_lib.read_mo_num(tf) + core = trexio_lib.read_mo_1e_int_core_hamiltonian(tf) + core = np.asarray(core) + if core.ndim == 1: + core = core.reshape(n, n) + + eri_size = trexio_lib.read_mo_2e_int_eri_size(tf) + assert eri_size == n ** 4 + offset = 0 + W_mo_phys = np.zeros((n, n, n, n)) + while offset < eri_size: + bufsize = min(BUFSIZE, eri_size - offset) + buf_idx, buf_val, icount, feof = trexio_lib.read_mo_2e_int_eri(tf, offset, bufsize) + offset += icount + if icount: + i, j, k, l = np.asarray(buf_idx, dtype=np.int64).T + W_mo_phys[i, j, k, l] = buf_val + if feof: + break + dens = trexio_lib.read_rdm_1e(tf, doReshape=True) + dens = np.asarray(dens) + if dens.ndim == 1: + dens = dens.reshape(n, n) + + spin_rdm = ( + trexio_lib.has_rdm_1e_up(tf) + and trexio_lib.has_rdm_2e_upup(tf) + and trexio_lib.has_rdm_2e_dndn(tf) + and trexio_lib.has_rdm_2e_updn(tf) + ) + + assert spin_rdm + + rdm2_size = trexio_lib.read_rdm_2e_upup_size(tf) + n_up = int(round(rdm2_size ** 0.25)) + assert rdm2_size == n_up ** 4 + + rdm2_size = trexio_lib.read_rdm_2e_dndn_size(tf) + n_dn = int(round(rdm2_size ** 0.25)) + assert rdm2_size == n_dn ** 4 + + rdm2_size = trexio_lib.read_rdm_2e_updn_size(tf) + assert rdm2_size == (n_up * n_dn) ** 2 + + offset = 0 + G_uu = np.zeros((n_up, n_up, n_up, n_up)) + while True: + buf_idx, buf_val, icount, feof = trexio_lib.read_rdm_2e_upup(tf, offset, BUFSIZE) + offset += icount + if icount: + i, j, k, l = np.asarray(buf_idx, dtype=np.int64).T + G_uu[i, j, k, l] = buf_val + if feof: + break + + offset = 0 + G_dd = np.zeros((n_dn, n_dn, n_dn, n_dn)) + while True: + buf_idx, buf_val, icount, feof = trexio_lib.read_rdm_2e_dndn(tf, offset, BUFSIZE) + offset += icount + if icount: + i, j, k, l = np.asarray(buf_idx, dtype=np.int64).T + G_dd[i, j, k, l] = buf_val + if feof: + break + + offset = 0 + G_ud = np.zeros((n_up, n_dn, n_up, n_dn)) + while True: + buf_idx, buf_val, icount, feof = trexio_lib.read_rdm_2e_updn(tf, offset, BUFSIZE) + offset += icount + if icount: + i, j, k, l = np.asarray(buf_idx, dtype=np.int64).T + G_ud[i, j, k, l] = buf_val + if feof: + break + + e_mo = e_nn + off = n_up + + W_uu = W_mo_phys[:n_up, :n_up, :n_up, :n_up] + W_dd = W_mo_phys[off:off + n_dn, off:off + n_dn, off:off + n_dn, off:off + n_dn] + W_ud = W_mo_phys[:n_up, off:off + n_dn, :n_up, off:off + n_dn] + + e_mo += np.einsum('pq,pq->', dens, core) + e_mo += 0.5 * ( + np.tensordot(G_uu, W_uu, axes=4) + + np.tensordot(G_dd, W_dd, axes=4) + + 2.0 * np.tensordot(G_ud, W_ud, axes=4) + ) + + assert abs(e_ao - mf0.e_tot) < 1e-8 + assert abs(e_mo - mf0.e_tot) < 1e-8 + + +@pytest.mark.parametrize("cart", [False, True], ids=["cart=false", "cart=true"]) +def test_energy_molecule_integrals_sym_s1_in_trexio_rhf_ecp(cart): + with tempfile.TemporaryDirectory() as d: + filename = os.path.join(d, 'mol_integrals_energy_ecp.h5') + + mol0 = pyscf.M( + atom='H 0 0 0; F 0 0 1', + basis='ccecp-ccpvdz', + ecp='ccecp', + cart=cart, + ) + mf0 = mol0.RHF().run() + + overlap = _hermitize(mf0.get_ovlp()) + kinetic = _hermitize(mol0.intor('int1e_kin')) + potential = _hermitize(mol0.intor('int1e_nuc')) + if mol0._ecp: + potential += _hermitize(mol0.intor('ECPscalar')) + core = kinetic + potential + + trexio.to_trexio(mf0, filename) + + trexio.write_scf_1e_int_eri(mf0, filename, basis='AO') + with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: + np.testing.assert_allclose(trexio_lib.read_ao_1e_int_overlap(tf), overlap, atol=DIFF_TOL) + np.testing.assert_allclose(trexio_lib.read_ao_1e_int_kinetic(tf), kinetic, atol=DIFF_TOL) + np.testing.assert_allclose(trexio_lib.read_ao_1e_int_potential_n_e(tf), potential, atol=DIFF_TOL) + np.testing.assert_allclose(trexio_lib.read_ao_1e_int_core_hamiltonian(tf), core, atol=DIFF_TOL) + + trexio.write_scf_2e_int_eri(mf0, filename, basis='AO') + with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: + assert trexio_lib.has_ao_2e_int_eri(tf) + size = trexio_lib.read_ao_2e_int_eri_size(tf) + idx, val, n_read, _ = trexio_lib.read_ao_2e_int_eri(tf, 0, size) + assert n_read == size + + trexio.write_scf_1e_int_eri(mf0, filename, basis='MO') + trexio.write_scf_2e_int_eri(mf0, filename, basis='MO') + trexio.write_scf_rdm_1e(mf0, filename) + trexio.write_scf_rdm_2e(mf0, filename) + + BUFSIZE = 100000 + with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: + e_nn = trexio_lib.read_nucleus_repulsion(tf) + + # AO energy reconstruction + nao = trexio_lib.read_ao_num(tf) + core_ao = trexio_lib.read_ao_1e_int_core_hamiltonian(tf) + core_ao = np.asarray(core_ao) + if core_ao.ndim == 1: + core_ao = core_ao.reshape(nao, nao) + + dm_ao = mf0.make_rdm1() + + ao_eri_size = trexio_lib.read_ao_2e_int_eri_size(tf) + assert ao_eri_size == nao ** 4 + offset = 0 + W_ao_phys = np.zeros((nao, nao, nao, nao)) # raw order as stored + while offset < ao_eri_size: + bufsize = min(BUFSIZE, ao_eri_size - offset) + buf_idx, buf_val, icount, feof = trexio_lib.read_ao_2e_int_eri(tf, offset, bufsize) + offset += icount + if icount: + p, r, q, s = np.asarray(buf_idx, dtype=np.int64).T # stored as (p,r,q,s) + W_ao_phys[p, r, q, s] = buf_val + if feof: + break + + W_ao_chem = W_ao_phys.transpose(0, 2, 1, 3) + + J = np.einsum('pqrs,rs->pq', W_ao_chem, dm_ao) + K = np.einsum('prqs,rs->pq', W_ao_chem, dm_ao) + vhf = J - 0.5 * K + e_ao = e_nn + e_ao += np.einsum('pq,pq->', dm_ao, core_ao) + e_ao += 0.5 * np.einsum('pq,pq->', dm_ao, vhf) + + # MO energy reconstruction + n = trexio_lib.read_mo_num(tf) + core = trexio_lib.read_mo_1e_int_core_hamiltonian(tf) + core = np.asarray(core) + if core.ndim == 1: + core = core.reshape(n, n) + + eri_size = trexio_lib.read_mo_2e_int_eri_size(tf) + assert eri_size == n ** 4 + offset = 0 + W_mo_phys = np.zeros((n, n, n, n)) + while offset < eri_size: + bufsize = min(BUFSIZE, eri_size - offset) + buf_idx, buf_val, icount, feof = trexio_lib.read_mo_2e_int_eri(tf, offset, bufsize) + offset += icount + if icount: + i, j, k, l = np.asarray(buf_idx, dtype=np.int64).T + W_mo_phys[i, j, k, l] = buf_val + if feof: + break + dens = trexio_lib.read_rdm_1e(tf, doReshape=True) + dens = np.asarray(dens) + if dens.ndim == 1: + dens = dens.reshape(n, n) + + offset = 0 + G_mo_phys = np.zeros((n, n, n, n)) + while True: + buf_idx, buf_val, icount, feof = trexio_lib.read_rdm_2e(tf, offset, BUFSIZE) + offset += icount + if icount: + i, j, k, l = np.asarray(buf_idx, dtype=np.int64).T + G_mo_phys[i, j, k, l] = buf_val + if feof: + break + + e_mo = e_nn + e_mo += np.einsum('pq,pq->', dens, core) + e_mo += 0.5 * np.tensordot(G_mo_phys, W_mo_phys, axes=4) + + assert abs(e_ao - mf0.e_tot) < 1e-8 + assert abs(e_mo - mf0.e_tot) < 1e-8 + + +@pytest.mark.parametrize("cart", [False, True], ids=["cart=false", "cart=true"]) +def test_energy_molecule_integrals_sym_s1_in_trexio_uhf_ecp(cart): + with tempfile.TemporaryDirectory() as d: + filename = os.path.join(d, 'mol_integrals_energy_uks_ecp.h5') + + mol0 = pyscf.M( + atom='H 0 0 0; F 0 0 1', + basis='ccecp-ccpvdz', + ecp='ccecp', + spin=2, + cart=cart, + ) + mf0 = mol0.UHF().run() + assert mf0.converged + + overlap = _hermitize(mf0.get_ovlp()) + kinetic = _hermitize(mol0.intor('int1e_kin')) + potential = _hermitize(mol0.intor('int1e_nuc')) + if mol0._ecp: + potential += _hermitize(mol0.intor('ECPscalar')) + core = kinetic + potential + + trexio.to_trexio(mf0, filename) + + trexio.write_scf_1e_int_eri(mf0, filename, basis='AO') + with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: + np.testing.assert_allclose(trexio_lib.read_ao_1e_int_overlap(tf), overlap, atol=DIFF_TOL) + np.testing.assert_allclose(trexio_lib.read_ao_1e_int_kinetic(tf), kinetic, atol=DIFF_TOL) + np.testing.assert_allclose(trexio_lib.read_ao_1e_int_potential_n_e(tf), potential, atol=DIFF_TOL) + np.testing.assert_allclose(trexio_lib.read_ao_1e_int_core_hamiltonian(tf), core, atol=DIFF_TOL) + + trexio.write_scf_2e_int_eri(mf0, filename, basis='AO') + with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: + assert trexio_lib.has_ao_2e_int_eri(tf) + size = trexio_lib.read_ao_2e_int_eri_size(tf) + nao = trexio_lib.read_ao_num(tf) + assert size == nao ** 4 + + trexio.write_scf_1e_int_eri(mf0, filename, basis='MO') + trexio.write_scf_2e_int_eri(mf0, filename, basis='MO') + trexio.write_scf_rdm_1e(mf0, filename) + trexio.write_scf_rdm_2e(mf0, filename) + + BUFSIZE = 100000 + with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: + e_nn = trexio_lib.read_nucleus_repulsion(tf) + + # AO energy reconstruction (UKS) + nao = trexio_lib.read_ao_num(tf) + core_ao = trexio_lib.read_ao_1e_int_core_hamiltonian(tf) + core_ao = np.asarray(core_ao) + if core_ao.ndim == 1: + core_ao = core_ao.reshape(nao, nao) + + dm_ao = mf0.make_rdm1() + dm_a = dm_b = None + if isinstance(dm_ao, (tuple, list)) and len(dm_ao) == 2: + dm_a, dm_b = dm_ao + dm_tot = dm_a + dm_b + elif isinstance(dm_ao, np.ndarray) and dm_ao.ndim == 3 and dm_ao.shape[0] == 2: + dm_a, dm_b = dm_ao[0], dm_ao[1] + dm_tot = dm_a + dm_b + else: + dm_tot = dm_ao + + ao_eri_size = trexio_lib.read_ao_2e_int_eri_size(tf) + assert ao_eri_size == nao ** 4 + offset = 0 + W_ao_phys = np.zeros((nao, nao, nao, nao)) # raw order as stored + while offset < ao_eri_size: + bufsize = min(BUFSIZE, ao_eri_size - offset) + buf_idx, buf_val, icount, feof = trexio_lib.read_ao_2e_int_eri(tf, offset, bufsize) + offset += icount + if icount: + p, r, q, s = np.asarray(buf_idx, dtype=np.int64).T # stored as (p,r,q,s) + W_ao_phys[p, r, q, s] = buf_val + if feof: + break + + W_ao_chem = W_ao_phys.transpose(0, 2, 1, 3) + + J = np.einsum('pqrs,rs->pq', W_ao_chem, dm_tot) + e_ao = e_nn + if dm_a is not None and dm_b is not None: + K_a = np.einsum('prqs,rs->pq', W_ao_chem, dm_a) + K_b = np.einsum('prqs,rs->pq', W_ao_chem, dm_b) + + e_ao += np.einsum('pq,pq->', dm_a, core_ao) + e_ao += np.einsum('pq,pq->', dm_b, core_ao) + e_ao += 0.5 * np.einsum('pq,pq->', dm_tot, J) + e_ao -= 0.5 * ( + np.einsum('pq,pq->', dm_a, K_a) + + np.einsum('pq,pq->', dm_b, K_b) + ) + else: + K = np.einsum('prqs,rs->pq', W_ao_chem, dm_tot) + + e_ao += np.einsum('pq,pq->', dm_tot, core_ao) + e_ao += 0.5 * np.einsum('pq,pq->', dm_tot, J) + e_ao -= 0.25 * np.einsum('pq,pq->', dm_tot, K) + + # MO energy reconstruction + n = trexio_lib.read_mo_num(tf) + core = trexio_lib.read_mo_1e_int_core_hamiltonian(tf) + core = np.asarray(core) + if core.ndim == 1: + core = core.reshape(n, n) + + eri_size = trexio_lib.read_mo_2e_int_eri_size(tf) + assert eri_size == n ** 4 + offset = 0 + W_mo_phys = np.zeros((n, n, n, n)) + while offset < eri_size: + bufsize = min(BUFSIZE, eri_size - offset) + buf_idx, buf_val, icount, feof = trexio_lib.read_mo_2e_int_eri(tf, offset, bufsize) + offset += icount + if icount: + i, j, k, l = np.asarray(buf_idx, dtype=np.int64).T + W_mo_phys[i, j, k, l] = buf_val + if feof: + break + dens = trexio_lib.read_rdm_1e(tf, doReshape=True) + dens = np.asarray(dens) + if dens.ndim == 1: + dens = dens.reshape(n, n) + + spin_rdm = ( + trexio_lib.has_rdm_1e_up(tf) + and trexio_lib.has_rdm_2e_upup(tf) + and trexio_lib.has_rdm_2e_dndn(tf) + and trexio_lib.has_rdm_2e_updn(tf) + ) + + assert spin_rdm + + rdm2_size = trexio_lib.read_rdm_2e_upup_size(tf) + n_up = int(round(rdm2_size ** 0.25)) + assert rdm2_size == n_up ** 4 + + rdm2_size = trexio_lib.read_rdm_2e_dndn_size(tf) + n_dn = int(round(rdm2_size ** 0.25)) + assert rdm2_size == n_dn ** 4 + + rdm2_size = trexio_lib.read_rdm_2e_updn_size(tf) + assert rdm2_size == (n_up * n_dn) ** 2 + + offset = 0 + G_uu = np.zeros((n_up, n_up, n_up, n_up)) + while True: + buf_idx, buf_val, icount, feof = trexio_lib.read_rdm_2e_upup(tf, offset, BUFSIZE) + offset += icount + if icount: + i, j, k, l = np.asarray(buf_idx, dtype=np.int64).T + G_uu[i, j, k, l] = buf_val + if feof: + break + + offset = 0 + G_dd = np.zeros((n_dn, n_dn, n_dn, n_dn)) + while True: + buf_idx, buf_val, icount, feof = trexio_lib.read_rdm_2e_dndn(tf, offset, BUFSIZE) + offset += icount + if icount: + i, j, k, l = np.asarray(buf_idx, dtype=np.int64).T + G_dd[i, j, k, l] = buf_val + if feof: + break + + offset = 0 + G_ud = np.zeros((n_up, n_dn, n_up, n_dn)) + while True: + buf_idx, buf_val, icount, feof = trexio_lib.read_rdm_2e_updn(tf, offset, BUFSIZE) + offset += icount + if icount: + i, j, k, l = np.asarray(buf_idx, dtype=np.int64).T + G_ud[i, j, k, l] = buf_val + if feof: + break + + e_mo = e_nn + off = n_up + + W_uu = W_mo_phys[:n_up, :n_up, :n_up, :n_up] + W_dd = W_mo_phys[off:off + n_dn, off:off + n_dn, off:off + n_dn, off:off + n_dn] + W_ud = W_mo_phys[:n_up, off:off + n_dn, :n_up, off:off + n_dn] + + e_mo += np.einsum('pq,pq->', dens, core) + e_mo += 0.5 * ( + np.tensordot(G_uu, W_uu, axes=4) + + np.tensordot(G_dd, W_dd, axes=4) + + 2.0 * np.tensordot(G_ud, W_ud, axes=4) + ) + + assert abs(e_ao - mf0.e_tot) < 1e-8 + assert abs(e_mo - mf0.e_tot) < 1e-8 + + +@pytest.mark.parametrize("cart", [False, True], ids=["cart=false", "cart=true"]) +def test_energy_molecule_integrals_sym_s4_in_trexio_rhf_ae(cart): + with tempfile.TemporaryDirectory() as d: + filename = os.path.join(d, 'mol_integrals_energy_s4.h5') + + mol0 = pyscf.M(atom='H 0 0 0; F 0 0 1', basis='6-31g*', cart=cart) + mf0 = mol0.RHF().run() + + overlap = _hermitize(mf0.get_ovlp()) + kinetic = _hermitize(mol0.intor('int1e_kin')) + potential = _hermitize(mol0.intor('int1e_nuc')) + if mol0._ecp: + from pyscf.gto import ecp + potential += _hermitize(ecp.ecp_int(mol0)) + core = kinetic + potential + + trexio.to_trexio(mf0, filename) + + trexio.write_scf_1e_int_eri(mf0, filename, basis='AO') + trexio.write_scf_2e_int_eri(mf0, filename, basis='AO', sym='s4') + trexio.write_scf_1e_int_eri(mf0, filename, basis='MO') + trexio.write_scf_2e_int_eri(mf0, filename, basis='MO', sym='s4') + trexio.write_scf_rdm_1e(mf0, filename) + trexio.write_scf_rdm_2e(mf0, filename) + + BUFSIZE = 100000 + with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: + e_nn = trexio_lib.read_nucleus_repulsion(tf) + + # AO energy reconstruction + nao = trexio_lib.read_ao_num(tf) + core_ao = trexio_lib.read_ao_1e_int_core_hamiltonian(tf) + core_ao = np.asarray(core_ao) + if core_ao.ndim == 1: + core_ao = core_ao.reshape(nao, nao) + + dm_ao = mf0.make_rdm1() + + ao_eri_size = trexio_lib.read_ao_2e_int_eri_size(tf) + idx, val, n_read, _ = trexio_lib.read_ao_2e_int_eri(tf, 0, ao_eri_size) + assert n_read == ao_eri_size + W_ao_phys = _expand_trexio_eri(idx, val, nao) + + W_ao_chem = W_ao_phys.transpose(0, 2, 1, 3) + + J = np.einsum('pqrs,rs->pq', W_ao_chem, dm_ao) + K = np.einsum('prqs,rs->pq', W_ao_chem, dm_ao) + vhf = J - 0.5 * K + e_ao = e_nn + e_ao += np.einsum('pq,pq->', dm_ao, core_ao) + e_ao += 0.5 * np.einsum('pq,pq->', dm_ao, vhf) + + # MO energy reconstruction + n = trexio_lib.read_mo_num(tf) + core = trexio_lib.read_mo_1e_int_core_hamiltonian(tf) + core = np.asarray(core) + if core.ndim == 1: + core = core.reshape(n, n) + + mo_eri_size = trexio_lib.read_mo_2e_int_eri_size(tf) + idx, val, n_read, _ = trexio_lib.read_mo_2e_int_eri(tf, 0, mo_eri_size) + assert n_read == mo_eri_size + W_mo_phys = _expand_trexio_eri(idx, val, n) + + dens = trexio_lib.read_rdm_1e(tf, doReshape=True) + dens = np.asarray(dens) + if dens.ndim == 1: + dens = dens.reshape(n, n) + + offset = 0 + G_mo_phys = np.zeros((n, n, n, n)) + while True: + buf_idx, buf_val, icount, feof = trexio_lib.read_rdm_2e(tf, offset, BUFSIZE) + offset += icount + if icount: + i, j, k, l = np.asarray(buf_idx, dtype=np.int64).T + G_mo_phys[i, j, k, l] = buf_val + if feof: + break + + e_mo = e_nn + e_mo += np.einsum('pq,pq->', dens, core) + e_mo += 0.5 * np.tensordot(G_mo_phys, W_mo_phys, axes=4) + + assert abs(e_ao - mf0.e_tot) < 1e-8 + assert abs(e_mo - mf0.e_tot) < 1e-8 + + +@pytest.mark.parametrize("cart", [False, True], ids=["cart=false", "cart=true"]) +def test_energy_molecule_integrals_sym_s4_in_trexio_uhf_ae(cart): + with tempfile.TemporaryDirectory() as d: + filename = os.path.join(d, 'mol_integrals_energy_uks_s4.h5') + + mol0 = pyscf.M(atom='O 0 0 0', basis='6-31g*', spin=2, cart=cart) + mf0 = mol0.UHF().run() + assert mf0.converged + + overlap = _hermitize(mf0.get_ovlp()) + kinetic = _hermitize(mol0.intor('int1e_kin')) + potential = _hermitize(mol0.intor('int1e_nuc')) + if mol0._ecp: + from pyscf.gto import ecp + potential += _hermitize(ecp.ecp_int(mol0)) + core = kinetic + potential + + trexio.to_trexio(mf0, filename) + + trexio.write_scf_1e_int_eri(mf0, filename, basis='AO') + trexio.write_scf_2e_int_eri(mf0, filename, basis='AO', sym='s4') + trexio.write_scf_1e_int_eri(mf0, filename, basis='MO') + trexio.write_scf_2e_int_eri(mf0, filename, basis='MO', sym='s4') + trexio.write_scf_rdm_1e(mf0, filename) + trexio.write_scf_rdm_2e(mf0, filename) + + BUFSIZE = 100000 + with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: + e_nn = trexio_lib.read_nucleus_repulsion(tf) + + # AO energy reconstruction (UKS) + nao = trexio_lib.read_ao_num(tf) + core_ao = trexio_lib.read_ao_1e_int_core_hamiltonian(tf) + core_ao = np.asarray(core_ao) + if core_ao.ndim == 1: + core_ao = core_ao.reshape(nao, nao) + + dm_ao = mf0.make_rdm1() + dm_a = dm_b = None + if isinstance(dm_ao, (tuple, list)) and len(dm_ao) == 2: + dm_a, dm_b = dm_ao + dm_tot = dm_a + dm_b + elif isinstance(dm_ao, np.ndarray) and dm_ao.ndim == 3 and dm_ao.shape[0] == 2: + dm_a, dm_b = dm_ao[0], dm_ao[1] + dm_tot = dm_a + dm_b + else: + dm_tot = dm_ao + + ao_eri_size = trexio_lib.read_ao_2e_int_eri_size(tf) + idx, val, n_read, _ = trexio_lib.read_ao_2e_int_eri(tf, 0, ao_eri_size) + assert n_read == ao_eri_size + W_ao_phys = _expand_trexio_eri(idx, val, nao) + W_ao_chem = W_ao_phys.transpose(0, 2, 1, 3) + + J = np.einsum('pqrs,rs->pq', W_ao_chem, dm_tot) + e_ao = e_nn + if dm_a is not None and dm_b is not None: + K_a = np.einsum('prqs,rs->pq', W_ao_chem, dm_a) + K_b = np.einsum('prqs,rs->pq', W_ao_chem, dm_b) + + e_ao += np.einsum('pq,pq->', dm_a, core_ao) + e_ao += np.einsum('pq,pq->', dm_b, core_ao) + e_ao += 0.5 * np.einsum('pq,pq->', dm_tot, J) + e_ao -= 0.5 * ( + np.einsum('pq,pq->', dm_a, K_a) + + np.einsum('pq,pq->', dm_b, K_b) + ) + else: + K = np.einsum('prqs,rs->pq', W_ao_chem, dm_tot) + + e_ao += np.einsum('pq,pq->', dm_tot, core_ao) + e_ao += 0.5 * np.einsum('pq,pq->', dm_tot, J) + e_ao -= 0.25 * np.einsum('pq,pq->', dm_tot, K) + + # MO energy reconstruction + n = trexio_lib.read_mo_num(tf) + core = trexio_lib.read_mo_1e_int_core_hamiltonian(tf) + core = np.asarray(core) + if core.ndim == 1: + core = core.reshape(n, n) + + mo_eri_size = trexio_lib.read_mo_2e_int_eri_size(tf) + idx, val, n_read, _ = trexio_lib.read_mo_2e_int_eri(tf, 0, mo_eri_size) + assert n_read == mo_eri_size + W_mo_phys = _expand_trexio_eri(idx, val, n) + + dens = trexio_lib.read_rdm_1e(tf, doReshape=True) + dens = np.asarray(dens) + if dens.ndim == 1: + dens = dens.reshape(n, n) + + spin_rdm = ( + trexio_lib.has_rdm_1e_up(tf) + and trexio_lib.has_rdm_2e_upup(tf) + and trexio_lib.has_rdm_2e_dndn(tf) + and trexio_lib.has_rdm_2e_updn(tf) + ) + + assert spin_rdm + + rdm2_size = trexio_lib.read_rdm_2e_upup_size(tf) + n_up = int(round(rdm2_size ** 0.25)) + assert rdm2_size == n_up ** 4 + + rdm2_size = trexio_lib.read_rdm_2e_dndn_size(tf) + n_dn = int(round(rdm2_size ** 0.25)) + assert rdm2_size == n_dn ** 4 + + rdm2_size = trexio_lib.read_rdm_2e_updn_size(tf) + assert rdm2_size == (n_up * n_dn) ** 2 + + offset = 0 + G_uu = np.zeros((n_up, n_up, n_up, n_up)) + while True: + buf_idx, buf_val, icount, feof = trexio_lib.read_rdm_2e_upup(tf, offset, BUFSIZE) + offset += icount + if icount: + i, j, k, l = np.asarray(buf_idx, dtype=np.int64).T + G_uu[i, j, k, l] = buf_val + if feof: + break + + offset = 0 + G_dd = np.zeros((n_dn, n_dn, n_dn, n_dn)) + while True: + buf_idx, buf_val, icount, feof = trexio_lib.read_rdm_2e_dndn(tf, offset, BUFSIZE) + offset += icount + if icount: + i, j, k, l = np.asarray(buf_idx, dtype=np.int64).T + G_dd[i, j, k, l] = buf_val + if feof: + break + + offset = 0 + G_ud = np.zeros((n_up, n_dn, n_up, n_dn)) + while True: + buf_idx, buf_val, icount, feof = trexio_lib.read_rdm_2e_updn(tf, offset, BUFSIZE) + offset += icount + if icount: + i, j, k, l = np.asarray(buf_idx, dtype=np.int64).T + G_ud[i, j, k, l] = buf_val + if feof: + break + + e_mo = e_nn + off = n_up + + W_uu = W_mo_phys[:n_up, :n_up, :n_up, :n_up] + W_dd = W_mo_phys[off:off + n_dn, off:off + n_dn, off:off + n_dn, off:off + n_dn] + W_ud = W_mo_phys[:n_up, off:off + n_dn, :n_up, off:off + n_dn] + + e_mo += np.einsum('pq,pq->', dens, core) + e_mo += 0.5 * ( + np.tensordot(G_uu, W_uu, axes=4) + + np.tensordot(G_dd, W_dd, axes=4) + + 2.0 * np.tensordot(G_ud, W_ud, axes=4) + ) + + assert abs(e_ao - mf0.e_tot) < 1e-8 + assert abs(e_mo - mf0.e_tot) < 1e-8 + + +@pytest.mark.parametrize("cart", [False, True], ids=["cart=false", "cart=true"]) +def test_energy_molecule_integrals_sym_s8_in_trexio_rhf_ae(cart): + with tempfile.TemporaryDirectory() as d: + filename = os.path.join(d, 'mol_integrals_energy_s8.h5') + + mol0 = pyscf.M(atom='H 0 0 0; F 0 0 1', basis='6-31g*', cart=cart) + mf0 = mol0.RHF().run() + + overlap = _hermitize(mf0.get_ovlp()) + kinetic = _hermitize(mol0.intor('int1e_kin')) + potential = _hermitize(mol0.intor('int1e_nuc')) + if mol0._ecp: + from pyscf.gto import ecp + potential += _hermitize(ecp.ecp_int(mol0)) + core = kinetic + potential + + trexio.to_trexio(mf0, filename) + + trexio.write_scf_1e_int_eri(mf0, filename, basis='AO') + trexio.write_scf_2e_int_eri(mf0, filename, basis='AO', sym='s8') + + BUFSIZE = 100000 + with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: + e_nn = trexio_lib.read_nucleus_repulsion(tf) + + # AO energy reconstruction + nao = trexio_lib.read_ao_num(tf) + core_ao = trexio_lib.read_ao_1e_int_core_hamiltonian(tf) + core_ao = np.asarray(core_ao) + if core_ao.ndim == 1: + core_ao = core_ao.reshape(nao, nao) + + dm_ao = mf0.make_rdm1() + + ao_eri_size = trexio_lib.read_ao_2e_int_eri_size(tf) + idx, val, n_read, _ = trexio_lib.read_ao_2e_int_eri(tf, 0, ao_eri_size) + assert n_read == ao_eri_size + W_ao_phys = _expand_trexio_eri(idx, val, nao) + + W_ao_chem = W_ao_phys.transpose(0, 2, 1, 3) + + J = np.einsum('pqrs,rs->pq', W_ao_chem, dm_ao) + K = np.einsum('prqs,rs->pq', W_ao_chem, dm_ao) + vhf = J - 0.5 * K + e_ao = e_nn + e_ao += np.einsum('pq,pq->', dm_ao, core_ao) + e_ao += 0.5 * np.einsum('pq,pq->', dm_ao, vhf) + + assert abs(e_ao - mf0.e_tot) < 1e-8 + + +@pytest.mark.parametrize("cart", [False, True], ids=["cart=false", "cart=true"]) +def test_energy_molecule_integrals_sym_s8_in_trexio_uhf_ae(cart): + with tempfile.TemporaryDirectory() as d: + filename = os.path.join(d, 'mol_integrals_energy_uks_s8.h5') + + mol0 = pyscf.M(atom='O 0 0 0', basis='6-31g*', spin=2, cart=cart) + mf0 = mol0.UHF().run() + assert mf0.converged + + overlap = _hermitize(mf0.get_ovlp()) + kinetic = _hermitize(mol0.intor('int1e_kin')) + potential = _hermitize(mol0.intor('int1e_nuc')) + if mol0._ecp: + from pyscf.gto import ecp + potential += _hermitize(ecp.ecp_int(mol0)) + core = kinetic + potential + + trexio.to_trexio(mf0, filename) + + trexio.write_scf_1e_int_eri(mf0, filename, basis='AO') + trexio.write_scf_2e_int_eri(mf0, filename, basis='AO', sym='s8') + + BUFSIZE = 100000 + with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: + e_nn = trexio_lib.read_nucleus_repulsion(tf) + + # AO energy reconstruction (UKS) + nao = trexio_lib.read_ao_num(tf) + core_ao = trexio_lib.read_ao_1e_int_core_hamiltonian(tf) + core_ao = np.asarray(core_ao) + if core_ao.ndim == 1: + core_ao = core_ao.reshape(nao, nao) + + dm_ao = mf0.make_rdm1() + dm_a = dm_b = None + if isinstance(dm_ao, (tuple, list)) and len(dm_ao) == 2: + dm_a, dm_b = dm_ao + dm_tot = dm_a + dm_b + elif isinstance(dm_ao, np.ndarray) and dm_ao.ndim == 3 and dm_ao.shape[0] == 2: + dm_a, dm_b = dm_ao[0], dm_ao[1] + dm_tot = dm_a + dm_b + else: + dm_tot = dm_ao + + ao_eri_size = trexio_lib.read_ao_2e_int_eri_size(tf) + idx, val, n_read, _ = trexio_lib.read_ao_2e_int_eri(tf, 0, ao_eri_size) + assert n_read == ao_eri_size + W_ao_phys = _expand_trexio_eri(idx, val, nao) + W_ao_chem = W_ao_phys.transpose(0, 2, 1, 3) + + J = np.einsum('pqrs,rs->pq', W_ao_chem, dm_tot) + e_ao = e_nn + if dm_a is not None and dm_b is not None: + K_a = np.einsum('prqs,rs->pq', W_ao_chem, dm_a) + K_b = np.einsum('prqs,rs->pq', W_ao_chem, dm_b) + + e_ao += np.einsum('pq,pq->', dm_a, core_ao) + e_ao += np.einsum('pq,pq->', dm_b, core_ao) + e_ao += 0.5 * np.einsum('pq,pq->', dm_tot, J) + e_ao -= 0.5 * ( + np.einsum('pq,pq->', dm_a, K_a) + + np.einsum('pq,pq->', dm_b, K_b) + ) + else: + K = np.einsum('prqs,rs->pq', W_ao_chem, dm_tot) + + e_ao += np.einsum('pq,pq->', dm_tot, core_ao) + e_ao += 0.5 * np.einsum('pq,pq->', dm_tot, J) + e_ao -= 0.25 * np.einsum('pq,pq->', dm_tot, K) + + assert abs(e_ao - mf0.e_tot) < 1e-8 diff --git a/pyscf/tools/trexio.py b/pyscf/tools/trexio.py index cdf8d81cd..86687c528 100644 --- a/pyscf/tools/trexio.py +++ b/pyscf/tools/trexio.py @@ -26,7 +26,7 @@ from pyscf import fci from pyscf.pbc import gto as pbcgto from pyscf import ao2mo -from pyscf.pbc import df as pbcdf +from pyscf.pbc import df as pbcdf, tools as pbctools import trexio @@ -81,6 +81,13 @@ def _mol_to_trexio(mol, trexio_file): trexio.write_electron_up_num(trexio_file, electron_up_num) trexio.write_electron_dn_num(trexio_file, electron_dn_num) + # 2.5 Nuclear repulsion energy + try: + trexio.write_nucleus_repulsion(trexio_file, mol.energy_nuc()) + except trexio.Error: + # TREXIO ver 2.6.0 bug. See #231. A negative value is not accepted. + trexio.write_nucleus_repulsion(trexio_file, 0.0) + # 3.1 Basis set trexio.write_basis_type(trexio_file, 'Gaussian') if any(mol._bas[:,gto.NCTR_OF] > 1): @@ -224,12 +231,14 @@ def _scf_to_trexio(mf, trexio_file): kpts = mol.get_scaled_kpts(kpts) nk = len(mf.kpts) weights = np.full(nk, 1.0/nk) + madelung = pbctools.pbc.madelung(mol, kpts) if nk == 1: # 2.3 Periodic boundary calculations (pbc group) trexio.write_pbc_k_point_num(trexio_file, 1) trexio.write_pbc_k_point(trexio_file, kpts) trexio.write_pbc_k_point_weight(trexio_file, weights[np.newaxis]) + trexio.write_pbc_madelung(trexio_file, madelung) if isinstance(mf, (pbc.scf.uhf.UHF, pbc.dft.uks.UKS, pbc.scf.kuhf.KUHF, pbc.dft.kuks.KUKS)): mo_type = 'UHF' @@ -306,6 +315,7 @@ def _scf_to_trexio(mf, trexio_file): trexio.write_pbc_k_point_num(trexio_file, nk) trexio.write_pbc_k_point(trexio_file, kpts) trexio.write_pbc_k_point_weight(trexio_file, weights) + trexio.write_pbc_madelung(trexio_file, madelung) # stack k-dependent molecular orbitals mo_k_point_pbc = [] @@ -779,8 +789,7 @@ def scf_from_trexio(filename): raise ValueError(f'Unknown spin multiplicity {uniq}') def write_scf_2e_int_eri( - mf, filename, backend='h5', basis='mo', mo_compact=True, aosym='s8', - df_engine='MDF', + mf, filename, backend='h5', basis='mo', df_engine='MDF', sym='s1', ): """ Write two-electron integrals (ERI) in AO or MO basis. @@ -794,6 +803,7 @@ def write_scf_2e_int_eri( * AO: spin-independent (single tensor). * MO: build a combined MO space by column-wise concatenation [alpha | beta] and compute full ERIs including cross-spin terms. + - Symmetry packing is supported for s1/s4 (MO) and s1/s4/s8 (AO). Notes ----- @@ -806,10 +816,14 @@ def write_scf_2e_int_eri( TOL_IMAG = 1e-12 # drop imag parts <= this; otherwise raise (real-only backend) basis = basis.upper() - if aosym not in ('s8', 's4', 's1'): - raise ValueError("aosym must be one of {'s8','s4','s1'}") + sym = sym.lower() is_pbc = hasattr(mf, 'cell') + if sym not in ('s1', 's4', 's8'): + raise ValueError("sym must be 's1', 's4', or 's8'") + if basis == 'MO' and sym == 's8': + raise NotImplementedError("MO ERI does not support s8 symmetry; use s1 or s4") + # ---------- helpers ---------- def _is_gamma_single_k(mf_obj) -> bool: """True if this PBC calculation is single-k at Gamma.""" @@ -874,6 +888,9 @@ def _get_uks_coeff_pair(mf_obj): if basis == 'MO': # Molecular if not is_pbc: + if sym == 's8': + raise NotImplementedError("MO ERI does not support s8 symmetry") + mo_compact = sym == 's4' if (isinstance(mf.mo_coeff, (list, tuple)) or (isinstance(mf.mo_coeff, np.ndarray) and mf.mo_coeff.ndim >= 3 and mf.mo_coeff.shape[0] == 2)): # UKS/UHF -> concatenate [alpha | beta], include cross-spin terms @@ -884,7 +901,11 @@ def _get_uks_coeff_pair(mf_obj): else: eri_mo = ao2mo.kernel(mf.mol, C, compact=mo_compact) eri_mo = _ensure_real(eri_mo) - _write_2e_int_eri(np.ascontiguousarray(eri_mo), filename, backend, 'MO') + nmo = C.shape[1] + if sym == 's1': + if eri_mo.ndim < 4: + eri_mo = ao2mo.restore(1, eri_mo, nmo) + _write_2e_int_eri(np.ascontiguousarray(eri_mo), filename, backend, 'MO', sym=sym) else: # RHF/RKS C = mf.mo_coeff if getattr(mf, '_eri', None) is not None: @@ -892,98 +913,462 @@ def _get_uks_coeff_pair(mf_obj): else: eri_mo = ao2mo.kernel(mf.mol, C, compact=mo_compact) eri_mo = _ensure_real(eri_mo) - _write_2e_int_eri(np.ascontiguousarray(eri_mo), filename, backend, 'MO') - return + nmo = C.shape[1] + if sym == 's1': + if eri_mo.ndim < 4: + eri_mo = ao2mo.restore(1, eri_mo, nmo) + _write_2e_int_eri(np.ascontiguousarray(eri_mo), filename, backend, 'MO', sym=sym) # PBC (Gamma only) - if not _is_gamma_single_k(mf): - raise NotImplementedError("PBC MO-ERI: non-Gamma k-points are not supported (real-only backend).") - dfobj = _df_obj() - - if (isinstance(mf.mo_coeff, (list, tuple)) or - (isinstance(mf.mo_coeff, np.ndarray) and mf.mo_coeff.ndim >= 3 and mf.mo_coeff.shape[0] == 2)): - # UKS/UHF @ Gamma: combined MO matrix [Ca | Cb] - Ca, Cb = _get_uks_coeff_pair(mf) - C = np.concatenate([Ca, Cb], axis=1) - eri_mo = dfobj.get_mo_eri((C, C, C, C)) - eri_mo = _ensure_real(eri_mo) - _write_2e_int_eri(np.ascontiguousarray(eri_mo), filename, backend, 'MO') - else: # RHF/RKS @ Gamma - C = mf.mo_coeff - if C.ndim == 3 and C.shape[0] == 1: # normalize (1,nao,nmo) -> (nao,nmo) - C = C[0] - eri_mo = dfobj.get_mo_eri((C, C, C, C)) - eri_mo = _ensure_real(eri_mo) - _write_2e_int_eri(np.ascontiguousarray(eri_mo), filename, backend, 'MO') - return + else: + if not _is_gamma_single_k(mf): + raise NotImplementedError("PBC MO-ERI: non-Gamma k-points are not supported (real-only backend).") + if sym == 's8': + raise NotImplementedError("PBC MO-ERI does not support s8 symmetry; use s1 or s4") + dfobj = _df_obj() + + if (isinstance(mf.mo_coeff, (list, tuple)) or + (isinstance(mf.mo_coeff, np.ndarray) and mf.mo_coeff.ndim >= 3 and mf.mo_coeff.shape[0] == 2)): + # UKS/UHF @ Gamma: combined MO matrix [Ca | Cb] + Ca, Cb = _get_uks_coeff_pair(mf) + C = np.concatenate([Ca, Cb], axis=1) + eri_mo = dfobj.get_mo_eri((C, C, C, C)) + eri_mo = _ensure_real(eri_mo) + nmo = C.shape[1] + if sym == 's1': + if eri_mo.ndim == 2: + eri_mo = ao2mo.restore(1, ao2mo.restore(4, eri_mo, nmo), nmo) + elif eri_mo.ndim < 4: + eri_mo = ao2mo.restore(1, eri_mo, nmo) + elif sym == 's4': + if eri_mo.ndim == 4: + eri_mo = ao2mo.restore(4, eri_mo, nmo) + _write_2e_int_eri(np.ascontiguousarray(eri_mo), filename, backend, 'MO', sym=sym) + else: # RHF/RKS @ Gamma + C = mf.mo_coeff + if C.ndim == 3 and C.shape[0] == 1: # normalize (1,nao,nmo) -> (nao,nmo) + C = C[0] + eri_mo = dfobj.get_mo_eri((C, C, C, C)) + eri_mo = _ensure_real(eri_mo) + nmo = C.shape[1] + if sym == 's1': + if eri_mo.ndim == 2: + eri_mo = ao2mo.restore(1, ao2mo.restore(4, eri_mo, nmo), nmo) + elif eri_mo.ndim < 4: + eri_mo = ao2mo.restore(1, eri_mo, nmo) + elif sym == 's4': + if eri_mo.ndim == 4: + eri_mo = ao2mo.restore(4, eri_mo, nmo) + _write_2e_int_eri(np.ascontiguousarray(eri_mo), filename, backend, 'MO', sym=sym) # --------------------- # AO-basis ERI writing (spin-independent even for UKS/UHF) # --------------------- - if is_pbc: - # PBC AO: Gamma only via DF (real-only) - if not _is_gamma_single_k(mf): - raise NotImplementedError("PBC AO-ERI: non-Gamma k-points are not supported (real-only backend).") - dfobj = _df_obj() - eri2 = pbcdf.df_ao2mo.get_eri(dfobj, compact=False) # (nao^2, nao^2) - nao = mf.cell.nao_nr() - if eri2.shape != (nao * nao, nao * nao): - raise RuntimeError(f"Unexpected ERI shape {eri2.shape}; expected ({nao*nao}, {nao*nao}) at Gamma.") - eri_ao = eri2.reshape(nao, nao, nao, nao) - eri_ao = _ensure_real(eri_ao) - _write_2e_int_eri(np.ascontiguousarray(eri_ao), filename, backend, 'AO') + else: # basis == 'AO' + if is_pbc: + # PBC AO: Gamma only via DF (real-only) + if not _is_gamma_single_k(mf): + raise NotImplementedError("PBC AO-ERI: non-Gamma k-points are not supported (real-only backend).") + dfobj = _df_obj() + eri2 = pbcdf.df_ao2mo.get_eri(dfobj, compact=False) # (nao^2, nao^2) + nao = mf.cell.nao_nr() + if eri2.shape != (nao * nao, nao * nao): + raise RuntimeError(f"Unexpected ERI shape {eri2.shape}; expected ({nao*nao}, {nao*nao}) at Gamma.") + eri_ao = eri2.reshape(nao, nao, nao, nao) + eri_ao = _ensure_real(eri_ao) + if sym == 's4': + eri_ao = ao2mo.restore(4, eri_ao, nao) + elif sym == 's8': + eri_ao = ao2mo.restore(8, eri_ao, nao) + _write_2e_int_eri(np.ascontiguousarray(eri_ao), filename, backend, 'AO', sym=sym) + else: + # Molecular AO + eri_ao = None + if sym == 's1': + eri_ao = getattr(mf, '_eri', None) + if eri_ao is None: + # 'int2e' follows mol.cart automatically (spherical vs Cartesian) + eri_ao = mf.mol.intor('int2e', aosym='s1') + else: + # _eri may be stored in s4/s8; expand to full tensor for TREXIO write + if eri_ao.ndim < 4: + n_ao = mf.mol.nao + eri_ao = ao2mo.restore(1, eri_ao, n_ao) + else: + eri_ao = mf.mol.intor('int2e', aosym=sym) + eri_ao = _ensure_real(eri_ao) + _write_2e_int_eri(np.ascontiguousarray(eri_ao), filename, backend, 'AO', sym=sym) + +def write_scf_rdm_1e(mf, filename, backend='h5'): + """Write one-electron RDM (density matrix) to TREXIO in MO basis. + + TREXIO RDM I/O assumes MO-basis data. This helper enforces MO and + builds a spin-summed density: diag(mo_occ) for RHF/RKS and + block-diagonal [diag(occ_a), diag(occ_b)] for UHF/UKS. PBC is limited + to single-k Gamma. + """ + + def _is_gamma_single_k(mf_obj) -> bool: + if not hasattr(mf_obj, 'cell'): + return False + if hasattr(mf_obj, 'kpt'): + return np.allclose(np.asarray(mf_obj.kpt), 0.0) + if hasattr(mf_obj, 'kpts'): + kpts = np.asarray(mf_obj.kpts) + return (kpts.shape[0] == 1) and np.allclose(kpts[0], 0.0) + return True + + def _ensure_real(x): + if np.iscomplexobj(x): + if np.all(np.abs(np.imag(x)) <= 1e-12): + x = np.real(x) + else: + raise NotImplementedError( + "Complex RDM encountered; real-only backend.") + return x + + is_pbc = hasattr(mf, 'cell') + if is_pbc and not _is_gamma_single_k(mf): + raise NotImplementedError("RDM write supports Gamma-point only for PBC.") + + is_uhf_like = isinstance( + mf, + ( + scf.uhf.UHF, + dft.uks.UKS, + pbc.scf.uhf.UHF, + pbc.dft.uks.UKS, + pbc.scf.kuhf.KUHF, + pbc.dft.kuks.KUKS, + ), + ) + + # MO-basis density is diagonal in canonical orbitals + if is_uhf_like and isinstance(mf.mo_occ, (tuple, list)) and len(mf.mo_occ) == 2: + occ_a, occ_b = mf.mo_occ + occ_a = np.asarray(occ_a) + occ_b = np.asarray(occ_b) + else: + occ = np.asarray(mf.mo_occ) + if is_uhf_like and occ.ndim == 2 and occ.shape[0] == 2: + occ_a, occ_b = occ[0], occ[1] + else: + occ_a = occ + occ_b = None + + if is_pbc and isinstance(occ_a, np.ndarray) and occ_a.ndim == 2: + if occ_a.shape[0] != 1: + raise NotImplementedError( + "PBC RDM write: only single-k Gamma supported for MO basis.") + occ_a = occ_a[0] + if occ_b is not None and occ_b.ndim == 2: + occ_b = occ_b[0] + + if occ_b is not None: + dm_a = np.diag(occ_a) + dm_b = np.diag(occ_b) + dm_a = _ensure_real(np.asarray(dm_a)) + dm_b = _ensure_real(np.asarray(dm_b)) + dm_a = np.ascontiguousarray(dm_a) + dm_b = np.ascontiguousarray(dm_b) + with trexio.File(filename, 'u', back_end=_mode(backend)) as tf: + nmo_up = dm_a.shape[0] + nmo_dn = dm_b.shape[0] + if nmo_dn != nmo_up: + raise ValueError("Alpha/beta MO sizes do not match for RDM write.") + nmo = nmo_up + nmo_dn + if not trexio.has_mo_num(tf): + trexio.write_mo_num(tf, nmo) + trexio.write_rdm_1e_up(tf, dm_a) + trexio.write_rdm_1e_dn(tf, dm_b) + dm_tot = np.zeros((nmo, nmo), dtype=dm_a.dtype) + dm_tot[:nmo_up, :nmo_up] = dm_a + dm_tot[nmo_up:, nmo_up:] = dm_b + trexio.write_rdm_1e(tf, dm_tot) + else: + dm = np.diag(occ_a) + dm = _ensure_real(np.asarray(dm)) + dm = np.ascontiguousarray(dm) + with trexio.File(filename, 'u', back_end=_mode(backend)) as tf: + nmo = dm.shape[0] + if not trexio.has_mo_num(tf): + trexio.write_mo_num(tf, nmo) + trexio.write_rdm_1e(tf, dm) + +def write_scf_rdm_2e(mf, filename, backend='h5', chunk_size=100000): + """Write spin-summed two-electron RDM to TREXIO in MO basis. + + Assumes MO-basis storage (TREXIO convention/physicist) and builds a Hartree–Fock-like + RDM using the spin-summed diagonal density ``dm = diag(mo_occ)``: + ``G[pqrs] = dm[p,r]*dm[q,s] - 0.5*dm[p,s]*dm[q,r]``. + + Notes + ----- + - PBC: only single-k Gamma is supported. + - Memory scales as O(n^4); keep `chunk_size` modest for the write loop. + """ + + def _is_gamma_single_k(mf_obj) -> bool: + if not hasattr(mf_obj, 'cell'): + return False + if hasattr(mf_obj, 'kpt'): + return np.allclose(np.asarray(mf_obj.kpt), 0.0) + if hasattr(mf_obj, 'kpts'): + kpts = np.asarray(mf_obj.kpts) + return (kpts.shape[0] == 1) and np.allclose(kpts[0], 0.0) + return True + + is_pbc = hasattr(mf, 'cell') + if is_pbc and not _is_gamma_single_k(mf): + raise NotImplementedError("RDM write supports Gamma-point only for PBC.") + + is_uhf_like = isinstance( + mf, + ( + scf.uhf.UHF, + dft.uks.UKS, + pbc.scf.uhf.UHF, + pbc.dft.uks.UKS, + pbc.scf.kuhf.KUHF, + pbc.dft.kuks.KUKS, + ), + ) + + # Spin-summed occupations or spin-separated for UHF/UKS + if is_uhf_like and isinstance(mf.mo_occ, (tuple, list)) and len(mf.mo_occ) == 2: + occ_a, occ_b = mf.mo_occ + occ_a = np.asarray(occ_a) + occ_b = np.asarray(occ_b) else: - # Molecular AO - eri_ao = getattr(mf, '_eri', None) - if eri_ao is None: - # 'int2e' follows mol.cart automatically (spherical vs Cartesian) - eri_ao = mf.mol.intor('int2e', aosym=aosym) - eri_ao = _ensure_real(eri_ao) - _write_2e_int_eri(np.ascontiguousarray(eri_ao), filename, backend, 'AO') - -def _write_2e_int_eri(eri, filename, backend='h5', basis='MO'): - assert basis.upper() in ['MO','AO'] - num_integrals = eri.size - if eri.ndim == 4: + occ = np.asarray(mf.mo_occ) + if is_uhf_like and occ.ndim == 2 and occ.shape[0] == 2: + occ_a, occ_b = occ[0], occ[1] + else: + occ_a = occ + occ_b = None + + if is_pbc and isinstance(occ_a, np.ndarray) and occ_a.ndim == 2: + if occ_a.shape[0] != 1: + raise NotImplementedError( + "PBC RDM write: only single-k Gamma supported for MO basis.") + occ_a = occ_a[0] + if occ_b is not None and occ_b.ndim == 2: + occ_b = occ_b[0] + + if occ_b is not None: + dm_a = np.diag(occ_a) + dm_b = np.diag(occ_b) + dm_a = np.ascontiguousarray(dm_a) + dm_b = np.ascontiguousarray(dm_b) + + g2_uu = ( + np.einsum('pr,qs->pqrs', dm_a, dm_a) + - np.einsum('ps,qr->pqrs', dm_a, dm_a) + ) + g2_dd = ( + np.einsum('pr,qs->pqrs', dm_b, dm_b) + - np.einsum('ps,qr->pqrs', dm_b, dm_b) + ) + g2_ud = np.einsum('pr,qs->pqrs', dm_a, dm_b) + + g2_uu = np.ascontiguousarray(g2_uu) + g2_dd = np.ascontiguousarray(g2_dd) + g2_ud = np.ascontiguousarray(g2_ud) + + nmo = dm_a.shape[0] + if dm_b.shape[0] != nmo: + raise ValueError("Alpha/beta MO sizes do not match for RDM write.") + idx = lib.cartesian_prod([np.arange(nmo, dtype=np.int32)] * 4) + flat_uu = g2_uu.reshape(-1) + flat_dd = g2_dd.reshape(-1) + flat_ud = g2_ud.reshape(-1) + + with trexio.File(filename, 'u', back_end=_mode(backend)) as tf: + if not trexio.has_mo_num(tf): + trexio.write_mo_num(tf, nmo) + + total = idx.shape[0] + offset = 0 + while offset < total: + end = min(offset + chunk_size, total) + trexio.write_rdm_2e_upup( + tf, + offset, + end - offset, + idx[offset:end], + flat_uu[offset:end], + ) + trexio.write_rdm_2e_dndn( + tf, + offset, + end - offset, + idx[offset:end], + flat_dd[offset:end], + ) + trexio.write_rdm_2e_updn( + tf, + offset, + end - offset, + idx[offset:end], + flat_ud[offset:end], + ) + offset = end + else: + dm = np.diag(occ_a) + dm = np.ascontiguousarray(dm) + + # Build dense spin-summed 2-RDM in MO basis + g2 = ( + np.einsum('pr,qs->pqrs', dm, dm) + - 0.5 * np.einsum('ps,qr->pqrs', dm, dm) + ) + g2 = np.ascontiguousarray(g2) + + nmo = dm.shape[0] + idx = lib.cartesian_prod([np.arange(nmo, dtype=np.int32)] * 4) + flat_g2 = g2.reshape(-1) + + with trexio.File(filename, 'u', back_end=_mode(backend)) as tf: + if not trexio.has_mo_num(tf): + trexio.write_mo_num(tf, nmo) + + total = idx.shape[0] + offset = 0 + while offset < total: + end = min(offset + chunk_size, total) + trexio.write_rdm_2e( + tf, + offset, + end - offset, + idx[offset:end], + flat_g2[offset:end], + ) + offset = end + +def _write_2e_int_eri(eri, filename, backend='h5', basis='MO', sym='s1'): + basis = basis.upper() + sym = sym.lower() + if basis not in ['MO', 'AO']: + raise ValueError("basis must be 'MO' or 'AO'") + if sym not in ('s1', 's4', 's8'): + raise ValueError("sym must be 's1', 's4', or 's8'") + + def _pair_from_tril(n): + i, j = np.tril_indices(n) + return i.astype(np.int32), j.astype(np.int32) + + if sym == 's1': + if eri.ndim != 4: + raise ValueError(f'ERI array must be a full 4D tensor (p,q,r,s); got ndim={eri.ndim}') + + num_integrals = eri.size n = eri.shape[0] - idx = lib.cartesian_prod([np.arange(n, dtype=np.int32)]*4) - elif eri.ndim == 2: # 4-fold symmetry + idx = lib.cartesian_prod([np.arange(n, dtype=np.int32)] * 4) + + # Physicist notation + idx=idx.reshape((num_integrals,4)) + for i in range(num_integrals): + idx[i,1],idx[i,2]=idx[i,2],idx[i,1] + idx=idx.flatten() + + # write ERI + with trexio.File(filename, 'u', back_end=_mode(backend)) as tf: + if basis == 'AO': + if not trexio.has_ao_num(tf): + trexio.write_ao_num(tf, n) + else: + if not trexio.has_mo_num(tf): + trexio.write_mo_num(tf, n) + if basis == 'MO': + trexio.write_mo_2e_int_eri(tf, 0, num_integrals, idx, eri.ravel()) + else: + trexio.write_ao_2e_int_eri(tf, 0, num_integrals, idx, eri.ravel()) + return + + if sym == 's4': + if eri.ndim != 2 or eri.shape[0] != eri.shape[1]: + raise ValueError("s4 ERI must be a square 2D array (npair, npair)") npair = eri.shape[0] - n = int((npair * 2)**.5) - idx_pair = np.argwhere(np.arange(n)[:,None] >= np.arange(n)) - idx = np.empty((npair, npair, 4), dtype=np.int32) - idx[:,:,:2] = idx_pair[:,None,:] - idx[:,:,2:] = idx_pair[None,:,:] - idx = idx.reshape(npair**2, 4) - elif eri.ndim == 1: # 8-fold symmetry - npair = int((eri.shape[0] * 2)**.5) - n = int((npair * 2)**.5) - idx_pair = np.argwhere(np.arange(n)[:,None] >= np.arange(n)) - idx = np.empty((npair, npair, 4), dtype=np.int32) - idx[:,:,:2] = idx_pair[:,None,:] - idx[:,:,2:] = idx_pair[None,:,:] - idx = idx[np.tril_indices(npair)] - else: - raise ValueError(f'ERI array must be 1, 2 or 4-dimensional, got {eri.ndim}') + n = int(round((np.sqrt(8 * npair + 1) - 1) / 2)) + if n * (n + 1) // 2 != npair: + raise ValueError('Invalid s4 ERI size for pair indexing') + + pair_i, pair_j = _pair_from_tril(n) + total = npair * npair + chunk = 100000 + + with trexio.File(filename, 'u', back_end=_mode(backend)) as tf: + if basis == 'AO': + if not trexio.has_ao_num(tf): + trexio.write_ao_num(tf, n) + else: + if not trexio.has_mo_num(tf): + trexio.write_mo_num(tf, n) + + offset = 0 + while offset < total: + end = min(offset + chunk, total) + t = np.arange(offset, end, dtype=np.int64) + ij = t // npair + kl = t % npair + + # Physicist notation + i = pair_i[ij] + j = pair_j[ij] + k = pair_i[kl] + l = pair_j[kl] + idx = np.stack([i, k, j, l], axis=1).astype(np.int32).ravel() + val = eri[ij, kl].ravel() + + if basis == 'MO': + trexio.write_mo_2e_int_eri(tf, offset, end - offset, idx, val) + else: + trexio.write_ao_2e_int_eri(tf, offset, end - offset, idx, val) + offset = end + return - # Physicist notation - idx=idx.reshape((num_integrals,4)) - for i in range(num_integrals): - idx[i,1],idx[i,2]=idx[i,2],idx[i,1] - idx=idx.flatten() + # sym == 's8' + if eri.ndim != 1: + raise ValueError('s8 ERI must be a 1D packed array') + total = eri.size + npair = int(round((np.sqrt(8 * total + 1) - 1) / 2)) + if npair * (npair + 1) // 2 != total: + raise ValueError('Invalid s8 ERI size for pair indexing') + n = int(round((np.sqrt(8 * npair + 1) - 1) / 2)) + if n * (n + 1) // 2 != npair: + raise ValueError('Invalid s8 ERI size for AO/MO indexing') + + pair_i, pair_j = _pair_from_tril(n) + tri_i, tri_j = np.tril_indices(npair) + chunk = 100000 - # write ERI with trexio.File(filename, 'u', back_end=_mode(backend)) as tf: - if basis.upper() == 'AO': + if basis == 'AO': if not trexio.has_ao_num(tf): trexio.write_ao_num(tf, n) else: if not trexio.has_mo_num(tf): trexio.write_mo_num(tf, n) - if basis.upper() == 'MO': - trexio.write_mo_2e_int_eri(tf, 0, num_integrals, idx, eri.ravel()) - else: - trexio.write_ao_2e_int_eri(tf, 0, num_integrals, idx, eri.ravel()) + + offset = 0 + while offset < total: + end = min(offset + chunk, total) + ij = tri_i[offset:end] + kl = tri_j[offset:end] + + # Physicist notation + i = pair_i[ij] + j = pair_j[ij] + k = pair_i[kl] + l = pair_j[kl] + idx = np.stack([i, k, j, l], axis=1).astype(np.int32).ravel() + val = eri[offset:end] + + if basis == 'MO': + trexio.write_mo_2e_int_eri(tf, offset, end - offset, idx, val) + else: + trexio.write_ao_2e_int_eri(tf, offset, end - offset, idx, val) + offset = end def write_scf_1e_int_eri( mf, filename, backend='h5', basis='AO', df_engine='MDF', @@ -1050,6 +1435,8 @@ def _hermitize(mat): is_pbc = hasattr(mf, 'cell') + ecp_mat = None + if is_pbc: if not _is_gamma_single_k(mf): raise NotImplementedError( @@ -1078,7 +1465,8 @@ def _hermitize(mat): if len(getattr(cell, '_ecpbas', [])) > 0: from pyscf.pbc.gto import ecp - potential += _as_matrix(ecp.ecp_int(cell), 'AO ECP potential') + ecp_mat = _as_matrix(ecp.ecp_int(cell), 'AO ECP potential') + potential += ecp_mat core = kinetic + potential else: @@ -1087,17 +1475,22 @@ def _hermitize(mat): kinetic = _as_matrix(mol.intor('int1e_kin'), 'AO kinetic') potential = _as_matrix(mol.intor('int1e_nuc'), 'AO potential') if mol._ecp: - from pyscf.gto import ecp - potential += _as_matrix(ecp.ecp_int(mol), 'AO ECP potential') + ecp_mat = _as_matrix(mol.intor('ECPscalar'), 'AO ECP potential') + potential += ecp_mat core = kinetic + potential overlap = np.ascontiguousarray(_ensure_real(_hermitize(overlap))) kinetic = np.ascontiguousarray(_ensure_real(_hermitize(kinetic))) potential = np.ascontiguousarray(_ensure_real(_hermitize(potential))) core = np.ascontiguousarray(_ensure_real(_hermitize(core))) + if ecp_mat is not None: + ecp_mat = np.ascontiguousarray(_ensure_real(_hermitize(ecp_mat))) if basis == 'AO': _write_1e_int_eri(overlap, kinetic, potential, core, filename, backend, 'AO') + if ecp_mat is not None: + with trexio.File(filename, 'u', back_end=_mode(backend)) as tf: + trexio.write_ao_1e_int_ecp(tf, ecp_mat) return def _get_uks_coeff_pair(mf_obj): @@ -1174,10 +1567,16 @@ def _ao_to_mo(mat, coeff): mo_kinetic = np.ascontiguousarray(_ensure_real(_ao_to_mo(kinetic, C))) mo_potential = np.ascontiguousarray(_ensure_real(_ao_to_mo(potential, C))) mo_core = np.ascontiguousarray(_ensure_real(_ao_to_mo(core, C))) + mo_ecp = None + if ecp_mat is not None: + mo_ecp = np.ascontiguousarray(_ensure_real(_ao_to_mo(ecp_mat, C))) _write_1e_int_eri( mo_overlap, mo_kinetic, mo_potential, mo_core, filename, backend, 'MO' ) + if mo_ecp is not None: + with trexio.File(filename, 'u', back_end=_mode(backend)) as tf: + trexio.write_mo_1e_int_ecp(tf, mo_ecp) def _write_1e_int_eri(overlap, kinetic, potential, core, filename, backend='h5', basis='AO'): basis = basis.upper() From dfe0009eab111d1e4bbcdc46dbfd4ba6f3d37869 Mon Sep 17 00:00:00 2001 From: kousuke-nakano <37653569+kousuke-nakano@users.noreply.github.com> Date: Sat, 31 Jan 2026 21:25:22 +0900 Subject: [PATCH 6/6] Cleaned up the code for pull-request. --- pyscf/tools/test/test_trexio.py | 228 ++++---- pyscf/tools/trexio.py | 933 ++++++++++++++++++-------------- 2 files changed, 640 insertions(+), 521 deletions(-) diff --git a/pyscf/tools/test/test_trexio.py b/pyscf/tools/test/test_trexio.py index 01bb47d47..3cf2c4937 100644 --- a/pyscf/tools/test/test_trexio.py +++ b/pyscf/tools/test/test_trexio.py @@ -102,12 +102,12 @@ def test_cell_k_gamma_ae_6_31g(cart): cell0.build(atom='H 0 0 0; H 0 0 1', basis='6-31g**', a=np.diag([3.0, 3.0, 5.0])) s0 = cell0.pbc_intor('int1e_ovlp', kpts=kpt) t0 = cell0.pbc_intor('int1e_kin', kpts=kpt) - v0 = cell0.pbc_intor('int1e_nuc', kpts=kpt) + #v0 = cell0.pbc_intor('int1e_nuc', kpts=kpt) trexio.to_trexio(cell0, filename) cell1 = trexio.mol_from_trexio(filename) s1 = cell1.pbc_intor('int1e_ovlp', kpts=kpt) t1 = cell1.pbc_intor('int1e_kin', kpts=kpt) - v1 = cell1.pbc_intor('int1e_nuc', kpts=kpt) + #v1 = cell1.pbc_intor('int1e_nuc', kpts=kpt) assert abs(s0 - s1).max() < DIFF_TOL assert abs(t0 - t1).max() < DIFF_TOL #assert abs(v0 - v1).max() < DIFF_TOL @@ -124,14 +124,14 @@ def test_cell_k_grid_ae_6_31g(cart): kpts0 = cell0.make_kpts(kmesh) s0 = np.asarray(cell0.pbc_intor('int1e_ovlp', kpts=kpts0)) t0 = np.asarray(cell0.pbc_intor('int1e_kin', kpts=kpts0)) - v0 = np.asarray(cell0.pbc_intor('int1e_nuc', kpts=kpts0)) + #v0 = np.asarray(cell0.pbc_intor('int1e_nuc', kpts=kpts0)) trexio.to_trexio(cell0, filename) cell1 = trexio.mol_from_trexio(filename) kpts1 = kpts0 #kpts1 = cell1.make_kpts(kmesh) s1 = np.asarray(cell1.pbc_intor('int1e_ovlp', kpts=kpts1)) t1 = np.asarray(cell1.pbc_intor('int1e_kin', kpts=kpts1)) - v1 = np.asarray(cell1.pbc_intor('int1e_nuc', kpts=kpts1)) + #v1 = np.asarray(cell1.pbc_intor('int1e_nuc', kpts=kpts1)) assert abs(s0 - s1).max() < DIFF_TOL assert abs(t0 - t1).max() < DIFF_TOL #assert abs(v0 - v1).max() < DIFF_TOL @@ -536,7 +536,7 @@ def test_mol_rhf_ccecp_ccpvqz(cart): ## molecule, segment contraction (ccecp-cc-pVQZ), ccecp, UHF @pytest.mark.parametrize("cart", [False, True], ids=["cart=false", "cart=true"]) -def test_mol_rhf_ccecp_ccpvqz(cart): +def test_mol_uhf_ccecp_ccpvqz(cart): with tempfile.TemporaryDirectory() as d: filename = os.path.join(d, 'test.h5') mol0 = pyscf.M(atom='H 0 0 0; F 0 0 1', basis='ccecp-ccpvdz', ecp='ccecp', spin=2, cart=cart) @@ -628,7 +628,7 @@ def test_write_molecule_integrals_sym_s1_to_trexio_rhf_ae(cart): trexio.to_trexio(mf0, filename) - trexio.write_scf_1e_int_eri(mf0, filename, basis='AO') + trexio.write_1e_eri(mf0, filename, basis='AO') with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: np.testing.assert_allclose(trexio_lib.read_ao_1e_int_overlap(tf), overlap, atol=DIFF_TOL) np.testing.assert_allclose(trexio_lib.read_ao_1e_int_kinetic(tf), kinetic, atol=DIFF_TOL) @@ -641,7 +641,7 @@ def test_write_molecule_integrals_sym_s1_to_trexio_rhf_ae(cart): mo_potential = _hermitize(coeff.conj().T @ potential @ coeff) mo_core = _hermitize(coeff.conj().T @ core @ coeff) - trexio.write_scf_1e_int_eri(mf0, filename, basis='MO') + trexio.write_1e_eri(mf0, filename, basis='MO') with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: np.testing.assert_allclose(trexio_lib.read_mo_1e_int_overlap(tf), mo_overlap, atol=DIFF_TOL) np.testing.assert_allclose(trexio_lib.read_mo_1e_int_kinetic(tf), mo_kinetic, atol=DIFF_TOL) @@ -650,7 +650,7 @@ def test_write_molecule_integrals_sym_s1_to_trexio_rhf_ae(cart): ao_eri = mol0.intor('int2e', aosym='s1') ao_idx_exp, ao_val_exp = _trexio_pack_eri(ao_eri, 'AO') - trexio.write_scf_2e_int_eri(mf0, filename, basis='AO') + trexio.write_2e_eri(mf0, filename, basis='AO') with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: assert trexio_lib.has_ao_2e_int_eri(tf) size = trexio_lib.read_ao_2e_int_eri_size(tf) @@ -663,7 +663,7 @@ def test_write_molecule_integrals_sym_s1_to_trexio_rhf_ae(cart): nmo = coeff.shape[1] mo_eri = mo_eri.reshape(nmo, nmo, nmo, nmo) mo_idx_exp, mo_val_exp = _trexio_pack_eri(mo_eri, 'MO') - trexio.write_scf_2e_int_eri(mf0, filename, basis='MO') + trexio.write_2e_eri(mf0, filename, basis='MO') with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: assert trexio_lib.has_mo_2e_int_eri(tf) size = trexio_lib.read_mo_2e_int_eri_size(tf) @@ -691,7 +691,7 @@ def test_write_molecule_integrals_sym_s1_to_trexio_uhf_ae(cart): trexio.to_trexio(mf0, filename) - trexio.write_scf_1e_int_eri(mf0, filename, basis='AO') + trexio.write_1e_eri(mf0, filename, basis='AO') with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: np.testing.assert_allclose(trexio_lib.read_ao_1e_int_overlap(tf), overlap, atol=DIFF_TOL) np.testing.assert_allclose(trexio_lib.read_ao_1e_int_kinetic(tf), kinetic, atol=DIFF_TOL) @@ -709,7 +709,7 @@ def test_write_molecule_integrals_sym_s1_to_trexio_uhf_ae(cart): mo_potential = _hermitize(coeff.conj().T @ potential @ coeff) mo_core = _hermitize(coeff.conj().T @ core @ coeff) - trexio.write_scf_1e_int_eri(mf0, filename, basis='MO') + trexio.write_1e_eri(mf0, filename, basis='MO') with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: np.testing.assert_allclose(trexio_lib.read_mo_1e_int_overlap(tf), mo_overlap, atol=DIFF_TOL) np.testing.assert_allclose(trexio_lib.read_mo_1e_int_kinetic(tf), mo_kinetic, atol=DIFF_TOL) @@ -722,7 +722,7 @@ def test_write_molecule_integrals_sym_s1_to_trexio_uhf_ae(cart): ao_eri = mol0.intor('int2e', aosym='s1') ao_idx_exp, ao_val_exp = _trexio_pack_eri(ao_eri, 'AO') - trexio.write_scf_2e_int_eri(mf0, filename, basis='AO') + trexio.write_2e_eri(mf0, filename, basis='AO') with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: assert trexio_lib.has_ao_2e_int_eri(tf) size = trexio_lib.read_ao_2e_int_eri_size(tf) @@ -735,7 +735,7 @@ def test_write_molecule_integrals_sym_s1_to_trexio_uhf_ae(cart): nmo = coeff.shape[1] mo_eri = mo_eri.reshape(nmo, nmo, nmo, nmo) mo_idx_exp, mo_val_exp = _trexio_pack_eri(mo_eri, 'MO') - trexio.write_scf_2e_int_eri(mf0, filename, basis='MO') + trexio.write_2e_eri(mf0, filename, basis='MO') with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: assert trexio_lib.has_mo_2e_int_eri(tf) size = trexio_lib.read_mo_2e_int_eri_size(tf) @@ -757,7 +757,7 @@ def test_write_molecule_integrals_sym_s4_to_trexio_rhf_ae(cart): ao_eri = mol0.intor('int2e', aosym='s4') ao_idx_exp, ao_val_exp = _trexio_pack_eri(ao_eri, 'AO', sym='s4') - trexio.write_scf_2e_int_eri(mf0, filename, basis='AO', sym='s4') + trexio.write_2e_eri(mf0, filename, basis='AO', sym='s4') with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: assert trexio_lib.has_ao_2e_int_eri(tf) size = trexio_lib.read_ao_2e_int_eri_size(tf) @@ -769,7 +769,7 @@ def test_write_molecule_integrals_sym_s4_to_trexio_rhf_ae(cart): coeff = mf0.mo_coeff mo_eri = ao2mo.kernel(mol0, coeff, compact=True) mo_idx_exp, mo_val_exp = _trexio_pack_eri(mo_eri, 'MO', sym='s4') - trexio.write_scf_2e_int_eri(mf0, filename, basis='MO', sym='s4') + trexio.write_2e_eri(mf0, filename, basis='MO', sym='s4') with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: assert trexio_lib.has_mo_2e_int_eri(tf) size = trexio_lib.read_mo_2e_int_eri_size(tf) @@ -792,7 +792,7 @@ def test_write_molecule_integrals_sym_s4_to_trexio_uhf_ae(cart): ao_eri = mol0.intor('int2e', aosym='s4') ao_idx_exp, ao_val_exp = _trexio_pack_eri(ao_eri, 'AO', sym='s4') - trexio.write_scf_2e_int_eri(mf0, filename, basis='AO', sym='s4') + trexio.write_2e_eri(mf0, filename, basis='AO', sym='s4') with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: assert trexio_lib.has_ao_2e_int_eri(tf) size = trexio_lib.read_ao_2e_int_eri_size(tf) @@ -805,7 +805,7 @@ def test_write_molecule_integrals_sym_s4_to_trexio_uhf_ae(cart): coeff = np.concatenate([coeff_alpha, coeff_beta], axis=1) mo_eri = ao2mo.kernel(mol0, coeff, compact=True) mo_idx_exp, mo_val_exp = _trexio_pack_eri(mo_eri, 'MO', sym='s4') - trexio.write_scf_2e_int_eri(mf0, filename, basis='MO', sym='s4') + trexio.write_2e_eri(mf0, filename, basis='MO', sym='s4') with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: assert trexio_lib.has_mo_2e_int_eri(tf) size = trexio_lib.read_mo_2e_int_eri_size(tf) @@ -827,7 +827,7 @@ def test_write_molecule_integrals_sym_s8_to_trexio_rhf_ae(cart): ao_eri = mol0.intor('int2e', aosym='s8') ao_idx_exp, ao_val_exp = _trexio_pack_eri(ao_eri, 'AO', sym='s8') - trexio.write_scf_2e_int_eri(mf0, filename, basis='AO', sym='s8') + trexio.write_2e_eri(mf0, filename, basis='AO', sym='s8') with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: assert trexio_lib.has_ao_2e_int_eri(tf) size = trexio_lib.read_ao_2e_int_eri_size(tf) @@ -850,7 +850,7 @@ def test_write_molecule_integrals_sym_s8_to_trexio_uhf_ae(cart): ao_eri = mol0.intor('int2e', aosym='s8') ao_idx_exp, ao_val_exp = _trexio_pack_eri(ao_eri, 'AO', sym='s8') - trexio.write_scf_2e_int_eri(mf0, filename, basis='AO', sym='s8') + trexio.write_2e_eri(mf0, filename, basis='AO', sym='s8') with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: assert trexio_lib.has_ao_2e_int_eri(tf) size = trexio_lib.read_ao_2e_int_eri_size(tf) @@ -888,7 +888,7 @@ def test_write_cell_gamma_integrals_sym_s1_to_trexio_rhf_ae(cart): trexio.to_trexio(mf0, filename) - trexio.write_scf_1e_int_eri(mf0, filename, basis='AO') + trexio.write_1e_eri(mf0, filename, basis='AO') with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: np.testing.assert_allclose(trexio_lib.read_ao_1e_int_overlap(tf), overlap, atol=DIFF_TOL) np.testing.assert_allclose(trexio_lib.read_ao_1e_int_kinetic(tf), kinetic, atol=DIFF_TOL) @@ -902,7 +902,7 @@ def test_write_cell_gamma_integrals_sym_s1_to_trexio_rhf_ae(cart): mo_potential = _hermitize(coeff.conj().T @ potential @ coeff) mo_core = _hermitize(coeff.conj().T @ core @ coeff) - trexio.write_scf_1e_int_eri(mf0, filename, basis='MO') + trexio.write_1e_eri(mf0, filename, basis='MO') with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: np.testing.assert_allclose(trexio_lib.read_mo_1e_int_overlap(tf), mo_overlap, atol=DIFF_TOL) np.testing.assert_allclose(trexio_lib.read_mo_1e_int_kinetic(tf), mo_kinetic, atol=DIFF_TOL) @@ -914,7 +914,7 @@ def test_write_cell_gamma_integrals_sym_s1_to_trexio_rhf_ae(cart): nao = cell0.nao_nr() eri_ao = eri_ao.reshape(nao, nao, nao, nao) ao_idx_exp, ao_val_exp = _trexio_pack_eri(eri_ao, 'AO') - trexio.write_scf_2e_int_eri(mf0, filename, basis='AO') + trexio.write_2e_eri(mf0, filename, basis='AO') with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: assert trexio_lib.has_ao_2e_int_eri(tf) size = trexio_lib.read_ao_2e_int_eri_size(tf) @@ -933,7 +933,7 @@ def test_write_cell_gamma_integrals_sym_s1_to_trexio_rhf_ae(cart): elif mo_eri.ndim < 4: mo_eri = ao2mo.restore(1, mo_eri, nmo) mo_idx_exp, mo_val_exp = _trexio_pack_eri(mo_eri, 'MO') - trexio.write_scf_2e_int_eri(mf0, filename, basis='MO') + trexio.write_2e_eri(mf0, filename, basis='MO') with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: assert trexio_lib.has_mo_2e_int_eri(tf) size = trexio_lib.read_mo_2e_int_eri_size(tf) @@ -972,7 +972,7 @@ def test_write_cell_gamma_integrals_sym_s1_to_trexio_uhf_ae(cart): trexio.to_trexio(mf0, filename) - trexio.write_scf_1e_int_eri(mf0, filename, basis='AO') + trexio.write_1e_eri(mf0, filename, basis='AO') with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: np.testing.assert_allclose(trexio_lib.read_ao_1e_int_overlap(tf), overlap, atol=DIFF_TOL) np.testing.assert_allclose(trexio_lib.read_ao_1e_int_kinetic(tf), kinetic, atol=DIFF_TOL) @@ -988,7 +988,7 @@ def test_write_cell_gamma_integrals_sym_s1_to_trexio_uhf_ae(cart): mo_potential = _hermitize(coeff.conj().T @ potential @ coeff) mo_core = _hermitize(coeff.conj().T @ core @ coeff) - trexio.write_scf_1e_int_eri(mf0, filename, basis='MO') + trexio.write_1e_eri(mf0, filename, basis='MO') with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: np.testing.assert_allclose(trexio_lib.read_mo_1e_int_overlap(tf), mo_overlap, atol=DIFF_TOL) np.testing.assert_allclose(trexio_lib.read_mo_1e_int_kinetic(tf), mo_kinetic, atol=DIFF_TOL) @@ -1000,7 +1000,7 @@ def test_write_cell_gamma_integrals_sym_s1_to_trexio_uhf_ae(cart): nao = cell0.nao_nr() eri_ao = eri_ao.reshape(nao, nao, nao, nao) ao_idx_exp, ao_val_exp = _trexio_pack_eri(eri_ao, 'AO') - trexio.write_scf_2e_int_eri(mf0, filename, basis='AO') + trexio.write_2e_eri(mf0, filename, basis='AO') with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: assert trexio_lib.has_ao_2e_int_eri(tf) size = trexio_lib.read_ao_2e_int_eri_size(tf) @@ -1018,7 +1018,7 @@ def test_write_cell_gamma_integrals_sym_s1_to_trexio_uhf_ae(cart): elif mo_eri.ndim < 4: mo_eri = ao2mo.restore(1, mo_eri, nmo) mo_idx_exp, mo_val_exp = _trexio_pack_eri(mo_eri, 'MO') - trexio.write_scf_2e_int_eri(mf0, filename, basis='MO') + trexio.write_2e_eri(mf0, filename, basis='MO') with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: assert trexio_lib.has_mo_2e_int_eri(tf) size = trexio_lib.read_mo_2e_int_eri_size(tf) @@ -1045,24 +1045,24 @@ def test_energy_molecule_integrals_sym_s1_in_trexio_rhf_ae(cart): trexio.to_trexio(mf0, filename) - trexio.write_scf_1e_int_eri(mf0, filename, basis='AO') + trexio.write_1e_eri(mf0, filename, basis='AO') with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: np.testing.assert_allclose(trexio_lib.read_ao_1e_int_overlap(tf), overlap, atol=DIFF_TOL) np.testing.assert_allclose(trexio_lib.read_ao_1e_int_kinetic(tf), kinetic, atol=DIFF_TOL) np.testing.assert_allclose(trexio_lib.read_ao_1e_int_potential_n_e(tf), potential, atol=DIFF_TOL) np.testing.assert_allclose(trexio_lib.read_ao_1e_int_core_hamiltonian(tf), core, atol=DIFF_TOL) - trexio.write_scf_2e_int_eri(mf0, filename, basis='AO') + trexio.write_2e_eri(mf0, filename, basis='AO') with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: assert trexio_lib.has_ao_2e_int_eri(tf) size = trexio_lib.read_ao_2e_int_eri_size(tf) idx, val, n_read, _ = trexio_lib.read_ao_2e_int_eri(tf, 0, size) assert n_read == size - trexio.write_scf_1e_int_eri(mf0, filename, basis='MO') - trexio.write_scf_2e_int_eri(mf0, filename, basis='MO') - trexio.write_scf_rdm_1e(mf0, filename) - trexio.write_scf_rdm_2e(mf0, filename) + trexio.write_1e_eri(mf0, filename, basis='MO') + trexio.write_2e_eri(mf0, filename, basis='MO') + trexio.write_1b_rdm(mf0, filename) + trexio.write_2b_rdm(mf0, filename) BUFSIZE = 100000 with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: @@ -1164,12 +1164,12 @@ def test_energy_crystal_integrals_sym_s1_in_trexio_rhf_ae(cart): assert mf0.converged trexio.to_trexio(mf0, filename) - trexio.write_scf_1e_int_eri(mf0, filename, basis='AO') - trexio.write_scf_2e_int_eri(mf0, filename, basis='AO') - trexio.write_scf_1e_int_eri(mf0, filename, basis='MO') - trexio.write_scf_2e_int_eri(mf0, filename, basis='MO') - trexio.write_scf_rdm_1e(mf0, filename) - trexio.write_scf_rdm_2e(mf0, filename) + trexio.write_1e_eri(mf0, filename, basis='AO') + trexio.write_2e_eri(mf0, filename, basis='AO') + trexio.write_1e_eri(mf0, filename, basis='MO') + trexio.write_2e_eri(mf0, filename, basis='MO') + trexio.write_1b_rdm(mf0, filename) + trexio.write_2b_rdm(mf0, filename) BUFSIZE = 100000 with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: @@ -1249,12 +1249,12 @@ def test_energy_crystal_integrals_sym_s1_in_trexio_uhf_ae(cart): assert mf0.converged trexio.to_trexio(mf0, filename) - trexio.write_scf_1e_int_eri(mf0, filename, basis='AO') - trexio.write_scf_2e_int_eri(mf0, filename, basis='AO') - trexio.write_scf_1e_int_eri(mf0, filename, basis='MO') - trexio.write_scf_2e_int_eri(mf0, filename, basis='MO') - trexio.write_scf_rdm_1e(mf0, filename) - trexio.write_scf_rdm_2e(mf0, filename) + trexio.write_1e_eri(mf0, filename, basis='AO') + trexio.write_2e_eri(mf0, filename, basis='AO') + trexio.write_1e_eri(mf0, filename, basis='MO') + trexio.write_2e_eri(mf0, filename, basis='MO') + trexio.write_1b_rdm(mf0, filename) + trexio.write_2b_rdm(mf0, filename) BUFSIZE = 100000 with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: @@ -1373,13 +1373,13 @@ def test_energy_crystal_integrals_sym_s1_in_trexio_rhf_ccecp(cart): cell = pyscf.pbc.gto.Cell() cell.cart = cart cell.unit = 'Bohr' + cell.exp_to_discard=0.2 cell.build( atom='H 0 0 0; H 0 0 2.6', basis='ccecp-ccpvdz', ecp='ccecp', a=np.diag([6.0, 6.0, 6.0]), ) - cell.exp_to_discard=0.2 mf0 = pyscf.pbc.scf.RHF(cell) mf0.with_df = pbcdf.MDF(cell).build() @@ -1387,12 +1387,12 @@ def test_energy_crystal_integrals_sym_s1_in_trexio_rhf_ccecp(cart): assert mf0.converged trexio.to_trexio(mf0, filename) - trexio.write_scf_1e_int_eri(mf0, filename, basis='AO') - trexio.write_scf_2e_int_eri(mf0, filename, basis='AO') - trexio.write_scf_1e_int_eri(mf0, filename, basis='MO') - trexio.write_scf_2e_int_eri(mf0, filename, basis='MO') - trexio.write_scf_rdm_1e(mf0, filename) - trexio.write_scf_rdm_2e(mf0, filename) + trexio.write_1e_eri(mf0, filename, basis='AO') + trexio.write_2e_eri(mf0, filename, basis='AO') + trexio.write_1e_eri(mf0, filename, basis='MO') + trexio.write_2e_eri(mf0, filename, basis='MO') + trexio.write_1b_rdm(mf0, filename) + trexio.write_2b_rdm(mf0, filename) BUFSIZE = 100000 with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: @@ -1460,13 +1460,13 @@ def test_energy_crystal_integrals_sym_s1_in_trexio_uhf_ccecp(cart): cell.cart = cart cell.unit = 'Bohr' cell.spin = 2 + cell.exp_to_discard=0.2 cell.build( atom='H 0 0 0; H 0 0 2.6', basis='ccecp-ccpvdz', ecp='ccecp', a=np.diag([6.0, 6.0, 6.0]), ) - cell.exp_to_discard=0.2 mf0 = pyscf.pbc.scf.UHF(cell) mf0.with_df = pbcdf.MDF(cell).build() @@ -1474,12 +1474,12 @@ def test_energy_crystal_integrals_sym_s1_in_trexio_uhf_ccecp(cart): assert mf0.converged trexio.to_trexio(mf0, filename) - trexio.write_scf_1e_int_eri(mf0, filename, basis='AO') - trexio.write_scf_2e_int_eri(mf0, filename, basis='AO') - trexio.write_scf_1e_int_eri(mf0, filename, basis='MO') - trexio.write_scf_2e_int_eri(mf0, filename, basis='MO') - trexio.write_scf_rdm_1e(mf0, filename) - trexio.write_scf_rdm_2e(mf0, filename) + trexio.write_1e_eri(mf0, filename, basis='AO') + trexio.write_2e_eri(mf0, filename, basis='AO') + trexio.write_1e_eri(mf0, filename, basis='MO') + trexio.write_2e_eri(mf0, filename, basis='MO') + trexio.write_1b_rdm(mf0, filename) + trexio.write_2b_rdm(mf0, filename) BUFSIZE = 100000 with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: @@ -1610,12 +1610,12 @@ def test_energy_crystal_integrals_sym_s4_in_trexio_rhf_ae(cart): assert mf0.converged trexio.to_trexio(mf0, filename) - trexio.write_scf_1e_int_eri(mf0, filename, basis='AO') - trexio.write_scf_2e_int_eri(mf0, filename, basis='AO', sym='s4') - trexio.write_scf_1e_int_eri(mf0, filename, basis='MO') - trexio.write_scf_2e_int_eri(mf0, filename, basis='MO', sym='s4') - trexio.write_scf_rdm_1e(mf0, filename) - trexio.write_scf_rdm_2e(mf0, filename) + trexio.write_1e_eri(mf0, filename, basis='AO') + trexio.write_2e_eri(mf0, filename, basis='AO', sym='s4') + trexio.write_1e_eri(mf0, filename, basis='MO') + trexio.write_2e_eri(mf0, filename, basis='MO', sym='s4') + trexio.write_1b_rdm(mf0, filename) + trexio.write_2b_rdm(mf0, filename) BUFSIZE = 100000 with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: @@ -1686,12 +1686,12 @@ def test_energy_crystal_integrals_sym_s4_in_trexio_uhf_ae(cart): assert mf0.converged trexio.to_trexio(mf0, filename) - trexio.write_scf_1e_int_eri(mf0, filename, basis='AO') - trexio.write_scf_2e_int_eri(mf0, filename, basis='AO', sym='s4') - trexio.write_scf_1e_int_eri(mf0, filename, basis='MO') - trexio.write_scf_2e_int_eri(mf0, filename, basis='MO', sym='s4') - trexio.write_scf_rdm_1e(mf0, filename) - trexio.write_scf_rdm_2e(mf0, filename) + trexio.write_1e_eri(mf0, filename, basis='AO') + trexio.write_2e_eri(mf0, filename, basis='AO', sym='s4') + trexio.write_1e_eri(mf0, filename, basis='MO') + trexio.write_2e_eri(mf0, filename, basis='MO', sym='s4') + trexio.write_1b_rdm(mf0, filename) + trexio.write_2b_rdm(mf0, filename) BUFSIZE = 100000 with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: @@ -1811,24 +1811,24 @@ def test_energy_molecule_integrals_sym_s1_in_trexio_uhf_ae(cart): trexio.to_trexio(mf0, filename) - trexio.write_scf_1e_int_eri(mf0, filename, basis='AO') + trexio.write_1e_eri(mf0, filename, basis='AO') with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: np.testing.assert_allclose(trexio_lib.read_ao_1e_int_overlap(tf), overlap, atol=DIFF_TOL) np.testing.assert_allclose(trexio_lib.read_ao_1e_int_kinetic(tf), kinetic, atol=DIFF_TOL) np.testing.assert_allclose(trexio_lib.read_ao_1e_int_potential_n_e(tf), potential, atol=DIFF_TOL) np.testing.assert_allclose(trexio_lib.read_ao_1e_int_core_hamiltonian(tf), core, atol=DIFF_TOL) - trexio.write_scf_2e_int_eri(mf0, filename, basis='AO') + trexio.write_2e_eri(mf0, filename, basis='AO') with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: assert trexio_lib.has_ao_2e_int_eri(tf) size = trexio_lib.read_ao_2e_int_eri_size(tf) nao = trexio_lib.read_ao_num(tf) assert size == nao ** 4 - trexio.write_scf_1e_int_eri(mf0, filename, basis='MO') - trexio.write_scf_2e_int_eri(mf0, filename, basis='MO') - trexio.write_scf_rdm_1e(mf0, filename) - trexio.write_scf_rdm_2e(mf0, filename) + trexio.write_1e_eri(mf0, filename, basis='MO') + trexio.write_2e_eri(mf0, filename, basis='MO') + trexio.write_1b_rdm(mf0, filename) + trexio.write_2b_rdm(mf0, filename) BUFSIZE = 100000 with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: @@ -2006,24 +2006,24 @@ def test_energy_molecule_integrals_sym_s1_in_trexio_rhf_ecp(cart): trexio.to_trexio(mf0, filename) - trexio.write_scf_1e_int_eri(mf0, filename, basis='AO') + trexio.write_1e_eri(mf0, filename, basis='AO') with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: np.testing.assert_allclose(trexio_lib.read_ao_1e_int_overlap(tf), overlap, atol=DIFF_TOL) np.testing.assert_allclose(trexio_lib.read_ao_1e_int_kinetic(tf), kinetic, atol=DIFF_TOL) np.testing.assert_allclose(trexio_lib.read_ao_1e_int_potential_n_e(tf), potential, atol=DIFF_TOL) np.testing.assert_allclose(trexio_lib.read_ao_1e_int_core_hamiltonian(tf), core, atol=DIFF_TOL) - trexio.write_scf_2e_int_eri(mf0, filename, basis='AO') + trexio.write_2e_eri(mf0, filename, basis='AO') with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: assert trexio_lib.has_ao_2e_int_eri(tf) size = trexio_lib.read_ao_2e_int_eri_size(tf) idx, val, n_read, _ = trexio_lib.read_ao_2e_int_eri(tf, 0, size) assert n_read == size - trexio.write_scf_1e_int_eri(mf0, filename, basis='MO') - trexio.write_scf_2e_int_eri(mf0, filename, basis='MO') - trexio.write_scf_rdm_1e(mf0, filename) - trexio.write_scf_rdm_2e(mf0, filename) + trexio.write_1e_eri(mf0, filename, basis='MO') + trexio.write_2e_eri(mf0, filename, basis='MO') + trexio.write_1b_rdm(mf0, filename) + trexio.write_2b_rdm(mf0, filename) BUFSIZE = 100000 with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: @@ -2129,24 +2129,24 @@ def test_energy_molecule_integrals_sym_s1_in_trexio_uhf_ecp(cart): trexio.to_trexio(mf0, filename) - trexio.write_scf_1e_int_eri(mf0, filename, basis='AO') + trexio.write_1e_eri(mf0, filename, basis='AO') with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: np.testing.assert_allclose(trexio_lib.read_ao_1e_int_overlap(tf), overlap, atol=DIFF_TOL) np.testing.assert_allclose(trexio_lib.read_ao_1e_int_kinetic(tf), kinetic, atol=DIFF_TOL) np.testing.assert_allclose(trexio_lib.read_ao_1e_int_potential_n_e(tf), potential, atol=DIFF_TOL) np.testing.assert_allclose(trexio_lib.read_ao_1e_int_core_hamiltonian(tf), core, atol=DIFF_TOL) - trexio.write_scf_2e_int_eri(mf0, filename, basis='AO') + trexio.write_2e_eri(mf0, filename, basis='AO') with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: assert trexio_lib.has_ao_2e_int_eri(tf) size = trexio_lib.read_ao_2e_int_eri_size(tf) nao = trexio_lib.read_ao_num(tf) assert size == nao ** 4 - trexio.write_scf_1e_int_eri(mf0, filename, basis='MO') - trexio.write_scf_2e_int_eri(mf0, filename, basis='MO') - trexio.write_scf_rdm_1e(mf0, filename) - trexio.write_scf_rdm_2e(mf0, filename) + trexio.write_1e_eri(mf0, filename, basis='MO') + trexio.write_2e_eri(mf0, filename, basis='MO') + trexio.write_1b_rdm(mf0, filename) + trexio.write_2b_rdm(mf0, filename) BUFSIZE = 100000 with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: @@ -2310,7 +2310,7 @@ def test_energy_molecule_integrals_sym_s4_in_trexio_rhf_ae(cart): mol0 = pyscf.M(atom='H 0 0 0; F 0 0 1', basis='6-31g*', cart=cart) mf0 = mol0.RHF().run() - overlap = _hermitize(mf0.get_ovlp()) + #overlap = _hermitize(mf0.get_ovlp()) kinetic = _hermitize(mol0.intor('int1e_kin')) potential = _hermitize(mol0.intor('int1e_nuc')) if mol0._ecp: @@ -2320,12 +2320,12 @@ def test_energy_molecule_integrals_sym_s4_in_trexio_rhf_ae(cart): trexio.to_trexio(mf0, filename) - trexio.write_scf_1e_int_eri(mf0, filename, basis='AO') - trexio.write_scf_2e_int_eri(mf0, filename, basis='AO', sym='s4') - trexio.write_scf_1e_int_eri(mf0, filename, basis='MO') - trexio.write_scf_2e_int_eri(mf0, filename, basis='MO', sym='s4') - trexio.write_scf_rdm_1e(mf0, filename) - trexio.write_scf_rdm_2e(mf0, filename) + trexio.write_1e_eri(mf0, filename, basis='AO') + trexio.write_2e_eri(mf0, filename, basis='AO', sym='s4') + trexio.write_1e_eri(mf0, filename, basis='MO') + trexio.write_2e_eri(mf0, filename, basis='MO', sym='s4') + trexio.write_1b_rdm(mf0, filename) + trexio.write_2b_rdm(mf0, filename) BUFSIZE = 100000 with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: @@ -2399,7 +2399,7 @@ def test_energy_molecule_integrals_sym_s4_in_trexio_uhf_ae(cart): mf0 = mol0.UHF().run() assert mf0.converged - overlap = _hermitize(mf0.get_ovlp()) + #overlap = _hermitize(mf0.get_ovlp()) kinetic = _hermitize(mol0.intor('int1e_kin')) potential = _hermitize(mol0.intor('int1e_nuc')) if mol0._ecp: @@ -2409,12 +2409,12 @@ def test_energy_molecule_integrals_sym_s4_in_trexio_uhf_ae(cart): trexio.to_trexio(mf0, filename) - trexio.write_scf_1e_int_eri(mf0, filename, basis='AO') - trexio.write_scf_2e_int_eri(mf0, filename, basis='AO', sym='s4') - trexio.write_scf_1e_int_eri(mf0, filename, basis='MO') - trexio.write_scf_2e_int_eri(mf0, filename, basis='MO', sym='s4') - trexio.write_scf_rdm_1e(mf0, filename) - trexio.write_scf_rdm_2e(mf0, filename) + trexio.write_1e_eri(mf0, filename, basis='AO') + trexio.write_2e_eri(mf0, filename, basis='AO', sym='s4') + trexio.write_1e_eri(mf0, filename, basis='MO') + trexio.write_2e_eri(mf0, filename, basis='MO', sym='s4') + trexio.write_1b_rdm(mf0, filename) + trexio.write_2b_rdm(mf0, filename) BUFSIZE = 100000 with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: @@ -2560,20 +2560,19 @@ def test_energy_molecule_integrals_sym_s8_in_trexio_rhf_ae(cart): mol0 = pyscf.M(atom='H 0 0 0; F 0 0 1', basis='6-31g*', cart=cart) mf0 = mol0.RHF().run() - overlap = _hermitize(mf0.get_ovlp()) - kinetic = _hermitize(mol0.intor('int1e_kin')) + #overlap = _hermitize(mf0.get_ovlp()) + #kinetic = _hermitize(mol0.intor('int1e_kin')) potential = _hermitize(mol0.intor('int1e_nuc')) if mol0._ecp: from pyscf.gto import ecp potential += _hermitize(ecp.ecp_int(mol0)) - core = kinetic + potential + #core = kinetic + potential trexio.to_trexio(mf0, filename) - trexio.write_scf_1e_int_eri(mf0, filename, basis='AO') - trexio.write_scf_2e_int_eri(mf0, filename, basis='AO', sym='s8') + trexio.write_1e_eri(mf0, filename, basis='AO') + trexio.write_2e_eri(mf0, filename, basis='AO', sym='s8') - BUFSIZE = 100000 with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: e_nn = trexio_lib.read_nucleus_repulsion(tf) @@ -2612,20 +2611,19 @@ def test_energy_molecule_integrals_sym_s8_in_trexio_uhf_ae(cart): mf0 = mol0.UHF().run() assert mf0.converged - overlap = _hermitize(mf0.get_ovlp()) - kinetic = _hermitize(mol0.intor('int1e_kin')) + #overlap = _hermitize(mf0.get_ovlp()) + #kinetic = _hermitize(mol0.intor('int1e_kin')) potential = _hermitize(mol0.intor('int1e_nuc')) if mol0._ecp: from pyscf.gto import ecp potential += _hermitize(ecp.ecp_int(mol0)) - core = kinetic + potential + #core = kinetic + potential trexio.to_trexio(mf0, filename) - trexio.write_scf_1e_int_eri(mf0, filename, basis='AO') - trexio.write_scf_2e_int_eri(mf0, filename, basis='AO', sym='s8') + trexio.write_1e_eri(mf0, filename, basis='AO') + trexio.write_2e_eri(mf0, filename, basis='AO', sym='s8') - BUFSIZE = 100000 with trexio_lib.File(filename, 'r', back_end=trexio_lib.TREXIO_AUTO) as tf: e_nn = trexio_lib.read_nucleus_repulsion(tf) diff --git a/pyscf/tools/trexio.py b/pyscf/tools/trexio.py index 86687c528..10f36b932 100644 --- a/pyscf/tools/trexio.py +++ b/pyscf/tools/trexio.py @@ -788,32 +788,107 @@ def scf_from_trexio(filename): else: raise ValueError(f'Unknown spin multiplicity {uniq}') -def write_scf_2e_int_eri( +_REAL_ONLY_TOL = 1e-12 + + +def _trexio_ensure_real(x, *, tol=_REAL_ONLY_TOL, what="Complex data encountered but the backend is real-only."): + if np.iscomplexobj(x): + if np.all(np.abs(np.imag(x)) <= tol): + x = np.real(x) + else: + raise NotImplementedError(what) + return x + + +def _trexio_is_gamma_single_k(obj) -> bool: + if not hasattr(obj, 'cell'): + return False + if hasattr(obj, 'kpt'): + return np.allclose(np.asarray(obj.kpt), 0.0) + if hasattr(obj, 'kpts'): + kpts = np.asarray(obj.kpts) + return (kpts.shape[0] == 1) and np.allclose(kpts[0], 0.0) + return True + + +def _trexio_get_uks_coeff_pair(mf_obj, *, expect_gamma=False): + C = mf_obj.mo_coeff + if isinstance(C, (list, tuple)) and len(C) == 2: + Ca, Cb = C + elif isinstance(C, np.ndarray) and C.ndim >= 3 and C.shape[0] == 2: + if C.ndim == 3: + Ca, Cb = C[0], C[1] + elif C.ndim == 4: + if expect_gamma and C.shape[1] == 1: + Ca, Cb = C[0, 0], C[1, 0] + else: + raise NotImplementedError("Only single-k UKS/UHF is supported.") + else: + raise ValueError(f"Unexpected mo_coeff shape: {C.shape}") + else: + raise TypeError("Not a UKS/UHF object or unsupported mo_coeff layout.") + + if expect_gamma: + if Ca.ndim == 3 and Ca.shape[0] == 1: + Ca = Ca[0] + if Cb.ndim == 3 and Cb.shape[0] == 1: + Cb = Cb[0] + if Ca.ndim != 2 or Cb.ndim != 2: + raise NotImplementedError( + "Only Gamma-point data are supported for combined MO coefficients." + ) + + if Ca.ndim != 2 or Cb.ndim != 2: + raise ValueError(f"Unexpected UKS/UHF mo_coeff shapes: Ca {Ca.shape}, Cb {Cb.shape}") + return Ca, Cb + + +def _trexio_concat_spin_coeff(Ca, Cb): + if Ca.ndim != 2 or Cb.ndim != 2: + raise ValueError(f"Unexpected UKS/UHF mo_coeff shapes: Ca {Ca.shape}, Cb {Cb.shape}") + return np.concatenate([Ca, Cb], axis=1) + +def write_2e_eri( mf, filename, backend='h5', basis='mo', df_engine='MDF', sym='s1', ): + """Write two-electron repulsion integrals to a TREXIO file. + + Parameters + ---------- + mf : SCF/KSCF object + Converged molecular or PBC mean-field object. PBC data must be + Gamma-only. + filename : str + Path to the TREXIO file to create or update. + backend : {'h5', 'text'}, optional + TREXIO backend selector passed through to ``trexio.File``. + basis : {'AO', 'MO'}, optional + Basis in which integrals are written. AO is always spin-independent. + MO concatenates alpha and beta MOs for UHF/UKS to include cross-spin + blocks. + df_engine : {'MDF', 'GDF'}, optional + Density-fitting engine used for PBC ERIs (Gamma-only). + sym : {'s1', 's4', 's8'}, optional + ERI symmetry/packing. MO supports ``s1``/``s4``; AO supports + ``s1``/``s4``/``s8``. + + Behavior + -------- + - Real-only backend: complex ERIs are rejected unless the imaginary part + is <=1e-12 (in which case it is discarded). + - PBC: only single-k Gamma is supported; non-Gamma raises + ``NotImplementedError``. + - MO + UHF/UKS: constructs the combined coefficient matrix ``[Ca | Cb]`` + and writes all spin blocks. + - Arrays are made C-contiguous before writing; dtype is preserved. + + Raises + ------ + ValueError + For invalid ``basis``/``sym`` combinations or unexpected shapes. + NotImplementedError + For complex ERIs, non-Gamma PBC data, or unsupported symmetry in MO. """ - Write two-electron integrals (ERI) in AO or MO basis. - - Rules - ----- - - Backend is real-only: complex ERIs are not supported. - - PBC: Only single-k Gamma is supported (non-Gamma -> NotImplementedError). - - Molecular RHF/RKS: AO and MO supported. - - Molecular/PBC UKS/UHF: - * AO: spin-independent (single tensor). - * MO: build a combined MO space by column-wise concatenation [alpha | beta] - and compute full ERIs including cross-spin terms. - - Symmetry packing is supported for s1/s4 (MO) and s1/s4/s8 (AO). - - Notes - ----- - - Immediately before writing, arrays are forced to be C-contiguous *without* - changing dtype: - np.ascontiguousarray(arr) - """ - - # ----- internal constants ----- - TOL_IMAG = 1e-12 # drop imag parts <= this; otherwise raise (real-only backend) basis = basis.upper() sym = sym.lower() @@ -825,30 +900,6 @@ def write_scf_2e_int_eri( raise NotImplementedError("MO ERI does not support s8 symmetry; use s1 or s4") # ---------- helpers ---------- - def _is_gamma_single_k(mf_obj) -> bool: - """True if this PBC calculation is single-k at Gamma.""" - if not hasattr(mf_obj, 'cell'): - return False - if hasattr(mf_obj, 'kpt'): - return np.allclose(np.asarray(mf_obj.kpt), 0.0) - if hasattr(mf_obj, 'kpts'): - kpts = np.asarray(mf_obj.kpts) - return (kpts.shape[0] == 1) and np.allclose(kpts[0], 0.0) - return True # treat no-kpt attribute as Gamma - - def _ensure_real(x): - """Ensure array is real; drop tiny imaginary part else raise (dtype preserved).""" - if np.iscomplexobj(x): - if np.all(np.abs(np.imag(x)) <= TOL_IMAG): - # np.real keeps precision: complex128->float64, complex64->float32 - x = np.real(x) - else: - raise NotImplementedError( - "Complex ERI encountered but the backend is real-only. " - "Use Gamma-point (k=0) or a complex-capable backend." - ) - return x - def _df_obj(): """Construct a DF engine for PBC Γ-point.""" if df_engine.upper() == 'MDF': @@ -857,31 +908,6 @@ def _df_obj(): return pbcdf.GDF(mf.cell).build() raise ValueError("df_engine must be 'MDF' or 'GDF'.") - def _get_uks_coeff_pair(mf_obj): - """ - Extract (Ca, Cb) as 2D arrays (nao, nalpha), (nao, nbeta), squeezing possible k-dims. - Supports: - - list/tuple: (Ca, Cb) - - ndarray shapes: (2, nao, nmo) or (2, 1, nao, nmo) - """ - C = mf_obj.mo_coeff - if isinstance(C, (list, tuple)) and len(C) == 2: - Ca, Cb = C - elif isinstance(C, np.ndarray) and C.ndim >= 3 and C.shape[0] == 2: - if C.ndim == 3: - Ca, Cb = C[0], C[1] - elif C.ndim == 4: - if C.shape[1] != 1: - raise NotImplementedError("Only single-k UKS/UHF is supported.") - Ca, Cb = C[0, 0], C[1, 0] - else: - raise ValueError(f"Unexpected mo_coeff shape: {C.shape}") - else: - raise TypeError("Not a UKS/UHF object or unsupported mo_coeff layout.") - if Ca.ndim != 2 or Cb.ndim != 2: - raise ValueError(f"Unexpected UKS/UHF mo_coeff shapes: Ca {Ca.shape}, Cb {Cb.shape}") - return Ca, Cb - # --------------------- # MO-basis ERI writing # --------------------- @@ -894,13 +920,19 @@ def _get_uks_coeff_pair(mf_obj): if (isinstance(mf.mo_coeff, (list, tuple)) or (isinstance(mf.mo_coeff, np.ndarray) and mf.mo_coeff.ndim >= 3 and mf.mo_coeff.shape[0] == 2)): # UKS/UHF -> concatenate [alpha | beta], include cross-spin terms - Ca, Cb = _get_uks_coeff_pair(mf) - C = np.concatenate([Ca, Cb], axis=1) # (nao, nalpha+nbeta) + Ca, Cb = _trexio_get_uks_coeff_pair(mf) + C = _trexio_concat_spin_coeff(Ca, Cb) # (nao, nalpha+nbeta) if getattr(mf, '_eri', None) is not None: eri_mo = ao2mo.incore.full(mf._eri, C, compact=mo_compact) else: eri_mo = ao2mo.kernel(mf.mol, C, compact=mo_compact) - eri_mo = _ensure_real(eri_mo) + eri_mo = _trexio_ensure_real( + eri_mo, + what=( + "Complex ERI encountered but the backend is real-only. " + "Use Gamma-point (k=0) or a complex-capable backend." + ), + ) nmo = C.shape[1] if sym == 's1': if eri_mo.ndim < 4: @@ -912,7 +944,13 @@ def _get_uks_coeff_pair(mf_obj): eri_mo = ao2mo.incore.full(mf._eri, C, compact=mo_compact) else: eri_mo = ao2mo.kernel(mf.mol, C, compact=mo_compact) - eri_mo = _ensure_real(eri_mo) + eri_mo = _trexio_ensure_real( + eri_mo, + what=( + "Complex ERI encountered but the backend is real-only. " + "Use Gamma-point (k=0) or a complex-capable backend." + ), + ) nmo = C.shape[1] if sym == 's1': if eri_mo.ndim < 4: @@ -921,7 +959,7 @@ def _get_uks_coeff_pair(mf_obj): # PBC (Gamma only) else: - if not _is_gamma_single_k(mf): + if not _trexio_is_gamma_single_k(mf): raise NotImplementedError("PBC MO-ERI: non-Gamma k-points are not supported (real-only backend).") if sym == 's8': raise NotImplementedError("PBC MO-ERI does not support s8 symmetry; use s1 or s4") @@ -930,10 +968,16 @@ def _get_uks_coeff_pair(mf_obj): if (isinstance(mf.mo_coeff, (list, tuple)) or (isinstance(mf.mo_coeff, np.ndarray) and mf.mo_coeff.ndim >= 3 and mf.mo_coeff.shape[0] == 2)): # UKS/UHF @ Gamma: combined MO matrix [Ca | Cb] - Ca, Cb = _get_uks_coeff_pair(mf) - C = np.concatenate([Ca, Cb], axis=1) + Ca, Cb = _trexio_get_uks_coeff_pair(mf, expect_gamma=True) + C = _trexio_concat_spin_coeff(Ca, Cb) eri_mo = dfobj.get_mo_eri((C, C, C, C)) - eri_mo = _ensure_real(eri_mo) + eri_mo = _trexio_ensure_real( + eri_mo, + what=( + "Complex ERI encountered but the backend is real-only. " + "Use Gamma-point (k=0) or a complex-capable backend." + ), + ) nmo = C.shape[1] if sym == 's1': if eri_mo.ndim == 2: @@ -949,7 +993,13 @@ def _get_uks_coeff_pair(mf_obj): if C.ndim == 3 and C.shape[0] == 1: # normalize (1,nao,nmo) -> (nao,nmo) C = C[0] eri_mo = dfobj.get_mo_eri((C, C, C, C)) - eri_mo = _ensure_real(eri_mo) + eri_mo = _trexio_ensure_real( + eri_mo, + what=( + "Complex ERI encountered but the backend is real-only. " + "Use Gamma-point (k=0) or a complex-capable backend." + ), + ) nmo = C.shape[1] if sym == 's1': if eri_mo.ndim == 2: @@ -967,7 +1017,7 @@ def _get_uks_coeff_pair(mf_obj): else: # basis == 'AO' if is_pbc: # PBC AO: Gamma only via DF (real-only) - if not _is_gamma_single_k(mf): + if not _trexio_is_gamma_single_k(mf): raise NotImplementedError("PBC AO-ERI: non-Gamma k-points are not supported (real-only backend).") dfobj = _df_obj() eri2 = pbcdf.df_ao2mo.get_eri(dfobj, compact=False) # (nao^2, nao^2) @@ -975,7 +1025,13 @@ def _get_uks_coeff_pair(mf_obj): if eri2.shape != (nao * nao, nao * nao): raise RuntimeError(f"Unexpected ERI shape {eri2.shape}; expected ({nao*nao}, {nao*nao}) at Gamma.") eri_ao = eri2.reshape(nao, nao, nao, nao) - eri_ao = _ensure_real(eri_ao) + eri_ao = _trexio_ensure_real( + eri_ao, + what=( + "Complex ERI encountered but the backend is real-only. " + "Use Gamma-point (k=0) or a complex-capable backend." + ), + ) if sym == 's4': eri_ao = ao2mo.restore(4, eri_ao, nao) elif sym == 's8': @@ -996,255 +1052,15 @@ def _get_uks_coeff_pair(mf_obj): eri_ao = ao2mo.restore(1, eri_ao, n_ao) else: eri_ao = mf.mol.intor('int2e', aosym=sym) - eri_ao = _ensure_real(eri_ao) + eri_ao = _trexio_ensure_real( + eri_ao, + what=( + "Complex ERI encountered but the backend is real-only. " + "Use Gamma-point (k=0) or a complex-capable backend." + ), + ) _write_2e_int_eri(np.ascontiguousarray(eri_ao), filename, backend, 'AO', sym=sym) -def write_scf_rdm_1e(mf, filename, backend='h5'): - """Write one-electron RDM (density matrix) to TREXIO in MO basis. - - TREXIO RDM I/O assumes MO-basis data. This helper enforces MO and - builds a spin-summed density: diag(mo_occ) for RHF/RKS and - block-diagonal [diag(occ_a), diag(occ_b)] for UHF/UKS. PBC is limited - to single-k Gamma. - """ - - def _is_gamma_single_k(mf_obj) -> bool: - if not hasattr(mf_obj, 'cell'): - return False - if hasattr(mf_obj, 'kpt'): - return np.allclose(np.asarray(mf_obj.kpt), 0.0) - if hasattr(mf_obj, 'kpts'): - kpts = np.asarray(mf_obj.kpts) - return (kpts.shape[0] == 1) and np.allclose(kpts[0], 0.0) - return True - - def _ensure_real(x): - if np.iscomplexobj(x): - if np.all(np.abs(np.imag(x)) <= 1e-12): - x = np.real(x) - else: - raise NotImplementedError( - "Complex RDM encountered; real-only backend.") - return x - - is_pbc = hasattr(mf, 'cell') - if is_pbc and not _is_gamma_single_k(mf): - raise NotImplementedError("RDM write supports Gamma-point only for PBC.") - - is_uhf_like = isinstance( - mf, - ( - scf.uhf.UHF, - dft.uks.UKS, - pbc.scf.uhf.UHF, - pbc.dft.uks.UKS, - pbc.scf.kuhf.KUHF, - pbc.dft.kuks.KUKS, - ), - ) - - # MO-basis density is diagonal in canonical orbitals - if is_uhf_like and isinstance(mf.mo_occ, (tuple, list)) and len(mf.mo_occ) == 2: - occ_a, occ_b = mf.mo_occ - occ_a = np.asarray(occ_a) - occ_b = np.asarray(occ_b) - else: - occ = np.asarray(mf.mo_occ) - if is_uhf_like and occ.ndim == 2 and occ.shape[0] == 2: - occ_a, occ_b = occ[0], occ[1] - else: - occ_a = occ - occ_b = None - - if is_pbc and isinstance(occ_a, np.ndarray) and occ_a.ndim == 2: - if occ_a.shape[0] != 1: - raise NotImplementedError( - "PBC RDM write: only single-k Gamma supported for MO basis.") - occ_a = occ_a[0] - if occ_b is not None and occ_b.ndim == 2: - occ_b = occ_b[0] - - if occ_b is not None: - dm_a = np.diag(occ_a) - dm_b = np.diag(occ_b) - dm_a = _ensure_real(np.asarray(dm_a)) - dm_b = _ensure_real(np.asarray(dm_b)) - dm_a = np.ascontiguousarray(dm_a) - dm_b = np.ascontiguousarray(dm_b) - with trexio.File(filename, 'u', back_end=_mode(backend)) as tf: - nmo_up = dm_a.shape[0] - nmo_dn = dm_b.shape[0] - if nmo_dn != nmo_up: - raise ValueError("Alpha/beta MO sizes do not match for RDM write.") - nmo = nmo_up + nmo_dn - if not trexio.has_mo_num(tf): - trexio.write_mo_num(tf, nmo) - trexio.write_rdm_1e_up(tf, dm_a) - trexio.write_rdm_1e_dn(tf, dm_b) - dm_tot = np.zeros((nmo, nmo), dtype=dm_a.dtype) - dm_tot[:nmo_up, :nmo_up] = dm_a - dm_tot[nmo_up:, nmo_up:] = dm_b - trexio.write_rdm_1e(tf, dm_tot) - else: - dm = np.diag(occ_a) - dm = _ensure_real(np.asarray(dm)) - dm = np.ascontiguousarray(dm) - with trexio.File(filename, 'u', back_end=_mode(backend)) as tf: - nmo = dm.shape[0] - if not trexio.has_mo_num(tf): - trexio.write_mo_num(tf, nmo) - trexio.write_rdm_1e(tf, dm) - -def write_scf_rdm_2e(mf, filename, backend='h5', chunk_size=100000): - """Write spin-summed two-electron RDM to TREXIO in MO basis. - - Assumes MO-basis storage (TREXIO convention/physicist) and builds a Hartree–Fock-like - RDM using the spin-summed diagonal density ``dm = diag(mo_occ)``: - ``G[pqrs] = dm[p,r]*dm[q,s] - 0.5*dm[p,s]*dm[q,r]``. - - Notes - ----- - - PBC: only single-k Gamma is supported. - - Memory scales as O(n^4); keep `chunk_size` modest for the write loop. - """ - - def _is_gamma_single_k(mf_obj) -> bool: - if not hasattr(mf_obj, 'cell'): - return False - if hasattr(mf_obj, 'kpt'): - return np.allclose(np.asarray(mf_obj.kpt), 0.0) - if hasattr(mf_obj, 'kpts'): - kpts = np.asarray(mf_obj.kpts) - return (kpts.shape[0] == 1) and np.allclose(kpts[0], 0.0) - return True - - is_pbc = hasattr(mf, 'cell') - if is_pbc and not _is_gamma_single_k(mf): - raise NotImplementedError("RDM write supports Gamma-point only for PBC.") - - is_uhf_like = isinstance( - mf, - ( - scf.uhf.UHF, - dft.uks.UKS, - pbc.scf.uhf.UHF, - pbc.dft.uks.UKS, - pbc.scf.kuhf.KUHF, - pbc.dft.kuks.KUKS, - ), - ) - - # Spin-summed occupations or spin-separated for UHF/UKS - if is_uhf_like and isinstance(mf.mo_occ, (tuple, list)) and len(mf.mo_occ) == 2: - occ_a, occ_b = mf.mo_occ - occ_a = np.asarray(occ_a) - occ_b = np.asarray(occ_b) - else: - occ = np.asarray(mf.mo_occ) - if is_uhf_like and occ.ndim == 2 and occ.shape[0] == 2: - occ_a, occ_b = occ[0], occ[1] - else: - occ_a = occ - occ_b = None - - if is_pbc and isinstance(occ_a, np.ndarray) and occ_a.ndim == 2: - if occ_a.shape[0] != 1: - raise NotImplementedError( - "PBC RDM write: only single-k Gamma supported for MO basis.") - occ_a = occ_a[0] - if occ_b is not None and occ_b.ndim == 2: - occ_b = occ_b[0] - - if occ_b is not None: - dm_a = np.diag(occ_a) - dm_b = np.diag(occ_b) - dm_a = np.ascontiguousarray(dm_a) - dm_b = np.ascontiguousarray(dm_b) - - g2_uu = ( - np.einsum('pr,qs->pqrs', dm_a, dm_a) - - np.einsum('ps,qr->pqrs', dm_a, dm_a) - ) - g2_dd = ( - np.einsum('pr,qs->pqrs', dm_b, dm_b) - - np.einsum('ps,qr->pqrs', dm_b, dm_b) - ) - g2_ud = np.einsum('pr,qs->pqrs', dm_a, dm_b) - - g2_uu = np.ascontiguousarray(g2_uu) - g2_dd = np.ascontiguousarray(g2_dd) - g2_ud = np.ascontiguousarray(g2_ud) - - nmo = dm_a.shape[0] - if dm_b.shape[0] != nmo: - raise ValueError("Alpha/beta MO sizes do not match for RDM write.") - idx = lib.cartesian_prod([np.arange(nmo, dtype=np.int32)] * 4) - flat_uu = g2_uu.reshape(-1) - flat_dd = g2_dd.reshape(-1) - flat_ud = g2_ud.reshape(-1) - - with trexio.File(filename, 'u', back_end=_mode(backend)) as tf: - if not trexio.has_mo_num(tf): - trexio.write_mo_num(tf, nmo) - - total = idx.shape[0] - offset = 0 - while offset < total: - end = min(offset + chunk_size, total) - trexio.write_rdm_2e_upup( - tf, - offset, - end - offset, - idx[offset:end], - flat_uu[offset:end], - ) - trexio.write_rdm_2e_dndn( - tf, - offset, - end - offset, - idx[offset:end], - flat_dd[offset:end], - ) - trexio.write_rdm_2e_updn( - tf, - offset, - end - offset, - idx[offset:end], - flat_ud[offset:end], - ) - offset = end - else: - dm = np.diag(occ_a) - dm = np.ascontiguousarray(dm) - - # Build dense spin-summed 2-RDM in MO basis - g2 = ( - np.einsum('pr,qs->pqrs', dm, dm) - - 0.5 * np.einsum('ps,qr->pqrs', dm, dm) - ) - g2 = np.ascontiguousarray(g2) - - nmo = dm.shape[0] - idx = lib.cartesian_prod([np.arange(nmo, dtype=np.int32)] * 4) - flat_g2 = g2.reshape(-1) - - with trexio.File(filename, 'u', back_end=_mode(backend)) as tf: - if not trexio.has_mo_num(tf): - trexio.write_mo_num(tf, nmo) - - total = idx.shape[0] - offset = 0 - while offset < total: - end = min(offset + chunk_size, total) - trexio.write_rdm_2e( - tf, - offset, - end - offset, - idx[offset:end], - flat_g2[offset:end], - ) - offset = end - def _write_2e_int_eri(eri, filename, backend='h5', basis='MO', sym='s1'): basis = basis.upper() sym = sym.lower() @@ -1370,46 +1186,49 @@ def _pair_from_tril(n): trexio.write_ao_2e_int_eri(tf, offset, end - offset, idx, val) offset = end -def write_scf_1e_int_eri( +def write_1e_eri( mf, filename, backend='h5', basis='AO', df_engine='MDF', ): - """ - Write one-electron integrals (overlap, kinetic, nuclear-electron potential, - and core Hamiltonian) in AO or MO basis. - - Rules - ----- - - Backend is real-only: negligible imaginary parts are discarded, otherwise - a NotImplementedError is raised. + """Write one-electron integrals to a TREXIO file. + + Stored quantities are overlap, kinetic, nuclear-electron potential, and + the core Hamiltonian (plus ECP if present) in either AO or MO basis. + + Parameters + ---------- + mf : SCF/KSCF object + Converged molecular or PBC mean-field object. PBC data must be + Gamma-only. + filename : str + Path to the TREXIO file to create or update. + backend : {'h5', 'text'}, optional + TREXIO backend selector passed through to ``trexio.File``. + basis : {'AO', 'MO'}, optional + Basis in which integrals are written. For MO + UHF/UKS, alpha and beta + coefficients are concatenated column-wise to form a single block. + df_engine : {'MDF', 'GDF'}, optional + Density-fitting engine for PBC integrals (Gamma-only). + + Behavior + -------- + - Real-only backend: imaginary parts larger than 1e-12 raise + ``NotImplementedError``; smaller parts are discarded. - PBC: only single-k Gamma calculations are supported. - - MO basis: UHF/UKS uses a concatenated MO matrix [alpha | beta]. + - All matrices are hermitized before writing and made C-contiguous; dtype + is preserved. + + Raises + ------ + ValueError + For invalid ``basis`` or unexpected shapes. + NotImplementedError + For complex data, non-Gamma PBC calculations, or unsupported MO layout. """ - TOL_IMAG = 1e-12 basis = basis.upper() if basis not in ('AO', 'MO'): raise ValueError("basis must be either 'AO' or 'MO'") - def _ensure_real(x): - if np.iscomplexobj(x): - if np.all(np.abs(np.imag(x)) <= TOL_IMAG): - x = np.real(x) - else: - raise NotImplementedError( - "Complex one-electron integrals encountered but the backend is real-only." - ) - return x - - def _is_gamma_single_k(mf_obj) -> bool: - if not hasattr(mf_obj, 'cell'): - return False - if hasattr(mf_obj, 'kpt'): - return np.allclose(np.asarray(mf_obj.kpt), 0.0) - if hasattr(mf_obj, 'kpts'): - kpts = np.asarray(mf_obj.kpts) - return (kpts.shape[0] == 1) and np.allclose(kpts[0], 0.0) - return True - def _as_matrix(mat, label): if isinstance(mat, (tuple, list)): if len(mat) == 0: @@ -1438,7 +1257,7 @@ def _hermitize(mat): ecp_mat = None if is_pbc: - if not _is_gamma_single_k(mf): + if not _trexio_is_gamma_single_k(mf): raise NotImplementedError( "PBC one-electron integrals are implemented for Gamma-point only." ) @@ -1479,12 +1298,37 @@ def _hermitize(mat): potential += ecp_mat core = kinetic + potential - overlap = np.ascontiguousarray(_ensure_real(_hermitize(overlap))) - kinetic = np.ascontiguousarray(_ensure_real(_hermitize(kinetic))) - potential = np.ascontiguousarray(_ensure_real(_hermitize(potential))) - core = np.ascontiguousarray(_ensure_real(_hermitize(core))) + overlap = np.ascontiguousarray( + _trexio_ensure_real( + _hermitize(overlap), + what="Complex one-electron integrals encountered but the backend is real-only.", + ) + ) + kinetic = np.ascontiguousarray( + _trexio_ensure_real( + _hermitize(kinetic), + what="Complex one-electron integrals encountered but the backend is real-only.", + ) + ) + potential = np.ascontiguousarray( + _trexio_ensure_real( + _hermitize(potential), + what="Complex one-electron integrals encountered but the backend is real-only.", + ) + ) + core = np.ascontiguousarray( + _trexio_ensure_real( + _hermitize(core), + what="Complex one-electron integrals encountered but the backend is real-only.", + ) + ) if ecp_mat is not None: - ecp_mat = np.ascontiguousarray(_ensure_real(_hermitize(ecp_mat))) + ecp_mat = np.ascontiguousarray( + _trexio_ensure_real( + _hermitize(ecp_mat), + what="Complex one-electron integrals encountered but the backend is real-only.", + ) + ) if basis == 'AO': _write_1e_int_eri(overlap, kinetic, potential, core, filename, backend, 'AO') @@ -1493,29 +1337,6 @@ def _hermitize(mat): trexio.write_ao_1e_int_ecp(tf, ecp_mat) return - def _get_uks_coeff_pair(mf_obj): - coeff = mf_obj.mo_coeff - if isinstance(coeff, (list, tuple)) and len(coeff) == 2: - Ca, Cb = coeff - elif isinstance(coeff, np.ndarray) and coeff.ndim >= 3 and coeff.shape[0] == 2: - if coeff.ndim == 3: - Ca, Cb = coeff[0], coeff[1] - elif coeff.ndim == 4: - if coeff.shape[1] != 1: - raise NotImplementedError( - "Only single-k UKS/UHF is supported for MO one-electron integrals." - ) - Ca, Cb = coeff[0, 0], coeff[1, 0] - else: - raise ValueError(f"Unexpected mo_coeff shape: {coeff.shape}") - else: - raise TypeError("Unsupported mo_coeff layout for UKS/UHF object") - if Ca.ndim != 2 or Cb.ndim != 2: - raise ValueError( - f"Unexpected UKS/UHF mo_coeff shapes: Ca {Ca.shape}, Cb {Cb.shape}" - ) - return Ca, Cb - def _get_rhf_coeff(mf_obj): coeff = mf_obj.mo_coeff if isinstance(coeff, np.ndarray): @@ -1536,17 +1357,8 @@ def _get_rhf_coeff(mf_obj): ) or ( isinstance(mf.mo_coeff, np.ndarray) and mf.mo_coeff.ndim >= 3 and mf.mo_coeff.shape[0] == 2 ): - Ca, Cb = _get_uks_coeff_pair(mf) - if is_pbc: - if Ca.ndim == 3 and Ca.shape[0] == 1: - Ca = Ca[0] - if Cb.ndim == 3 and Cb.shape[0] == 1: - Cb = Cb[0] - if Ca.ndim != 2 or Cb.ndim != 2: - raise NotImplementedError( - "Only Gamma-point data are supported for MO one-electron integrals" - ) - C = np.concatenate([Ca, Cb], axis=1) + Ca, Cb = _trexio_get_uks_coeff_pair(mf, expect_gamma=is_pbc) + C = _trexio_concat_spin_coeff(Ca, Cb) else: C = _get_rhf_coeff(mf) @@ -1563,13 +1375,38 @@ def _get_rhf_coeff(mf_obj): def _ao_to_mo(mat, coeff): return _hermitize(coeff.conj().T @ mat @ coeff) - mo_overlap = np.ascontiguousarray(_ensure_real(_ao_to_mo(overlap, C))) - mo_kinetic = np.ascontiguousarray(_ensure_real(_ao_to_mo(kinetic, C))) - mo_potential = np.ascontiguousarray(_ensure_real(_ao_to_mo(potential, C))) - mo_core = np.ascontiguousarray(_ensure_real(_ao_to_mo(core, C))) + mo_overlap = np.ascontiguousarray( + _trexio_ensure_real( + _ao_to_mo(overlap, C), + what="Complex one-electron integrals encountered but the backend is real-only.", + ) + ) + mo_kinetic = np.ascontiguousarray( + _trexio_ensure_real( + _ao_to_mo(kinetic, C), + what="Complex one-electron integrals encountered but the backend is real-only.", + ) + ) + mo_potential = np.ascontiguousarray( + _trexio_ensure_real( + _ao_to_mo(potential, C), + what="Complex one-electron integrals encountered but the backend is real-only.", + ) + ) + mo_core = np.ascontiguousarray( + _trexio_ensure_real( + _ao_to_mo(core, C), + what="Complex one-electron integrals encountered but the backend is real-only.", + ) + ) mo_ecp = None if ecp_mat is not None: - mo_ecp = np.ascontiguousarray(_ensure_real(_ao_to_mo(ecp_mat, C))) + mo_ecp = np.ascontiguousarray( + _trexio_ensure_real( + _ao_to_mo(ecp_mat, C), + what="Complex one-electron integrals encountered but the backend is real-only.", + ) + ) _write_1e_int_eri( mo_overlap, mo_kinetic, mo_potential, mo_core, filename, backend, 'MO' @@ -1606,6 +1443,290 @@ def _write_1e_int_eri(overlap, kinetic, potential, core, filename, backend='h5', trexio.write_mo_1e_int_potential_n_e(tf, potential) trexio.write_mo_1e_int_core_hamiltonian(tf, core) +def write_1b_rdm(mf, filename, backend='h5'): + """Write a one-body reduced density matrix in MO basis to TREXIO. + + Parameters + ---------- + mf : SCF/KSCF object + Converged molecular or PBC mean-field object. Uses ``mo_occ`` to + build diagonal MO densities. PBC data must be Gamma-only. + filename : str + Path to the TREXIO file to create or update. + backend : {'h5', 'text'}, optional + TREXIO backend selector passed through to ``trexio.File``. + + Behavior + -------- + - MO basis is enforced (TREXIO convention). + - RHF/RKS: writes a spin-summed diagonal density ``diag(mo_occ)``. + - UHF/UKS: writes spin-blocked ``rdm_1e_up``/``rdm_1e_dn`` and a + block-diagonal spin-summed matrix. Alpha/beta MO counts must match. + - PBC: only single-k Gamma is supported; k-resolved occupations are + squeezed to 1D first. + - Real-only backend: imaginary parts larger than 1e-12 raise + ``NotImplementedError``. + + Raises + ------ + ValueError + When alpha/beta dimensions differ or input shapes are inconsistent. + NotImplementedError + For complex densities or non-Gamma PBC calculations. + """ + + is_pbc = hasattr(mf, 'cell') + if is_pbc and not _trexio_is_gamma_single_k(mf): + raise NotImplementedError("RDM write supports Gamma-point only for PBC.") + + is_uhf_like = isinstance( + mf, + ( + scf.uhf.UHF, + dft.uks.UKS, + pbc.scf.uhf.UHF, + pbc.dft.uks.UKS, + pbc.scf.kuhf.KUHF, + pbc.dft.kuks.KUKS, + ), + ) + + # MO-basis density is diagonal in canonical orbitals + if is_uhf_like and isinstance(mf.mo_occ, (tuple, list)) and len(mf.mo_occ) == 2: + occ_a, occ_b = mf.mo_occ + occ_a = np.asarray(occ_a) + occ_b = np.asarray(occ_b) + else: + occ = np.asarray(mf.mo_occ) + if is_uhf_like and occ.ndim == 2 and occ.shape[0] == 2: + occ_a, occ_b = occ[0], occ[1] + else: + occ_a = occ + occ_b = None + + if is_pbc and isinstance(occ_a, np.ndarray) and occ_a.ndim == 2: + if occ_a.shape[0] != 1: + raise NotImplementedError( + "PBC RDM write: only single-k Gamma supported for MO basis.") + occ_a = occ_a[0] + if occ_b is not None and occ_b.ndim == 2: + occ_b = occ_b[0] + + if occ_b is not None: + dm_a = np.diag(occ_a) + dm_b = np.diag(occ_b) + dm_a = _trexio_ensure_real( + np.asarray(dm_a), + what="Complex RDM encountered; real-only backend.", + ) + dm_b = _trexio_ensure_real( + np.asarray(dm_b), + what="Complex RDM encountered; real-only backend.", + ) + dm_a = np.ascontiguousarray(dm_a) + dm_b = np.ascontiguousarray(dm_b) + with trexio.File(filename, 'u', back_end=_mode(backend)) as tf: + nmo_up = dm_a.shape[0] + nmo_dn = dm_b.shape[0] + if nmo_dn != nmo_up: + raise ValueError("Alpha/beta MO sizes do not match for RDM write.") + nmo = nmo_up + nmo_dn + if not trexio.has_mo_num(tf): + trexio.write_mo_num(tf, nmo) + trexio.write_rdm_1e_up(tf, dm_a) + trexio.write_rdm_1e_dn(tf, dm_b) + dm_tot = np.zeros((nmo, nmo), dtype=dm_a.dtype) + dm_tot[:nmo_up, :nmo_up] = dm_a + dm_tot[nmo_up:, nmo_up:] = dm_b + trexio.write_rdm_1e(tf, dm_tot) + else: + dm = np.diag(occ_a) + dm = _trexio_ensure_real( + np.asarray(dm), + what="Complex RDM encountered; real-only backend.", + ) + dm = np.ascontiguousarray(dm) + with trexio.File(filename, 'u', back_end=_mode(backend)) as tf: + nmo = dm.shape[0] + if not trexio.has_mo_num(tf): + trexio.write_mo_num(tf, nmo) + trexio.write_rdm_1e(tf, dm) + +def write_2b_rdm(mf, filename, backend='h5', chunk_size=100000): + """Write a two-body reduced density matrix in MO basis to TREXIO. + + Parameters + ---------- + mf : SCF/KSCF object + Converged molecular or PBC mean-field object. Uses ``mo_occ`` to + build diagonal MO densities. PBC data must be Gamma-only. + filename : str + Path to the TREXIO file to create or update. + backend : {'h5', 'text'}, optional + TREXIO backend selector passed through to ``trexio.File``. + chunk_size : int, optional + Number of elements streamed per block when writing flattened ERIs. + + Behavior + -------- + - MO basis is enforced (TREXIO convention). + - RHF/RKS: writes spin-summed 2-RDM using ``dm = diag(mo_occ)`` and + ``G[pqrs] = dm[p,r]*dm[q,s] - 0.5*dm[p,s]*dm[q,r]``. + - UHF/UKS: writes spin-resolved blocks ``G_uu``, ``G_dd``, and ``G_ud`` + constructed from ``diag(occ_a)`` and ``diag(occ_b)``; alpha/beta sizes + must match. + - PBC: only single-k Gamma is supported; k-resolved occupations are + squeezed to 1D first. + - Data are streamed in chunks to avoid holding the entire flattened array + in memory at once; memory still scales as O(n^4). + + Raises + ------ + ValueError + When alpha/beta dimensions differ or input shapes are inconsistent. + NotImplementedError + For non-Gamma PBC calculations. + """ + + is_pbc = hasattr(mf, 'cell') + if is_pbc and not _trexio_is_gamma_single_k(mf): + raise NotImplementedError("RDM write supports Gamma-point only for PBC.") + + is_uhf_like = isinstance( + mf, + ( + scf.uhf.UHF, + dft.uks.UKS, + pbc.scf.uhf.UHF, + pbc.dft.uks.UKS, + pbc.scf.kuhf.KUHF, + pbc.dft.kuks.KUKS, + ), + ) + + # Spin-summed occupations or spin-separated for UHF/UKS + if is_uhf_like and isinstance(mf.mo_occ, (tuple, list)) and len(mf.mo_occ) == 2: + occ_a, occ_b = mf.mo_occ + occ_a = np.asarray(occ_a) + occ_b = np.asarray(occ_b) + else: + occ = np.asarray(mf.mo_occ) + if is_uhf_like and occ.ndim == 2 and occ.shape[0] == 2: + occ_a, occ_b = occ[0], occ[1] + else: + occ_a = occ + occ_b = None + + if is_pbc and isinstance(occ_a, np.ndarray) and occ_a.ndim == 2: + if occ_a.shape[0] != 1: + raise NotImplementedError( + "PBC RDM write: only single-k Gamma supported for MO basis.") + occ_a = occ_a[0] + if occ_b is not None and occ_b.ndim == 2: + occ_b = occ_b[0] + + if occ_b is not None: + dm_a = np.diag(occ_a) + dm_b = np.diag(occ_b) + dm_a = _trexio_ensure_real( + np.asarray(dm_a), + what="Complex RDM encountered; real-only backend.", + ) + dm_b = _trexio_ensure_real( + np.asarray(dm_b), + what="Complex RDM encountered; real-only backend.", + ) + dm_a = np.ascontiguousarray(dm_a) + dm_b = np.ascontiguousarray(dm_b) + + g2_uu = ( + np.einsum('pr,qs->pqrs', dm_a, dm_a) + - np.einsum('ps,qr->pqrs', dm_a, dm_a) + ) + g2_dd = ( + np.einsum('pr,qs->pqrs', dm_b, dm_b) + - np.einsum('ps,qr->pqrs', dm_b, dm_b) + ) + g2_ud = np.einsum('pr,qs->pqrs', dm_a, dm_b) + + g2_uu = np.ascontiguousarray(g2_uu) + g2_dd = np.ascontiguousarray(g2_dd) + g2_ud = np.ascontiguousarray(g2_ud) + + nmo = dm_a.shape[0] + if dm_b.shape[0] != nmo: + raise ValueError("Alpha/beta MO sizes do not match for RDM write.") + idx = lib.cartesian_prod([np.arange(nmo, dtype=np.int32)] * 4) + flat_uu = g2_uu.reshape(-1) + flat_dd = g2_dd.reshape(-1) + flat_ud = g2_ud.reshape(-1) + + with trexio.File(filename, 'u', back_end=_mode(backend)) as tf: + if not trexio.has_mo_num(tf): + trexio.write_mo_num(tf, nmo) + + total = idx.shape[0] + offset = 0 + while offset < total: + end = min(offset + chunk_size, total) + trexio.write_rdm_2e_upup( + tf, + offset, + end - offset, + idx[offset:end], + flat_uu[offset:end], + ) + trexio.write_rdm_2e_dndn( + tf, + offset, + end - offset, + idx[offset:end], + flat_dd[offset:end], + ) + trexio.write_rdm_2e_updn( + tf, + offset, + end - offset, + idx[offset:end], + flat_ud[offset:end], + ) + offset = end + else: + dm = np.diag(occ_a) + dm = _trexio_ensure_real( + np.asarray(dm), + what="Complex RDM encountered; real-only backend.", + ) + dm = np.ascontiguousarray(dm) + + # Build dense spin-summed 2-RDM in MO basis + g2 = ( + np.einsum('pr,qs->pqrs', dm, dm) + - 0.5 * np.einsum('ps,qr->pqrs', dm, dm) + ) + g2 = np.ascontiguousarray(g2) + + nmo = dm.shape[0] + idx = lib.cartesian_prod([np.arange(nmo, dtype=np.int32)] * 4) + flat_g2 = g2.reshape(-1) + + with trexio.File(filename, 'u', back_end=_mode(backend)) as tf: + if not trexio.has_mo_num(tf): + trexio.write_mo_num(tf, nmo) + + total = idx.shape[0] + offset = 0 + while offset < total: + end = min(offset + chunk_size, total) + trexio.write_rdm_2e( + tf, + offset, + end - offset, + idx[offset:end], + flat_g2[offset:end], + ) + offset = end + def _order_ao_index(mol): if mol.cart: return np.arange(mol.nao)