diff --git a/examples/mcscf/90_gasci_spin_supergroup.py b/examples/mcscf/90_gasci_spin_supergroup.py new file mode 100644 index 000000000..959609478 --- /dev/null +++ b/examples/mcscf/90_gasci_spin_supergroup.py @@ -0,0 +1,43 @@ +#!/usr/bin/env python +# +# Author: Yi Deng +# + +"""Fixed-orbital GASCI with an explicit spin-supergroup set D.""" + +import numpy + +from pyscf import gto +from pyscf import scf +from pyscf.mcscf import gasci + + +mol = gto.M( + atom="N 0 0 0; N 0 0 1.10", + basis="sto-3g", + spin=0, + verbose=0, +) +mf = scf.RHF(mol).run() + +# Each row is (alpha occupations, beta occupations). Rows are normalized to +# lexicographic order and duplicates are removed before entering the C kernel. +gas_restr = numpy.asarray([ + [0, 2, 2, 0], + [1, 1, 1, 1], + [2, 0, 0, 2], +], dtype=numpy.int32) + +mc = gasci.GASCI( + mf, 4, (2, 2), ncore=5, + gas_orbs=(2, 2), + gas_restr=gas_restr, + gas_restr_type="spin-supergroup", +) +mc.verbose = 4 +e_tot, e_gas, ci, _, _ = mc.kernel() + +print("GASCI total energy = %.12f" % e_tot) +print("GASCI active-space energy = %.12f" % e_gas) +print("GAS determinant count = %d" % ci.size) +print("GAS space information = %s" % mc.gas_space_info()) diff --git a/examples/mcscf/91_gas_restriction_types.py b/examples/mcscf/91_gas_restriction_types.py new file mode 100644 index 000000000..d4dfbeff0 --- /dev/null +++ b/examples/mcscf/91_gas_restriction_types.py @@ -0,0 +1,38 @@ +#!/usr/bin/env python +# +# Author: Yi Deng +# + +"""Normalize the four supported GAS restriction input formats.""" + +import numpy + +from pyscf.mcscf import addons_gas + + +gas_orbs = (2, 2) +nelec = (2, 2) + +spin_supergroups = numpy.asarray([ + [2, 0, 0, 2], + [1, 1, 1, 1], + [0, 2, 2, 0], + [1, 1, 1, 1], # Duplicate rows are removed. +]) +supergroups = numpy.asarray([[2, 2]]) +cumulative_occ = numpy.asarray([[2, 2], [4, 4]]) + +for restriction_type, restriction in ( + ("spin-supergroup", spin_supergroups), + ("supergroup", supergroups), + ("cumulative-occ", cumulative_occ)): + blocks = addons_gas.normalize_gas_restr( + gas_orbs, nelec, restriction, restriction_type) + print("%-17s ->\n%s" % (restriction_type, blocks)) + +ras_blocks = addons_gas.normalize_gas_restr( + (1, 2, 1), nelec, + {"max_holes": 1, "max_particles": 1}, + "ras", +) +print("ras ->\n%s" % ras_blocks) diff --git a/examples/mcscf/92_gasci_multiroot_rdm.py b/examples/mcscf/92_gasci_multiroot_rdm.py new file mode 100644 index 000000000..2b3503128 --- /dev/null +++ b/examples/mcscf/92_gasci_multiroot_rdm.py @@ -0,0 +1,48 @@ +#!/usr/bin/env python +# +# Author: Yi Deng +# + +"""Multiroot GASCI with cumulative restrictions and state-resolved RDMs.""" + +import numpy + +from pyscf import gto +from pyscf import scf +from pyscf.mcscf import gasci + + +mol = gto.M( + atom="N 0 0 0; N 0 0 1.10", basis="sto-3g", spin=0, verbose=0) +mf = scf.RHF(mol).run() + +# Each row gives the lower and upper cumulative electron occupation after +# adding the corresponding GAS subspace. +mc = gasci.GASCI( + mf, 4, (2, 2), ncore=5, + gas_orbs=(2, 2), + gas_restr=[[1, 3], [4, 4]], + gas_restr_type="cumulative-occ", +) +mc.fcisolver.nroots = 3 +mc.verbose = 4 +energies, _, _, _, _ = mc.kernel() + +for state, energy in enumerate(numpy.asarray(energies).reshape(-1)): + dm1, dm2 = mc.make_gasdm12(state=state) + ss, multiplicity = mc.spin_square(state=state) + pair_trace = numpy.einsum("pprr", dm2).real + print("state %d E = %.12f Tr(D1) = %.8f S^2 = %.8f mult = %.4f" % + (state, energy, numpy.trace(dm1), ss, multiplicity)) + print("state %d Tr_pair(D2) = %.8f" % (state, pair_trace)) + +tdm1, tdm2 = mc.trans_gasdm12(bra_state=0, ket_state=2) +print("transition 0 -> 2: ||D1|| = %.8f ||D2|| = %.8f" % + (numpy.linalg.norm(tdm1), numpy.linalg.norm(tdm2))) + +_, natural_occ = mc.get_gas_natorb(state=0) +_, pseudo_occ = mc.get_gas_pseudo_natorb(state=0) +print("natural occupations = %s" % natural_occ) +for igas, occupations in enumerate(pseudo_occ): + print("GAS%d pseudo-natural occupations = %s" % + (igas + 1, occupations)) diff --git a/examples/mcscf/gasci_openmolcas/01_o2_multiroot.py b/examples/mcscf/gasci_openmolcas/01_o2_multiroot.py new file mode 100644 index 000000000..168fe4395 --- /dev/null +++ b/examples/mcscf/gasci_openmolcas/01_o2_multiroot.py @@ -0,0 +1,252 @@ +#!/usr/bin/env python +# +# Author: Yi Deng +# + +"""Fixed-orbital multiroot GASCI example using OpenMolcas orbitals.""" + +from pathlib import Path + +import numpy + +from pyscf import gto +from pyscf import scf +from pyscf.mcscf import gasci + + +# OpenMolcas reference input: +# +# &GATEWAY +# Title = O2 GASCI 2S+1=3 +# Coord = o2.xyz +# Basis = def2-TZVP +# Group = C1 +# noCD +# +# &SEWARD +# +# &RASSCF +# Title = O2 GASCI(12e,8o), GAS=2+4+2, triplet +# +# CIONLY +# FILEORB = o2_guessorb.INPORB +# +# Spin = 3 +# Symmetry = 1 +# +# Nactel = 12 0 0 +# Inactive = 2 +# +# GASSCF +# 3 +# 2 +# 2 4 +# 4 +# 8 10 +# 2 +# 12 12 +# +# CIRoot = 20 20 1 +# CIMX = 200 +# +# Notes: +# - OpenMolcas Spin and PySCF spin use multiplicity 2S+1 and 2S, +# respectively. +# - CIONLY keeps the molecular orbitals fixed. +# - noCD avoids the Cholesky approximation to the ERIs. +# +# The exact PySCF molecule and fixed MO coefficients were imported from +# the converged OpenMolcas HDF5 file as follows: +# +# import numpy +# from mrh.my_pyscf.tools.molcas2pyscf import get_mo_from_h5 +# from mrh.my_pyscf.tools.molcas2pyscf import get_mol_from_h5 +# +# h5_file = "o2_m3.rasscf.h5" +# mol = get_mol_from_h5( +# h5_file, +# charge=0, +# spin=2, +# ) +# mo_coeff = get_mo_from_h5(mol, h5_file) +# numpy.savez_compressed( +# "data/o2_mo.npz", +# mol=mol.dumps(), +# mo_coeff=mo_coeff[:, :10], +# ) +# +# OpenMolcas output: +# +# :: RASSCF root number 1 Total energy: -149.76505363 +# :: RASSCF root number 2 Total energy: -149.53703734 +# :: RASSCF root number 3 Total energy: -149.53703734 +# :: RASSCF root number 4 Total energy: -149.53215375 +# :: RASSCF root number 5 Total energy: -149.38576564 +# :: RASSCF root number 6 Total energy: -149.28581775 +# :: RASSCF root number 7 Total energy: -149.28581722 +# :: RASSCF root number 8 Total energy: -149.15394925 +# :: RASSCF root number 9 Total energy: -149.11022925 +# :: RASSCF root number 10 Total energy: -149.11022908 +# :: RASSCF root number 11 Total energy: -149.07315914 +# :: RASSCF root number 12 Total energy: -149.07315679 +# :: RASSCF root number 13 Total energy: -149.02767903 +# :: RASSCF root number 14 Total energy: -149.02767903 +# :: RASSCF root number 15 Total energy: -149.01720780 +# :: RASSCF root number 16 Total energy: -149.01720383 +# :: RASSCF root number 17 Total energy: -149.00345234 +# :: RASSCF root number 18 Total energy: -149.00345234 +# :: RASSCF root number 19 Total energy: -148.99747740 +# :: RASSCF root number 20 Total energy: -148.99747411 + + +openmolcas_energies = numpy.asarray([ + -149.76505363, + -149.53703734, + -149.53703734, + -149.53215375, + -149.38576564, + -149.28581775, + -149.28581722, + -149.15394925, + -149.11022925, + -149.11022908, + -149.07315914, + -149.07315679, + -149.02767903, + -149.02767903, + -149.01720780, + -149.01720383, + -149.00345234, + -149.00345234, + -148.99747740, + -148.99747411, +]) + +here = Path(__file__).resolve().parent +data_file = here / "data" / "o2_mo.npz" + +# Load the exact PySCF molecule and fixed OpenMolcas MO coefficients. +with numpy.load(data_file, allow_pickle=False) as data: + mol = gto.loads(data["mol"].item()) + mo_coeff = data["mo_coeff"] + +mol.max_memory = 8000 # MB + +# One-cycle ROHF for mf generation. +mf = scf.ROHF(mol) +mf.max_cycle = 1 +mf.kernel() +mf.mo_coeff = mo_coeff + +# Active space +ncore = 2 +ncas = 8 +nelecas = (7, 5) + +gas_orbs = (2, 4, 2) + +# OpenMolcas cumulative GAS constraints. +gas_restr = ( + (2, 4), + (8, 10), + (12, 12), +) +gas_restr_type = "cumulative-occ" + +# GASCI +mc = gasci.GASCI( + mf, + ncas, + nelecas, + ncore=ncore, + gas_orbs=gas_orbs, + gas_restr=gas_restr, + gas_restr_type=gas_restr_type, +) + +nroots = 30 + +mc.verbose = 4 +mc.canonicalization = False +mc.fcisolver.spin = 2 +mc.fcisolver.nroots = nroots +mc.fcisolver.max_cycle = 300 +mc.fcisolver.conv_tol = 1e-10 + +mc.kernel() + +# The PySCF M_S = 1 space contains states with different total spin. +# Select the triplet roots using = 2. +pyscf_triplet_energies = [] + +print() +print(f"PySCF triplet roots among the lowest {nroots} states") +print("state E(PySCF) / Eh ") + +for state, energy in enumerate(numpy.asarray(mc.e_tot).reshape(-1)): + ss, _ = mc.spin_square(state=state) + if abs(ss - 2.0) < 1e-6: + pyscf_triplet_energies.append(energy) + print(f"{state + 1:5d} {energy:22.12f} {ss:12.8f}") + +ncompare = min( + len(pyscf_triplet_energies), + len(openmolcas_energies), +) +if ncompare == 0: + raise RuntimeError( + f"No triplet root was found among the lowest {nroots} states" + ) + +pyscf_triplet_energies = numpy.asarray( + pyscf_triplet_energies[:ncompare] +) +reference_energies = openmolcas_energies[:ncompare] +differences = pyscf_triplet_energies - reference_energies + +print() +print("OpenMolcas/PySCF GASCI comparison") +print("root E(OpenMolcas) / Eh E(PySCF) / Eh diff / Eh") + +for root, (e_ref, e_pyscf, diff) in enumerate( + zip(reference_energies, pyscf_triplet_energies, differences), + start=1, +): + print(f"{root:4d} {e_ref:22.12f} {e_pyscf:22.12f} {diff: .3e}") + +max_diff = numpy.max(numpy.abs(differences)) +print() +print(f"Compared roots : {ncompare}") +print(f"max |difference| : {max_diff:.3e} Eh") + +# Example output (final comparison section only): +# +# Numerical values in the last few digits may depend on the platform. +# +# GASCI converged +# +# OpenMolcas/PySCF GASCI comparison +# root E(OpenMolcas) / Eh E(PySCF) / Eh diff / Eh +# 1 -149.765053630000 -149.765053626447 3.553e-09 +# 2 -149.537037340000 -149.537037344145 -4.145e-09 +# 3 -149.537037340000 -149.537037343830 -3.830e-09 +# 4 -149.532153750000 -149.532153753145 -3.145e-09 +# 5 -149.385765640000 -149.385765640052 -5.247e-11 +# 6 -149.285817750000 -149.285817749263 7.374e-10 +# 7 -149.285817220000 -149.285817220731 -7.307e-10 +# 8 -149.153949250000 -149.153949254242 -4.242e-09 +# 9 -149.110229250000 -149.110229248524 1.476e-09 +# 10 -149.110229080000 -149.110229082511 -2.511e-09 +# 11 -149.073159140000 -149.073159141160 -1.160e-09 +# 12 -149.073156790000 -149.073156787594 2.406e-09 +# 13 -149.027679030000 -149.027679032858 -2.858e-09 +# 14 -149.027679030000 -149.027679032856 -2.856e-09 +# 15 -149.017207800000 -149.017207802629 -2.629e-09 +# 16 -149.017203830000 -149.017203833468 -3.468e-09 +# 17 -149.003452340000 -149.003452343080 -3.080e-09 +# 18 -149.003452340000 -149.003452342985 -2.985e-09 +# 19 -148.997477400000 -148.997477403623 -3.623e-09 +# 20 -148.997474110000 -148.997474109395 6.051e-10 +# +# Compared roots : 20 +# max |difference| : 4.242e-09 Eh diff --git a/examples/mcscf/gasci_openmolcas/02_alfe2o4_plus_single_root.py b/examples/mcscf/gasci_openmolcas/02_alfe2o4_plus_single_root.py new file mode 100644 index 000000000..10d271cae --- /dev/null +++ b/examples/mcscf/gasci_openmolcas/02_alfe2o4_plus_single_root.py @@ -0,0 +1,174 @@ +#!/usr/bin/env python +# +# Author: Yi Deng +# + +"""Fixed-orbital single-root GASCI example for undecet AlFe2O4+.""" + +from pathlib import Path + +import numpy + +from pyscf import gto +from pyscf import scf +from pyscf.mcscf import gasci + + +# OpenMolcas reference input: +# +# &GATEWAY +# Title = AlFe2O4+ GASCI 2S+1=11 +# Coord = alfe2o4.xyz +# Basis = cc-pVDZ +# Group = C1 +# noCD +# +# &SEWARD +# +# &RASSCF +# Title = AlFe2O4+ GASCI(34e,34o), O2p/Fe3d/O3p, undecet +# +# CIONLY +# FILEORB = alfe2o4_guessorb.INPORB +# +# Charge = 1 +# Spin = 11 +# Symmetry = 1 +# +# Nactel = 34 0 0 +# Inactive = 31 +# +# ! GAS1: MO32-43, O 2p, 12 orbitals +# ! GAS2: MO44-53, Fe 3d, 10 orbitals +# ! GAS3: MO54-65, O 3p, 12 orbitals +# ! +# ! Constraints: +# ! 23 <= N(O 2p) <= 24 +# ! 33 <= N(O 2p + Fe 3d) <= 34 +# ! 34 = N(O 2p + Fe 3d + O 3p) = 34 +# +# GASSCF +# 3 +# 12 +# 23 24 +# 10 +# 33 34 +# 12 +# 34 34 +# +# CIRoot = 1 1 1 +# CIMX = 200 +# ITERations = 300 300 +# +# Notes: +# - OpenMolcas Spin = 11 means multiplicity 2S+1 = 11. +# - PySCF spin = 2S = 10 for the same undecet calculation. +# - CIONLY keeps the molecular orbitals fixed. +# - noCD avoids the Cholesky approximation to the ERIs. +# +# The exact PySCF molecule and fixed MO coefficients were imported from +# the converged OpenMolcas HDF5 file as follows: +# +# import numpy +# from mrh.my_pyscf.tools.molcas2pyscf import get_mo_from_h5 +# from mrh.my_pyscf.tools.molcas2pyscf import get_mol_from_h5 +# +# h5_file = "fe2_d2d_m11.rasscf.h5" +# mol = get_mol_from_h5( +# h5_file, +# charge=1, +# spin=10, +# ) +# mo_coeff = get_mo_from_h5(mol, h5_file) +# numpy.savez_compressed( +# "data/alfe2o4_plus_mo.npz", +# mol=mol.dumps(), +# mo_coeff=mo_coeff[:, :65], +# ) +# +# OpenMolcas output: +# +# :: RASSCF root number 1 Total energy: -3064.32294249 + + +openmolcas_energy = -3064.32294249 + +here = Path(__file__).resolve().parent +data_file = here / "data" / "alfe2o4_plus_mo.npz" + +# Load the exact PySCF molecule and fixed OpenMolcas MO coefficients. +with numpy.load(data_file, allow_pickle=False) as data: + mol = gto.loads(data["mol"].item()) + mo_coeff = data["mo_coeff"] + +mol.max_memory = 8000 # MB + +# One-cycle ROHF for mf generation. +mf = scf.ROHF(mol) +mf.max_cycle = 1 +mf.kernel() +mf.mo_coeff = mo_coeff + +# Active space: GASCI(34e,34o) +ncore = 31 +ncas = 34 +nelecas = (22, 12) + +gas_orbs = (12, 10, 12) + +# OpenMolcas cumulative GAS constraints. +gas_restr = ( + (23, 24), + (33, 34), + (34, 34), +) +gas_restr_type = "cumulative-occ" + +# GASCI +mc = gasci.GASCI( + mf, + ncas, + nelecas, + ncore=ncore, + gas_orbs=gas_orbs, + gas_restr=gas_restr, + gas_restr_type=gas_restr_type, +) + +mc.verbose = 4 +mc.canonicalization = False +mc.fcisolver.spin = 10 +mc.fcisolver.nroots = 1 +mc.fcisolver.max_cycle = 300 +mc.fcisolver.max_space = 30 +mc.fcisolver.conv_tol = 1e-10 +mc.fcisolver.conv_tol_residual = 1e-6 + +mc.kernel() + +pyscf_energy = float(numpy.asarray(mc.e_tot).reshape(-1)[0]) +ss, _ = mc.spin_square(state=0) +difference = pyscf_energy - openmolcas_energy + +print() +print("OpenMolcas/PySCF GASCI comparison") +print("E(OpenMolcas) / Eh E(PySCF) / Eh diff / Eh ") +print( + f"{openmolcas_energy:20.12f} " + f"{pyscf_energy:20.12f} " + f"{difference: .3e} " + f"{ss:12.8f}" +) + +# Example output (final comparison section only): +# +# Numerical values in the last few digits may depend on the platform. +# +# GASCI converged +# GASCI E = -3064.32294249369 +# E(CI) = -179.331711827401 +# S^2 = 30.0000000 +# +# OpenMolcas/PySCF GASCI comparison +# E(OpenMolcas) / Eh E(PySCF) / Eh diff / Eh +# -3064.322942490000 -3064.322942493695 -3.695e-09 30.00000000 diff --git a/examples/mcscf/gasci_openmolcas/03_al2feo4_plus_rasci.py b/examples/mcscf/gasci_openmolcas/03_al2feo4_plus_rasci.py new file mode 100644 index 000000000..7ff1b9bd0 --- /dev/null +++ b/examples/mcscf/gasci_openmolcas/03_al2feo4_plus_rasci.py @@ -0,0 +1,168 @@ +#!/usr/bin/env python +# +# Author: Yi Deng +# + +"""Fixed-orbital single-root RASCI example for sextet Al2FeO4+.""" + +from pathlib import Path + +import numpy + +from pyscf import gto +from pyscf import scf +from pyscf.mcscf import gasci + + +# OpenMolcas reference input: +# +# &GATEWAY +# Title = Al2FeO4+ RASCI 2S+1=6 +# Coord = al2feo4.xyz +# Basis = cc-pVTZ +# Group = C1 +# noCD +# +# &SEWARD +# +# &RASSCF +# Title = Al2FeO4+ RASCI(29e,22o), O2p/Fe3d/Fe4d, sextet +# +# CIONLY +# FILEORB = al2feo4_guessorb.INPORB +# +# Charge = 1 +# Spin = 6 +# Symmetry = 1 +# +# Nactel = 29 2 2 +# Inactive = 27 +# +# ! RAS1: MO28-39, O 2p, 12 orbitals, maximum 2 holes +# ! RAS2: MO40-44, Fe 3d, 5 orbitals +# ! RAS3: MO45-49, Fe 4d, 5 orbitals, maximum 2 electrons +# +# RAS1 +# 12 +# +# RAS2 +# 5 +# +# RAS3 +# 5 +# +# CIRoot = 1 1 1 +# CIMX = 200 +# ITERations = 300 300 +# +# Notes: +# - OpenMolcas Spin = 6 means multiplicity 2S+1 = 6. +# - PySCF spin = 2S = 5 for the same sextet calculation. +# - CIONLY keeps the molecular orbitals fixed. +# - noCD avoids the Cholesky approximation to the ERIs. +# +# The exact PySCF molecule and fixed MO coefficients were imported from +# the converged OpenMolcas HDF5 file as follows: +# +# import numpy +# from mrh.my_pyscf.tools.molcas2pyscf import get_mo_from_h5 +# from mrh.my_pyscf.tools.molcas2pyscf import get_mol_from_h5 +# +# h5_file = "al2feo4_rasci.rasscf.h5" +# mol = get_mol_from_h5( +# h5_file, +# charge=1, +# spin=5, +# ) +# mo_coeff = get_mo_from_h5(mol, h5_file) +# numpy.savez_compressed( +# "data/al2feo4_plus_rasci_mo.npz", +# mol=mol.dumps(), +# mo_coeff=mo_coeff[:, :49], +# ) +# +# OpenMolcas output: +# +# :: RASSCF root number 1 Total energy: -2045.01359885 + + +openmolcas_energy = -2045.01359885 + +here = Path(__file__).resolve().parent +data_file = here / "data" / "al2feo4_plus_rasci_mo.npz" + +# Load the exact PySCF molecule and fixed OpenMolcas MO coefficients. +with numpy.load(data_file, allow_pickle=False) as data: + mol = gto.loads(data["mol"].item()) + mo_coeff = data["mo_coeff"] + +mol.max_memory = 8000 # MB + +# One-cycle ROHF for mf generation. +mf = scf.ROHF(mol) +mf.max_cycle = 1 +mf.kernel() +mf.mo_coeff = mo_coeff + +# Active space: RASCI(29e,22o) +ncore = 27 +ncas = 22 +nelecas = (17, 12) + +gas_orbs = (12, 5, 5) + +# OpenMolcas RAS constraints. +gas_restr = { + "max_holes": 2, + "max_particles": 2, +} +gas_restr_type = "ras" + +# GASCI +mc = gasci.GASCI( + mf, + ncas, + nelecas, + ncore=ncore, + gas_orbs=gas_orbs, + gas_restr=gas_restr, + gas_restr_type=gas_restr_type, +) + +mc.verbose = 4 +mc.canonicalization = False +mc.fcisolver.spin = 5 +mc.fcisolver.nroots = 1 +mc.fcisolver.max_cycle = 300 +mc.fcisolver.max_space = 30 +mc.fcisolver.conv_tol = 1e-10 +mc.fcisolver.conv_tol_residual = 1e-6 + +mc.kernel() + +pyscf_energy = float(numpy.asarray(mc.e_tot).reshape(-1)[0]) +ss, _ = mc.spin_square(state=0) +difference = pyscf_energy - openmolcas_energy + +print() +print("OpenMolcas/PySCF RASCI comparison") +print("E(OpenMolcas) / Eh E(PySCF) / Eh diff / Eh ") +print( + f"{openmolcas_energy:20.12f} " + f"{pyscf_energy:20.12f} " + f"{difference: .3e} " + f"{ss:12.8f}" +) + +# Example output (final comparison section only): +# +# Numerical values in the last few digits may depend on the platform. +# +# GASCI converged +# GASCI E = -2045.01359884909 +# E(CI) = -363.115477794565 +# S^2 = 8.7500000 +# +# OpenMolcas/PySCF RASCI comparison +# E(OpenMolcas) / Eh E(PySCF) / Eh diff / Eh +# -2045.013598850000 -2045.013598849092 9.081e-10 8.75000000 diff --git a/examples/mcscf/gasci_openmolcas/data/al2feo4.xyz b/examples/mcscf/gasci_openmolcas/data/al2feo4.xyz new file mode 100644 index 000000000..f7a838fe8 --- /dev/null +++ b/examples/mcscf/gasci_openmolcas/data/al2feo4.xyz @@ -0,0 +1,9 @@ +7 +al2feo4 geo +Fe -1.3429818 -0.8116774 -0.0042259 +Al 1.2206276 -0.6483692 0.1680069 +Al -0.0105184 1.3891192 0.0899530 +O 0.1154860 -1.7981817 -0.4918142 +O 1.5591198 0.9238418 -0.4704792 +O -1.5205804 0.9083794 -0.5938707 +O -0.0211529 0.0368879 1.3024301 diff --git a/examples/mcscf/gasci_openmolcas/data/al2feo4_plus_rasci_mo.npz b/examples/mcscf/gasci_openmolcas/data/al2feo4_plus_rasci_mo.npz new file mode 100644 index 000000000..1cb9c5334 Binary files /dev/null and b/examples/mcscf/gasci_openmolcas/data/al2feo4_plus_rasci_mo.npz differ diff --git a/examples/mcscf/gasci_openmolcas/data/alfe2o4.xyz b/examples/mcscf/gasci_openmolcas/data/alfe2o4.xyz new file mode 100644 index 000000000..72668e31f --- /dev/null +++ b/examples/mcscf/gasci_openmolcas/data/alfe2o4.xyz @@ -0,0 +1,9 @@ +7 +alfe2o4 geo +Fe 2.5044793 0.0014445 0.0005249 +Fe -2.5043955 -0.0013998 0.0002220 +Al -0.0000312 -0.0000463 -0.0003452 +O -1.2849760 -0.9231676 0.8771297 +O 1.2862771 -0.8760934 -0.9223767 +O 1.2848112 0.8779657 0.9223203 +O -1.2861647 0.9212968 -0.8774752 diff --git a/examples/mcscf/gasci_openmolcas/data/alfe2o4_plus_mo.npz b/examples/mcscf/gasci_openmolcas/data/alfe2o4_plus_mo.npz new file mode 100644 index 000000000..7a5b90133 Binary files /dev/null and b/examples/mcscf/gasci_openmolcas/data/alfe2o4_plus_mo.npz differ diff --git a/examples/mcscf/gasci_openmolcas/data/o2.xyz b/examples/mcscf/gasci_openmolcas/data/o2.xyz new file mode 100644 index 000000000..91e38d589 --- /dev/null +++ b/examples/mcscf/gasci_openmolcas/data/o2.xyz @@ -0,0 +1,4 @@ +2 +o2 geo +O 0 0 0 +O 0 0 1.2075 diff --git a/examples/mcscf/gasci_openmolcas/data/o2_mo.npz b/examples/mcscf/gasci_openmolcas/data/o2_mo.npz new file mode 100644 index 000000000..acc0793eb Binary files /dev/null and b/examples/mcscf/gasci_openmolcas/data/o2_mo.npz differ diff --git a/pyscf/lib/CMakeLists.txt b/pyscf/lib/CMakeLists.txt index 2268c7cb4..79c9966d2 100644 --- a/pyscf/lib/CMakeLists.txt +++ b/pyscf/lib/CMakeLists.txt @@ -1,6 +1,19 @@ cmake_minimum_required (VERSION 3.5) project (pyscf_forge) +# Build only the GAS FCI kernel library. This mode is intended for +# reproducible local kernel builds and deliberately bypasses unrelated +# dependencies such as BLAS, FFTW, and the other PySCF-forge extensions. +option( + BUILD_GAS_FCI_ONLY + "Configure only the GAS FCI shared library" + OFF +) + +option(ENABLE_OPENMP "Compiling C extensions with openmp" ON) + +if(NOT BUILD_GAS_FCI_ONLY) + set(CMAKE_VERBOSE_MAKEFILE OFF) if (NOT CMAKE_BUILD_TYPE) set(CMAKE_BUILD_TYPE RELWITHDEBINFO) @@ -71,7 +84,6 @@ find_package(BLAS REQUIRED) #find_package (LAPACK REQUIRED) # OpenMP -option(ENABLE_OPENMP "Compiling C extensions with openmp" ON) set(OPENMP_C_PROPERTIES "") if(ENABLE_OPENMP) find_package(OpenMP) @@ -200,3 +212,23 @@ set_target_properties (clib_csf PROPERTIES add_subdirectory(gbci) add_subdirectory(pwscf) + +else() + + if(NOT ENABLE_OPENMP) + message(FATAL_ERROR + "BUILD_GAS_FCI_ONLY=ON requires ENABLE_OPENMP=ON" + ) + endif() + + find_package(OpenMP REQUIRED COMPONENTS C) + + if(NOT TARGET OpenMP::OpenMP_C) + message(FATAL_ERROR + "OpenMP was reported as found, but OpenMP::OpenMP_C is unavailable" + ) + endif() + +endif() + +add_subdirectory(mcscf) diff --git a/pyscf/lib/mcscf/CMakeLists.txt b/pyscf/lib/mcscf/CMakeLists.txt new file mode 100644 index 000000000..5302442e4 --- /dev/null +++ b/pyscf/lib/mcscf/CMakeLists.txt @@ -0,0 +1,52 @@ +# Copyright 2026 The PySCF Developers. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +set(GAS_FCI_SOURCE_FILES + fci_string_gas.c + fci_contract_gas.c + fci_hdiag_gas.c + fci_rdm_gas.c +) + +add_library(clib_fci_gas SHARED ${GAS_FCI_SOURCE_FILES}) + +target_include_directories(clib_fci_gas PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR} +) + +target_compile_features(clib_fci_gas PRIVATE c_std_11) + +if(CMAKE_C_COMPILER_ID MATCHES "GNU|Clang") + target_compile_options(clib_fci_gas PRIVATE + -Wall + -Wextra + -Wpedantic + ) +endif() + +target_link_libraries(clib_fci_gas PRIVATE m) + +if(ENABLE_OPENMP AND TARGET OpenMP::OpenMP_C) + target_link_libraries(clib_fci_gas PRIVATE OpenMP::OpenMP_C) +endif() + +set_target_properties(clib_fci_gas PROPERTIES + LINKER_LANGUAGE C + C_STANDARD 11 + C_STANDARD_REQUIRED ON + C_EXTENSIONS OFF + CLEAN_DIRECT_OUTPUT 1 + LIBRARY_OUTPUT_DIRECTORY ${PROJECT_SOURCE_DIR} + OUTPUT_NAME "fci_gas" +) diff --git a/pyscf/lib/mcscf/fci_contract_gas.c b/pyscf/lib/mcscf/fci_contract_gas.c new file mode 100644 index 000000000..29a095ece --- /dev/null +++ b/pyscf/lib/mcscf/fci_contract_gas.c @@ -0,0 +1,2487 @@ +/* Copyright 2026 The PySCF Developers. All Rights Reserved. + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. + + * + * Author: Yi Deng + */ + +#include "fci_gas.h" + +#include +#include + +#if defined(__AVX2__) +#include +#endif + +#ifdef _OPENMP +#include +#endif + +/* ========================================================================== */ +/* 1. Configuration and low-level utilities */ +/* ========================================================================== */ + +enum { + BB_ALPHA_BLOCK = 64, + BB_TRANSPOSE_MIN_DIM = 4, + BB_TRANSPOSE_TILE = 32, + ABBA_WRITEBACK_TILE = 32, + OMP_TASK_FACTOR = 32, + OMP_MAX_ALPHA_SPLIT = 64, + OMP_MIN_ALPHA_TILE = 4 +}; + +static const size_t BB_TRANSPOSE_MAX_ELEMS = 67108864u; +static const size_t ABBA_WORKSPACE_MAX_ELEMS = 67108864u; +static size_t gas_abba_t1_target_bytes = 64u * 1024u * 1024u; + +void fci_contract_gas_set_abba_t1_bytes(uint64_t bytes) +{ + gas_abba_t1_target_bytes = (size_t)bytes; +} + +static uint32_t abba_beta_tile_size(uint32_t nb0, uint32_t nop, uint32_t na1) +{ + if (nb0 == 0u || nop == 0u || na1 == 0u || + gas_abba_t1_target_bytes == 0u) { + return nb0; + } + + size_t denom = (size_t)nop * (size_t)na1 * sizeof(double); + size_t tile = denom == 0u ? nb0 : gas_abba_t1_target_bytes / denom; + if (tile < 1u) { + tile = 1u; + } + if (tile > nb0) { + tile = nb0; + } + return (uint32_t)tile; +} + +static inline double abba_gvalue(const double *restrict g, + uint32_t nnorb, uint16_t opb, uint16_t opa) +{ + return g[(uint64_t)opb * nnorb + opa]; +} + +/* Consecutive beta link tables b0->b1 and b1->b2. */ +typedef struct { + gas_tid_t first; + gas_tid_t second; +} gas_tid_pair_t; + +/* + * Alpha link landing in the current destination tile [a0,a1). + * rel_addr is relative to a0; op_index selects the active alpha operator. + */ +typedef struct { + uint32_t src_row; + uint32_t rel_addr; + uint16_t op_index; + int8_t sign; + uint8_t pad; +} gas_alpha_hit_t; + +typedef struct { + double *arena; + size_t arena_cap; + gas_tid_pair_t *pair; + size_t pair_cap; + int32_t *op_map; + size_t op_map_cap; + uint16_t *op; + size_t op_cap; + + size_t arena_zeroed_upto; + + /* A hit is an alpha link whose destination address lies inside the + * current alpha tile [a0,a1). + * Building the list once per ABBA batch lets each beta tile reuse it + * instead of rescanning mostly-missing alpha links. */ + gas_alpha_hit_t *alpha_hit; + size_t alpha_hit_cap; + uint32_t alpha_hit_n; +} gas_contract_ws_t; + +typedef enum { + GAS_BATCH_DONE = 0, + GAS_BATCH_FALLBACK = 1 +} gas_batch_status_t; + +#ifdef _OPENMP +typedef struct { + uint64_t cost; + gas_sid_t adst; + uint32_t q0; + uint32_t q1; + uint32_t a0; + uint32_t a1; +} gas_omp_task_t; + +static int cmp_omp_task_desc(const void *pa, const void *pb) +{ + const gas_omp_task_t *a = (const gas_omp_task_t *)pa; + const gas_omp_task_t *b = (const gas_omp_task_t *)pb; + if (a->cost != b->cost) { + return a->cost < b->cost ? 1 : -1; + } + return (a->adst > b->adst) - (a->adst < b->adst); +} +#endif + +struct gas_contract_plan { + const gas_space_t *gas; + const double *eri; + const double *gos; + size_t abba_t1_target_bytes; + gas_contract_ws_t *workspace; + uint32_t nworkspace; +#ifdef _OPENMP + gas_omp_task_t *task; + uint32_t ntask; + uint32_t task_threads; +#endif +}; + +static inline void contract_zero(double *x, uint32_t n) +{ + memset(x, 0, (size_t)n * sizeof(*x)); +} + +static inline void contract_axpy(double *restrict y, + const double *restrict x, + double a, uint32_t n) +{ +#ifdef _OPENMP +#pragma omp simd +#endif + for (uint32_t i = 0; i < n; i++) { + y[i] += a * x[i]; + } +} + +/* One AA source row updates several distinct destination rows. Loading x once + * for four destinations reduces the dominant source-row traffic without + * changing the arithmetic order within any destination row. */ +static inline void contract_axpy4(double *restrict y0, + double *restrict y1, + double *restrict y2, + double *restrict y3, + const double *restrict x, + double a0, double a1, + double a2, double a3, + uint32_t n) +{ + uint32_t i = 0; +#if defined(__AVX2__) + const __m256d va0 = _mm256_set1_pd(a0); + const __m256d va1 = _mm256_set1_pd(a1); + const __m256d va2 = _mm256_set1_pd(a2); + const __m256d va3 = _mm256_set1_pd(a3); + for (; i + 4u <= n; i += 4u) { + const __m256d vx = _mm256_loadu_pd(x + i); + __m256d vy0 = _mm256_loadu_pd(y0 + i); + __m256d vy1 = _mm256_loadu_pd(y1 + i); + __m256d vy2 = _mm256_loadu_pd(y2 + i); + __m256d vy3 = _mm256_loadu_pd(y3 + i); + vy0 = _mm256_add_pd(vy0, _mm256_mul_pd(va0, vx)); + vy1 = _mm256_add_pd(vy1, _mm256_mul_pd(va1, vx)); + vy2 = _mm256_add_pd(vy2, _mm256_mul_pd(va2, vx)); + vy3 = _mm256_add_pd(vy3, _mm256_mul_pd(va3, vx)); + _mm256_storeu_pd(y0 + i, vy0); + _mm256_storeu_pd(y1 + i, vy1); + _mm256_storeu_pd(y2 + i, vy2); + _mm256_storeu_pd(y3 + i, vy3); + } +#endif + for (; i < n; i++) { + const double xi = x[i]; + y0[i] += a0 * xi; + y1[i] += a1 * xi; + y2[i] += a2 * xi; + y3[i] += a3 * xi; + } +} + +static inline void contract_axpy2(double *restrict y0, + double *restrict y1, + const double *restrict x, + double a0, double a1, + uint32_t n) +{ + uint32_t i = 0; +#if defined(__AVX2__) + const __m256d va0 = _mm256_set1_pd(a0); + const __m256d va1 = _mm256_set1_pd(a1); + for (; i + 4u <= n; i += 4u) { + const __m256d vx = _mm256_loadu_pd(x + i); + __m256d vy0 = _mm256_loadu_pd(y0 + i); + __m256d vy1 = _mm256_loadu_pd(y1 + i); + vy0 = _mm256_add_pd(vy0, _mm256_mul_pd(va0, vx)); + vy1 = _mm256_add_pd(vy1, _mm256_mul_pd(va1, vx)); + _mm256_storeu_pd(y0 + i, vy0); + _mm256_storeu_pd(y1 + i, vy1); + } +#endif + for (; i < n; i++) { + const double xi = x[i]; + y0[i] += a0 * xi; + y1[i] += a1 * xi; + } +} + +/* ========================================================================== */ +/* 2. ABBA microkernels and workspace helpers */ +/* ========================================================================== */ + +static inline void abba_beta_kernel1(double *restrict y, + const double *restrict t0, + double g0, uint32_t n) +{ + for (uint32_t i = 0; i < n; i++) { + y[i] += g0 * t0[i]; + } +} + +static inline void abba_beta_kernel2(double *restrict y, + const double *restrict t0, + const double *restrict t1, + double g0, double g1, + uint32_t n) +{ + for (uint32_t i = 0; i < n; i++) { + y[i] += g0 * t0[i] + g1 * t1[i]; + } +} + +static inline void abba_beta_kernel4_scalar(double *restrict y, + const double *restrict t0, + const double *restrict t1, + const double *restrict t2, + const double *restrict t3, + double g0, double g1, + double g2, double g3, + uint32_t n) +{ + for (uint32_t i = 0; i < n; i++) { + y[i] += g0 * t0[i] + g1 * t1[i] + + g2 * t2[i] + g3 * t3[i]; + } +} + +#if defined(__AVX2__) +static inline void abba_beta_kernel4_avx2(double *restrict y, + const double *restrict t0, + const double *restrict t1, + const double *restrict t2, + const double *restrict t3, + double g0, double g1, + double g2, double g3, + uint32_t n) +{ + const __m256d vg0 = _mm256_set1_pd(g0); + const __m256d vg1 = _mm256_set1_pd(g1); + const __m256d vg2 = _mm256_set1_pd(g2); + const __m256d vg3 = _mm256_set1_pd(g3); + uint32_t i = 0; + for (; i + 4u <= n; i += 4u) { + __m256d vy = _mm256_loadu_pd(y + i); + __m256d v0 = _mm256_loadu_pd(t0 + i); + __m256d v1 = _mm256_loadu_pd(t1 + i); + __m256d v2 = _mm256_loadu_pd(t2 + i); + __m256d v3 = _mm256_loadu_pd(t3 + i); + vy = _mm256_add_pd(vy, _mm256_mul_pd(vg0, v0)); + vy = _mm256_add_pd(vy, _mm256_mul_pd(vg1, v1)); + vy = _mm256_add_pd(vy, _mm256_mul_pd(vg2, v2)); + vy = _mm256_add_pd(vy, _mm256_mul_pd(vg3, v3)); + _mm256_storeu_pd(y + i, vy); + } + for (; i < n; i++) { + y[i] += g0 * t0[i] + g1 * t1[i] + + g2 * t2[i] + g3 * t3[i]; + } +} +#endif + +static inline void abba_beta_kernel4(double *restrict y, + const double *restrict t0, + const double *restrict t1, + const double *restrict t2, + const double *restrict t3, + double g0, double g1, + double g2, double g3, + uint32_t n) +{ +#if defined(__AVX2__) + abba_beta_kernel4_avx2(y, t0, t1, t2, t3, g0, g1, g2, g3, n); +#else + abba_beta_kernel4_scalar(y, t0, t1, t2, t3, g0, g1, g2, g3, n); +#endif +} + +/* Update two non-aliasing destination rows from the same four T1 streams. + * Keeping four operators leaves enough AVX2 registers for both coefficient + * sets and both Y accumulators. The multiply/add order within each Y row is + * identical to the first/second half of the established width8 kernel. */ +static inline void abba_beta_kernel2dst4( + double *restrict y0, double *restrict y1, + const double *restrict t0, const double *restrict t1, + const double *restrict t2, const double *restrict t3, + double g00, double g01, double g02, double g03, + double g10, double g11, double g12, double g13, + uint32_t n) +{ +#if defined(__AVX2__) + const __m256d vg00 = _mm256_set1_pd(g00); + const __m256d vg01 = _mm256_set1_pd(g01); + const __m256d vg02 = _mm256_set1_pd(g02); + const __m256d vg03 = _mm256_set1_pd(g03); + const __m256d vg10 = _mm256_set1_pd(g10); + const __m256d vg11 = _mm256_set1_pd(g11); + const __m256d vg12 = _mm256_set1_pd(g12); + const __m256d vg13 = _mm256_set1_pd(g13); + uint32_t i = 0u; + for (; i + 4u <= n; i += 4u) { + __m256d vy0 = _mm256_loadu_pd(y0 + i); + __m256d vy1 = _mm256_loadu_pd(y1 + i); + __m256d vt = _mm256_loadu_pd(t0 + i); + vy0 = _mm256_add_pd(vy0, _mm256_mul_pd(vg00, vt)); + vy1 = _mm256_add_pd(vy1, _mm256_mul_pd(vg10, vt)); + vt = _mm256_loadu_pd(t1 + i); + vy0 = _mm256_add_pd(vy0, _mm256_mul_pd(vg01, vt)); + vy1 = _mm256_add_pd(vy1, _mm256_mul_pd(vg11, vt)); + vt = _mm256_loadu_pd(t2 + i); + vy0 = _mm256_add_pd(vy0, _mm256_mul_pd(vg02, vt)); + vy1 = _mm256_add_pd(vy1, _mm256_mul_pd(vg12, vt)); + vt = _mm256_loadu_pd(t3 + i); + vy0 = _mm256_add_pd(vy0, _mm256_mul_pd(vg03, vt)); + vy1 = _mm256_add_pd(vy1, _mm256_mul_pd(vg13, vt)); + _mm256_storeu_pd(y0 + i, vy0); + _mm256_storeu_pd(y1 + i, vy1); + } + for (; i < n; i++) { + y0[i] += g00 * t0[i] + g01 * t1[i] + + g02 * t2[i] + g03 * t3[i]; + y1[i] += g10 * t0[i] + g11 * t1[i] + + g12 * t2[i] + g13 * t3[i]; + } +#else + for (uint32_t i = 0; i < n; i++) { + y0[i] += g00 * t0[i] + g01 * t1[i] + + g02 * t2[i] + g03 * t3[i]; + y1[i] += g10 * t0[i] + g11 * t1[i] + + g12 * t2[i] + g13 * t3[i]; + } +#endif +} + +static inline void abba_beta_kernel8(double *restrict y, + const double *restrict t0, + const double *restrict t1, + const double *restrict t2, + const double *restrict t3, + const double *restrict t4, + const double *restrict t5, + const double *restrict t6, + const double *restrict t7, + double g0, double g1, double g2, double g3, + double g4, double g5, double g6, double g7, + uint32_t n) +{ +#if defined(__AVX2__) + const __m256d vg0 = _mm256_set1_pd(g0); + const __m256d vg1 = _mm256_set1_pd(g1); + const __m256d vg2 = _mm256_set1_pd(g2); + const __m256d vg3 = _mm256_set1_pd(g3); + const __m256d vg4 = _mm256_set1_pd(g4); + const __m256d vg5 = _mm256_set1_pd(g5); + const __m256d vg6 = _mm256_set1_pd(g6); + const __m256d vg7 = _mm256_set1_pd(g7); + uint32_t i = 0; + for (; i + 4u <= n; i += 4u) { + __m256d vy = _mm256_loadu_pd(y + i); + __m256d v0 = _mm256_loadu_pd(t0 + i); + __m256d v1 = _mm256_loadu_pd(t1 + i); + __m256d v2 = _mm256_loadu_pd(t2 + i); + __m256d v3 = _mm256_loadu_pd(t3 + i); + __m256d v4 = _mm256_loadu_pd(t4 + i); + __m256d v5 = _mm256_loadu_pd(t5 + i); + __m256d v6 = _mm256_loadu_pd(t6 + i); + __m256d v7 = _mm256_loadu_pd(t7 + i); + vy = _mm256_add_pd(vy, _mm256_mul_pd(vg0, v0)); + vy = _mm256_add_pd(vy, _mm256_mul_pd(vg1, v1)); + vy = _mm256_add_pd(vy, _mm256_mul_pd(vg2, v2)); + vy = _mm256_add_pd(vy, _mm256_mul_pd(vg3, v3)); + vy = _mm256_add_pd(vy, _mm256_mul_pd(vg4, v4)); + vy = _mm256_add_pd(vy, _mm256_mul_pd(vg5, v5)); + vy = _mm256_add_pd(vy, _mm256_mul_pd(vg6, v6)); + vy = _mm256_add_pd(vy, _mm256_mul_pd(vg7, v7)); + _mm256_storeu_pd(y + i, vy); + } + for (; i < n; i++) { + y[i] += g0 * t0[i] + g1 * t1[i] + + g2 * t2[i] + g3 * t3[i] + + g4 * t4[i] + g5 * t5[i] + + g6 * t6[i] + g7 * t7[i]; + } +#else + for (uint32_t i = 0; i < n; i++) { + y[i] += g0 * t0[i] + g1 * t1[i] + + g2 * t2[i] + g3 * t3[i] + + g4 * t4[i] + g5 * t5[i] + + g6 * t6[i] + g7 * t7[i]; + } +#endif +} + +static inline void abba_fused_ops(double *restrict y, + const double *restrict tbase, + size_t op_stride, + const double *restrict gos, + uint32_t nnorb, uint16_t opb, + const uint16_t *restrict opa, + uint32_t nop, double sbeta, uint32_t n) +{ + uint32_t io = 0; + for (; io + 8u <= nop; io += 8u) { + const double *restrict t0 = tbase + (size_t)(io + 0u) * op_stride; + const double *restrict t1 = tbase + (size_t)(io + 1u) * op_stride; + const double *restrict t2 = tbase + (size_t)(io + 2u) * op_stride; + const double *restrict t3 = tbase + (size_t)(io + 3u) * op_stride; + const double *restrict t4 = tbase + (size_t)(io + 4u) * op_stride; + const double *restrict t5 = tbase + (size_t)(io + 5u) * op_stride; + const double *restrict t6 = tbase + (size_t)(io + 6u) * op_stride; + const double *restrict t7 = tbase + (size_t)(io + 7u) * op_stride; + double g0 = sbeta * abba_gvalue(gos, nnorb, opb, opa[io + 0u]); + double g1 = sbeta * abba_gvalue(gos, nnorb, opb, opa[io + 1u]); + double g2 = sbeta * abba_gvalue(gos, nnorb, opb, opa[io + 2u]); + double g3 = sbeta * abba_gvalue(gos, nnorb, opb, opa[io + 3u]); + double g4 = sbeta * abba_gvalue(gos, nnorb, opb, opa[io + 4u]); + double g5 = sbeta * abba_gvalue(gos, nnorb, opb, opa[io + 5u]); + double g6 = sbeta * abba_gvalue(gos, nnorb, opb, opa[io + 6u]); + double g7 = sbeta * abba_gvalue(gos, nnorb, opb, opa[io + 7u]); + abba_beta_kernel8(y, t0, t1, t2, t3, t4, t5, t6, t7, + g0, g1, g2, g3, g4, g5, g6, g7, n); + } + for (; io + 4u <= nop; io += 4u) { + const double *restrict t0 = tbase + (size_t)(io + 0u) * op_stride; + const double *restrict t1 = tbase + (size_t)(io + 1u) * op_stride; + const double *restrict t2 = tbase + (size_t)(io + 2u) * op_stride; + const double *restrict t3 = tbase + (size_t)(io + 3u) * op_stride; + double g0 = sbeta * abba_gvalue(gos, nnorb, opb, opa[io + 0u]); + double g1 = sbeta * abba_gvalue(gos, nnorb, opb, opa[io + 1u]); + double g2 = sbeta * abba_gvalue(gos, nnorb, opb, opa[io + 2u]); + double g3 = sbeta * abba_gvalue(gos, nnorb, opb, opa[io + 3u]); + abba_beta_kernel4(y, t0, t1, t2, t3, g0, g1, g2, g3, n); + } + for (; io + 2u <= nop; io += 2u) { + const double *restrict t0 = tbase + (size_t)(io + 0u) * op_stride; + const double *restrict t1 = tbase + (size_t)(io + 1u) * op_stride; + double g0 = sbeta * abba_gvalue(gos, nnorb, opb, opa[io + 0u]); + double g1 = sbeta * abba_gvalue(gos, nnorb, opb, opa[io + 1u]); + abba_beta_kernel2(y, t0, t1, g0, g1, n); + } + for (; io < nop; io++) { + const double *restrict t0 = tbase + (size_t)io * op_stride; + double g0 = sbeta * abba_gvalue(gos, nnorb, opb, opa[io]); + abba_beta_kernel1(y, t0, g0, n); + } +} + +static inline void abba_fused_ops_pair2x4( + double *restrict y0, double *restrict y1, + const double *restrict tbase, size_t op_stride, + const double *restrict gos, uint32_t nnorb, + uint16_t opb0, uint16_t opb1, + const uint16_t *restrict opa, uint32_t nop, + double sbeta0, double sbeta1, uint32_t n) +{ + uint32_t io = 0u; + for (; io + 4u <= nop; io += 4u) { + const double *restrict t0 = tbase + (size_t)(io + 0u) * op_stride; + const double *restrict t1 = tbase + (size_t)(io + 1u) * op_stride; + const double *restrict t2 = tbase + (size_t)(io + 2u) * op_stride; + const double *restrict t3 = tbase + (size_t)(io + 3u) * op_stride; + double g00 = sbeta0 * abba_gvalue(gos, nnorb, opb0, opa[io + 0u]); + double g01 = sbeta0 * abba_gvalue(gos, nnorb, opb0, opa[io + 1u]); + double g02 = sbeta0 * abba_gvalue(gos, nnorb, opb0, opa[io + 2u]); + double g03 = sbeta0 * abba_gvalue(gos, nnorb, opb0, opa[io + 3u]); + double g10 = sbeta1 * abba_gvalue(gos, nnorb, opb1, opa[io + 0u]); + double g11 = sbeta1 * abba_gvalue(gos, nnorb, opb1, opa[io + 1u]); + double g12 = sbeta1 * abba_gvalue(gos, nnorb, opb1, opa[io + 2u]); + double g13 = sbeta1 * abba_gvalue(gos, nnorb, opb1, opa[io + 3u]); + abba_beta_kernel2dst4(y0, y1, t0, t1, t2, t3, + g00, g01, g02, g03, + g10, g11, g12, g13, n); + } + for (; io + 2u <= nop; io += 2u) { + const double *restrict t0 = tbase + (size_t)(io + 0u) * op_stride; + const double *restrict t1 = tbase + (size_t)(io + 1u) * op_stride; + double g00 = sbeta0 * abba_gvalue(gos, nnorb, opb0, opa[io + 0u]); + double g01 = sbeta0 * abba_gvalue(gos, nnorb, opb0, opa[io + 1u]); + double g10 = sbeta1 * abba_gvalue(gos, nnorb, opb1, opa[io + 0u]); + double g11 = sbeta1 * abba_gvalue(gos, nnorb, opb1, opa[io + 1u]); + abba_beta_kernel2(y0, t0, t1, g00, g01, n); + abba_beta_kernel2(y1, t0, t1, g10, g11, n); + } + for (; io < nop; io++) { + const double *restrict t0 = tbase + (size_t)io * op_stride; + double g0 = sbeta0 * abba_gvalue(gos, nnorb, opb0, opa[io]); + double g1 = sbeta1 * abba_gvalue(gos, nnorb, opb1, opa[io]); + abba_beta_kernel1(y0, t0, g0, n); + abba_beta_kernel1(y1, t0, g1, n); + } +} + +/* Blocked transpose from C[alpha,beta] to T[beta,alpha]. */ +static void transpose_pack(double *restrict t, + const double *restrict c, uint32_t c_nb, + uint32_t na, uint32_t nb, uint32_t c_a0) +{ + for (uint32_t a0 = 0; a0 < na; a0 += BB_TRANSPOSE_TILE) { + uint32_t a1 = na - a0 < BB_TRANSPOSE_TILE ? + na : a0 + BB_TRANSPOSE_TILE; + for (uint32_t b0 = 0; b0 < nb; b0 += BB_TRANSPOSE_TILE) { + uint32_t b1 = nb - b0 < BB_TRANSPOSE_TILE ? + nb : b0 + BB_TRANSPOSE_TILE; + for (uint32_t b = b0; b < b1; b++) { + double *restrict tr = t + (uint64_t)b * na + a0; +#ifdef _OPENMP +#pragma omp simd +#endif + for (uint32_t a = a0; a < a1; a++) { + tr[a - a0] = c[(uint64_t)(c_a0 + a) * c_nb + b]; + } + } + } + } +} + +/* Add T[beta,alpha] into the row-major block C[alpha,beta]. */ +static void transpose_add(double *restrict c, uint32_t c_nb, + const double *restrict t, + uint32_t na, uint32_t nb, uint32_t c_a0, + uint32_t tile) +{ + for (uint32_t a0 = 0; a0 < na; a0 += tile) { + uint32_t a1 = na - a0 < tile ? na : a0 + tile; + for (uint32_t b0 = 0; b0 < nb; b0 += tile) { + uint32_t b1 = nb - b0 < tile ? nb : b0 + tile; + for (uint32_t a = a0; a < a1; a++) { + double *restrict cr = c + + (uint64_t)(c_a0 + a) * c_nb + b0; +#ifdef _OPENMP +#pragma omp simd +#endif + for (uint32_t b = b0; b < b1; b++) { + cr[b - b0] += t[(uint64_t)b * na + a]; + } + } + } + } +} + +static inline void abba_transpose_add(double *restrict dst, uint32_t dst_nb, + const double *restrict y, + uint32_t na, uint32_t nb, + uint32_t dst_a0) +{ + transpose_add(dst, dst_nb, y, na, nb, dst_a0, ABBA_WRITEBACK_TILE); +} + +static inline void contract_ws_init(gas_contract_ws_t *ws) +{ + memset(ws, 0, sizeof(*ws)); +} + +static inline void contract_ws_free(gas_contract_ws_t *ws) +{ + free(ws->arena); + free(ws->pair); + free(ws->op_map); + free(ws->op); + free(ws->alpha_hit); + memset(ws, 0, sizeof(*ws)); +} + +static inline uint64_t contract_ws_bytes(const gas_contract_ws_t *ws) +{ + return (uint64_t)ws->arena_cap * sizeof(*ws->arena) + + (uint64_t)ws->pair_cap * sizeof(*ws->pair) + + (uint64_t)ws->op_map_cap * sizeof(*ws->op_map) + + (uint64_t)ws->op_cap * sizeof(*ws->op) + + (uint64_t)ws->alpha_hit_cap * sizeof(*ws->alpha_hit); +} + +static size_t contract_grow_capacity(size_t cap, size_t need) +{ + size_t next = cap ? cap : 64u; + while (next < need) { + if (next > SIZE_MAX / 2u) { + return need; + } + next *= 2u; + } + return next; +} + +static int reserve_arena(gas_contract_ws_t *ws, size_t n) +{ + if (n <= ws->arena_cap) { + return GAS_SUCCESS; + } + size_t cap = contract_grow_capacity(ws->arena_cap, n); + double *p = realloc(ws->arena, cap * sizeof(*p)); + if (p == 0) { + return GAS_ERR_MEMORY; + } + ws->arena = p; + ws->arena_cap = cap; + if (ws->arena_zeroed_upto > cap) { + ws->arena_zeroed_upto = cap; + } + return GAS_SUCCESS; +} + +static int reserve_alpha_hitlist(gas_contract_ws_t *ws, size_t n) +{ + if (n > UINT32_MAX) { + return GAS_ERR_MEMORY; + } + if (n <= ws->alpha_hit_cap) { + return GAS_SUCCESS; + } + size_t cap = contract_grow_capacity(ws->alpha_hit_cap, n); + gas_alpha_hit_t *p = realloc(ws->alpha_hit, cap * sizeof(*p)); + if (p == 0) { + return GAS_ERR_MEMORY; + } + ws->alpha_hit = p; + ws->alpha_hit_cap = cap; + return GAS_SUCCESS; +} + +static int reserve_pairs(gas_contract_ws_t *ws, size_t n) +{ + if (n <= ws->pair_cap) { + return GAS_SUCCESS; + } + size_t cap = contract_grow_capacity(ws->pair_cap, n); + gas_tid_pair_t *p = realloc(ws->pair, cap * sizeof(*p)); + if (p == 0) { + return GAS_ERR_MEMORY; + } + ws->pair = p; + ws->pair_cap = cap; + return GAS_SUCCESS; +} + +static int reserve_op_map(gas_contract_ws_t *ws, size_t n) +{ + if (n <= ws->op_map_cap) { + return GAS_SUCCESS; + } + size_t cap = contract_grow_capacity(ws->op_map_cap, n); + int32_t *p = realloc(ws->op_map, cap * sizeof(*p)); + if (p == 0) { + return GAS_ERR_MEMORY; + } + ws->op_map = p; + ws->op_map_cap = cap; + return GAS_SUCCESS; +} + +static int reserve_ops(gas_contract_ws_t *ws, size_t n) +{ + if (n <= ws->op_cap) { + return GAS_SUCCESS; + } + size_t cap = contract_grow_capacity(ws->op_cap, n); + uint16_t *p = realloc(ws->op, cap * sizeof(*p)); + if (p == 0) { + return GAS_ERR_MEMORY; + } + ws->op = p; + ws->op_cap = cap; + return GAS_SUCCESS; +} + +#ifdef _OPENMP +static gas_tid_t find_table_in_row(gas_sid_t dst, + const gas_sid_t *restrict sid, + gas_tid_t off, uint32_t n) +{ + if (n <= 8u) { + for (uint32_t i = 0; i < n; i++) { + if (sid[i] == dst) { + return off + i; + } + } + return GAS_INVALID_TID; + } + + uint32_t lo = 0; + uint32_t hi = n; + while (lo < hi) { + uint32_t mid = lo + ((hi - lo) >> 1); + if (sid[mid] < dst) { + lo = mid + 1u; + } else { + hi = mid; + } + } + return lo < n && sid[lo] == dst ? off + lo : GAS_INVALID_TID; +} +#endif + +/* ========================================================================== */ +/* 3. Two-electron path kernels */ +/* ========================================================================== */ + +static void contract_aa_path_rows(const gas_space_t *gas, + const double *restrict eri, + const double *restrict ci0, + double *restrict ci1, + uint32_t nnorb, + gas_bid_t bsrc, gas_bid_t bdst, + gas_tid_t tid1, gas_tid_t tid2, + uint32_t a0, uint32_t a1) +{ + const gas_block_t *bs = gas->block + bsrc; + const gas_block_t *bd = gas->block + bdst; + const gas_link_table_t *t1 = gas->table + tid1; + const gas_link_table_t *t2 = gas->table + tid2; + const uint32_t nb = gas->sector_nstr[bs->sb]; + const uint32_t na1 = gas->sector_nstr[bd->sa]; + const double *src = ci0 + bs->offset; + double *dst = ci1 + bd->offset; + + if (a0 > na1) a0 = na1; + if (a1 > na1) a1 = na1; + if (a1 <= a0) return; + + for (uint32_t ia0 = 0; ia0 < t1->nsrc; ia0++) { + const double *src_row = src + (uint64_t)ia0 * nb; + const gas_link_entry_t *r1 = t1->link + (uint64_t)ia0 * t1->nlink; + for (uint32_t k1 = 0; k1 < t1->nlink; k1++) { + const gas_link_entry_t *e1 = r1 + k1; + const gas_link_entry_t *r2 = t2->link + + (uint64_t)e1->addr * t2->nlink; + uint16_t op1 = gas_link_ia(e1); + for (uint32_t k2 = 0; k2 < t2->nlink; k2++) { + const gas_link_entry_t *e2 = r2 + k2; + if (e2->addr < a0 || e2->addr >= a1) continue; + double fac = (double)(e1->sign * e2->sign) * + eri[(uint64_t)gas_link_ia(e2) * nnorb + op1]; + double *dst_row = dst + (uint64_t)e2->addr * nb; + contract_axpy(dst_row, src_row, fac, nb); + } + } + } +} + +static void contract_aa_path(const gas_space_t *gas, + const double *restrict eri, + const double *restrict ci0, + double *restrict ci1, + uint32_t nnorb, + gas_bid_t bsrc, gas_bid_t bdst, + gas_tid_t tid1, gas_tid_t tid2) +{ + contract_aa_path_rows(gas, eri, ci0, ci1, nnorb, bsrc, bdst, + tid1, tid2, 0u, + gas->sector_nstr[gas->block[bdst].sa]); +} + +#ifdef _OPENMP +static int contract_aa_path_rows_accum(const gas_space_t *gas, + const double *restrict eri, + const double *restrict ci0, + double *restrict ci1, + uint32_t nnorb, + gas_bid_t bsrc, gas_bid_t bdst, + gas_tid_t tid1, gas_tid_t tid2, + uint32_t a0, uint32_t a1, + gas_contract_ws_t *ws) +{ + const gas_block_t *bs = gas->block + bsrc; + const gas_block_t *bd = gas->block + bdst; + const gas_link_table_t *t1 = gas->table + tid1; + const gas_link_table_t *t2 = gas->table + tid2; + const uint32_t nb = gas->sector_nstr[bs->sb]; + const uint32_t na1 = gas->sector_nstr[bd->sa]; + const double *src = ci0 + bs->offset; + double *dst = ci1 + bd->offset; + + if (a0 > na1) a0 = na1; + if (a1 > na1) a1 = na1; + if (a1 <= a0) return GAS_SUCCESS; + + if (na1 > UINT16_MAX || + reserve_op_map(ws, na1) != GAS_SUCCESS || + reserve_ops(ws, na1) != GAS_SUCCESS || + reserve_arena(ws, na1) != GAS_SUCCESS || + ws->op_map == 0 || ws->op == 0 || ws->arena == 0) { + return GAS_ERR_MEMORY; + } + for (uint32_t i = 0; i < na1; i++) { + ws->op_map[i] = -1; + } + double *acc = ws->arena; + + /* Large-block AA kernel. For a fixed source alpha string, the two + * one-electron alpha paths can reach the same destination alpha string + * through different intermediate strings/orderings. A direct kernel would do + * one full beta-row AXPY per path. Here all coefficients landing on the + * same destination alpha row are accumulated first, then one AXPY is + * issued. This mirrors the BB coefficient-accumulation kernel and + * targets the large dense beta rows in grouped2x8. */ + for (uint32_t ia0 = 0; ia0 < t1->nsrc; ia0++) { + const double *src_row = src + (uint64_t)ia0 * nb; + const gas_link_entry_t *r1 = t1->link + (uint64_t)ia0 * t1->nlink; + uint32_t ntouch = 0; + + for (uint32_t k1 = 0; k1 < t1->nlink; k1++) { + const gas_link_entry_t *e1 = r1 + k1; + const gas_link_entry_t *r2 = t2->link + + (uint64_t)e1->addr * t2->nlink; + uint16_t op1 = gas_link_ia(e1); + int s1 = e1->sign; + for (uint32_t k2 = 0; k2 < t2->nlink; k2++) { + const gas_link_entry_t *e2 = r2 + k2; + uint32_t ia2 = e2->addr; + if (ia2 < a0 || ia2 >= a1) continue; + double fac = (double)(s1 * e2->sign) * + eri[(uint64_t)gas_link_ia(e2) * nnorb + op1]; + int32_t pos = ws->op_map[ia2]; + if (pos < 0) { + pos = (int32_t)ntouch; + ws->op_map[ia2] = pos; + ws->op[ntouch++] = (uint16_t)ia2; + acc[ia2] = fac; + } else { + acc[ia2] += fac; + } + } + } + + uint32_t it = 0; + for (; it + 4u <= ntouch; it += 4u) { + uint32_t ia20 = ws->op[it + 0u]; + uint32_t ia21 = ws->op[it + 1u]; + uint32_t ia22 = ws->op[it + 2u]; + uint32_t ia23 = ws->op[it + 3u]; + contract_axpy4(dst + (uint64_t)ia20 * nb, + dst + (uint64_t)ia21 * nb, + dst + (uint64_t)ia22 * nb, + dst + (uint64_t)ia23 * nb, + src_row, acc[ia20], acc[ia21], + acc[ia22], acc[ia23], nb); + ws->op_map[ia20] = -1; + ws->op_map[ia21] = -1; + ws->op_map[ia22] = -1; + ws->op_map[ia23] = -1; + } + if (it + 2u <= ntouch) { + uint32_t ia20 = ws->op[it + 0u]; + uint32_t ia21 = ws->op[it + 1u]; + contract_axpy2(dst + (uint64_t)ia20 * nb, + dst + (uint64_t)ia21 * nb, + src_row, acc[ia20], acc[ia21], nb); + ws->op_map[ia20] = -1; + ws->op_map[ia21] = -1; + it += 2u; + } + for (; it < ntouch; it++) { + uint32_t ia2 = ws->op[it]; + contract_axpy(dst + (uint64_t)ia2 * nb, + src_row, acc[ia2], nb); + ws->op_map[ia2] = -1; + } + } + return GAS_SUCCESS; +} +#endif + +static void contract_bb_path_rows(const gas_space_t *gas, + const double *restrict eri, + const double *restrict ci0, + double *restrict ci1, + uint32_t nnorb, + gas_bid_t bsrc, gas_bid_t bdst, + gas_tid_t tid1, gas_tid_t tid2, + uint32_t a0, uint32_t a1) +{ + const gas_block_t *bs = gas->block + bsrc; + const gas_block_t *bd = gas->block + bdst; + const gas_link_table_t *t1 = gas->table + tid1; + const gas_link_table_t *t2 = gas->table + tid2; + const uint32_t na = gas->sector_nstr[bs->sa]; + const uint32_t nb0 = gas->sector_nstr[bs->sb]; + const uint32_t nb2 = gas->sector_nstr[bd->sb]; + const double *src = ci0 + bs->offset; + double *dst = ci1 + bd->offset; + + if (a0 > na) a0 = na; + if (a1 > na) a1 = na; + if (a1 <= a0) return; + + for (uint32_t aa0 = a0; aa0 < a1; aa0 += BB_ALPHA_BLOCK) { + uint32_t aa1 = a1 - aa0 < BB_ALPHA_BLOCK ? a1 : aa0 + BB_ALPHA_BLOCK; + for (uint32_t ib0 = 0; ib0 < t1->nsrc; ib0++) { + const gas_link_entry_t *r1 = t1->link + + (uint64_t)ib0 * t1->nlink; + for (uint32_t k1 = 0; k1 < t1->nlink; k1++) { + const gas_link_entry_t *e1 = r1 + k1; + const gas_link_entry_t *r2 = t2->link + + (uint64_t)e1->addr * t2->nlink; + uint16_t op1 = gas_link_ia(e1); + int s1 = e1->sign; + for (uint32_t k2 = 0; k2 < t2->nlink; k2++) { + const gas_link_entry_t *e2 = r2 + k2; + double fac = (double)(s1 * e2->sign) * + eri[(uint64_t)gas_link_ia(e2) * nnorb + op1]; + uint32_t ib2 = e2->addr; +#ifdef _OPENMP +#pragma omp simd +#endif + for (uint32_t ia = aa0; ia < aa1; ia++) { + dst[(uint64_t)ia * nb2 + ib2] += + fac * src[(uint64_t)ia * nb0 + ib0]; + } + } + } + } + } +} + +static void contract_bb_path(const gas_space_t *gas, + const double *restrict eri, + const double *restrict ci0, + double *restrict ci1, + uint32_t nnorb, + gas_bid_t bsrc, gas_bid_t bdst, + gas_tid_t tid1, gas_tid_t tid2) +{ + contract_bb_path_rows(gas, eri, ci0, ci1, nnorb, bsrc, bdst, + tid1, tid2, 0u, + gas->sector_nstr[gas->block[bsrc].sa]); +} + +static int use_bb_transpose(uint32_t na, uint32_t nb0, + uint32_t nb2, uint32_t npair) +{ + size_t elems = (size_t)na * nb0 + (size_t)na * nb2; + return npair != 0 && na >= BB_TRANSPOSE_MIN_DIM && + nb0 >= BB_TRANSPOSE_MIN_DIM && nb2 >= BB_TRANSPOSE_MIN_DIM && + elems <= BB_TRANSPOSE_MAX_ELEMS; +} + +static int contract_bb_blockpair_transpose_rows( + const gas_space_t *gas, + const double *restrict eri, + const double *restrict ci0, + double *restrict ci1, + uint32_t nnorb, + gas_bid_t bsrc, gas_bid_t bdst, + const gas_tid_pair_t *pair, + uint32_t npair, + gas_contract_ws_t *ws, + uint32_t a0, uint32_t a1) +{ + const gas_block_t *bs = gas->block + bsrc; + const gas_block_t *bd = gas->block + bdst; + const uint32_t na = gas->sector_nstr[bs->sa]; + const uint32_t nb0 = gas->sector_nstr[bs->sb]; + const uint32_t nb2 = gas->sector_nstr[bd->sb]; + if (a0 > na) a0 = na; + if (a1 > na) a1 = na; + if (a1 <= a0) return GAS_SUCCESS; + const uint32_t nat = a1 - a0; + const size_t nsrc = (size_t)nat * nb0; + const size_t ndst = (size_t)nat * nb2; + const double *src = ci0 + bs->offset; + double *dst = ci1 + bd->offset; + int use_accum = 0; + if (nb2 <= UINT16_MAX && + reserve_op_map(ws, nb2) == GAS_SUCCESS && + reserve_ops(ws, nb2) == GAS_SUCCESS && + reserve_arena(ws, nsrc + ndst + (size_t)nb2) == GAS_SUCCESS) { + use_accum = 1; + for (uint32_t i = 0; i < nb2; i++) { + ws->op_map[i] = -1; + } + } else if (reserve_arena(ws, nsrc + ndst) != GAS_SUCCESS) { + return GAS_ERR_MEMORY; + } + double *src_t = ws->arena; + double *dst_t = ws->arena + nsrc; + double *acc = use_accum ? ws->arena + nsrc + ndst : 0; + + transpose_pack(src_t, src, nb0, nat, nb0, a0); + memset(dst_t, 0, ndst * sizeof(*dst_t)); + if (use_accum) { + /* Large-block BB kernel: for each source beta string, combine all + * path coefficients that land on the same destination beta string. + * This turns many repeated AXPYs over the same alpha tile into one + * AXPY with the accumulated coefficient. The touched list keeps + * reset cost proportional to the number of destinations reached, + * not to nb2. */ + for (uint32_t ib0 = 0; ib0 < nb0; ib0++) { + const double *src_vec = src_t + (uint64_t)ib0 * nat; + uint32_t ntouch = 0; + + for (uint32_t ip = 0; ip < npair; ip++) { + const gas_link_table_t *t1 = gas->table + pair[ip].first; + const gas_link_table_t *t2 = gas->table + pair[ip].second; + if (ib0 >= t1->nsrc) { + continue; + } + const gas_link_entry_t *r1 = t1->link + + (uint64_t)ib0 * t1->nlink; + for (uint32_t k1 = 0; k1 < t1->nlink; k1++) { + const gas_link_entry_t *e1 = r1 + k1; + const gas_link_entry_t *r2 = t2->link + + (uint64_t)e1->addr * t2->nlink; + uint16_t op1 = gas_link_ia(e1); + int s1 = e1->sign; + for (uint32_t k2 = 0; k2 < t2->nlink; k2++) { + const gas_link_entry_t *e2 = r2 + k2; + uint32_t ib2 = e2->addr; + double fac = (double)(s1 * e2->sign) * + eri[(uint64_t)gas_link_ia(e2) * + nnorb + op1]; + int32_t pos = ws->op_map[ib2]; + if (pos < 0) { + pos = (int32_t)ntouch; + ws->op_map[ib2] = pos; + ws->op[ntouch++] = (uint16_t)ib2; + acc[ib2] = fac; + } else { + acc[ib2] += fac; + } + } + } + } + + for (uint32_t it = 0; it < ntouch; it++) { + uint32_t ib2 = ws->op[it]; + contract_axpy(dst_t + (uint64_t)ib2 * nat, + src_vec, acc[ib2], nat); + ws->op_map[ib2] = -1; + } + } + } else { + for (uint32_t ip = 0; ip < npair; ip++) { + const gas_link_table_t *t1 = gas->table + pair[ip].first; + const gas_link_table_t *t2 = gas->table + pair[ip].second; + for (uint32_t ib0 = 0; ib0 < t1->nsrc; ib0++) { + const double *src_vec = src_t + (uint64_t)ib0 * nat; + const gas_link_entry_t *r1 = t1->link + + (uint64_t)ib0 * t1->nlink; + for (uint32_t k1 = 0; k1 < t1->nlink; k1++) { + const gas_link_entry_t *e1 = r1 + k1; + const gas_link_entry_t *r2 = t2->link + + (uint64_t)e1->addr * t2->nlink; + uint16_t op1 = gas_link_ia(e1); + int s1 = e1->sign; + for (uint32_t k2 = 0; k2 < t2->nlink; k2++) { + const gas_link_entry_t *e2 = r2 + k2; + double fac = (double)(s1 * e2->sign) * + eri[(uint64_t)gas_link_ia(e2) * + nnorb + op1]; + double *dst_vec = dst_t + + (uint64_t)e2->addr * nat; + contract_axpy(dst_vec, src_vec, fac, nat); + } + } + } + } + } + transpose_add(dst, nb2, dst_t, nat, nb2, a0, BB_TRANSPOSE_TILE); + + return GAS_SUCCESS; +} + +static void contract_bb_path_list_rows(const gas_space_t *gas, + const double *restrict eri, + const double *restrict ci0, + double *restrict ci1, + uint32_t nnorb, + gas_bid_t bsrc, gas_bid_t bdst, + const gas_tid_pair_t *pair, + uint32_t npair, + uint32_t a0, uint32_t a1) +{ + for (uint32_t i = 0; i < npair; i++) { + contract_bb_path_rows(gas, eri, ci0, ci1, nnorb, bsrc, bdst, + pair[i].first, pair[i].second, a0, a1); + } +} + +static void contract_bb_blockpair_rows(const gas_space_t *gas, + const double *restrict eri, + const double *restrict ci0, + double *restrict ci1, + uint32_t nnorb, + gas_bid_t bsrc, gas_bid_t bdst, + const gas_tid_pair_t *pair, + uint32_t npair, + gas_contract_ws_t *ws, + uint32_t a0, uint32_t a1) +{ + const gas_block_t *bs = gas->block + bsrc; + const gas_block_t *bd = gas->block + bdst; + uint32_t na = gas->sector_nstr[bs->sa]; + uint32_t nb0 = gas->sector_nstr[bs->sb]; + uint32_t nb2 = gas->sector_nstr[bd->sb]; + if (a0 > na) a0 = na; + if (a1 > na) a1 = na; + uint32_t nat = a1 > a0 ? a1 - a0 : 0u; + + if (use_bb_transpose(nat, nb0, nb2, npair) && + contract_bb_blockpair_transpose_rows(gas, eri, ci0, ci1, nnorb, + bsrc, bdst, pair, npair, ws, + a0, a1) == GAS_SUCCESS) { + return; + } + contract_bb_path_list_rows(gas, eri, ci0, ci1, nnorb, + bsrc, bdst, pair, npair, a0, a1); +} + +static void contract_bb_blockpair(const gas_space_t *gas, + const double *restrict eri, + const double *restrict ci0, + double *restrict ci1, + uint32_t nnorb, + gas_bid_t bsrc, gas_bid_t bdst, + const gas_tid_pair_t *pair, uint32_t npair, + gas_contract_ws_t *ws) +{ + contract_bb_blockpair_rows(gas, eri, ci0, ci1, nnorb, + bsrc, bdst, pair, npair, ws, 0u, + gas->sector_nstr[gas->block[bsrc].sa]); +} + +static void contract_abba_path_rows(const gas_space_t *gas, + const double *restrict eri, + const double *restrict ci0, + double *restrict ci1, + uint32_t nnorb, + gas_bid_t bsrc, gas_bid_t bdst, + gas_tid_t tida, gas_tid_t tidb, + uint32_t a0, uint32_t a1) +{ + const gas_block_t *bs = gas->block + bsrc; + const gas_block_t *bd = gas->block + bdst; + const gas_link_table_t *ta = gas->table + tida; + const gas_link_table_t *tb = gas->table + tidb; + const uint32_t nb0 = gas->sector_nstr[bs->sb]; + const uint32_t nb1 = gas->sector_nstr[bd->sb]; + const uint32_t na1 = gas->sector_nstr[bd->sa]; + const double *src = ci0 + bs->offset; + double *dst = ci1 + bd->offset; + if (a0 > na1) a0 = na1; + if (a1 > na1) a1 = na1; + if (a1 <= a0) return; + + for (uint32_t ia0 = 0; ia0 < ta->nsrc; ia0++) { + const double *src_row = src + (uint64_t)ia0 * nb0; + const gas_link_entry_t *ra = ta->link + (uint64_t)ia0 * ta->nlink; + for (uint32_t ka = 0; ka < ta->nlink; ka++) { + const gas_link_entry_t *ea = ra + ka; + if (ea->addr < a0 || ea->addr >= a1) continue; + double *dst_row = dst + (uint64_t)ea->addr * nb1; + uint16_t opa = gas_link_ia(ea); + int sa = ea->sign; + for (uint32_t ib0 = 0; ib0 < tb->nsrc; ib0++) { + double c = src_row[ib0]; + const gas_link_entry_t *rb = tb->link + + (uint64_t)ib0 * tb->nlink; + for (uint32_t kb = 0; kb < tb->nlink; kb++) { + const gas_link_entry_t *eb = rb + kb; + uint16_t opb = gas_link_ia(eb); + dst_row[eb->addr] += (double)(sa * eb->sign) * c * + (eri[(uint64_t)opb * nnorb + opa] + + eri[(uint64_t)opa * nnorb + opb]); + } + } + } + } +} + +static void contract_abba_path(const gas_space_t *gas, + const double *restrict eri, + const double *restrict ci0, + double *restrict ci1, + uint32_t nnorb, + gas_bid_t bsrc, gas_bid_t bdst, + gas_tid_t tida, gas_tid_t tidb) +{ + contract_abba_path_rows(gas, eri, ci0, ci1, nnorb, bsrc, bdst, + tida, tidb, 0u, + gas->sector_nstr[gas->block[bdst].sa]); +} + +/* ========================================================================== */ +/* 4. Destination-driven two-electron traversal */ +/* ========================================================================== */ + +static uint32_t aa_forward_scan_cost(const gas_space_t *gas, + const gas_sid_t *as, uint32_t n) +{ + uint32_t cost = 0; + for (uint32_t i = 0; i < n; i++) { + cost += gas->T.row[as[i]].n; + } + return cost; +} + +static void contract_aa_forward(const gas_space_t *gas, + const double *restrict eri, + const double *restrict ci0, + double *restrict ci1, + uint32_t nnorb, gas_bid_t bsrc) +{ + const gas_block_t *blk = gas->block + bsrc; + gas_row_t ar = gas->T.row[blk->sa]; + const gas_sid_t *as = gas->T.dst + ar.off; + + for (uint32_t i = 0; i < ar.n; i++) { + gas_row_t row = gas->T.row[as[i]]; + const gas_sid_t *dst = gas->T.dst + row.off; + for (uint32_t j = 0; j < row.n; j++) { + gas_bid_t bd = gas_find_block(gas, dst[j], blk->sb); + if (bd != GAS_INVALID_BID) { + contract_aa_path(gas, eri, ci0, ci1, nnorb, + bsrc, bd, ar.off + i, row.off + j); + } + } + } +} + +static uint32_t bb_forward_scan_cost(const gas_space_t *gas, + const gas_sid_t *bs, uint32_t n) +{ + uint32_t cost = 0; + for (uint32_t i = 0; i < n; i++) { + cost += gas->T.row[bs[i]].n; + } + return cost; +} + +static void contract_bb_forward(const gas_space_t *gas, + const double *restrict eri, + const double *restrict ci0, + double *restrict ci1, + uint32_t nnorb, gas_bid_t bsrc) +{ + const gas_block_t *blk = gas->block + bsrc; + gas_row_t br = gas->T.row[blk->sb]; + const gas_sid_t *bs = gas->T.dst + br.off; + + for (uint32_t i = 0; i < br.n; i++) { + gas_row_t row = gas->T.row[bs[i]]; + const gas_sid_t *dst = gas->T.dst + row.off; + for (uint32_t j = 0; j < row.n; j++) { + gas_bid_t bd = gas_find_block(gas, blk->sa, dst[j]); + if (bd != GAS_INVALID_BID) { + contract_bb_path(gas, eri, ci0, ci1, nnorb, + bsrc, bd, br.off + i, row.off + j); + } + } + } +} + +static void contract_aa_by_dst(const gas_space_t *gas, + const double *restrict eri, + const double *restrict ci0, + double *restrict ci1, + uint32_t nnorb, gas_bid_t bsrc) +{ + const gas_block_t *blk = gas->block + bsrc; + gas_row_t ar = gas->T.row[blk->sa]; + const gas_sid_t *as = gas->T.dst + ar.off; + gas_row_t qrow = gas->D.by_beta_row[blk->sb]; + const gas_sid_t *qa = gas->D.by_beta_sid + qrow.off; + const gas_bid_t *qb = gas->D.by_beta_bid + qrow.off; + + for (uint32_t q = 0; q < qrow.n; q++) { + gas_row_t rr = gas->R.row[qa[q]]; + const gas_sid_t *rs = gas->R.src + rr.off; + const gas_tid_t *rt = gas->R.tid + rr.off; + uint32_t i = 0; + uint32_t j = 0; + while (i < ar.n && j < rr.n) { + if (as[i] < rs[j]) { + i++; + } else if (as[i] > rs[j]) { + j++; + } else { + contract_aa_path(gas, eri, ci0, ci1, nnorb, + bsrc, qb[q], ar.off + i, rt[j]); + i++; + j++; + } + } + } +} + +static void contract_bb_by_dst(const gas_space_t *gas, + const double *restrict eri, + const double *restrict ci0, + double *restrict ci1, + uint32_t nnorb, gas_bid_t bsrc, + gas_contract_ws_t *ws) +{ + const gas_block_t *blk = gas->block + bsrc; + uint32_t na = gas->sector_nstr[blk->sa]; + uint32_t nb0 = gas->sector_nstr[blk->sb]; + gas_row_t br = gas->T.row[blk->sb]; + const gas_sid_t *bs = gas->T.dst + br.off; + gas_row_t qrow = gas->D.by_alpha_row[blk->sa]; + + for (uint32_t q = 0; q < qrow.n; q++) { + gas_bid_t bdst = qrow.off + q; + gas_sid_t b2 = gas->block[bdst].sb; + gas_row_t rr = gas->R.row[b2]; + const gas_sid_t *rs = gas->R.src + rr.off; + const gas_tid_t *rt = gas->R.tid + rr.off; + uint32_t i = 0; + uint32_t j = 0; + uint32_t maxpair = br.n < rr.n ? br.n : rr.n; + uint32_t nb2 = gas->sector_nstr[b2]; + + if (use_bb_transpose(na, nb0, nb2, maxpair) && + reserve_pairs(ws, maxpair) == GAS_SUCCESS) { + uint32_t npair = 0; + while (i < br.n && j < rr.n) { + if (bs[i] < rs[j]) { + i++; + } else if (bs[i] > rs[j]) { + j++; + } else { + ws->pair[npair].first = br.off + i; + ws->pair[npair].second = rt[j]; + npair++; + i++; + j++; + } + } + if (npair != 0) { + contract_bb_blockpair(gas, eri, ci0, ci1, nnorb, + bsrc, bdst, ws->pair, npair, ws); + } + continue; + } + + while (i < br.n && j < rr.n) { + if (bs[i] < rs[j]) { + i++; + } else if (bs[i] > rs[j]) { + j++; + } else { + contract_bb_path(gas, eri, ci0, ci1, nnorb, + bsrc, bdst, br.off + i, rt[j]); + i++; + j++; + } + } + } +} + +static int collect_alpha_ops(const gas_link_table_t *ta, + uint32_t nnorb, gas_contract_ws_t *ws, + uint32_t *nop_out) +{ + uint32_t nop = ta->nop; + if (reserve_op_map(ws, nnorb) != GAS_SUCCESS) { + return GAS_ERR_MEMORY; + } + for (uint32_t i = 0; i < nnorb; i++) { + ws->op_map[i] = -1; + } + if (reserve_ops(ws, nop) != GAS_SUCCESS) { + return GAS_ERR_MEMORY; + } + if (nop != 0u && ta->active_op == 0) { + return GAS_ERR_INVALID; + } + for (uint32_t i = 0; i < nop; i++) { + uint16_t op = ta->active_op[i]; + ws->op_map[op] = (int32_t)i; + ws->op[i] = op; + } + *nop_out = nop; + return GAS_SUCCESS; +} + +static int build_alpha_hitlist(const gas_link_table_t *ta, + uint32_t a0, uint32_t a1, + gas_contract_ws_t *ws) +{ + uint32_t nhit = 0u; + + for (uint32_t ia0 = 0; ia0 < ta->nsrc; ia0++) { + const gas_link_entry_t *restrict ra = ta->link + + (uint64_t)ia0 * ta->nlink; + for (uint32_t ka = 0; ka < ta->nlink; ka++) { + const gas_link_entry_t *ea = ra + ka; + if (ea->addr >= a0 && ea->addr < a1) { + nhit++; + } + } + } + + if (reserve_alpha_hitlist(ws, nhit) != GAS_SUCCESS) { + return GAS_ERR_MEMORY; + } + + uint32_t out = 0u; + for (uint32_t ia0 = 0; ia0 < ta->nsrc; ia0++) { + const gas_link_entry_t *restrict ra = ta->link + + (uint64_t)ia0 * ta->nlink; + for (uint32_t ka = 0; ka < ta->nlink; ka++) { + const gas_link_entry_t *ea = ra + ka; + if (ea->addr < a0 || ea->addr >= a1) continue; + int32_t io = ws->op_map[gas_link_ia(ea)]; + if (io < 0) continue; + ws->alpha_hit[out].src_row = ia0; + ws->alpha_hit[out].rel_addr = ea->addr - a0; + ws->alpha_hit[out].op_index = (uint16_t)io; + ws->alpha_hit[out].sign = ea->sign; + ws->alpha_hit[out].pad = 0u; + out++; + } + } + ws->alpha_hit_n = out; + + return GAS_SUCCESS; +} + +static int collect_abba_beta_pairs_range(const gas_space_t *gas, + gas_sid_t adst, + uint32_t q0, uint32_t q1, + const gas_sid_t *restrict beta_sid, + gas_tid_t beta_off, uint32_t nbeta, + gas_contract_ws_t *ws, + uint32_t *npair_out) +{ + gas_row_t brow = gas->D.by_alpha_row[adst]; + uint32_t nrange; + uint32_t maxpair; + uint32_t q = q0; + uint32_t j = 0; + uint32_t npair = 0; + + if (q0 > brow.n) q0 = brow.n; + if (q1 > brow.n) q1 = brow.n; + if (q1 < q0) q1 = q0; + q = q0; + nrange = q1 - q0; + maxpair = nrange < nbeta ? nrange : nbeta; + + if (reserve_pairs(ws, maxpair) != GAS_SUCCESS) { + return GAS_ERR_MEMORY; + } + + while (q < q1 && j < nbeta) { + gas_bid_t bdst = brow.off + q; + gas_sid_t b1 = gas->block[bdst].sb; + gas_sid_t b2 = beta_sid[j]; + + if (b1 < b2) { + q++; + } else if (b1 > b2) { + j++; + } else { + ws->pair[npair].first = bdst; + ws->pair[npair].second = beta_off + j; + npair++; + q++; + j++; + } + } + + *npair_out = npair; + return GAS_SUCCESS; +} + +static gas_batch_status_t contract_abba_batch_alpha_tile( + const gas_space_t *gas, + const double *restrict gos, + const double *restrict ci0, + double *restrict ci1, + uint32_t nnorb, gas_bid_t bsrc, + gas_sid_t adst, gas_tid_t tida, + uint32_t q0, uint32_t q1, + uint32_t a0, uint32_t a1, + const gas_sid_t *restrict beta_sid, + gas_tid_t beta_off, uint32_t nbeta, + gas_contract_ws_t *ws) +{ + const gas_block_t *bs = gas->block + bsrc; + const gas_link_table_t *ta = gas->table + tida; + uint32_t nb0 = gas->sector_nstr[bs->sb]; + uint32_t na_full = gas->sector_nstr[adst]; + const double *src = ci0 + bs->offset; + uint32_t nop = 0; + uint32_t npair = 0; + if (a0 > na_full) a0 = na_full; + if (a1 > na_full) a1 = na_full; + if (a1 <= a0) return GAS_BATCH_DONE; + uint32_t nat = a1 - a0; + + if (collect_alpha_ops(ta, nnorb, ws, &nop) != GAS_SUCCESS) { + return GAS_BATCH_FALLBACK; + } + + if (nop == 0u) { + return GAS_BATCH_DONE; + } + if (collect_abba_beta_pairs_range(gas, adst, q0, q1, beta_sid, beta_off, + nbeta, ws, &npair) != GAS_SUCCESS) { + return GAS_BATCH_FALLBACK; + } + if (npair == 0u) { + return GAS_BATCH_DONE; + } + + uint32_t btile = abba_beta_tile_size(nb0, nop, nat); + if (btile == 0u) { + return GAS_BATCH_DONE; + } + + if (build_alpha_hitlist(ta, a0, a1, ws) != GAS_SUCCESS) { + return GAS_BATCH_FALLBACK; + } + + for (uint32_t b0 = 0; b0 < nb0; b0 += btile) { + uint32_t b1lim = nb0 - b0 < btile ? nb0 : b0 + btile; + /* Keep one invariant [op][full-btile][alpha] layout even for + * the final partial beta tile. Unused beta rows are never + * read, while stable strides make earlier zero state reusable. */ + uint32_t t1_nbt = btile; + size_t dense_t1_elems = (size_t)nop * t1_nbt * nat; + size_t t1_elems = dense_t1_elems; + if (t1_elems > ABBA_WORKSPACE_MAX_ELEMS || + reserve_arena(ws, t1_elems) != GAS_SUCCESS) { + return GAS_BATCH_FALLBACK; + } + if (b0 == 0u) { + memset(ws->arena, 0, t1_elems * sizeof(*ws->arena)); + ws->arena_zeroed_upto = t1_elems; + } + for (uint32_t ih = 0; ih < ws->alpha_hit_n; ih++) { + const gas_alpha_hit_t *h = ws->alpha_hit + ih; + const double *restrict src_row = src + (uint64_t)h->src_row * nb0; + double *t = ws->arena + + ((size_t)h->op_index * t1_nbt * nat + h->rel_addr); + double sgn = (double)h->sign; + for (uint32_t ib0 = b0; ib0 < b1lim; ib0++) { + t[(size_t)(ib0 - b0) * nat] = sgn * src_row[ib0]; + } + } + + for (uint32_t ip = 0; ip < npair; ip++) { + gas_bid_t bdst = (gas_bid_t)ws->pair[ip].first; + gas_tid_t tidb = ws->pair[ip].second; + const gas_link_table_t *tb = gas->table + tidb; + gas_sid_t bdst_sid = gas->block[bdst].sb; + uint32_t nb1 = gas->sector_nstr[bdst_sid]; + double *dst = ci1 + gas->block[bdst].offset; + size_t y_elems = (size_t)nb1 * nat; + int use_tile = t1_elems + y_elems <= ABBA_WORKSPACE_MAX_ELEMS && + reserve_arena(ws, t1_elems + y_elems) == GAS_SUCCESS; + double *t1 = ws->arena; + + if (use_tile) { + double *y = ws->arena + t1_elems; + memset(y, 0, y_elems * sizeof(*y)); + for (uint32_t ib0 = b0; ib0 < b1lim; ib0++) { + const gas_link_entry_t *rb = tb->link + + (uint64_t)ib0 * tb->nlink; + const double *tbase = t1 + + (size_t)(ib0 - b0) * nat; + uint32_t kb = 0u; + while (kb < tb->nlink) { + const gas_link_entry_t *eb0 = + rb + kb; + if (kb + 1u < tb->nlink) { + const gas_link_entry_t *eb1 = + rb + kb + 1u; + if (eb0->addr != eb1->addr) { + double *restrict y0 = y + + (size_t)eb0->addr * nat; + double *restrict y1 = y + + (size_t)eb1->addr * nat; + abba_fused_ops_pair2x4( + y0, y1, tbase, + (size_t)t1_nbt * nat, + gos, nnorb, + gas_link_ia(eb0), + gas_link_ia(eb1), + ws->op, nop, + (double)eb0->sign, + (double)eb1->sign, + nat); + kb += 2u; + continue; + } + } + double *restrict y0 = y + (size_t)eb0->addr * nat; + abba_fused_ops( + y0, tbase, (size_t)t1_nbt * nat, + gos, nnorb, gas_link_ia(eb0), + ws->op, nop, (double)eb0->sign, + nat); + kb++; + } + } + abba_transpose_add(dst, nb1, y, nat, nb1, a0); + continue; + } + + for (uint32_t ib0 = b0; ib0 < b1lim; ib0++) { + const gas_link_entry_t *rb = tb->link + + (uint64_t)ib0 * tb->nlink; + for (uint32_t kb = 0; kb < tb->nlink; kb++) { + const gas_link_entry_t *eb = rb + kb; + uint16_t opb = gas_link_ia(eb); + double sbeta = (double)eb->sign; + for (uint32_t io = 0; io < nop; io++) { + uint16_t opa = ws->op[io]; + double g = sbeta * + abba_gvalue(gos, nnorb, opb, opa); + const double *t = t1 + + ((size_t)io * t1_nbt + + (ib0 - b0)) * nat; + for (uint32_t ia = 0; ia < nat; ia++) { + dst[(uint64_t)(a0 + ia) * nb1 + + eb->addr] += g * t[ia]; + } + } + } + } + + } + } + return GAS_BATCH_DONE; +} + +static gas_batch_status_t contract_abba_batch_alpha_range( + const gas_space_t *gas, + const double *restrict gos, + const double *restrict ci0, + double *restrict ci1, + uint32_t nnorb, gas_bid_t bsrc, + gas_sid_t adst, gas_tid_t tida, + uint32_t q0, uint32_t q1, + const gas_sid_t *restrict beta_sid, + gas_tid_t beta_off, uint32_t nbeta, + gas_contract_ws_t *ws) +{ + return contract_abba_batch_alpha_tile( + gas, gos, ci0, ci1, nnorb, bsrc, adst, tida, + q0, q1, 0u, gas->sector_nstr[adst], + beta_sid, beta_off, nbeta, ws); +} + +static gas_batch_status_t contract_abba_batch_alpha( + const gas_space_t *gas, + const double *restrict gos, + const double *restrict ci0, + double *restrict ci1, + uint32_t nnorb, gas_bid_t bsrc, + gas_sid_t adst, gas_tid_t tida, + const gas_sid_t *restrict beta_sid, + gas_tid_t beta_off, uint32_t nbeta, + gas_contract_ws_t *ws) +{ + gas_row_t brow = gas->D.by_alpha_row[adst]; + return contract_abba_batch_alpha_range(gas, gos, ci0, ci1, nnorb, + bsrc, adst, tida, 0, brow.n, + beta_sid, beta_off, nbeta, ws); +} + +static void contract_abba_alpha_fallback( + const gas_space_t *gas, + const double *restrict eri, + const double *restrict ci0, + double *restrict ci1, + uint32_t nnorb, gas_bid_t bsrc, + gas_sid_t adst, gas_tid_t tida, + const gas_sid_t *restrict beta_sid, + gas_tid_t beta_off, uint32_t nbeta) +{ + for (uint32_t j = 0; j < nbeta; j++) { + gas_bid_t bdst = gas_find_block(gas, adst, beta_sid[j]); + if (bdst != GAS_INVALID_BID) { + contract_abba_path(gas, eri, ci0, ci1, nnorb, + bsrc, bdst, tida, beta_off + j); + } + } +} + +static void contract_abba_source_batch(const gas_space_t *gas, + const double *restrict eri, + const double *restrict gos, + const double *restrict ci0, + double *restrict ci1, + uint32_t nnorb, gas_bid_t bsrc, + gas_contract_ws_t *ws) +{ + const gas_block_t *blk = gas->block + bsrc; + gas_row_t ar = gas->T.row[blk->sa]; + gas_row_t br = gas->T.row[blk->sb]; + const gas_sid_t *as = gas->T.dst + ar.off; + const gas_sid_t *bs = gas->T.dst + br.off; + + for (uint32_t i = 0; i < ar.n; i++) { + gas_tid_t tida = ar.off + i; + gas_batch_status_t status = contract_abba_batch_alpha( + gas, gos, ci0, ci1, nnorb, bsrc, + as[i], tida, bs, br.off, br.n, ws); + if (status == GAS_BATCH_FALLBACK) { + contract_abba_alpha_fallback(gas, eri, ci0, ci1, nnorb, + bsrc, as[i], tida, + bs, br.off, br.n); + } + } +} + +static void contract_2e_block(const gas_space_t *gas, + const double *restrict eri, + const double *restrict gos, + const double *restrict ci0, + double *restrict ci1, + uint32_t nnorb, gas_bid_t b, + gas_contract_ws_t *ws) +{ + const gas_block_t *blk = gas->block + b; + gas_row_t ar = gas->T.row[blk->sa]; + gas_row_t br = gas->T.row[blk->sb]; + const gas_sid_t *as = gas->T.dst + ar.off; + const gas_sid_t *bs = gas->T.dst + br.off; + + if (gas->D.by_beta_row[blk->sb].n < aa_forward_scan_cost(gas, as, ar.n)) { + contract_aa_by_dst(gas, eri, ci0, ci1, nnorb, b); + } else { + contract_aa_forward(gas, eri, ci0, ci1, nnorb, b); + } + + if (gas->D.by_alpha_row[blk->sa].n < bb_forward_scan_cost(gas, bs, br.n)) { + contract_bb_by_dst(gas, eri, ci0, ci1, nnorb, b, ws); + } else { + contract_bb_forward(gas, eri, ci0, ci1, nnorb, b); + } + + contract_abba_source_batch(gas, eri, gos, ci0, ci1, nnorb, b, ws); +} + +/* ========================================================================== */ +/* 5. Destination alpha-sector traversal */ +/* ========================================================================== */ + +uint32_t fci_contract_gas_parallel_units(const gas_space_t *gas) +{ + uint32_t n = 0; + + if (gas == 0) { + return 0; + } + for (gas_sid_t s = 0; s < gas->nsector; s++) { + n += gas->D.by_alpha_row[s].n != 0; + } + return n; +} + +#ifdef _OPENMP +static uint64_t cost_add_sat(uint64_t a, uint64_t b) +{ + return UINT64_MAX - a < b ? UINT64_MAX : a + b; +} + +static uint64_t cost_mul_sat(uint64_t a, uint64_t b) +{ + return a != 0u && b > UINT64_MAX / a ? UINT64_MAX : a * b; +} + +static int append_omp_task(gas_omp_task_t **task, uint32_t *n, uint32_t *cap, + gas_sid_t adst, uint32_t q0, uint32_t q1, + uint32_t a0, uint32_t a1, uint64_t cost) +{ + if (*n == *cap) { + uint32_t next = *cap ? *cap * 2u : 1024u; + if (next < *cap) return GAS_ERR_MEMORY; + gas_omp_task_t *p = realloc( + *task, (size_t)next * sizeof(*p)); + if (p == 0) return GAS_ERR_MEMORY; + *task = p; + *cap = next; + } + (*task)[*n].adst = adst; + (*task)[*n].q0 = q0; + (*task)[*n].q1 = q1; + (*task)[*n].a0 = a0; + (*task)[*n].a1 = a1; + (*task)[*n].cost = cost; + (*n)++; + return GAS_SUCCESS; +} + +static uint64_t destination_block_cost(const gas_space_t *gas, gas_bid_t bdst) +{ + const gas_block_t *bd = gas->block + bdst; + const uint64_t na1 = gas->sector_nstr[bd->sa]; + const uint64_t nb1 = gas->sector_nstr[bd->sb]; + uint64_t cost = cost_mul_sat(na1, nb1); /* destination zero/write */ + + /* AA: all legal sources with the same beta sector, intersected by the + * directed intermediate alpha sector. */ + gas_row_t arr = gas->R.row[bd->sa]; + gas_row_t asrc_row = gas->D.by_beta_row[bd->sb]; + for (uint32_t q = 0; q < asrc_row.n; q++) { + gas_sid_t asrc = gas->D.by_beta_sid[asrc_row.off + q]; + gas_row_t tr = gas->T.row[asrc]; + uint32_t i = 0, j = 0; + while (i < tr.n && j < arr.n) { + gas_sid_t x = gas->T.dst[tr.off + i]; + gas_sid_t y = gas->R.src[arr.off + j]; + if (x < y) { + i++; + } else if (x > y) { + j++; + } else { + const gas_link_table_t *t1 = gas->table + tr.off + i; + const gas_link_table_t *t2 = gas->table + gas->R.tid[arr.off + j]; + uint64_t w = cost_mul_sat(t1->nsrc, t1->nlink); + w = cost_mul_sat(w, t2->nlink); + w = cost_mul_sat(w, nb1); + cost = cost_add_sat(cost, w); + i++; + j++; + } + } + } + + /* BB: same alpha sector, beta common-middle intersections. */ + gas_row_t brr = gas->R.row[bd->sb]; + gas_row_t same_alpha = gas->D.by_alpha_row[bd->sa]; + for (uint32_t q = 0; q < same_alpha.n; q++) { + gas_sid_t bsrc = gas->block[same_alpha.off + q].sb; + gas_row_t tr = gas->T.row[bsrc]; + uint32_t i = 0, j = 0; + while (i < tr.n && j < brr.n) { + gas_sid_t x = gas->T.dst[tr.off + i]; + gas_sid_t y = gas->R.src[brr.off + j]; + if (x < y) { + i++; + } else if (x > y) { + j++; + } else { + const gas_link_table_t *t1 = gas->table + tr.off + i; + const gas_link_table_t *t2 = gas->table + gas->R.tid[brr.off + j]; + uint64_t w = cost_mul_sat(t1->nsrc, t1->nlink); + w = cost_mul_sat(w, t2->nlink); + w = cost_mul_sat(w, na1); + cost = cost_add_sat(cost, w); + i++; + j++; + } + } + } + + /* AB/BA: direct alpha and beta sector transfers. The alpha gather is + * included in each destination weight; this deliberately errs on the + * side of splitting expensive, high-degree destination blocks. */ + for (uint32_t i = 0; i < arr.n; i++) { + gas_sid_t asrc = gas->R.src[arr.off + i]; + const gas_link_table_t *ta = gas->table + gas->R.tid[arr.off + i]; + gas_row_t srow = gas->D.by_alpha_row[asrc]; + for (uint32_t q = 0; q < srow.n; q++) { + gas_sid_t bsrc = gas->block[srow.off + q].sb; + gas_tid_t tidb = gas_find_table(gas, bsrc, bd->sb); + if (tidb != GAS_INVALID_TID) { + const gas_link_table_t *tb = gas->table + tidb; + uint64_t nb0 = gas->sector_nstr[bsrc]; + uint64_t gather = cost_mul_sat(ta->nsrc, ta->nlink); + gather = cost_mul_sat(gather, nb0); + uint64_t scatter = cost_mul_sat(tb->nsrc, tb->nlink); + scatter = cost_mul_sat(scatter, ta->nop ? ta->nop : 1u); + scatter = cost_mul_sat(scatter, na1); + cost = cost_add_sat(cost, gather); + cost = cost_add_sat(cost, scatter); + cost = cost_add_sat(cost, cost_mul_sat(na1, nb1)); + } + } + } + return cost; +} + +static gas_omp_task_t *build_omp_tasks(const gas_space_t *gas, uint32_t *ntask) +{ + gas_omp_task_t *task = 0; + uint32_t n = 0; + uint32_t cap = 0; + uint64_t total_cost = 0; + uint32_t nthread = (uint32_t)omp_get_max_threads(); + uint64_t *block_cost = malloc( + (size_t)gas->nblock * sizeof(*block_cost)); + + if (block_cost == 0) { + *ntask = 0; + return 0; + } + for (gas_bid_t b = 0; b < gas->nblock; b++) { + block_cost[b] = destination_block_cost(gas, b); + total_cost = cost_add_sat(total_cost, block_cost[b]); + } + uint64_t target = total_cost / + ((uint64_t)(nthread ? nthread : 1u) * OMP_TASK_FACTOR); + if (target < 1u) target = 1u; + + for (gas_sid_t s = 0; s < gas->nsector; s++) { + gas_row_t row = gas->D.by_alpha_row[s]; + uint32_t na = gas->sector_nstr[s]; + if (row.n == 0 || na == 0u) continue; + uint32_t q0 = 0; + uint64_t chunk = 0; + + for (uint32_t q = 0; q < row.n; q++) { + uint64_t w = block_cost[row.off + q]; + uint64_t want64 = (w + target - 1u) / target; + uint32_t by_rows = + (na + OMP_MIN_ALPHA_TILE - 1u) / + OMP_MIN_ALPHA_TILE; + uint32_t split = want64 > UINT32_MAX ? + UINT32_MAX : (uint32_t)want64; + if (split > OMP_MAX_ALPHA_SPLIT) { + split = OMP_MAX_ALPHA_SPLIT; + } + if (split > by_rows) split = by_rows; + + if (split > 1u) { + if (q > q0 && append_omp_task( + &task, &n, &cap, s, q0, q, 0u, na, + chunk) != GAS_SUCCESS) { + goto fail; + } + uint32_t a0 = 0; + for (uint32_t part = 0; part < split; part++) { + uint32_t left = na - a0; + uint32_t remain = split - part; + uint32_t width = + (left + remain - 1u) / remain; + uint32_t a1 = a0 + width; + uint64_t task_cost = + cost_mul_sat(w, width); + task_cost = na ? task_cost / na : w; + if (task_cost == 0u) task_cost = 1u; + if (append_omp_task( + &task, &n, &cap, s, + q, q + 1u, a0, a1, + task_cost) != GAS_SUCCESS) { + goto fail; + } + a0 = a1; + } + q0 = q + 1u; + chunk = 0; + continue; + } + + if (q > q0 && chunk != 0u && + cost_add_sat(chunk, w) > target) { + if (append_omp_task(&task, &n, &cap, + s, q0, q, 0u, na, + chunk) != GAS_SUCCESS) { + goto fail; + } + q0 = q; + chunk = 0; + } + chunk = cost_add_sat(chunk, w); + } + if (q0 < row.n && append_omp_task( + &task, &n, &cap, s, q0, row.n, 0u, na, + chunk) != GAS_SUCCESS) { + goto fail; + } + } + free(block_cost); + if (n > 1u) { + qsort(task, n, sizeof(*task), cmp_omp_task_desc); + } + *ntask = n; + return task; + +fail: + free(block_cost); + free(task); + *ntask = 0; + return 0; +} + +uint32_t fci_contract_gas_omp_task_count(const gas_space_t *gas) +{ + uint32_t n = 0; +#ifdef _OPENMP + gas_omp_task_t *task; + if (gas == 0) return 0; + task = build_omp_tasks(gas, &n); + free(task); +#else + (void)gas; +#endif + return n; +} +static void zero_alpha_range(const gas_space_t *gas, double *ci1, + gas_sid_t adst, uint32_t q0, uint32_t q1, + uint32_t a0, uint32_t a1) +{ + gas_row_t row = gas->D.by_alpha_row[adst]; + uint32_t na = gas->sector_nstr[adst]; + if (q0 > row.n) q0 = row.n; + if (q1 > row.n) q1 = row.n; + if (q1 < q0) q1 = q0; + if (a0 > na) a0 = na; + if (a1 > na) a1 = na; + if (a1 <= a0) return; + for (uint32_t q = q0; q < q1; q++) { + gas_bid_t b = row.off + q; + uint32_t nb = gas->sector_nstr[gas->block[b].sb]; + contract_zero(ci1 + gas->block[b].offset + (uint64_t)a0 * nb, + (a1 - a0) * nb); + } +} + +static void contract_aa_alpha_range(const gas_space_t *gas, + const double *restrict eri, + const double *restrict ci0, + double *restrict ci1, + uint32_t nnorb, gas_sid_t adst, + uint32_t q0, uint32_t q1, + uint32_t a0, uint32_t a1, + gas_contract_ws_t *ws) +{ + gas_row_t drow = gas->D.by_alpha_row[adst]; + gas_row_t rr = gas->R.row[adst]; + const gas_sid_t *middle_rev = gas->R.src + rr.off; + const gas_tid_t *second_tid = gas->R.tid + rr.off; + uint32_t na = gas->sector_nstr[adst]; + + if (q0 > drow.n) q0 = drow.n; + if (q1 > drow.n) q1 = drow.n; + if (q1 < q0) q1 = q0; + if (a0 > na) a0 = na; + if (a1 > na) a1 = na; + + for (uint32_t d = q0; d < q1; d++) { + gas_bid_t bdst = drow.off + d; + gas_sid_t beta = gas->block[bdst].sb; + gas_row_t srow = gas->D.by_beta_row[beta]; + const gas_sid_t *alpha_src = gas->D.by_beta_sid + srow.off; + const gas_bid_t *source_block = gas->D.by_beta_bid + srow.off; + + for (uint32_t q = 0; q < srow.n; q++) { + gas_row_t tr = gas->T.row[alpha_src[q]]; + const gas_sid_t *middle_fwd = gas->T.dst + tr.off; + uint32_t i = 0; + uint32_t j = 0; + + while (i < tr.n && j < rr.n) { + if (middle_fwd[i] < middle_rev[j]) { + i++; + } else if (middle_fwd[i] > middle_rev[j]) { + j++; + } else { + if (ws == 0 || + contract_aa_path_rows_accum( + gas, eri, ci0, ci1, nnorb, + source_block[q], bdst, + tr.off + i, second_tid[j], + a0, a1, ws) != GAS_SUCCESS) { + contract_aa_path_rows( + gas, eri, ci0, ci1, nnorb, + source_block[q], bdst, + tr.off + i, second_tid[j], a0, a1); + } + i++; + j++; + } + } + } + } +} + +static void contract_bb_alpha_range(const gas_space_t *gas, + const double *restrict eri, + const double *restrict ci0, + double *restrict ci1, + uint32_t nnorb, gas_sid_t alpha, + uint32_t q0, uint32_t q1, + uint32_t a0, uint32_t a1, + gas_contract_ws_t *ws) +{ + gas_row_t row = gas->D.by_alpha_row[alpha]; + uint32_t na = gas->sector_nstr[alpha]; + if (q0 > row.n) q0 = row.n; + if (q1 > row.n) q1 = row.n; + if (q1 < q0) q1 = q0; + if (a0 > na) a0 = na; + if (a1 > na) a1 = na; + + for (uint32_t d = q0; d < q1; d++) { + gas_bid_t bdst = row.off + d; + gas_sid_t b2 = gas->block[bdst].sb; + gas_row_t rr = gas->R.row[b2]; + const gas_sid_t *rs = gas->R.src + rr.off; + const gas_tid_t *rt = gas->R.tid + rr.off; + uint32_t nb2 = gas->sector_nstr[b2]; + + for (uint32_t q = 0; q < row.n; q++) { + gas_bid_t bsrc = row.off + q; + gas_sid_t b0 = gas->block[bsrc].sb; + gas_row_t br = gas->T.row[b0]; + const gas_sid_t *bs = gas->T.dst + br.off; + uint32_t nb0 = gas->sector_nstr[b0]; + uint32_t maxpair = br.n < rr.n ? br.n : rr.n; + uint32_t i = 0; + uint32_t j = 0; + uint32_t nat = a1 > a0 ? a1 - a0 : 0u; + + if (use_bb_transpose(nat, nb0, nb2, maxpair) && + reserve_pairs(ws, maxpair) == GAS_SUCCESS) { + uint32_t npair = 0; + while (i < br.n && j < rr.n) { + if (bs[i] < rs[j]) { + i++; + } else if (bs[i] > rs[j]) { + j++; + } else { + ws->pair[npair].first = br.off + i; + ws->pair[npair].second = rt[j]; + npair++; + i++; + j++; + } + } + if (npair != 0) { + contract_bb_blockpair_rows( + gas, eri, ci0, ci1, nnorb, + bsrc, bdst, ws->pair, npair, ws, + a0, a1); + } + continue; + } + + while (i < br.n && j < rr.n) { + if (bs[i] < rs[j]) { + i++; + } else if (bs[i] > rs[j]) { + j++; + } else { + contract_bb_path_rows( + gas, eri, ci0, ci1, nnorb, + bsrc, bdst, br.off + i, rt[j], a0, a1); + i++; + j++; + } + } + } + } +} + +static void contract_abba_alpha_fallback_range( + const gas_space_t *gas, + const double *restrict eri, + const double *restrict ci0, + double *restrict ci1, + uint32_t nnorb, gas_bid_t bsrc, + gas_sid_t adst, gas_tid_t tida, + uint32_t q0, uint32_t q1, + uint32_t a0, uint32_t a1, + const gas_sid_t *restrict beta_sid, + gas_tid_t beta_off, uint32_t nbeta) +{ + gas_row_t brow = gas->D.by_alpha_row[adst]; + if (q0 > brow.n) q0 = brow.n; + if (q1 > brow.n) q1 = brow.n; + if (q1 < q0) q1 = q0; + for (uint32_t q = q0; q < q1; q++) { + gas_bid_t bdst = brow.off + q; + gas_sid_t b1 = gas->block[bdst].sb; + gas_tid_t tidb = find_table_in_row(b1, beta_sid, beta_off, nbeta); + if (tidb != GAS_INVALID_TID) { + contract_abba_path_rows(gas, eri, ci0, ci1, nnorb, + bsrc, bdst, tida, tidb, a0, a1); + } + } +} + +static void contract_abba_alpha_range(const gas_space_t *gas, + const double *restrict eri, + const double *restrict gos, + const double *restrict ci0, + double *restrict ci1, + uint32_t nnorb, gas_sid_t adst, + uint32_t q0, uint32_t q1, + uint32_t a0, uint32_t a1, + gas_contract_ws_t *ws) +{ + gas_row_t ar = gas->R.row[adst]; + const gas_sid_t *alpha_src = gas->R.src + ar.off; + const gas_tid_t *alpha_tid = gas->R.tid + ar.off; + gas_row_t drow = gas->D.by_alpha_row[adst]; + uint32_t na = gas->sector_nstr[adst]; + + if (q0 > drow.n) q0 = drow.n; + if (q1 > drow.n) q1 = drow.n; + if (q1 < q0) q1 = q0; + if (a0 > na) a0 = na; + if (a1 > na) a1 = na; + + for (uint32_t i = 0; i < ar.n; i++) { + gas_row_t srow = gas->D.by_alpha_row[alpha_src[i]]; + + for (uint32_t q = 0; q < srow.n; q++) { + gas_bid_t bsrc = srow.off + q; + gas_sid_t beta0 = gas->block[bsrc].sb; + gas_row_t br = gas->T.row[beta0]; + const gas_sid_t *beta_dst = gas->T.dst + br.off; + gas_batch_status_t status; + + status = contract_abba_batch_alpha_tile( + gas, gos, ci0, ci1, nnorb, bsrc, + adst, alpha_tid[i], q0, q1, a0, a1, + beta_dst, br.off, br.n, ws); + + if (status == GAS_BATCH_FALLBACK) { + contract_abba_alpha_fallback_range( + gas, eri, ci0, ci1, nnorb, bsrc, + adst, alpha_tid[i], q0, q1, a0, a1, + beta_dst, br.off, br.n); + } + } + } +} + +static void contract_2e_alpha_range(const gas_space_t *gas, + const double *restrict eri, + const double *restrict gos, + const double *restrict ci0, + double *restrict ci1, + uint32_t nnorb, gas_sid_t adst, + uint32_t q0, uint32_t q1, + uint32_t a0, uint32_t a1, + gas_contract_ws_t *ws) +{ + contract_aa_alpha_range(gas, eri, ci0, ci1, nnorb, + adst, q0, q1, a0, a1, ws); + contract_bb_alpha_range(gas, eri, ci0, ci1, nnorb, + adst, q0, q1, a0, a1, ws); + contract_abba_alpha_range(gas, eri, gos, ci0, ci1, nnorb, + adst, q0, q1, a0, a1, ws); +} + +#endif + +#ifndef _OPENMP +uint32_t fci_contract_gas_omp_task_count(const gas_space_t *gas) +{ + (void)gas; + return 1u; +} +#endif + +/* ========================================================================== */ +/* 6. Reusable contraction plan */ +/* ========================================================================== */ + +static int contract_plan_config_matches(const gas_contract_plan_t *plan) +{ + return plan->abba_t1_target_bytes == gas_abba_t1_target_bytes; +} + +static int contract_plan_resize_workspaces(gas_contract_plan_t *plan, + uint32_t need) +{ + if (need <= plan->nworkspace) { + return GAS_SUCCESS; + } + gas_contract_ws_t *p = realloc( + plan->workspace, (size_t)need * sizeof(*p)); + if (p == 0) { + return GAS_ERR_MEMORY; + } + plan->workspace = p; + for (uint32_t i = plan->nworkspace; i < need; i++) { + contract_ws_init(plan->workspace + i); + } + plan->nworkspace = need; + return GAS_SUCCESS; +} + +#ifdef _OPENMP +static int contract_plan_rebuild_tasks(gas_contract_plan_t *plan) +{ + uint32_t nthread = (uint32_t)omp_get_max_threads(); + if (nthread == 0u) nthread = 1u; + if (contract_plan_resize_workspaces(plan, nthread) != GAS_SUCCESS) { + return GAS_ERR_MEMORY; + } + if (plan->task != 0 && plan->task_threads == nthread) { + return GAS_SUCCESS; + } + + free(plan->task); + plan->task = build_omp_tasks(plan->gas, &plan->ntask); + plan->task_threads = nthread; + if (plan->task == 0 && plan->gas->nblock != 0u) { + plan->ntask = 0u; + return GAS_ERR_MEMORY; + } + return GAS_SUCCESS; +} + +static void contract_gas_2e_parallel_plan(gas_contract_plan_t *plan, + const double *ci0, double *ci1, + uint32_t nnorb) +{ + const gas_space_t *gas = plan->gas; + const double *abba_eri = plan->gos; + +#pragma omp parallel + { + uint32_t tid = (uint32_t)omp_get_thread_num(); + gas_contract_ws_t *ws = plan->workspace + tid; + +#pragma omp for schedule(dynamic, 1) + for (int64_t tt = 0; tt < (int64_t)plan->ntask; tt++) { + gas_sid_t adst = plan->task[tt].adst; + uint32_t q0 = plan->task[tt].q0; + uint32_t q1 = plan->task[tt].q1; + uint32_t a0 = plan->task[tt].a0; + uint32_t a1 = plan->task[tt].a1; + zero_alpha_range(gas, ci1, adst, q0, q1, a0, a1); + contract_2e_alpha_range(gas, plan->eri, abba_eri, + ci0, ci1, nnorb, adst, + q0, q1, a0, a1, ws); + } + } +} +#endif + +int fci_contract_gas_plan_create(gas_contract_plan_t **out, + const gas_space_t *gas, + const double *eri, + const double *gos) +{ + if (out == 0 || gas == 0 || eri == 0 || gos == 0 || + gas->link_format != GAS_LINK_COMPRESSED) { + return GAS_ERR_INVALID; + } + *out = 0; + + gas_contract_plan_t *plan = calloc(1, sizeof(*plan)); + if (plan == 0) { + return GAS_ERR_MEMORY; + } + plan->gas = gas; + plan->eri = eri; + plan->gos = gos; + plan->abba_t1_target_bytes = gas_abba_t1_target_bytes; + +#ifdef _OPENMP + if (contract_plan_rebuild_tasks(plan) != GAS_SUCCESS) { + fci_contract_gas_plan_free(plan); + return GAS_ERR_MEMORY; + } +#else + if (contract_plan_resize_workspaces(plan, 1u) != GAS_SUCCESS) { + fci_contract_gas_plan_free(plan); + return GAS_ERR_MEMORY; + } +#endif + *out = plan; + return GAS_SUCCESS; +} + +int fci_contract_gas_plan_execute(gas_contract_plan_t *plan, + const double *ci0, double *ci1) +{ + if (plan == 0 || ci0 == 0 || ci1 == 0 || ci0 == ci1 || + !contract_plan_config_matches(plan)) { + return GAS_ERR_INVALID; + } + const gas_space_t *gas = plan->gas; + uint32_t nnorb = (uint32_t)(gas->norb_tot * (gas->norb_tot + 1) / 2); + const double *abba_eri = plan->gos; + +#ifdef _OPENMP + if (contract_plan_rebuild_tasks(plan) != GAS_SUCCESS) { + return GAS_ERR_MEMORY; + } + if (omp_get_max_threads() > 1 && plan->ntask > 1u) { + contract_gas_2e_parallel_plan(plan, ci0, ci1, nnorb); + return GAS_SUCCESS; + } +#endif + if (contract_plan_resize_workspaces(plan, 1u) != GAS_SUCCESS) { + return GAS_ERR_MEMORY; + } + contract_zero(ci1, gas->ndet); + for (gas_bid_t b = 0; b < gas->nblock; b++) { + contract_2e_block(gas, plan->eri, abba_eri, + ci0, ci1, nnorb, b, plan->workspace); + } + return GAS_SUCCESS; +} + +void fci_contract_gas_plan_free(gas_contract_plan_t *plan) +{ + if (plan == 0) { + return; + } + for (uint32_t i = 0; i < plan->nworkspace; i++) { + contract_ws_free(plan->workspace + i); + } + free(plan->workspace); +#ifdef _OPENMP + free(plan->task); +#endif + free(plan); +} + +uint32_t fci_contract_gas_plan_task_count(const gas_contract_plan_t *plan) +{ + if (plan == 0) return 0u; +#ifdef _OPENMP + return plan->ntask; +#else + return 1u; +#endif +} + +uint64_t fci_contract_gas_plan_workspace_bytes(const gas_contract_plan_t *plan) +{ + uint64_t n = 0; + if (plan == 0) return 0u; + for (uint32_t i = 0; i < plan->nworkspace; i++) { + n += contract_ws_bytes(plan->workspace + i); + } + return n; +} diff --git a/pyscf/lib/mcscf/fci_gas.h b/pyscf/lib/mcscf/fci_gas.h new file mode 100644 index 000000000..67cf99dc6 --- /dev/null +++ b/pyscf/lib/mcscf/fci_gas.h @@ -0,0 +1,294 @@ +/* Copyright 2026 The PySCF Developers. All Rights Reserved. + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. + + * + * Author: Yi Deng + */ + +#ifndef FCI_GAS_H +#define FCI_GAS_H + +#include + +enum { + GAS_MAX_ORB = 63, + GAS_MAX_LOCAL_ORB = 31, + GAS_MAX_NGAS = 15 +}; + +#define GAS_INVALID_SID UINT16_MAX +#define GAS_INVALID_BID UINT32_MAX +#define GAS_INVALID_TID UINT32_MAX + +typedef uint16_t gas_sid_t; +typedef uint32_t gas_bid_t; +typedef uint32_t gas_tid_t; + +typedef enum { + GAS_SUCCESS = 0, + GAS_ERR_INVALID = -1, + GAS_ERR_MEMORY = -2 +} gas_status_t; + +typedef enum { + GAS_LINK_RAW = 0, + GAS_LINK_COMPRESSED = 1 +} gas_link_format_t; + +/* + * A legal determinant block couples one alpha sector to one beta sector. + * Determinants occupy one contiguous alpha-major rectangle in the flattened + * CI vector: offset + ia * nstr(beta) + ib. + */ +typedef struct { + uint32_t offset; + gas_sid_t sa; + gas_sid_t sb; +} gas_block_t; + +/* Slice [off, off+n) in a packed index array. */ +typedef struct { + uint32_t off; + uint32_t n; +} gas_row_t; + +/* + * Index D for the legal block set. Alpha rows address gas->block directly; + * beta rows use the packed sector-ID and block-ID arrays because canonical + * block order is alpha-major. + */ +typedef struct { + gas_row_t *by_alpha_row; + gas_row_t *by_beta_row; + gas_sid_t *by_beta_sid; + gas_bid_t *by_beta_bid; +} gas_block_index_t; + +/* Forward table index T: source sector -> packed destination sectors. */ +typedef struct { + gas_row_t *row; + gas_sid_t *dst; +} gas_table_index_t; + +/* + * Reverse table index R: destination sector -> packed (source sector, + * link-table ID) entries. + */ +typedef struct { + gas_row_t *row; + gas_sid_t *src; + gas_tid_t *tid; +} gas_table_rev_index_t; + +/* + * One spin-string excitation. addr is the destination string address and + * sign is the fermionic phase. In a raw table, op packs creation orbital p + * in the low byte and destruction orbital q in the high byte. In a + * compressed table, op is the lower-triangular pair index for (p,q). + */ +typedef struct { + uint32_t addr; + uint16_t op; + int8_t sign; + uint8_t padding; /* completes the 8-byte entry */ +} gas_link_entry_t; + +/* + * Directed source-sector -> destination-sector link table. active_op is the + * sorted list of triangular operator-pair indices present in the table. + */ +typedef struct { + gas_link_entry_t *link; + uint16_t *active_op; + uint32_t nsrc; + uint32_t nlink; + uint16_t nop; +} gas_link_table_t; + +/* + * Reusable H*c execution plan. A plan caches Hamiltonian-dependent data, + * OpenMP tasks and per-thread workspaces across Davidson iterations. + * The referenced gas space, eri array and gos array must outlive the plan. + * A single plan must not be executed concurrently by multiple host threads. + */ +typedef struct gas_contract_plan gas_contract_plan_t; +typedef struct gas_rdm_plan gas_rdm_plan_t; + +typedef struct { + int ngas; + int norb_tot; + int *norb; + int *start; + int na; + int nb; + uint32_t ndet; + uint16_t nsector; + uint16_t sector_padding; /* 2 + 2 bytes align sector_nstr */ + uint32_t *sector_nstr; + uint8_t *sector_occ; + uint32_t *sector_stride; + + uint32_t nblock; + gas_block_t *block; + gas_block_index_t D; /* legal determinant-block index */ + + uint32_t ntable; + uint8_t link_format; + uint8_t link_format_padding[3]; /* 1 + 3 bytes align table */ + gas_link_table_t *table; + gas_table_index_t T; /* forward link-table index */ + gas_table_rev_index_t R; /* reverse link-table index */ +} gas_space_t; + +typedef struct { + uint64_t metadata; + uint64_t sector; + uint64_t block; + uint64_t block_index; + uint64_t link_table; + uint64_t link_table_index; + uint64_t total; +} gas_memory_report_t; + +_Static_assert(sizeof(void *) == 8, "fci_gas requires a 64-bit platform"); +_Static_assert(sizeof(gas_link_entry_t) == 8, "gas_link_entry_t must be 8 bytes"); +_Static_assert(sizeof(gas_block_t) == 8, "gas_block_t must be 8 bytes"); +_Static_assert(sizeof(gas_row_t) == 8, "gas_row_t must be 8 bytes"); +_Static_assert(sizeof(gas_link_table_t) == 32, "gas_link_table_t must be 32 bytes"); +_Static_assert(sizeof(gas_space_t) == 168, "gas_space_t must be 168 bytes"); + +static inline uint16_t gas_link_pair_index(int p, int q) +{ + uint32_t hi = p > q ? (uint32_t)p : (uint32_t)q; + uint32_t lo = p > q ? (uint32_t)q : (uint32_t)p; + return (uint16_t)(hi * (hi + 1u) / 2u + lo); +} + +static inline void gas_link_set_raw(gas_link_entry_t *e, int p, int q) +{ + uint32_t packed = (uint32_t)p | ((uint32_t)q << 8); + e->op = (uint16_t)packed; +} + +static inline uint8_t gas_link_cre(const gas_link_entry_t *e) +{ + return (uint8_t)(e->op & 0xffu); +} + +static inline uint8_t gas_link_des(const gas_link_entry_t *e) +{ + return (uint8_t)(e->op >> 8); +} + +static inline void gas_link_set_ia(gas_link_entry_t *e, uint16_t ia) +{ + e->op = ia; +} + +static inline uint16_t gas_link_ia(const gas_link_entry_t *e) +{ + return e->op; +} + +uint32_t gas_str2addr_sector(const gas_space_t *gas, gas_sid_t s, uint64_t str); +uint64_t gas_addr2str_sector(const gas_space_t *gas, gas_sid_t s, uint32_t addr); +uint32_t gas_block_ndet(const gas_space_t *gas, gas_bid_t b); +uint32_t gas_det_addr(const gas_space_t *gas, gas_bid_t b, uint32_t ia, uint32_t ib); +gas_bid_t gas_find_block(const gas_space_t *gas, gas_sid_t sa, gas_sid_t sb); +gas_tid_t gas_find_table(const gas_space_t *gas, gas_sid_t src, gas_sid_t dst); +const gas_link_table_t *gas_get_link_table(const gas_space_t *gas, + gas_sid_t src, gas_sid_t dst); + +int gas_space_from_blocks(gas_space_t *gas, int ngas, const int *norb, + int na, int nb, int nblock, const int *block_occ); +int gas_space_compress_links(gas_space_t *gas); +int gas_space_links_are_compressed(const gas_space_t *gas); +void gas_space_free(gas_space_t *gas); + +void gas_memory_report(const gas_space_t *gas, gas_memory_report_t *report); +uint64_t gas_memory_bytes(const gas_space_t *gas); + +/* + * Diagonal of the physical GAS Hamiltonian in the determinant order defined + * by gas->block. h1e is a row-major (norb,norb) matrix. eri is the original + * (unabsorbed) chemist-notation two-electron integral tensor restored to the + * four-fold pair-matrix layout (npair,npair), where + * npair = norb * (norb + 1) / 2. Raw and compressed link spaces are accepted. + */ +int fci_make_hdiag_gas(const gas_space_t *gas, + const double *h1e, const double *eri, + double *hdiag); + +/* ABBA workspace target; the production default is 64 MiB. */ +void fci_contract_gas_set_abba_t1_bytes(uint64_t bytes); +uint32_t fci_contract_gas_omp_task_count(const gas_space_t *gas); + +/* + * Contract the spin-free Hamiltonian after absorbing h1e into the + * two-electron tensor. eri_tril is a row-major (npair,npair) array in + * PySCF's four-fold pair-matrix layout; each orbital pair uses + * gas_link_pair_index. ci0 and ci1 follow gas->block determinant order. + * Contraction requires compressed link tables. + */ +int fci_contract_gas_plan_create(gas_contract_plan_t **out, + const gas_space_t *gas, + const double *eri_tril, + const double *gos); +int fci_contract_gas_plan_execute(gas_contract_plan_t *plan, + const double *ci0, double *ci1); +void fci_contract_gas_plan_free(gas_contract_plan_t *plan); +uint32_t fci_contract_gas_plan_task_count(const gas_contract_plan_t *plan); +uint64_t fci_contract_gas_plan_workspace_bytes(const gas_contract_plan_t *plan); +uint32_t fci_contract_gas_parallel_units(const gas_space_t *gas); + +/* + * GAS reduced density matrices. These entry points require raw link tables; + * rebuild the gas space after the Davidson contraction stage rather than + * passing a space that has been compressed for Hamiltonian contraction. + * + * All arrays are row-major. The conventions match PySCF's public RDMs: + * + * dm1[p,q] = + * dm2[p,q,r,s] = + * + * dm2ab has alpha spin on the (p,q) operator pair and beta spin on + * the (r,s) pair. Spin-free RDMs are assembled from these arrays in + * the Python layer. + */ + +/* + * Target size for each thread's tiled opposite-spin T1 workspace + * (16 MiB production default). + */ +void fci_rdm_gas_set_abba_t1_bytes(uint64_t bytes); + +/* + * Reusable RDM plan. The referenced raw-link gas space must outlive the plan. + * A plan caches destination-alpha range tasks and grows thread-private RDM + * workspaces as needed. It must not be executed concurrently by multiple + * host threads. Changing omp_get_max_threads() between calls is supported. + */ +int fci_rdm_gas_plan_create(gas_rdm_plan_t **out, const gas_space_t *gas); +int fci_rdm_gas_plan_make_rdm1s(gas_rdm_plan_t *plan, + const double *bra, const double *ket, + double *dm1a, double *dm1b); +int fci_rdm_gas_plan_make_rdm12s(gas_rdm_plan_t *plan, + const double *bra, const double *ket, + double *dm1a, double *dm1b, + double *dm2aa, double *dm2ab, double *dm2bb); +void fci_rdm_gas_plan_free(gas_rdm_plan_t *plan); +uint32_t fci_rdm_gas_plan_task_count(const gas_rdm_plan_t *plan); +uint64_t fci_rdm_gas_plan_workspace_bytes(const gas_rdm_plan_t *plan); + +#endif /* FCI_GAS_H */ diff --git a/pyscf/lib/mcscf/fci_hdiag_gas.c b/pyscf/lib/mcscf/fci_hdiag_gas.c new file mode 100644 index 000000000..413c7fb45 --- /dev/null +++ b/pyscf/lib/mcscf/fci_hdiag_gas.c @@ -0,0 +1,709 @@ +/* Copyright 2026 The PySCF Developers. All Rights Reserved. + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. + + * + * Author: Yi Deng + */ + +#include "fci_gas.h" + +#include +#include + +#if defined(__AVX2__) +#include +#endif + +#ifdef _OPENMP +#include +#endif + +enum { + HDIAG_POTENTIAL_BYTES = 64u * 1024u * 1024u, + HDIAG_TASK_FACTOR = 32u, + HDIAG_CACHE_CHUNK = 128u +}; + +typedef struct { + uint64_t *mask; + double *spin_diag; + double *beta_potential; +} hdiag_sector_cache_t; + +typedef struct { + gas_bid_t block; + uint32_t alpha_begin; + uint32_t alpha_end; +} hdiag_task_t; + +typedef struct { + gas_sid_t sector; + uint16_t orbital; + uint32_t begin; + uint32_t end; +} hdiag_cache_task_t; + +typedef struct { + hdiag_cache_task_t *mask_task; + hdiag_cache_task_t *potential_task; + uint32_t mask_task_count; + uint32_t potential_task_count; +} hdiag_cache_work_t; + +static inline int ctz64(uint64_t x) +{ + return __builtin_ctzll(x); +} + +static inline double coulomb_value(const double *restrict eri, + uint32_t npair, int p, int q) +{ + uint16_t pp = gas_link_pair_index(p, p); + uint16_t qq = gas_link_pair_index(q, q); + return eri[(size_t)pp * npair + qq]; +} + +static inline double exchange_value(const double *restrict eri, + uint32_t npair, int p, int q) +{ + uint16_t pq = gas_link_pair_index(p, q); + return eri[(size_t)pq * npair + pq]; +} + +static int memory_ranges_overlap(const void *left, size_t left_bytes, + const void *right, size_t right_bytes) +{ + uintptr_t a = (uintptr_t)left; + uintptr_t b = (uintptr_t)right; + if (left_bytes > UINTPTR_MAX - a || right_bytes > UINTPTR_MAX - b) { + return 1; + } + return a < b + right_bytes && b < a + left_bytes; +} + +static int hdiag_validate(const gas_space_t *gas, const double *h1e, + const double *eri, const double *hdiag) +{ + if (gas == 0 || h1e == 0 || eri == 0 || hdiag == 0) { + return GAS_ERR_INVALID; + } + if (gas->ngas <= 0 || gas->ngas > GAS_MAX_NGAS || + gas->norb_tot <= 0 || gas->norb_tot > GAS_MAX_ORB || + gas->nsector == 0u || gas->nblock == 0u || gas->ndet == 0u || + gas->norb == 0 || gas->start == 0 || + gas->sector_nstr == 0 || gas->sector_occ == 0 || + gas->sector_stride == 0 || gas->block == 0) { + return GAS_ERR_INVALID; + } + if (gas->link_format != GAS_LINK_RAW && + gas->link_format != GAS_LINK_COMPRESSED) { + return GAS_ERR_INVALID; + } + + size_t norb = (size_t)gas->norb_tot; + size_t npair = norb * (norb + 1u) / 2u; + size_t output_bytes = (size_t)gas->ndet * sizeof(*hdiag); + if (memory_ranges_overlap(hdiag, output_bytes, + h1e, norb * norb * sizeof(*h1e)) || + memory_ranges_overlap(hdiag, output_bytes, + eri, npair * npair * sizeof(*eri))) { + return GAS_ERR_INVALID; + } + return GAS_SUCCESS; +} + +static void free_sector_cache(hdiag_sector_cache_t *cache, uint32_t nsector) +{ + if (cache == 0) return; + for (uint32_t s = 0; s < nsector; s++) { + free(cache[s].mask); + free(cache[s].spin_diag); + free(cache[s].beta_potential); + } + free(cache); +} + +static void free_cache_work(hdiag_cache_work_t *work) +{ + free(work->mask_task); + free(work->potential_task); + memset(work, 0, sizeof(*work)); +} + +static double spin_diagonal(uint64_t mask, const double *restrict h1e, + const double *restrict eri, + uint32_t norb, uint32_t npair) +{ + double value = 0.0; + uint64_t remaining = mask; + while (remaining != 0u) { + int p = ctz64(remaining); + remaining &= remaining - 1u; + value += h1e[(size_t)p * norb + (uint32_t)p]; + + uint64_t partners = remaining; + while (partners != 0u) { + int q = ctz64(partners); + partners &= partners - 1u; + value += coulomb_value(eri, npair, p, q) + - exchange_value(eri, npair, p, q); + } + } + return value; +} + +static int prepare_sector_cache(const gas_space_t *gas, + hdiag_sector_cache_t **out, + hdiag_cache_work_t *work) +{ + uint32_t nsector = gas->nsector; + uint32_t norb = (uint32_t)gas->norb_tot; + unsigned char *used = calloc((size_t)nsector, sizeof(*used)); + hdiag_sector_cache_t *cache = calloc((size_t)nsector, sizeof(*cache)); + uint64_t nmask64 = 0u; + uint64_t npotential64 = 0u; + int status = GAS_ERR_MEMORY; + if (used == 0 || cache == 0) goto done; + + for (gas_bid_t b = 0; b < gas->nblock; b++) { + used[gas->block[b].sa] |= 1u; + used[gas->block[b].sb] |= 2u; + } + + size_t budget = (size_t)HDIAG_POTENTIAL_BYTES; + for (uint32_t s = 0; s < nsector; s++) { + if (used[s] == 0u) continue; + uint32_t nstr = gas->sector_nstr[s]; + cache[s].mask = malloc((size_t)nstr * sizeof(*cache[s].mask)); + if (cache[s].mask == 0) goto done; + cache[s].spin_diag = malloc( + (size_t)nstr * sizeof(*cache[s].spin_diag)); + if (cache[s].spin_diag == 0) goto done; + + uint64_t chunks = ((uint64_t)nstr + HDIAG_CACHE_CHUNK - 1u) / + HDIAG_CACHE_CHUNK; + nmask64 += chunks; + if ((used[s] & 2u) == 0u || + nstr > SIZE_MAX / norb / sizeof(double)) { + continue; + } + + size_t bytes = (size_t)norb * nstr * sizeof(double); + if (bytes > budget) continue; + cache[s].beta_potential = malloc(bytes); + if (cache[s].beta_potential == 0) continue; + budget -= bytes; + npotential64 += chunks * norb; + } + + if (nmask64 == 0u || nmask64 > UINT32_MAX || + npotential64 > UINT32_MAX || + nmask64 > SIZE_MAX / sizeof(*work->mask_task) || + npotential64 > SIZE_MAX / sizeof(*work->potential_task)) { + goto done; + } + + work->mask_task = malloc((size_t)nmask64 * sizeof(*work->mask_task)); + if (npotential64 != 0u) { + work->potential_task = malloc( + (size_t)npotential64 * sizeof(*work->potential_task)); + } + if (work->mask_task == 0 || + (npotential64 != 0u && work->potential_task == 0)) { + goto done; + } + + uint32_t nmask = 0u; + uint32_t npotential = 0u; + for (uint32_t s = 0; s < nsector; s++) { + if (used[s] == 0u) continue; + uint32_t nstr = gas->sector_nstr[s]; + for (uint32_t begin = 0u; begin < nstr; + begin += HDIAG_CACHE_CHUNK) { + uint32_t end = nstr - begin < HDIAG_CACHE_CHUNK ? + nstr : begin + HDIAG_CACHE_CHUNK; + work->mask_task[nmask++] = (hdiag_cache_task_t){ + (gas_sid_t)s, 0u, begin, end}; + if (cache[s].beta_potential == 0) continue; + for (uint32_t p = 0u; p < norb; p++) { + work->potential_task[npotential++] = + (hdiag_cache_task_t){ + (gas_sid_t)s, (uint16_t)p, + begin, end}; + } + } + } + + work->mask_task_count = nmask; + work->potential_task_count = npotential; + *out = cache; + cache = 0; + status = GAS_SUCCESS; + +done: + free(used); + if (status != GAS_SUCCESS) free_cache_work(work); + free_sector_cache(cache, nsector); + return status; +} + +static void execute_mask_cache_task(const gas_space_t *gas, + const double *restrict h1e, + const double *restrict eri, + hdiag_sector_cache_t *cache, + const hdiag_cache_task_t *task) +{ + uint32_t norb = (uint32_t)gas->norb_tot; + uint32_t npair = norb * (norb + 1u) / 2u; + hdiag_sector_cache_t *sector = cache + task->sector; + for (uint32_t i = task->begin; i < task->end; i++) { + uint64_t mask = gas_addr2str_sector(gas, task->sector, i); + sector->mask[i] = mask; + sector->spin_diag[i] = spin_diagonal( + mask, h1e, eri, norb, npair); + } +} + +static void execute_potential_cache_task(const gas_space_t *gas, + const double *restrict eri, + hdiag_sector_cache_t *cache, + const hdiag_cache_task_t *task) +{ + uint32_t norb = (uint32_t)gas->norb_tot; + uint32_t npair = norb * (norb + 1u) / 2u; + hdiag_sector_cache_t *sector = cache + task->sector; + uint32_t nstr = gas->sector_nstr[task->sector]; + double *potential = sector->beta_potential + + (size_t)task->orbital * nstr; + for (uint32_t i = task->begin; i < task->end; i++) { + uint64_t mask = sector->mask[i]; + double value = 0.0; + while (mask != 0u) { + int q = ctz64(mask); + mask &= mask - 1u; + value += coulomb_value(eri, npair, task->orbital, q); + } + potential[i] = value; + } +} + +static void execute_cache_work(const gas_space_t *gas, + const double *restrict h1e, + const double *restrict eri, + hdiag_sector_cache_t *cache, + const hdiag_cache_work_t *work) +{ +#ifdef _OPENMP +#pragma omp parallel + { +#pragma omp for schedule(static) + for (int64_t k = 0; + k < (int64_t)work->mask_task_count; k++) { + execute_mask_cache_task( + gas, h1e, eri, cache, work->mask_task + k); + } +#pragma omp for schedule(static) + for (int64_t k = 0; + k < (int64_t)work->potential_task_count; k++) { + execute_potential_cache_task( + gas, eri, cache, work->potential_task + k); + } + } +#else + for (uint32_t k = 0; k < work->mask_task_count; k++) { + execute_mask_cache_task( + gas, h1e, eri, cache, work->mask_task + k); + } + for (uint32_t k = 0; k < work->potential_task_count; k++) { + execute_potential_cache_task( + gas, eri, cache, work->potential_task + k); + } +#endif +} + +static int build_sector_cache(const gas_space_t *gas, + const double *restrict h1e, + const double *restrict eri, + hdiag_sector_cache_t **out) +{ + hdiag_cache_work_t work = {0}; + int status = prepare_sector_cache(gas, out, &work); + if (status == GAS_SUCCESS) { + execute_cache_work(gas, h1e, eri, *out, &work); + } + free_cache_work(&work); + return status; +} + +static int build_tasks(const gas_space_t *gas, + hdiag_task_t **out, uint32_t *task_count) +{ +#ifdef _OPENMP + uint32_t threads = (uint32_t)omp_get_max_threads(); +#else + uint32_t threads = 1u; +#endif + uint64_t target_tasks = (uint64_t)threads * HDIAG_TASK_FACTOR; + uint64_t target_dets = ((uint64_t)gas->ndet + target_tasks - 1u) / + target_tasks; + if (target_dets == 0u) target_dets = 1u; + + uint64_t count = 0u; + for (gas_bid_t b = 0; b < gas->nblock; b++) { + const gas_block_t *block = gas->block + b; + uint32_t na = gas->sector_nstr[block->sa]; + uint32_t nb = gas->sector_nstr[block->sb]; + uint64_t rows64 = target_dets / nb; + uint32_t rows = rows64 == 0u ? 1u : + rows64 > na ? na : (uint32_t)rows64; + count += ((uint64_t)na + rows - 1u) / rows; + } + if (count == 0u || count > UINT32_MAX || + count > SIZE_MAX / sizeof(hdiag_task_t)) { + return GAS_ERR_MEMORY; + } + + hdiag_task_t *task = malloc((size_t)count * sizeof(*task)); + if (task == 0) return GAS_ERR_MEMORY; + + uint32_t n = 0u; + for (gas_bid_t b = 0; b < gas->nblock; b++) { + const gas_block_t *block = gas->block + b; + uint32_t na = gas->sector_nstr[block->sa]; + uint32_t nb = gas->sector_nstr[block->sb]; + uint64_t rows64 = target_dets / nb; + uint32_t rows = rows64 == 0u ? 1u : + rows64 > na ? na : (uint32_t)rows64; + for (uint32_t a0 = 0; a0 < na; a0 += rows) { + uint32_t a1 = na - a0 < rows ? na : a0 + rows; + task[n].block = b; + task[n].alpha_begin = a0; + task[n].alpha_end = a1; + n++; + } + } + + *out = task; + *task_count = n; + return GAS_SUCCESS; +} + +static inline void initialize_row(double *restrict y, + const double *restrict beta, + double alpha, uint32_t n) +{ + uint32_t i = 0u; +#if defined(__AVX2__) + __m256d va = _mm256_set1_pd(alpha); + for (; i + 4u <= n; i += 4u) { + __m256d vb = _mm256_loadu_pd(beta + i); + _mm256_storeu_pd(y + i, _mm256_add_pd(va, vb)); + } +#endif + uint32_t tail = i; +#ifdef _OPENMP +#pragma omp simd +#endif + for (uint32_t j = tail; j < n; j++) y[j] = alpha + beta[j]; +} + +static inline void add_rows8(double *restrict y, + const double *restrict x0, + const double *restrict x1, + const double *restrict x2, + const double *restrict x3, + const double *restrict x4, + const double *restrict x5, + const double *restrict x6, + const double *restrict x7, + uint32_t n) +{ + uint32_t i = 0u; +#if defined(__AVX2__) + for (; i + 4u <= n; i += 4u) { + __m256d v = _mm256_loadu_pd(y + i); + v = _mm256_add_pd(v, _mm256_loadu_pd(x0 + i)); + v = _mm256_add_pd(v, _mm256_loadu_pd(x1 + i)); + v = _mm256_add_pd(v, _mm256_loadu_pd(x2 + i)); + v = _mm256_add_pd(v, _mm256_loadu_pd(x3 + i)); + v = _mm256_add_pd(v, _mm256_loadu_pd(x4 + i)); + v = _mm256_add_pd(v, _mm256_loadu_pd(x5 + i)); + v = _mm256_add_pd(v, _mm256_loadu_pd(x6 + i)); + v = _mm256_add_pd(v, _mm256_loadu_pd(x7 + i)); + _mm256_storeu_pd(y + i, v); + } +#endif + uint32_t tail = i; +#ifdef _OPENMP +#pragma omp simd +#endif + for (uint32_t j = tail; j < n; j++) { + y[j] += x0[j] + x1[j] + x2[j] + x3[j] + + x4[j] + x5[j] + x6[j] + x7[j]; + } +} + +static inline void add_rows4(double *restrict y, + const double *restrict x0, + const double *restrict x1, + const double *restrict x2, + const double *restrict x3, + uint32_t n) +{ + uint32_t i = 0u; +#if defined(__AVX2__) + for (; i + 4u <= n; i += 4u) { + __m256d v = _mm256_loadu_pd(y + i); + v = _mm256_add_pd(v, _mm256_loadu_pd(x0 + i)); + v = _mm256_add_pd(v, _mm256_loadu_pd(x1 + i)); + v = _mm256_add_pd(v, _mm256_loadu_pd(x2 + i)); + v = _mm256_add_pd(v, _mm256_loadu_pd(x3 + i)); + _mm256_storeu_pd(y + i, v); + } +#endif + uint32_t tail = i; +#ifdef _OPENMP +#pragma omp simd +#endif + for (uint32_t j = tail; j < n; j++) { + y[j] += x0[j] + x1[j] + x2[j] + x3[j]; + } +} + +static inline void add_rows2(double *restrict y, + const double *restrict x0, + const double *restrict x1, uint32_t n) +{ +#ifdef _OPENMP +#pragma omp simd +#endif + for (uint32_t i = 0; i < n; i++) y[i] += x0[i] + x1[i]; +} + +static inline void add_row(double *restrict y, + const double *restrict x, uint32_t n) +{ +#ifdef _OPENMP +#pragma omp simd +#endif + for (uint32_t i = 0; i < n; i++) y[i] += x[i]; +} + +static inline void set_fused_rows(double *restrict y, + const double *restrict beta, + double alpha, + const double *const restrict *row, + uint32_t nrow, uint32_t n) +{ + switch (nrow) { + case 8: +#ifdef _OPENMP +#pragma omp simd +#endif + for (uint32_t i = 0; i < n; i++) { + y[i] = alpha + beta[i] + row[0][i] + row[1][i] + + row[2][i] + row[3][i] + row[4][i] + row[5][i] + + row[6][i] + row[7][i]; + } + break; + case 7: +#ifdef _OPENMP +#pragma omp simd +#endif + for (uint32_t i = 0; i < n; i++) { + y[i] = alpha + beta[i] + row[0][i] + row[1][i] + + row[2][i] + row[3][i] + row[4][i] + row[5][i] + + row[6][i]; + } + break; + case 6: +#ifdef _OPENMP +#pragma omp simd +#endif + for (uint32_t i = 0; i < n; i++) { + y[i] = alpha + beta[i] + row[0][i] + row[1][i] + + row[2][i] + row[3][i] + row[4][i] + row[5][i]; + } + break; + case 5: +#ifdef _OPENMP +#pragma omp simd +#endif + for (uint32_t i = 0; i < n; i++) { + y[i] = alpha + beta[i] + row[0][i] + row[1][i] + + row[2][i] + row[3][i] + row[4][i]; + } + break; + case 4: +#ifdef _OPENMP +#pragma omp simd +#endif + for (uint32_t i = 0; i < n; i++) { + y[i] = alpha + beta[i] + row[0][i] + row[1][i] + + row[2][i] + row[3][i]; + } + break; + case 3: +#ifdef _OPENMP +#pragma omp simd +#endif + for (uint32_t i = 0; i < n; i++) { + y[i] = alpha + beta[i] + row[0][i] + row[1][i] + + row[2][i]; + } + break; + case 2: +#ifdef _OPENMP +#pragma omp simd +#endif + for (uint32_t i = 0; i < n; i++) { + y[i] = alpha + beta[i] + row[0][i] + row[1][i]; + } + break; + case 1: +#ifdef _OPENMP +#pragma omp simd +#endif + for (uint32_t i = 0; i < n; i++) { + y[i] = alpha + beta[i] + row[0][i]; + } + break; + default: + initialize_row(y, beta, alpha, n); + break; + } +} + +static void fill_alpha_row_potential(const gas_space_t *gas, + const hdiag_sector_cache_t *cache, + const gas_block_t *block, + uint32_t ia, double *restrict y) +{ + uint32_t nb = gas->sector_nstr[block->sb]; + const hdiag_sector_cache_t *alpha = cache + block->sa; + const hdiag_sector_cache_t *beta = cache + block->sb; + const double *row[GAS_MAX_ORB]; + uint32_t nrow = 0u; + uint64_t mask = alpha->mask[ia]; + while (mask != 0u) { + int p = ctz64(mask); + mask &= mask - 1u; + row[nrow++] = beta->beta_potential + (size_t)p * nb; + } + + uint32_t i = nrow < 8u ? nrow : 8u; + set_fused_rows(y, beta->spin_diag, alpha->spin_diag[ia], row, i, nb); + for (; i + 8u <= nrow; i += 8u) { + add_rows8(y, row[i], row[i + 1u], row[i + 2u], row[i + 3u], + row[i + 4u], row[i + 5u], row[i + 6u], + row[i + 7u], nb); + } + if (i + 4u <= nrow) { + add_rows4(y, row[i], row[i + 1u], row[i + 2u], row[i + 3u], + nb); + i += 4u; + } + if (i + 2u <= nrow) { + add_rows2(y, row[i], row[i + 1u], nb); + i += 2u; + } + if (i < nrow) add_row(y, row[i], nb); +} + +static void fill_alpha_row_direct(const gas_space_t *gas, + const double *restrict eri, + const hdiag_sector_cache_t *cache, + const gas_block_t *block, + uint32_t ia, double *restrict y) +{ + uint32_t norb = (uint32_t)gas->norb_tot; + uint32_t npair = norb * (norb + 1u) / 2u; + uint32_t nb = gas->sector_nstr[block->sb]; + const hdiag_sector_cache_t *alpha = cache + block->sa; + const hdiag_sector_cache_t *beta = cache + block->sb; + double potential[GAS_MAX_ORB] = {0.0}; + + uint64_t alpha_mask = alpha->mask[ia]; + while (alpha_mask != 0u) { + int p = ctz64(alpha_mask); + alpha_mask &= alpha_mask - 1u; + for (uint32_t q = 0; q < norb; q++) { + potential[q] += coulomb_value(eri, npair, p, (int)q); + } + } + + double alpha_value = alpha->spin_diag[ia]; + for (uint32_t ib = 0; ib < nb; ib++) { + double value = alpha_value + beta->spin_diag[ib]; + uint64_t beta_mask = beta->mask[ib]; + while (beta_mask != 0u) { + int q = ctz64(beta_mask); + beta_mask &= beta_mask - 1u; + value += potential[q]; + } + y[ib] = value; + } +} + +static void execute_task(const gas_space_t *gas, + const double *restrict eri, + const hdiag_sector_cache_t *cache, + const hdiag_task_t *task, + double *restrict hdiag) +{ + const gas_block_t *block = gas->block + task->block; + uint32_t nb = gas->sector_nstr[block->sb]; + int use_potential = cache[block->sb].beta_potential != 0; + for (uint32_t ia = task->alpha_begin; ia < task->alpha_end; ia++) { + double *row = hdiag + block->offset + (size_t)ia * nb; + if (use_potential) { + fill_alpha_row_potential(gas, cache, block, ia, row); + } else { + fill_alpha_row_direct(gas, eri, cache, block, ia, row); + } + } +} + +int fci_make_hdiag_gas(const gas_space_t *gas, + const double *h1e, const double *eri, + double *hdiag) +{ + int status = hdiag_validate(gas, h1e, eri, hdiag); + if (status != GAS_SUCCESS) return status; + + hdiag_sector_cache_t *cache = 0; + hdiag_task_t *task = 0; + uint32_t task_count = 0u; + status = build_sector_cache(gas, h1e, eri, &cache); + if (status != GAS_SUCCESS) goto done; + + status = build_tasks(gas, &task, &task_count); + if (status != GAS_SUCCESS) goto done; + +#ifdef _OPENMP +#pragma omp parallel for schedule(static) + for (int64_t i = 0; i < (int64_t)task_count; i++) { + execute_task(gas, eri, cache, task + i, hdiag); + } +#else + for (uint32_t i = 0; i < task_count; i++) { + execute_task(gas, eri, cache, task + i, hdiag); + } +#endif + +done: + free(task); + free_sector_cache(cache, gas->nsector); + return status; +} diff --git a/pyscf/lib/mcscf/fci_rdm_gas.c b/pyscf/lib/mcscf/fci_rdm_gas.c new file mode 100644 index 000000000..c8a32ea08 --- /dev/null +++ b/pyscf/lib/mcscf/fci_rdm_gas.c @@ -0,0 +1,2286 @@ +/* Copyright 2026 The PySCF Developers. All Rights Reserved. + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. + + * + * Author: Yi Deng + */ + +#include "fci_gas.h" + +#include +#include + +#if defined(__AVX2__) +#include +#endif + +#ifdef _OPENMP +#include +#endif + +enum { + RDM_TASK_FACTOR = 16, + RDM_MAX_ALPHA_SPLIT = 64, + RDM_MIN_ALPHA_TILE = 4, + RDM_TRANSPOSE_TILE = 32, + RDM_BB_TRANSPOSE_MIN_ALPHA = 4, + RDM_AA_GROUP_SIZE = 8 +}; + +static const size_t RDM_WORKSPACE_MAX_ELEMS = 67108864u; +static size_t rdm_abba_t1_target_bytes = 16u * 1024u * 1024u; + +/* + * BB kernels store consecutive link-table IDs. AB kernels reuse the buffer + * for (destination block ID, beta link-table ID). + */ +typedef struct { + gas_tid_t first; + gas_tid_t second; +} rdm_tid_pair_t; + +/* + * Alpha link landing in the current destination tile [a0,a1). + * rel_addr is relative to a0; op_index selects the plan's active operator. + */ +typedef struct { + uint32_t src_row; + uint32_t rel_addr; + uint16_t op_index; + int8_t sign; + uint8_t padding; +} rdm_alpha_hit_t; + +typedef struct { + uint32_t off; + uint16_t n; + uint16_t padding; +} rdm_table_ops_t; + +typedef struct { + double *data; + size_t capacity; + double *arena; + size_t arena_capacity; + rdm_tid_pair_t *pair; + size_t pair_capacity; + rdm_alpha_hit_t *alpha_hit; + size_t alpha_hit_capacity; + uint32_t alpha_hit_n; + int32_t *op_map; + size_t op_map_capacity; +} rdm_workspace_t; + +typedef struct { + uint64_t cost; + gas_sid_t adst; + uint32_t q0; + uint32_t q1; + uint32_t a0; + uint32_t a1; +} rdm_task_t; + +struct gas_rdm_plan { + const gas_space_t *gas; + rdm_workspace_t *workspace; + uint32_t nworkspace; + rdm_task_t *task; + uint32_t ntask; + uint32_t task_threads; + rdm_table_ops_t *table_ops; + uint16_t *active_op; + uint32_t *active_row; + size_t abba_t1_target_bytes; +}; + +static int rdm_arena_reserve(rdm_workspace_t *workspace, size_t need); +static int rdm_pair_reserve(rdm_workspace_t *workspace, size_t need); +static int rdm_hit_reserve(rdm_workspace_t *workspace, size_t need); +static int rdm_op_map_reserve(rdm_workspace_t *workspace, size_t need); + +void fci_rdm_gas_set_abba_t1_bytes(uint64_t bytes) +{ + rdm_abba_t1_target_bytes = (size_t)bytes; +} + +static inline uint32_t rdm_op(const gas_link_entry_t *e, uint32_t norb) +{ + return (uint32_t)gas_link_cre(e) * norb + gas_link_des(e); +} + +static inline size_t rdm2_index(uint32_t p, uint32_t q, + uint32_t r, uint32_t s, uint32_t norb) +{ + return (((size_t)p * norb + q) * norb + r) * norb + s; +} + +static inline double rdm_dot(const double *x, const double *y, uint32_t n) +{ + double value = 0.0; + for (uint32_t i = 0; i < n; i++) { + value += x[i] * y[i]; + } + return value; +} + +static inline void rdm_dot4(const double *restrict x, + const double *restrict y0, + const double *restrict y1, + const double *restrict y2, + const double *restrict y3, + uint32_t n, double value[4]) +{ + double s0 = 0.0, s1 = 0.0, s2 = 0.0, s3 = 0.0; + uint32_t i = 0; +#if defined(__AVX2__) + __m256d v0 = _mm256_setzero_pd(); + __m256d v1 = _mm256_setzero_pd(); + __m256d v2 = _mm256_setzero_pd(); + __m256d v3 = _mm256_setzero_pd(); + for (; i + 4u <= n; i += 4u) { + __m256d vx = _mm256_loadu_pd(x + i); + v0 = _mm256_add_pd(v0, _mm256_mul_pd(vx, _mm256_loadu_pd(y0 + i))); + v1 = _mm256_add_pd(v1, _mm256_mul_pd(vx, _mm256_loadu_pd(y1 + i))); + v2 = _mm256_add_pd(v2, _mm256_mul_pd(vx, _mm256_loadu_pd(y2 + i))); + v3 = _mm256_add_pd(v3, _mm256_mul_pd(vx, _mm256_loadu_pd(y3 + i))); + } + double lane[4]; + _mm256_storeu_pd(lane, v0); + s0 = lane[0] + lane[1] + lane[2] + lane[3]; + _mm256_storeu_pd(lane, v1); + s1 = lane[0] + lane[1] + lane[2] + lane[3]; + _mm256_storeu_pd(lane, v2); + s2 = lane[0] + lane[1] + lane[2] + lane[3]; + _mm256_storeu_pd(lane, v3); + s3 = lane[0] + lane[1] + lane[2] + lane[3]; +#endif + for (; i < n; i++) { + double xi = x[i]; + s0 += xi * y0[i]; + s1 += xi * y1[i]; + s2 += xi * y2[i]; + s3 += xi * y3[i]; + } + value[0] = s0; + value[1] = s1; + value[2] = s2; + value[3] = s3; +} + +static inline void rdm_dot2(const double *restrict x, + const double *restrict y0, + const double *restrict y1, + uint32_t n, double value[2]) +{ + double s0 = 0.0, s1 = 0.0; + uint32_t i = 0u; +#if defined(__AVX2__) + __m256d v0 = _mm256_setzero_pd(); + __m256d v1 = _mm256_setzero_pd(); + for (; i + 4u <= n; i += 4u) { + __m256d vx = _mm256_loadu_pd(x + i); + v0 = _mm256_add_pd(v0, _mm256_mul_pd( + vx, _mm256_loadu_pd(y0 + i))); + v1 = _mm256_add_pd(v1, _mm256_mul_pd( + vx, _mm256_loadu_pd(y1 + i))); + } + double lane[4]; + _mm256_storeu_pd(lane, v0); + s0 = lane[0] + lane[1] + lane[2] + lane[3]; + _mm256_storeu_pd(lane, v1); + s1 = lane[0] + lane[1] + lane[2] + lane[3]; +#endif + for (; i < n; i++) { + double xi = x[i]; + s0 += xi * y0[i]; + s1 += xi * y1[i]; + } + value[0] = s0; + value[1] = s1; +} + +static inline void rdm_dot3(const double *restrict x, + const double *restrict y0, + const double *restrict y1, + const double *restrict y2, + uint32_t n, double value[3]) +{ + double s0 = 0.0, s1 = 0.0, s2 = 0.0; + uint32_t i = 0u; +#if defined(__AVX2__) + __m256d v0 = _mm256_setzero_pd(); + __m256d v1 = _mm256_setzero_pd(); + __m256d v2 = _mm256_setzero_pd(); + for (; i + 4u <= n; i += 4u) { + __m256d vx = _mm256_loadu_pd(x + i); + v0 = _mm256_add_pd(v0, _mm256_mul_pd( + vx, _mm256_loadu_pd(y0 + i))); + v1 = _mm256_add_pd(v1, _mm256_mul_pd( + vx, _mm256_loadu_pd(y1 + i))); + v2 = _mm256_add_pd(v2, _mm256_mul_pd( + vx, _mm256_loadu_pd(y2 + i))); + } + double lane[4]; + _mm256_storeu_pd(lane, v0); + s0 = lane[0] + lane[1] + lane[2] + lane[3]; + _mm256_storeu_pd(lane, v1); + s1 = lane[0] + lane[1] + lane[2] + lane[3]; + _mm256_storeu_pd(lane, v2); + s2 = lane[0] + lane[1] + lane[2] + lane[3]; +#endif + for (; i < n; i++) { + double xi = x[i]; + s0 += xi * y0[i]; + s1 += xi * y1[i]; + s2 += xi * y2[i]; + } + value[0] = s0; + value[1] = s1; + value[2] = s2; +} + +static inline void rdm_dot8_ptrs(const double *restrict x, + const double *const y[8], + uint32_t n, double value[8]) +{ + for (uint32_t j = 0; j < 8u; j++) value[j] = 0.0; + uint32_t i = 0u; +#if defined(__AVX2__) + __m256d sum[8]; + for (uint32_t j = 0; j < 8u; j++) sum[j] = _mm256_setzero_pd(); + for (; i + 4u <= n; i += 4u) { + __m256d vx = _mm256_loadu_pd(x + i); + for (uint32_t j = 0; j < 8u; j++) { + sum[j] = _mm256_add_pd( + sum[j], _mm256_mul_pd( + vx, _mm256_loadu_pd(y[j] + i))); + } + } + for (uint32_t j = 0; j < 8u; j++) { + double lane[4]; + _mm256_storeu_pd(lane, sum[j]); + value[j] = lane[0] + lane[1] + lane[2] + lane[3]; + } +#endif + for (; i < n; i++) { + double xi = x[i]; + for (uint32_t j = 0; j < 8u; j++) value[j] += xi * y[j][i]; + } +} + +static void rdm_transpose_pack(double *restrict packed, + const double *restrict block, + uint32_t block_nb, uint32_t na, + uint32_t nb, uint32_t block_a0) +{ + for (uint32_t a0 = 0; a0 < na; a0 += RDM_TRANSPOSE_TILE) { + uint32_t a1 = na - a0 < RDM_TRANSPOSE_TILE ? + na : a0 + RDM_TRANSPOSE_TILE; + for (uint32_t b0 = 0; b0 < nb; b0 += RDM_TRANSPOSE_TILE) { + uint32_t b1 = nb - b0 < RDM_TRANSPOSE_TILE ? + nb : b0 + RDM_TRANSPOSE_TILE; + for (uint32_t b = b0; b < b1; b++) { + double *dst = packed + (size_t)b * na + a0; +#ifdef _OPENMP +#pragma omp simd +#endif + for (uint32_t a = a0; a < a1; a++) { + dst[a - a0] = + block[(size_t)(block_a0 + a) * + block_nb + b]; + } + } + } + } +} + +static int rdm_validate(const gas_space_t *gas, + const double *bra, const double *ket) +{ + if (gas == 0 || bra == 0 || ket == 0 || + gas->link_format != GAS_LINK_RAW || gas->norb_tot <= 0) { + return GAS_ERR_INVALID; + } + return GAS_SUCCESS; +} + +static void rdm1_alpha(const gas_space_t *gas, + const double *bra, const double *ket, double *dm1a) +{ + uint32_t norb = (uint32_t)gas->norb_tot; + + for (gas_bid_t bdst = 0; bdst < gas->nblock; bdst++) { + const gas_block_t *bd = gas->block + bdst; + gas_row_t incoming = gas->R.row[bd->sa]; + uint32_t nb = gas->sector_nstr[bd->sb]; + const double *dst = bra + bd->offset; + + for (uint32_t it = 0; it < incoming.n; it++) { + gas_sid_t asrc = gas->R.src[incoming.off + it]; + gas_tid_t tid = gas->R.tid[incoming.off + it]; + gas_bid_t bsrc = gas_find_block(gas, asrc, bd->sb); + if (bsrc == GAS_INVALID_BID) continue; + + const gas_link_table_t *table = gas->table + tid; + const double *src = ket + gas->block[bsrc].offset; + for (uint32_t ia0 = 0; ia0 < table->nsrc; ia0++) { + const double *src_row = src + (size_t)ia0 * nb; + const gas_link_entry_t *row = table->link + + (size_t)ia0 * table->nlink; + for (uint32_t k = 0; k < table->nlink; k++) { + const gas_link_entry_t *e = row + k; + const double *dst_row = dst + (size_t)e->addr * nb; + uint32_t p = gas_link_cre(e); + uint32_t q = gas_link_des(e); + dm1a[(size_t)q * norb + p] += + (double)e->sign * + rdm_dot(dst_row, src_row, nb); + } + } + } + } +} + +static void rdm1_beta(const gas_space_t *gas, + const double *bra, const double *ket, double *dm1b) +{ + uint32_t norb = (uint32_t)gas->norb_tot; + + for (gas_bid_t bdst = 0; bdst < gas->nblock; bdst++) { + const gas_block_t *bd = gas->block + bdst; + gas_row_t incoming = gas->R.row[bd->sb]; + uint32_t na = gas->sector_nstr[bd->sa]; + uint32_t nb1 = gas->sector_nstr[bd->sb]; + const double *dst = bra + bd->offset; + + for (uint32_t it = 0; it < incoming.n; it++) { + gas_sid_t bsrc_sid = gas->R.src[incoming.off + it]; + gas_tid_t tid = gas->R.tid[incoming.off + it]; + gas_bid_t bsrc = gas_find_block(gas, bd->sa, bsrc_sid); + if (bsrc == GAS_INVALID_BID) continue; + + const gas_link_table_t *table = gas->table + tid; + uint32_t nb0 = gas->sector_nstr[bsrc_sid]; + const double *src = ket + gas->block[bsrc].offset; + for (uint32_t ib0 = 0; ib0 < table->nsrc; ib0++) { + const gas_link_entry_t *row = table->link + + (size_t)ib0 * table->nlink; + for (uint32_t k = 0; k < table->nlink; k++) { + const gas_link_entry_t *e = row + k; + double value = 0.0; + for (uint32_t ia = 0; ia < na; ia++) { + value += dst[(size_t)ia * nb1 + e->addr] * + src[(size_t)ia * nb0 + ib0]; + } + uint32_t p = gas_link_cre(e); + uint32_t q = gas_link_des(e); + dm1b[(size_t)q * norb + p] += + (double)e->sign * value; + } + } + } + } +} + +static void rdm2_aa_path(const gas_space_t *gas, + const double *bra, const double *ket, + gas_bid_t bsrc, gas_bid_t bdst, + gas_tid_t tid1, gas_tid_t tid2, double *ph2aa) +{ + const gas_block_t *bs = gas->block + bsrc; + const gas_block_t *bd = gas->block + bdst; + const gas_link_table_t *t1 = gas->table + tid1; + const gas_link_table_t *t2 = gas->table + tid2; + uint32_t norb = (uint32_t)gas->norb_tot; + uint32_t nop = norb * norb; + uint32_t nb = gas->sector_nstr[bs->sb]; + const double *src = ket + bs->offset; + const double *dst = bra + bd->offset; + + for (uint32_t ia0 = 0; ia0 < t1->nsrc; ia0++) { + const double *src_row = src + (size_t)ia0 * nb; + const gas_link_entry_t *r1 = t1->link + + (size_t)ia0 * t1->nlink; + for (uint32_t k1 = 0; k1 < t1->nlink; k1++) { + const gas_link_entry_t *e1 = r1 + k1; + const gas_link_entry_t *r2 = t2->link + + (size_t)e1->addr * t2->nlink; + uint32_t op1 = rdm_op(e1, norb); + for (uint32_t k2 = 0; k2 < t2->nlink; k2++) { + const gas_link_entry_t *e2 = r2 + k2; + const double *dst_row = dst + (size_t)e2->addr * nb; + uint32_t op2 = rdm_op(e2, norb); + ph2aa[(size_t)op2 * nop + op1] += + (double)(e1->sign * e2->sign) * + rdm_dot(dst_row, src_row, nb); + } + } + } +} + +static void rdm2_aa(const gas_space_t *gas, + const double *bra, const double *ket, double *ph2aa) +{ + for (gas_bid_t bdst = 0; bdst < gas->nblock; bdst++) { + const gas_block_t *bd = gas->block + bdst; + gas_row_t reverse = gas->R.row[bd->sa]; + gas_row_t sources = gas->D.by_beta_row[bd->sb]; + const gas_sid_t *middle_rev = gas->R.src + reverse.off; + const gas_tid_t *second_tid = gas->R.tid + reverse.off; + + for (uint32_t q = 0; q < sources.n; q++) { + gas_sid_t asrc = gas->D.by_beta_sid[sources.off + q]; + gas_bid_t bsrc = gas->D.by_beta_bid[sources.off + q]; + gas_row_t forward = gas->T.row[asrc]; + const gas_sid_t *middle_fwd = gas->T.dst + forward.off; + uint32_t i = 0; + uint32_t j = 0; + + while (i < forward.n && j < reverse.n) { + if (middle_fwd[i] < middle_rev[j]) { + i++; + } else if (middle_fwd[i] > middle_rev[j]) { + j++; + } else { + rdm2_aa_path(gas, bra, ket, bsrc, bdst, + forward.off + i, second_tid[j], + ph2aa); + i++; + j++; + } + } + } + } +} + +static void rdm2_bb_path(const gas_space_t *gas, + const double *bra, const double *ket, + gas_bid_t bsrc, gas_bid_t bdst, + gas_tid_t tid1, gas_tid_t tid2, double *ph2bb) +{ + const gas_block_t *bs = gas->block + bsrc; + const gas_block_t *bd = gas->block + bdst; + const gas_link_table_t *t1 = gas->table + tid1; + const gas_link_table_t *t2 = gas->table + tid2; + uint32_t norb = (uint32_t)gas->norb_tot; + uint32_t nop = norb * norb; + uint32_t na = gas->sector_nstr[bs->sa]; + uint32_t nb0 = gas->sector_nstr[bs->sb]; + uint32_t nb2 = gas->sector_nstr[bd->sb]; + const double *src = ket + bs->offset; + const double *dst = bra + bd->offset; + + for (uint32_t ib0 = 0; ib0 < t1->nsrc; ib0++) { + const gas_link_entry_t *r1 = t1->link + + (size_t)ib0 * t1->nlink; + for (uint32_t k1 = 0; k1 < t1->nlink; k1++) { + const gas_link_entry_t *e1 = r1 + k1; + const gas_link_entry_t *r2 = t2->link + + (size_t)e1->addr * t2->nlink; + uint32_t op1 = rdm_op(e1, norb); + for (uint32_t k2 = 0; k2 < t2->nlink; k2++) { + const gas_link_entry_t *e2 = r2 + k2; + double value = 0.0; + for (uint32_t ia = 0; ia < na; ia++) { + value += dst[(size_t)ia * nb2 + e2->addr] * + src[(size_t)ia * nb0 + ib0]; + } + uint32_t op2 = rdm_op(e2, norb); + ph2bb[(size_t)op2 * nop + op1] += + (double)(e1->sign * e2->sign) * value; + } + } + } +} + +static void rdm2_bb(const gas_space_t *gas, + const double *bra, const double *ket, double *ph2bb) +{ + for (gas_bid_t bdst = 0; bdst < gas->nblock; bdst++) { + const gas_block_t *bd = gas->block + bdst; + gas_row_t reverse = gas->R.row[bd->sb]; + gas_row_t sources = gas->D.by_alpha_row[bd->sa]; + const gas_sid_t *middle_rev = gas->R.src + reverse.off; + const gas_tid_t *second_tid = gas->R.tid + reverse.off; + + for (uint32_t q = 0; q < sources.n; q++) { + gas_bid_t bsrc = sources.off + q; + gas_sid_t bsrc_sid = gas->block[bsrc].sb; + gas_row_t forward = gas->T.row[bsrc_sid]; + const gas_sid_t *middle_fwd = gas->T.dst + forward.off; + uint32_t i = 0; + uint32_t j = 0; + + while (i < forward.n && j < reverse.n) { + if (middle_fwd[i] < middle_rev[j]) { + i++; + } else if (middle_fwd[i] > middle_rev[j]) { + j++; + } else { + rdm2_bb_path(gas, bra, ket, bsrc, bdst, + forward.off + i, second_tid[j], + ph2bb); + i++; + j++; + } + } + } + } +} + +static void rdm2_ab_path(const gas_space_t *gas, + const double *bra, const double *ket, + gas_bid_t bsrc, gas_bid_t bdst, + gas_tid_t tida, gas_tid_t tidb, double *dm2ab) +{ + const gas_block_t *bs = gas->block + bsrc; + const gas_block_t *bd = gas->block + bdst; + const gas_link_table_t *ta = gas->table + tida; + const gas_link_table_t *tb = gas->table + tidb; + uint32_t norb = (uint32_t)gas->norb_tot; + uint32_t nb0 = gas->sector_nstr[bs->sb]; + uint32_t nb1 = gas->sector_nstr[bd->sb]; + const double *src = ket + bs->offset; + const double *dst = bra + bd->offset; + + for (uint32_t ia0 = 0; ia0 < ta->nsrc; ia0++) { + const gas_link_entry_t *ra = ta->link + + (size_t)ia0 * ta->nlink; + for (uint32_t ka = 0; ka < ta->nlink; ka++) { + const gas_link_entry_t *ea = ra + ka; + uint32_t p = gas_link_cre(ea); + uint32_t q = gas_link_des(ea); + for (uint32_t ib0 = 0; ib0 < tb->nsrc; ib0++) { + double c = src[(size_t)ia0 * nb0 + ib0]; + const gas_link_entry_t *rb = tb->link + + (size_t)ib0 * tb->nlink; + for (uint32_t kb = 0; kb < tb->nlink; kb++) { + const gas_link_entry_t *eb = rb + kb; + uint32_t r = gas_link_cre(eb); + uint32_t s = gas_link_des(eb); + dm2ab[rdm2_index(p, q, r, s, norb)] += + (double)(ea->sign * eb->sign) * + dst[(size_t)ea->addr * nb1 + eb->addr] * c; + } + } + } + } +} + +static void rdm2_ab(const gas_space_t *gas, + const double *bra, const double *ket, double *dm2ab) +{ + for (gas_bid_t bdst = 0; bdst < gas->nblock; bdst++) { + const gas_block_t *bd = gas->block + bdst; + gas_row_t alpha_in = gas->R.row[bd->sa]; + + for (uint32_t i = 0; i < alpha_in.n; i++) { + gas_sid_t asrc = gas->R.src[alpha_in.off + i]; + gas_tid_t tida = gas->R.tid[alpha_in.off + i]; + gas_row_t source_blocks = gas->D.by_alpha_row[asrc]; + + for (uint32_t q = 0; q < source_blocks.n; q++) { + gas_bid_t bsrc = source_blocks.off + q; + gas_sid_t bsrc_sid = gas->block[bsrc].sb; + gas_tid_t tidb = gas_find_table(gas, bsrc_sid, bd->sb); + if (tidb != GAS_INVALID_TID) { + rdm2_ab_path(gas, bra, ket, bsrc, bdst, + tida, tidb, dm2ab); + } + } + } + } +} + +static void rdm_task_bounds(const gas_space_t *gas, const rdm_task_t *task, + gas_row_t *row, uint32_t *q0, uint32_t *q1, + uint32_t *a0, uint32_t *a1) +{ + *row = gas->D.by_alpha_row[task->adst]; + uint32_t na = gas->sector_nstr[task->adst]; + *q0 = task->q0 < row->n ? task->q0 : row->n; + *q1 = task->q1 < row->n ? task->q1 : row->n; + if (*q1 < *q0) *q1 = *q0; + *a0 = task->a0 < na ? task->a0 : na; + *a1 = task->a1 < na ? task->a1 : na; + if (*a1 < *a0) *a1 = *a0; +} + +static void rdm1_alpha_task(const gas_space_t *gas, + const double *bra, const double *ket, + const rdm_task_t *task, double *dm1a) +{ + uint32_t norb = (uint32_t)gas->norb_tot; + gas_row_t blocks; + uint32_t q0, q1, a0, a1; + rdm_task_bounds(gas, task, &blocks, &q0, &q1, &a0, &a1); + + for (uint32_t qb = q0; qb < q1; qb++) { + gas_bid_t bdst = blocks.off + qb; + const gas_block_t *bd = gas->block + bdst; + gas_row_t incoming = gas->R.row[bd->sa]; + uint32_t nb = gas->sector_nstr[bd->sb]; + const double *dst = bra + bd->offset; + + for (uint32_t it = 0; it < incoming.n; it++) { + gas_sid_t asrc = gas->R.src[incoming.off + it]; + gas_tid_t tid = gas->R.tid[incoming.off + it]; + gas_bid_t bsrc = gas_find_block(gas, asrc, bd->sb); + if (bsrc == GAS_INVALID_BID) continue; + const gas_link_table_t *table = gas->table + tid; + const double *src = ket + gas->block[bsrc].offset; + + for (uint32_t ia0 = 0; ia0 < table->nsrc; ia0++) { + const double *src_row = src + (size_t)ia0 * nb; + const gas_link_entry_t *link = table->link + + (size_t)ia0 * table->nlink; + for (uint32_t k = 0; k < table->nlink; k++) { + const gas_link_entry_t *e = link + k; + if (e->addr < a0 || e->addr >= a1) continue; + const double *dst_row = dst + (size_t)e->addr * nb; + uint32_t p = gas_link_cre(e); + uint32_t q = gas_link_des(e); + dm1a[(size_t)q * norb + p] += + (double)e->sign * + rdm_dot(dst_row, src_row, nb); + } + } + } + } +} + +static void rdm1_beta_task(const gas_space_t *gas, + const double *bra, const double *ket, + const rdm_task_t *task, double *dm1b) +{ + uint32_t norb = (uint32_t)gas->norb_tot; + gas_row_t blocks; + uint32_t q0, q1, a0, a1; + rdm_task_bounds(gas, task, &blocks, &q0, &q1, &a0, &a1); + + for (uint32_t qb = q0; qb < q1; qb++) { + gas_bid_t bdst = blocks.off + qb; + const gas_block_t *bd = gas->block + bdst; + gas_row_t incoming = gas->R.row[bd->sb]; + uint32_t nb1 = gas->sector_nstr[bd->sb]; + const double *dst = bra + bd->offset; + + for (uint32_t it = 0; it < incoming.n; it++) { + gas_sid_t bsrc_sid = gas->R.src[incoming.off + it]; + gas_tid_t tid = gas->R.tid[incoming.off + it]; + gas_bid_t bsrc = gas_find_block(gas, bd->sa, bsrc_sid); + if (bsrc == GAS_INVALID_BID) continue; + const gas_link_table_t *table = gas->table + tid; + uint32_t nb0 = gas->sector_nstr[bsrc_sid]; + const double *src = ket + gas->block[bsrc].offset; + + for (uint32_t ib0 = 0; ib0 < table->nsrc; ib0++) { + const gas_link_entry_t *link = table->link + + (size_t)ib0 * table->nlink; + for (uint32_t k = 0; k < table->nlink; k++) { + const gas_link_entry_t *e = link + k; + double value = 0.0; + for (uint32_t ia = a0; ia < a1; ia++) { + value += dst[(size_t)ia * nb1 + e->addr] * + src[(size_t)ia * nb0 + ib0]; + } + uint32_t p = gas_link_cre(e); + uint32_t q = gas_link_des(e); + dm1b[(size_t)q * norb + p] += + (double)e->sign * value; + } + } + } + } +} + +static void rdm2_aa_path_task(const gas_space_t *gas, + const double *bra, const double *ket, + gas_bid_t bsrc, gas_bid_t bdst, + gas_tid_t tid1, gas_tid_t tid2, + uint32_t a0, uint32_t a1, double *ph2aa) +{ + const gas_block_t *bs = gas->block + bsrc; + const gas_block_t *bd = gas->block + bdst; + const gas_link_table_t *t1 = gas->table + tid1; + const gas_link_table_t *t2 = gas->table + tid2; + uint32_t norb = (uint32_t)gas->norb_tot; + uint32_t nop = norb * norb; + uint32_t nb = gas->sector_nstr[bs->sb]; + const double *src = ket + bs->offset; + const double *dst = bra + bd->offset; + + for (uint32_t ia0 = 0; ia0 < t1->nsrc; ia0++) { + const double *src_row = src + (size_t)ia0 * nb; + const gas_link_entry_t *r1 = t1->link + + (size_t)ia0 * t1->nlink; + for (uint32_t k1 = 0; k1 < t1->nlink; k1++) { + const gas_link_entry_t *e1 = r1 + k1; + const gas_link_entry_t *r2 = t2->link + + (size_t)e1->addr * t2->nlink; + uint32_t op1 = rdm_op(e1, norb); + const gas_link_entry_t *touched[8]; + uint32_t ntouched = 0; + for (uint32_t k2 = 0; k2 < t2->nlink; k2++) { + const gas_link_entry_t *e2 = r2 + k2; + if (e2->addr < a0 || e2->addr >= a1) continue; + touched[ntouched++] = e2; + if (ntouched == RDM_AA_GROUP_SIZE) { + const double *row[8]; + double value[8]; + for (uint32_t u = 0; u < 8u; u++) { + row[u] = dst + + (size_t)touched[u]->addr * nb; + } + rdm_dot8_ptrs(src_row, row, nb, value); + for (uint32_t u = 0; u < 8u; u++) { + uint32_t op2 = rdm_op(touched[u], norb); + ph2aa[(size_t)op2 * nop + op1] += + (double)(e1->sign * + touched[u]->sign) * value[u]; + } + ntouched = 0u; + } + } + if (ntouched >= 4u) { + double value[4]; + rdm_dot4(src_row, + dst + (size_t)touched[0]->addr * nb, + dst + (size_t)touched[1]->addr * nb, + dst + (size_t)touched[2]->addr * nb, + dst + (size_t)touched[3]->addr * nb, + nb, value); + for (uint32_t u = 0; u < 4u; u++) { + uint32_t op2 = rdm_op(touched[u], norb); + ph2aa[(size_t)op2 * nop + op1] += + (double)(e1->sign * + touched[u]->sign) * value[u]; + } + for (uint32_t u = 4u; u < ntouched; u++) { + touched[u - 4u] = touched[u]; + } + ntouched -= 4u; + } + if (ntouched == 3u) { + double value[3]; + rdm_dot3(src_row, + dst + (size_t)touched[0]->addr * nb, + dst + (size_t)touched[1]->addr * nb, + dst + (size_t)touched[2]->addr * nb, + nb, value); + for (uint32_t u = 0; u < 3u; u++) { + uint32_t op2 = rdm_op(touched[u], norb); + ph2aa[(size_t)op2 * nop + op1] += + (double)(e1->sign * + touched[u]->sign) * value[u]; + } + ntouched = 0u; + } else if (ntouched == 2u) { + double value[2]; + rdm_dot2(src_row, + dst + (size_t)touched[0]->addr * nb, + dst + (size_t)touched[1]->addr * nb, + nb, value); + for (uint32_t u = 0; u < 2u; u++) { + uint32_t op2 = rdm_op(touched[u], norb); + ph2aa[(size_t)op2 * nop + op1] += + (double)(e1->sign * + touched[u]->sign) * value[u]; + } + ntouched = 0u; + } + for (uint32_t u = 0; u < ntouched; u++) { + const gas_link_entry_t *e2 = touched[u]; + const double *dst_row = dst + (size_t)e2->addr * nb; + uint32_t op2 = rdm_op(e2, norb); + ph2aa[(size_t)op2 * nop + op1] += + (double)(e1->sign * e2->sign) * + rdm_dot(dst_row, src_row, nb); + } + } + } +} + +static void rdm2_aa_task(const gas_space_t *gas, + const double *bra, const double *ket, + const rdm_task_t *task, double *ph2aa) +{ + gas_row_t blocks; + uint32_t q0, q1, a0, a1; + rdm_task_bounds(gas, task, &blocks, &q0, &q1, &a0, &a1); + + for (uint32_t qb = q0; qb < q1; qb++) { + gas_bid_t bdst = blocks.off + qb; + const gas_block_t *bd = gas->block + bdst; + gas_row_t reverse = gas->R.row[bd->sa]; + gas_row_t sources = gas->D.by_beta_row[bd->sb]; + const gas_sid_t *middle_rev = gas->R.src + reverse.off; + const gas_tid_t *second_tid = gas->R.tid + reverse.off; + + for (uint32_t q = 0; q < sources.n; q++) { + gas_sid_t asrc = gas->D.by_beta_sid[sources.off + q]; + gas_bid_t bsrc = gas->D.by_beta_bid[sources.off + q]; + gas_row_t forward = gas->T.row[asrc]; + const gas_sid_t *middle_fwd = gas->T.dst + forward.off; + uint32_t i = 0, j = 0; + while (i < forward.n && j < reverse.n) { + if (middle_fwd[i] < middle_rev[j]) { + i++; + } else if (middle_fwd[i] > middle_rev[j]) { + j++; + } else { + rdm2_aa_path_task(gas, bra, ket, bsrc, bdst, + forward.off + i, + second_tid[j], a0, a1, + ph2aa); + i++; + j++; + } + } + } + } +} + +static void rdm2_bb_path_task(const gas_space_t *gas, + const double *bra, const double *ket, + gas_bid_t bsrc, gas_bid_t bdst, + gas_tid_t tid1, gas_tid_t tid2, + uint32_t a0, uint32_t a1, double *ph2bb) +{ + const gas_block_t *bs = gas->block + bsrc; + const gas_block_t *bd = gas->block + bdst; + const gas_link_table_t *t1 = gas->table + tid1; + const gas_link_table_t *t2 = gas->table + tid2; + uint32_t norb = (uint32_t)gas->norb_tot; + uint32_t nop = norb * norb; + uint32_t nb0 = gas->sector_nstr[bs->sb]; + uint32_t nb2 = gas->sector_nstr[bd->sb]; + const double *src = ket + bs->offset; + const double *dst = bra + bd->offset; + + for (uint32_t ib0 = 0; ib0 < t1->nsrc; ib0++) { + const gas_link_entry_t *r1 = t1->link + + (size_t)ib0 * t1->nlink; + for (uint32_t k1 = 0; k1 < t1->nlink; k1++) { + const gas_link_entry_t *e1 = r1 + k1; + const gas_link_entry_t *r2 = t2->link + + (size_t)e1->addr * t2->nlink; + uint32_t op1 = rdm_op(e1, norb); + for (uint32_t k2 = 0; k2 < t2->nlink; k2++) { + const gas_link_entry_t *e2 = r2 + k2; + double value = 0.0; + for (uint32_t ia = a0; ia < a1; ia++) { + value += dst[(size_t)ia * nb2 + e2->addr] * + src[(size_t)ia * nb0 + ib0]; + } + uint32_t op2 = rdm_op(e2, norb); + ph2bb[(size_t)op2 * nop + op1] += + (double)(e1->sign * e2->sign) * value; + } + } + } +} + +static void rdm2_bb_pair_list_task(const gas_space_t *gas, + const double *bra, const double *ket, + gas_bid_t bsrc, gas_bid_t bdst, + const rdm_tid_pair_t *pair, uint32_t npair, + uint32_t a0, uint32_t a1, double *ph2bb, + rdm_workspace_t *workspace) +{ + const gas_block_t *bs = gas->block + bsrc; + const gas_block_t *bd = gas->block + bdst; + uint32_t nb0 = gas->sector_nstr[bs->sb]; + uint32_t nb2 = gas->sector_nstr[bd->sb]; + uint32_t nat = a1 > a0 ? a1 - a0 : 0u; + size_t src_elems = (size_t)nb0 * nat; + size_t dst_elems = (size_t)nb2 * nat; + + if (nat >= RDM_BB_TRANSPOSE_MIN_ALPHA && + src_elems + dst_elems <= RDM_WORKSPACE_MAX_ELEMS && + rdm_arena_reserve(workspace, src_elems + dst_elems) == GAS_SUCCESS) { + double *src_t = workspace->arena; + double *dst_t = src_t + src_elems; + rdm_transpose_pack(src_t, ket + bs->offset, nb0, + nat, nb0, a0); + rdm_transpose_pack(dst_t, bra + bd->offset, nb2, + nat, nb2, a0); + uint32_t norb = (uint32_t)gas->norb_tot; + uint32_t nop = norb * norb; + for (uint32_t ip = 0; ip < npair; ip++) { + const gas_link_table_t *t1 = gas->table + pair[ip].first; + const gas_link_table_t *t2 = gas->table + pair[ip].second; + for (uint32_t ib0 = 0; ib0 < t1->nsrc; ib0++) { + const double *src_vec = src_t + (size_t)ib0 * nat; + const gas_link_entry_t *r1 = t1->link + + (size_t)ib0 * t1->nlink; + for (uint32_t k1 = 0; k1 < t1->nlink; k1++) { + const gas_link_entry_t *e1 = r1 + k1; + const gas_link_entry_t *r2 = t2->link + + (size_t)e1->addr * t2->nlink; + uint32_t op1 = rdm_op(e1, norb); + uint32_t k2 = 0u; + for (; k2 + 8u <= t2->nlink; k2 += 8u) { + const double *row[8]; + double value[8]; + for (uint32_t u = 0; u < 8u; u++) { + row[u] = dst_t + + (size_t)r2[k2 + u].addr * nat; + } + rdm_dot8_ptrs(src_vec, row, nat, value); + for (uint32_t u = 0; u < 8u; u++) { + const gas_link_entry_t *e2 = r2 + k2 + u; + uint32_t op2 = rdm_op(e2, norb); + ph2bb[(size_t)op2 * nop + op1] += + (double)(e1->sign * e2->sign) * + value[u]; + } + } + if (k2 + 4u <= t2->nlink) { + double value[4]; + rdm_dot4(src_vec, + dst_t + (size_t)r2[k2].addr * nat, + dst_t + (size_t)r2[k2 + 1u].addr * nat, + dst_t + (size_t)r2[k2 + 2u].addr * nat, + dst_t + (size_t)r2[k2 + 3u].addr * nat, + nat, value); + for (uint32_t u = 0; u < 4u; u++) { + const gas_link_entry_t *e2 = r2 + k2 + u; + uint32_t op2 = rdm_op(e2, norb); + ph2bb[(size_t)op2 * nop + op1] += + (double)(e1->sign * e2->sign) * + value[u]; + } + k2 += 4u; + } + for (; k2 < t2->nlink; k2++) { + const gas_link_entry_t *e2 = r2 + k2; + uint32_t op2 = rdm_op(e2, norb); + double value = rdm_dot( + dst_t + (size_t)e2->addr * nat, + src_vec, nat); + ph2bb[(size_t)op2 * nop + op1] += + (double)(e1->sign * e2->sign) * + value; + } + } + } + } + return; + } + + for (uint32_t ip = 0; ip < npair; ip++) { + rdm2_bb_path_task(gas, bra, ket, bsrc, bdst, + pair[ip].first, pair[ip].second, + a0, a1, ph2bb); + } +} + +static void rdm2_bb_task(const gas_space_t *gas, + const double *bra, const double *ket, + const rdm_task_t *task, double *ph2bb, + rdm_workspace_t *workspace) +{ + gas_row_t blocks; + uint32_t q0, q1, a0, a1; + rdm_task_bounds(gas, task, &blocks, &q0, &q1, &a0, &a1); + + for (uint32_t qb = q0; qb < q1; qb++) { + gas_bid_t bdst = blocks.off + qb; + const gas_block_t *bd = gas->block + bdst; + gas_row_t reverse = gas->R.row[bd->sb]; + gas_row_t sources = gas->D.by_alpha_row[bd->sa]; + const gas_sid_t *middle_rev = gas->R.src + reverse.off; + const gas_tid_t *second_tid = gas->R.tid + reverse.off; + + for (uint32_t q = 0; q < sources.n; q++) { + gas_bid_t bsrc = sources.off + q; + gas_sid_t bsrc_sid = gas->block[bsrc].sb; + gas_row_t forward = gas->T.row[bsrc_sid]; + const gas_sid_t *middle_fwd = gas->T.dst + forward.off; + uint32_t i = 0, j = 0; + uint32_t maxpair = forward.n < reverse.n ? + forward.n : reverse.n; + if (rdm_pair_reserve(workspace, maxpair) != GAS_SUCCESS) { + maxpair = 0u; + } + uint32_t npair = 0; + while (i < forward.n && j < reverse.n) { + if (middle_fwd[i] < middle_rev[j]) { + i++; + } else if (middle_fwd[i] > middle_rev[j]) { + j++; + } else { + if (npair < maxpair) { + workspace->pair[npair].first = + forward.off + i; + workspace->pair[npair].second = + second_tid[j]; + npair++; + } else { + rdm2_bb_path_task( + gas, bra, ket, bsrc, bdst, + forward.off + i, + second_tid[j], a0, a1, + ph2bb); + } + i++; + j++; + } + } + if (npair != 0u) { + rdm2_bb_pair_list_task(gas, bra, ket, bsrc, bdst, + workspace->pair, npair, + a0, a1, ph2bb, workspace); + } + } + } +} + +static void rdm2_ab_path_fallback(const gas_space_t *gas, + const double *bra, const double *ket, + gas_bid_t bsrc, gas_bid_t bdst, + gas_tid_t tida, gas_tid_t tidb, + uint32_t a0, uint32_t a1, double *dm2ab) +{ + const gas_block_t *bs = gas->block + bsrc; + const gas_block_t *bd = gas->block + bdst; + const gas_link_table_t *ta = gas->table + tida; + const gas_link_table_t *tb = gas->table + tidb; + uint32_t norb = (uint32_t)gas->norb_tot; + uint32_t nb0 = gas->sector_nstr[bs->sb]; + uint32_t nb1 = gas->sector_nstr[bd->sb]; + const double *src = ket + bs->offset; + const double *dst = bra + bd->offset; + + for (uint32_t ia0 = 0; ia0 < ta->nsrc; ia0++) { + const gas_link_entry_t *ra = ta->link + + (size_t)ia0 * ta->nlink; + for (uint32_t ka = 0; ka < ta->nlink; ka++) { + const gas_link_entry_t *ea = ra + ka; + if (ea->addr < a0 || ea->addr >= a1) continue; + uint32_t p = gas_link_cre(ea); + uint32_t q = gas_link_des(ea); + for (uint32_t ib0 = 0; ib0 < tb->nsrc; ib0++) { + double c = src[(size_t)ia0 * nb0 + ib0]; + const gas_link_entry_t *rb = tb->link + + (size_t)ib0 * tb->nlink; + for (uint32_t kb = 0; kb < tb->nlink; kb++) { + const gas_link_entry_t *eb = rb + kb; + uint32_t r = gas_link_cre(eb); + uint32_t s = gas_link_des(eb); + dm2ab[rdm2_index(p, q, r, s, norb)] += + (double)(ea->sign * eb->sign) * + dst[(size_t)ea->addr * nb1 + eb->addr] * c; + } + } + } + } +} + +#if defined(__AVX2__) +static inline double rdm_hsum256(__m256d value) +{ + double lane[4]; + _mm256_storeu_pd(lane, value); + return lane[0] + lane[1] + lane[2] + lane[3]; +} +#endif + +static void rdm_dot8_streams(const double *restrict y, + const double *restrict base, size_t stride, + uint32_t n, double value[8]) +{ + for (uint32_t j = 0; j < 8u; j++) value[j] = 0.0; + uint32_t i = 0; +#if defined(__AVX2__) + __m256d sum[8]; + for (uint32_t j = 0; j < 8u; j++) sum[j] = _mm256_setzero_pd(); + for (; i + 4u <= n; i += 4u) { + __m256d vy = _mm256_loadu_pd(y + i); + for (uint32_t j = 0; j < 8u; j++) { + __m256d vt = _mm256_loadu_pd(base + (size_t)j * stride + i); + sum[j] = _mm256_add_pd(sum[j], _mm256_mul_pd(vy, vt)); + } + } + for (uint32_t j = 0; j < 8u; j++) value[j] = rdm_hsum256(sum[j]); +#endif + for (; i < n; i++) { + double yi = y[i]; + for (uint32_t j = 0; j < 8u; j++) { + value[j] += yi * base[(size_t)j * stride + i]; + } + } +} + +static void rdm_dot2x4_streams(const double *restrict y0, + const double *restrict y1, + const double *restrict base, size_t stride, + uint32_t n, double value0[4], double value1[4]) +{ + for (uint32_t j = 0; j < 4u; j++) { + value0[j] = 0.0; + value1[j] = 0.0; + } + uint32_t i = 0; +#if defined(__AVX2__) + __m256d sum0[4], sum1[4]; + for (uint32_t j = 0; j < 4u; j++) { + sum0[j] = _mm256_setzero_pd(); + sum1[j] = _mm256_setzero_pd(); + } + for (; i + 4u <= n; i += 4u) { + __m256d vy0 = _mm256_loadu_pd(y0 + i); + __m256d vy1 = _mm256_loadu_pd(y1 + i); + for (uint32_t j = 0; j < 4u; j++) { + __m256d vt = _mm256_loadu_pd(base + (size_t)j * stride + i); + sum0[j] = _mm256_add_pd(sum0[j], _mm256_mul_pd(vy0, vt)); + sum1[j] = _mm256_add_pd(sum1[j], _mm256_mul_pd(vy1, vt)); + } + } + for (uint32_t j = 0; j < 4u; j++) { + value0[j] = rdm_hsum256(sum0[j]); + value1[j] = rdm_hsum256(sum1[j]); + } +#endif + for (; i < n; i++) { + double a = y0[i], b = y1[i]; + for (uint32_t j = 0; j < 4u; j++) { + double t = base[(size_t)j * stride + i]; + value0[j] += a * t; + value1[j] += b * t; + } + } +} + +static inline size_t rdm_abba_row_offset(const uint32_t *active_row, + uint32_t i) +{ + return active_row[i]; +} + +static void rdm_abba_fused_ops(double *dm2ab, + const uint32_t *active_row, uint32_t nactive, + uint32_t opb, double sign, + const double *y, const double *tbase, + size_t op_stride, uint32_t nat) +{ + uint32_t io = 0; + for (; io + 8u <= nactive; io += 8u) { + double value[8]; + rdm_dot8_streams(y, tbase + (size_t)io * op_stride, + op_stride, nat, value); + for (uint32_t j = 0; j < 8u; j++) { + dm2ab[rdm_abba_row_offset(active_row, io + j) + opb] += + sign * value[j]; + } + } + for (; io + 4u <= nactive; io += 4u) { + double value[4]; + const double *base = tbase + (size_t)io * op_stride; + rdm_dot4(y, base, base + op_stride, + base + 2u * op_stride, base + 3u * op_stride, + nat, value); + for (uint32_t j = 0; j < 4u; j++) { + dm2ab[rdm_abba_row_offset(active_row, io + j) + opb] += + sign * value[j]; + } + } + for (; io < nactive; io++) { + double value = rdm_dot(y, tbase + (size_t)io * op_stride, nat); + dm2ab[rdm_abba_row_offset(active_row, io) + opb] += sign * value; + } +} + +static void rdm_abba_pair2x4(double *dm2ab, + const uint32_t *active_row, uint32_t nactive, + uint32_t opb0, uint32_t opb1, + double sign0, double sign1, + const double *y0, const double *y1, + const double *tbase, size_t op_stride, + uint32_t nat) +{ + uint32_t io = 0; + for (; io + 4u <= nactive; io += 4u) { + double value0[4], value1[4]; + rdm_dot2x4_streams(y0, y1, + tbase + (size_t)io * op_stride, + op_stride, nat, value0, value1); + for (uint32_t j = 0; j < 4u; j++) { + size_t row = rdm_abba_row_offset(active_row, io + j); + dm2ab[row + opb0] += sign0 * value0[j]; + dm2ab[row + opb1] += sign1 * value1[j]; + } + } + if (io < nactive) { + rdm_abba_fused_ops(dm2ab, active_row + io, + nactive - io, opb0, sign0, y0, + tbase + (size_t)io * op_stride, + op_stride, nat); + rdm_abba_fused_ops(dm2ab, active_row + io, + nactive - io, opb1, sign1, y1, + tbase + (size_t)io * op_stride, + op_stride, nat); + } +} + +static uint32_t rdm_abba_beta_tile(uint32_t nb0, uint32_t nactive, + uint32_t nat, size_t target_bytes) +{ + if (nb0 == 0u || nactive == 0u || nat == 0u || target_bytes == 0u) { + return nb0; + } + if ((size_t)nactive > SIZE_MAX / nat / sizeof(double)) return 1u; + size_t denom = (size_t)nactive * nat * sizeof(double); + size_t tile = denom ? target_bytes / denom : nb0; + if (tile < 1u) tile = 1u; + if (tile > nb0) tile = nb0; + return (uint32_t)tile; +} + +static int rdm_build_alpha_hitlist(const gas_space_t *gas, + const gas_rdm_plan_t *plan, + gas_tid_t tida, uint32_t a0, uint32_t a1, + rdm_workspace_t *workspace) +{ + const gas_link_table_t *ta = gas->table + tida; + const rdm_table_ops_t *meta = plan->table_ops + tida; + uint32_t nopair = (uint32_t)gas->norb_tot * gas->norb_tot; + if (rdm_op_map_reserve(workspace, nopair) != GAS_SUCCESS) { + return GAS_ERR_MEMORY; + } + for (uint32_t i = 0; i < nopair; i++) workspace->op_map[i] = -1; + const uint16_t *active = plan->active_op + meta->off; + for (uint32_t i = 0; i < meta->n; i++) workspace->op_map[active[i]] = (int32_t)i; + + uint32_t nhit = 0; + for (uint32_t ia0 = 0; ia0 < ta->nsrc; ia0++) { + const gas_link_entry_t *row = ta->link + (size_t)ia0 * ta->nlink; + for (uint32_t k = 0; k < ta->nlink; k++) { + if (row[k].addr >= a0 && row[k].addr < a1) nhit++; + } + } + if (rdm_hit_reserve(workspace, nhit) != GAS_SUCCESS) { + return GAS_ERR_MEMORY; + } + uint32_t out = 0; + for (uint32_t ia0 = 0; ia0 < ta->nsrc; ia0++) { + const gas_link_entry_t *row = ta->link + (size_t)ia0 * ta->nlink; + for (uint32_t k = 0; k < ta->nlink; k++) { + const gas_link_entry_t *e = row + k; + if (e->addr < a0 || e->addr >= a1) continue; + int32_t io = workspace->op_map[rdm_op(e, + (uint32_t)gas->norb_tot)]; + if (io < 0) continue; + workspace->alpha_hit[out].src_row = ia0; + workspace->alpha_hit[out].rel_addr = e->addr - a0; + workspace->alpha_hit[out].op_index = (uint16_t)io; + workspace->alpha_hit[out].sign = e->sign; + workspace->alpha_hit[out].padding = 0u; + out++; + } + } + workspace->alpha_hit_n = out; + return out == nhit ? GAS_SUCCESS : GAS_ERR_INVALID; +} + +static int rdm_collect_ab_pairs(const gas_space_t *gas, gas_sid_t adst, + uint32_t q0, uint32_t q1, + gas_row_t beta_tables, + rdm_workspace_t *workspace, + uint32_t *npair_out) +{ + gas_row_t blocks = gas->D.by_alpha_row[adst]; + if (q0 > blocks.n) q0 = blocks.n; + if (q1 > blocks.n) q1 = blocks.n; + if (q1 < q0) q1 = q0; + uint32_t maxpair = q1 - q0 < beta_tables.n ? + q1 - q0 : beta_tables.n; + if (rdm_pair_reserve(workspace, maxpair) != GAS_SUCCESS) { + return GAS_ERR_MEMORY; + } + uint32_t q = q0, j = 0u, npair = 0u; + while (q < q1 && j < beta_tables.n) { + gas_bid_t bdst = blocks.off + q; + gas_sid_t bdst_beta = gas->block[bdst].sb; + gas_sid_t table_beta = gas->T.dst[beta_tables.off + j]; + if (bdst_beta < table_beta) { + q++; + } else if (bdst_beta > table_beta) { + j++; + } else { + workspace->pair[npair].first = bdst; + workspace->pair[npair].second = beta_tables.off + j; + npair++; + q++; + j++; + } + } + *npair_out = npair; + return GAS_SUCCESS; +} + +static int rdm2_ab_batch_tiled(const gas_rdm_plan_t *plan, + const double *bra, const double *ket, + gas_bid_t bsrc, gas_tid_t tida, + const rdm_tid_pair_t *pair, uint32_t npair, + uint32_t a0, uint32_t a1, double *dm2ab, + rdm_workspace_t *workspace) +{ + const gas_space_t *gas = plan->gas; + const gas_block_t *bs = gas->block + bsrc; + const rdm_table_ops_t *meta = plan->table_ops + tida; + uint32_t nactive = meta->n; + uint32_t nat = a1 > a0 ? a1 - a0 : 0u; + uint32_t nb0 = gas->sector_nstr[bs->sb]; + if (nactive == 0u || nat == 0u || nb0 == 0u || npair == 0u) { + return GAS_SUCCESS; + } + if (rdm_build_alpha_hitlist(gas, plan, tida, a0, a1, + workspace) != GAS_SUCCESS) { + return GAS_ERR_MEMORY; + } + + uint32_t btile = rdm_abba_beta_tile(nb0, nactive, nat, + plan->abba_t1_target_bytes); + if (btile == 0u) return GAS_SUCCESS; + if ((size_t)btile > SIZE_MAX / nat) return GAS_ERR_MEMORY; + uint32_t max_nb1 = 0u; + size_t all_packed_elems = 0u; + for (uint32_t ip = 0; ip < npair; ip++) { + const gas_block_t *bd = gas->block + pair[ip].first; + uint32_t nb1 = gas->sector_nstr[bd->sb]; + if (nb1 > max_nb1) max_nb1 = nb1; + if ((size_t)nb1 > SIZE_MAX / nat || + all_packed_elems > SIZE_MAX - (size_t)nb1 * nat) { + return GAS_ERR_MEMORY; + } + all_packed_elems += (size_t)nb1 * nat; + } + if ((size_t)max_nb1 > SIZE_MAX / nat) return GAS_ERR_MEMORY; + size_t op_stride = (size_t)btile * nat; + size_t one_packed_elems = (size_t)max_nb1 * nat; + if ((size_t)nactive > SIZE_MAX / op_stride) return GAS_ERR_MEMORY; + size_t t1_elems = (size_t)nactive * op_stride; + int pack_once = 0; + if (t1_elems <= RDM_WORKSPACE_MAX_ELEMS && + all_packed_elems <= RDM_WORKSPACE_MAX_ELEMS - t1_elems && + rdm_arena_reserve(workspace, t1_elems + all_packed_elems) == + GAS_SUCCESS) { + pack_once = 1; + } + if (!pack_once && + (t1_elems > RDM_WORKSPACE_MAX_ELEMS || + one_packed_elems > RDM_WORKSPACE_MAX_ELEMS - t1_elems || + rdm_arena_reserve(workspace, t1_elems + one_packed_elems) != + GAS_SUCCESS)) { + return GAS_ERR_MEMORY; + } + double *t1 = workspace->arena; + double *packed = t1 + t1_elems; + memset(t1, 0, t1_elems * sizeof(*t1)); + const double *src = ket + bs->offset; + const uint32_t *active_row = plan->active_row + meta->off; + + if (pack_once) { + size_t off = 0u; + for (uint32_t ip = 0; ip < npair; ip++) { + const gas_block_t *bd = gas->block + pair[ip].first; + uint32_t nb1 = gas->sector_nstr[bd->sb]; + rdm_transpose_pack(packed + off, bra + bd->offset, nb1, + nat, nb1, a0); + off += (size_t)nb1 * nat; + } + } + + for (uint32_t b0 = 0; b0 < nb0; b0 += btile) { + uint32_t b1 = nb0 - b0 < btile ? nb0 : b0 + btile; + for (uint32_t ih = 0; ih < workspace->alpha_hit_n; ih++) { + const rdm_alpha_hit_t *h = workspace->alpha_hit + ih; + const double *src_row = src + (size_t)h->src_row * nb0; + double *stream = t1 + (size_t)h->op_index * op_stride + + h->rel_addr; + double sign = (double)h->sign; + for (uint32_t ib = b0; ib < b1; ib++) { + stream[(size_t)(ib - b0) * nat] = sign * src_row[ib]; + } + } + + size_t packed_off = 0u; + for (uint32_t ip = 0; ip < npair; ip++) { + gas_bid_t bdst = pair[ip].first; + gas_tid_t tidb = pair[ip].second; + const gas_block_t *bd = gas->block + bdst; + const gas_link_table_t *tb = gas->table + tidb; + uint32_t nb1 = gas->sector_nstr[bd->sb]; + double *dst_t = pack_once ? packed + packed_off : packed; + if (!pack_once) { + rdm_transpose_pack(dst_t, bra + bd->offset, nb1, + nat, nb1, a0); + } + for (uint32_t ib = b0; ib < b1; ib++) { + const gas_link_entry_t *row = tb->link + + (size_t)ib * tb->nlink; + const double *tbase = t1 + + (size_t)(ib - b0) * nat; + uint32_t k = 0u; + for (; k + 1u < tb->nlink; k += 2u) { + const gas_link_entry_t *e0 = row + k; + const gas_link_entry_t *e1 = row + k + 1u; + rdm_abba_pair2x4( + dm2ab, active_row, nactive, + rdm_op(e0, (uint32_t)gas->norb_tot), + rdm_op(e1, (uint32_t)gas->norb_tot), + (double)e0->sign, + (double)e1->sign, + dst_t + (size_t)e0->addr * nat, + dst_t + (size_t)e1->addr * nat, + tbase, op_stride, nat); + } + for (; k < tb->nlink; k++) { + const gas_link_entry_t *e = row + k; + rdm_abba_fused_ops( + dm2ab, active_row, nactive, + rdm_op(e, (uint32_t)gas->norb_tot), + (double)e->sign, + dst_t + (size_t)e->addr * nat, + tbase, op_stride, nat); + } + } + packed_off += (size_t)nb1 * nat; + } + } + return GAS_SUCCESS; +} + +static void rdm2_ab_task(const gas_rdm_plan_t *plan, + const double *bra, const double *ket, + const rdm_task_t *task, double *dm2ab, + rdm_workspace_t *workspace) +{ + const gas_space_t *gas = plan->gas; + gas_row_t blocks; + uint32_t q0, q1, a0, a1; + rdm_task_bounds(gas, task, &blocks, &q0, &q1, &a0, &a1); + + gas_row_t alpha_in = gas->R.row[task->adst]; + for (uint32_t i = 0; i < alpha_in.n; i++) { + gas_sid_t asrc = gas->R.src[alpha_in.off + i]; + gas_tid_t tida = gas->R.tid[alpha_in.off + i]; + gas_row_t source_blocks = gas->D.by_alpha_row[asrc]; + for (uint32_t q = 0; q < source_blocks.n; q++) { + gas_bid_t bsrc = source_blocks.off + q; + gas_sid_t beta0 = gas->block[bsrc].sb; + gas_row_t beta_tables = gas->T.row[beta0]; + uint32_t npair = 0u; + if (rdm_collect_ab_pairs(gas, task->adst, q0, q1, + beta_tables, workspace, &npair) != GAS_SUCCESS) { + for (uint32_t qb = q0; qb < q1; qb++) { + gas_bid_t bdst = blocks.off + qb; + gas_tid_t tidb = gas_find_table( + gas, beta0, gas->block[bdst].sb); + if (tidb == GAS_INVALID_TID) continue; + rdm2_ab_path_fallback( + gas, bra, ket, bsrc, bdst, + tida, tidb, a0, a1, dm2ab); + } + continue; + } + if (npair == 0u) continue; + if (rdm2_ab_batch_tiled( + plan, bra, ket, bsrc, tida, + workspace->pair, npair, a0, a1, + dm2ab, workspace) != GAS_SUCCESS) { + for (uint32_t ip = 0; ip < npair; ip++) { + rdm2_ab_path_fallback( + gas, bra, ket, bsrc, + workspace->pair[ip].first, + tida, workspace->pair[ip].second, + a0, a1, dm2ab); + } + } + } + } +} + +static void rdm12_task_accumulate(const gas_rdm_plan_t *plan, + const double *bra, const double *ket, + const rdm_task_t *task, + rdm_workspace_t *workspace, + size_t n1, size_t n2) +{ + const gas_space_t *gas = plan->gas; + double *work = workspace->data; + double *dm1a = work; + double *dm1b = dm1a + n1; + double *dm2aa = dm1b + n1; + double *dm2ab = dm2aa + n2; + double *dm2bb = dm2ab + n2; + rdm1_alpha_task(gas, bra, ket, task, dm1a); + rdm1_beta_task(gas, bra, ket, task, dm1b); + rdm2_aa_task(gas, bra, ket, task, dm2aa); + rdm2_ab_task(plan, bra, ket, task, dm2ab, workspace); + rdm2_bb_task(gas, bra, ket, task, dm2bb, workspace); +} + +static uint64_t rdm_cost_add(uint64_t a, uint64_t b) +{ + return UINT64_MAX - a < b ? UINT64_MAX : a + b; +} + +static uint64_t rdm_cost_mul(uint64_t a, uint64_t b) +{ + return a != 0u && b > UINT64_MAX / a ? UINT64_MAX : a * b; +} + +static uint64_t rdm_destination_block_cost(const gas_space_t *gas, + gas_bid_t bdst) +{ + const gas_block_t *bd = gas->block + bdst; + uint64_t na1 = gas->sector_nstr[bd->sa]; + uint64_t nb1 = gas->sector_nstr[bd->sb]; + uint64_t cost = 0u; + gas_row_t alpha_reverse = gas->R.row[bd->sa]; + gas_row_t beta_reverse = gas->R.row[bd->sb]; + + /* One-body alpha and beta links. */ + for (uint32_t i = 0; i < alpha_reverse.n; i++) { + gas_sid_t asrc = gas->R.src[alpha_reverse.off + i]; + if (gas_find_block(gas, asrc, bd->sb) != GAS_INVALID_BID) { + const gas_link_table_t *t = gas->table + + gas->R.tid[alpha_reverse.off + i]; + cost = rdm_cost_add(cost, rdm_cost_mul( + rdm_cost_mul(t->nsrc, t->nlink), nb1)); + } + } + for (uint32_t i = 0; i < beta_reverse.n; i++) { + gas_sid_t bsrc = gas->R.src[beta_reverse.off + i]; + if (gas_find_block(gas, bd->sa, bsrc) != GAS_INVALID_BID) { + const gas_link_table_t *t = gas->table + + gas->R.tid[beta_reverse.off + i]; + cost = rdm_cost_add(cost, rdm_cost_mul( + rdm_cost_mul(t->nsrc, t->nlink), na1)); + } + } + + /* Same-spin alpha paths through a common middle sector. */ + gas_row_t alpha_sources = gas->D.by_beta_row[bd->sb]; + const gas_sid_t *alpha_middle_reverse = + gas->R.src + alpha_reverse.off; + for (uint32_t q = 0; q < alpha_sources.n; q++) { + gas_sid_t asrc = gas->D.by_beta_sid[alpha_sources.off + q]; + gas_row_t forward = gas->T.row[asrc]; + const gas_sid_t *middle = gas->T.dst + forward.off; + uint32_t i = 0, j = 0; + while (i < forward.n && j < alpha_reverse.n) { + if (middle[i] < alpha_middle_reverse[j]) { + i++; + } else if (middle[i] > alpha_middle_reverse[j]) { + j++; + } else { + const gas_link_table_t *t1 = gas->table + forward.off + i; + const gas_link_table_t *t2 = gas->table + + gas->R.tid[alpha_reverse.off + j]; + uint64_t w = rdm_cost_mul(t1->nsrc, t1->nlink); + w = rdm_cost_mul(w, t2->nlink); + w = rdm_cost_mul(w, nb1); + cost = rdm_cost_add(cost, w); + i++; + j++; + } + } + } + + /* Same-spin beta paths through a common middle sector. */ + gas_row_t beta_sources = gas->D.by_alpha_row[bd->sa]; + const gas_sid_t *beta_middle_reverse = gas->R.src + beta_reverse.off; + for (uint32_t q = 0; q < beta_sources.n; q++) { + gas_bid_t bsrc = beta_sources.off + q; + gas_sid_t bsrc_sid = gas->block[bsrc].sb; + gas_row_t forward = gas->T.row[bsrc_sid]; + const gas_sid_t *middle = gas->T.dst + forward.off; + uint32_t i = 0, j = 0; + while (i < forward.n && j < beta_reverse.n) { + if (middle[i] < beta_middle_reverse[j]) { + i++; + } else if (middle[i] > beta_middle_reverse[j]) { + j++; + } else { + const gas_link_table_t *t1 = gas->table + forward.off + i; + const gas_link_table_t *t2 = gas->table + + gas->R.tid[beta_reverse.off + j]; + uint64_t w = rdm_cost_mul(t1->nsrc, t1->nlink); + w = rdm_cost_mul(w, t2->nlink); + w = rdm_cost_mul(w, na1); + cost = rdm_cost_add(cost, w); + i++; + j++; + } + } + } + + /* Opposite-spin direct alpha/beta link products. */ + for (uint32_t i = 0; i < alpha_reverse.n; i++) { + gas_sid_t asrc = gas->R.src[alpha_reverse.off + i]; + const gas_link_table_t *ta = gas->table + + gas->R.tid[alpha_reverse.off + i]; + gas_row_t source_blocks = gas->D.by_alpha_row[asrc]; + for (uint32_t q = 0; q < source_blocks.n; q++) { + gas_sid_t bsrc = gas->block[source_blocks.off + q].sb; + gas_tid_t tidb = gas_find_table(gas, bsrc, bd->sb); + if (tidb != GAS_INVALID_TID) { + const gas_link_table_t *tb = gas->table + tidb; + uint64_t wa = rdm_cost_mul(ta->nsrc, ta->nlink); + uint64_t wb = rdm_cost_mul(tb->nsrc, tb->nlink); + cost = rdm_cost_add(cost, rdm_cost_mul(wa, wb)); + } + } + } + return cost ? cost : 1u; +} + +static int rdm_task_compare(const void *pa, const void *pb) +{ + const rdm_task_t *a = pa; + const rdm_task_t *b = pb; + if (a->cost != b->cost) return a->cost < b->cost ? 1 : -1; + if (a->adst != b->adst) return (a->adst > b->adst) - (a->adst < b->adst); + if (a->q0 != b->q0) return (a->q0 > b->q0) - (a->q0 < b->q0); + return (a->a0 > b->a0) - (a->a0 < b->a0); +} + +static int rdm_append_task(rdm_task_t **task, uint32_t *n, uint32_t *capacity, + gas_sid_t adst, uint32_t q0, uint32_t q1, + uint32_t a0, uint32_t a1, uint64_t cost) +{ + if (*n == *capacity) { + uint32_t next = *capacity ? 2u * *capacity : 256u; + if (next < *capacity) return GAS_ERR_MEMORY; + rdm_task_t *p = realloc(*task, (size_t)next * sizeof(*p)); + if (p == 0) return GAS_ERR_MEMORY; + *task = p; + *capacity = next; + } + (*task)[*n].cost = cost ? cost : 1u; + (*task)[*n].adst = adst; + (*task)[*n].q0 = q0; + (*task)[*n].q1 = q1; + (*task)[*n].a0 = a0; + (*task)[*n].a1 = a1; + (*n)++; + return GAS_SUCCESS; +} + +static uint32_t rdm_max_threads(void) +{ +#ifdef _OPENMP + int n = omp_get_max_threads(); + return n > 0 ? (uint32_t)n : 1u; +#else + return 1u; +#endif +} + +static rdm_task_t *rdm_build_tasks(const gas_space_t *gas, uint32_t nthread, + uint32_t *ntask) +{ + rdm_task_t *task = 0; + uint32_t n = 0, capacity = 0; + uint64_t total = 0; + uint64_t *block_cost = malloc((size_t)gas->nblock * sizeof(*block_cost)); + if (block_cost == 0) return 0; + + for (gas_bid_t b = 0; b < gas->nblock; b++) { + block_cost[b] = rdm_destination_block_cost(gas, b); + total = rdm_cost_add(total, block_cost[b]); + } + uint64_t target = total / + ((uint64_t)(nthread ? nthread : 1u) * RDM_TASK_FACTOR); + if (target == 0u) target = 1u; + + for (gas_sid_t s = 0; s < gas->nsector; s++) { + gas_row_t row = gas->D.by_alpha_row[s]; + uint32_t na = gas->sector_nstr[s]; + if (row.n == 0u || na == 0u) continue; + uint32_t q0 = 0; + uint64_t chunk = 0; + + for (uint32_t q = 0; q < row.n; q++) { + uint64_t w = block_cost[row.off + q]; + uint64_t want64 = (w + target - 1u) / target; + uint32_t by_rows = (na + RDM_MIN_ALPHA_TILE - 1u) / + RDM_MIN_ALPHA_TILE; + uint32_t split = want64 > UINT32_MAX ? + UINT32_MAX : (uint32_t)want64; + if (split > RDM_MAX_ALPHA_SPLIT) { + split = RDM_MAX_ALPHA_SPLIT; + } + if (split > by_rows) split = by_rows; + + if (split > 1u) { + if (q > q0 && rdm_append_task(&task, &n, &capacity, + s, q0, q, 0u, na, chunk) != GAS_SUCCESS) { + free(block_cost); + free(task); + return 0; + } + uint32_t a0 = 0; + for (uint32_t part = 0; part < split; part++) { + uint32_t left = na - a0; + uint32_t remain = split - part; + uint32_t width = (left + remain - 1u) / remain; + uint32_t a1 = a0 + width; + uint64_t part_cost = rdm_cost_mul(w, width) / na; + if (rdm_append_task(&task, &n, &capacity, + s, q, q + 1u, a0, a1, + part_cost) != GAS_SUCCESS) { + free(block_cost); + free(task); + return 0; + } + a0 = a1; + } + q0 = q + 1u; + chunk = 0; + continue; + } + if (q > q0 && chunk != 0u && + rdm_cost_add(chunk, w) > target) { + if (rdm_append_task(&task, &n, &capacity, + s, q0, q, 0u, na, chunk) != GAS_SUCCESS) { + free(block_cost); + free(task); + return 0; + } + q0 = q; + chunk = 0; + } + chunk = rdm_cost_add(chunk, w); + } + if (q0 < row.n && rdm_append_task(&task, &n, &capacity, + s, q0, row.n, 0u, na, chunk) != GAS_SUCCESS) { + free(block_cost); + free(task); + return 0; + } + } + free(block_cost); + qsort(task, n, sizeof(*task), rdm_task_compare); + *ntask = n; + return task; +} + +/* Convert to for one spin species. */ +static void reorder_same_spin(double *dm2, const double *dm1, uint32_t norb) +{ + for (uint32_t p = 0; p < norb; p++) { + for (uint32_t r = 0; r < norb; r++) { + for (uint32_t q = 0; q < norb; q++) { + size_t d = rdm2_index(p, q, r, q, norb); + double old = dm2[d]; + double delta = q == r ? dm1[(size_t)q * norb + p] : 0.0; + dm2[d] = delta - old; + for (uint32_t s = q + 1u; s < norb; s++) { + size_t i = rdm2_index(p, q, r, s, norb); + size_t j = rdm2_index(p, s, r, q, norb); + double x = dm2[i]; + double y = dm2[j]; + double di = s == r ? + dm1[(size_t)q * norb + p] : 0.0; + double dj = q == r ? + dm1[(size_t)s * norb + p] : 0.0; + dm2[i] = di - y; + dm2[j] = dj - x; + } + } + } + } +} + +static int rdm_output_validate(const gas_space_t *gas, + const double *bra, const double *ket, + const double *dm1a, const double *dm1b, + const double *dm2aa, const double *dm2ab, + const double *dm2bb) +{ + if (rdm_validate(gas, bra, ket) != GAS_SUCCESS || + dm1a == 0 || dm1b == 0 || dm1a == dm1b || + dm1a == bra || dm1a == ket || dm1b == bra || dm1b == ket) { + return GAS_ERR_INVALID; + } + if (dm2aa != 0 || dm2ab != 0 || dm2bb != 0) { + if (dm2aa == 0 || dm2ab == 0 || dm2bb == 0 || + dm2aa == dm2ab || dm2aa == dm2bb || dm2ab == dm2bb) { + return GAS_ERR_INVALID; + } + if (dm2aa == bra || dm2aa == ket || + dm2ab == bra || dm2ab == ket || + dm2bb == bra || dm2bb == ket || + dm1a == dm2aa || dm1a == dm2ab || dm1a == dm2bb || + dm1b == dm2aa || dm1b == dm2ab || dm1b == dm2bb) { + return GAS_ERR_INVALID; + } + } + return GAS_SUCCESS; +} + +static int rdm_workspace_reserve(rdm_workspace_t *workspace, size_t need) +{ + if (need <= workspace->capacity) return GAS_SUCCESS; + if (need > SIZE_MAX / sizeof(*workspace->data)) return GAS_ERR_MEMORY; + double *p = realloc(workspace->data, need * sizeof(*p)); + if (p == 0) return GAS_ERR_MEMORY; + workspace->data = p; + workspace->capacity = need; + return GAS_SUCCESS; +} + +static int rdm_arena_reserve(rdm_workspace_t *workspace, size_t need) +{ + if (need <= workspace->arena_capacity) return GAS_SUCCESS; + if (need > SIZE_MAX / sizeof(*workspace->arena)) return GAS_ERR_MEMORY; + double *p = realloc(workspace->arena, need * sizeof(*p)); + if (p == 0) return GAS_ERR_MEMORY; + workspace->arena = p; + workspace->arena_capacity = need; + return GAS_SUCCESS; +} + +static int rdm_pair_reserve(rdm_workspace_t *workspace, size_t need) +{ + if (need <= workspace->pair_capacity) return GAS_SUCCESS; + if (need > SIZE_MAX / sizeof(*workspace->pair)) return GAS_ERR_MEMORY; + rdm_tid_pair_t *p = realloc(workspace->pair, need * sizeof(*p)); + if (p == 0) return GAS_ERR_MEMORY; + workspace->pair = p; + workspace->pair_capacity = need; + return GAS_SUCCESS; +} + +static int rdm_hit_reserve(rdm_workspace_t *workspace, size_t need) +{ + if (need <= workspace->alpha_hit_capacity) return GAS_SUCCESS; + if (need > SIZE_MAX / sizeof(*workspace->alpha_hit)) return GAS_ERR_MEMORY; + rdm_alpha_hit_t *p = realloc(workspace->alpha_hit, need * sizeof(*p)); + if (p == 0) return GAS_ERR_MEMORY; + workspace->alpha_hit = p; + workspace->alpha_hit_capacity = need; + return GAS_SUCCESS; +} + +static int rdm_op_map_reserve(rdm_workspace_t *workspace, size_t need) +{ + if (need <= workspace->op_map_capacity) return GAS_SUCCESS; + if (need > SIZE_MAX / sizeof(*workspace->op_map)) return GAS_ERR_MEMORY; + int32_t *p = realloc(workspace->op_map, need * sizeof(*p)); + if (p == 0) return GAS_ERR_MEMORY; + workspace->op_map = p; + workspace->op_map_capacity = need; + return GAS_SUCCESS; +} + +static int rdm_plan_resize_workspaces(gas_rdm_plan_t *plan, + uint32_t nthread, size_t size) +{ + if (nthread > plan->nworkspace) { + rdm_workspace_t *p = realloc( + plan->workspace, (size_t)nthread * sizeof(*p)); + if (p == 0) return GAS_ERR_MEMORY; + plan->workspace = p; + for (uint32_t i = plan->nworkspace; i < nthread; i++) { + memset(plan->workspace + i, 0, sizeof(plan->workspace[i])); + } + plan->nworkspace = nthread; + } + for (uint32_t i = 0; i < nthread; i++) { + if (rdm_workspace_reserve(plan->workspace + i, size) != GAS_SUCCESS) { + return GAS_ERR_MEMORY; + } + } + return GAS_SUCCESS; +} + +static int rdm_plan_build_active_ops(gas_rdm_plan_t *plan) +{ + const gas_space_t *gas = plan->gas; + uint32_t nopair = (uint32_t)gas->norb_tot * gas->norb_tot; + uint8_t *seen = calloc(nopair, sizeof(*seen)); + rdm_table_ops_t *meta = calloc(gas->ntable, sizeof(*meta)); + if ((seen == 0 && nopair != 0u) || (meta == 0 && gas->ntable != 0u)) { + free(seen); + free(meta); + return GAS_ERR_MEMORY; + } + + uint64_t total = 0u; + for (gas_tid_t tid = 0; tid < gas->ntable; tid++) { + const gas_link_table_t *table = gas->table + tid; + if (nopair != 0u) { + memset(seen, 0, nopair * sizeof(*seen)); + } + uint32_t count = 0u; + size_t nentry = (size_t)table->nsrc * table->nlink; + for (size_t i = 0; i < nentry; i++) { + uint32_t op = rdm_op(table->link + i, + (uint32_t)gas->norb_tot); + if (!seen[op]) { + seen[op] = 1u; + count++; + } + } + if (count > UINT16_MAX || total + count > UINT32_MAX) { + free(seen); + free(meta); + return GAS_ERR_INVALID; + } + meta[tid].off = (uint32_t)total; + meta[tid].n = (uint16_t)count; + total += count; + } + + if (total == 0u) { + free(seen); + plan->table_ops = meta; + return GAS_SUCCESS; + } + uint16_t *active = malloc((size_t)total * sizeof(*active)); + uint32_t *active_row = malloc((size_t)total * sizeof(*active_row)); + if (active == 0 || active_row == 0) { + free(seen); + free(meta); + free(active); + free(active_row); + return GAS_ERR_MEMORY; + } + for (gas_tid_t tid = 0; tid < gas->ntable; tid++) { + const gas_link_table_t *table = gas->table + tid; + if (nopair != 0u) { + memset(seen, 0, nopair * sizeof(*seen)); + } + size_t nentry = (size_t)table->nsrc * table->nlink; + for (size_t i = 0; i < nentry; i++) { + seen[rdm_op(table->link + i, + (uint32_t)gas->norb_tot)] = 1u; + } + uint32_t out = meta[tid].off; + for (uint32_t op = 0; op < nopair; op++) { + if (seen[op]) { + active[out] = (uint16_t)op; + active_row[out] = + (uint32_t)((uint64_t)op * nopair); + out++; + } + } + } + free(seen); + plan->table_ops = meta; + plan->active_op = active; + plan->active_row = active_row; + return GAS_SUCCESS; +} + +static int rdm_plan_rebuild_tasks(gas_rdm_plan_t *plan) +{ + uint32_t nthread = rdm_max_threads(); + if (plan->task != 0 && plan->task_threads == nthread) { + return GAS_SUCCESS; + } + uint32_t ntask = 0; + rdm_task_t *task = rdm_build_tasks(plan->gas, nthread, &ntask); + if (task == 0 && plan->gas->nblock != 0u) return GAS_ERR_MEMORY; + free(plan->task); + plan->task = task; + plan->ntask = ntask; + plan->task_threads = nthread; + return GAS_SUCCESS; +} + +static void rdm1_serial(const gas_space_t *gas, + const double *bra, const double *ket, + double *dm1a, double *dm1b) +{ + size_t n1 = (size_t)gas->norb_tot * gas->norb_tot; + memset(dm1a, 0, n1 * sizeof(*dm1a)); + memset(dm1b, 0, n1 * sizeof(*dm1b)); + rdm1_alpha(gas, bra, ket, dm1a); + rdm1_beta(gas, bra, ket, dm1b); +} + +static void rdm12_serial(const gas_space_t *gas, + const double *bra, const double *ket, + double *dm1a, double *dm1b, + double *dm2aa, double *dm2ab, double *dm2bb) +{ + uint32_t norb = (uint32_t)gas->norb_tot; + size_t n1 = (size_t)norb * norb; + size_t n2 = n1 * n1; + memset(dm1a, 0, n1 * sizeof(*dm1a)); + memset(dm1b, 0, n1 * sizeof(*dm1b)); + memset(dm2aa, 0, n2 * sizeof(*dm2aa)); + memset(dm2ab, 0, n2 * sizeof(*dm2ab)); + memset(dm2bb, 0, n2 * sizeof(*dm2bb)); + rdm1_alpha(gas, bra, ket, dm1a); + rdm1_beta(gas, bra, ket, dm1b); + rdm2_aa(gas, bra, ket, dm2aa); + rdm2_ab(gas, bra, ket, dm2ab); + rdm2_bb(gas, bra, ket, dm2bb); + reorder_same_spin(dm2aa, dm1a, norb); + reorder_same_spin(dm2bb, dm1b, norb); +} + +static int rdm12_tasks_serial(gas_rdm_plan_t *plan, + const double *bra, const double *ket, + double *dm1a, double *dm1b, + double *dm2aa, double *dm2ab, double *dm2bb, + size_t n1, size_t n2) +{ + size_t local_size = 2u * n1 + 3u * n2; + if (rdm_plan_resize_workspaces(plan, 1u, local_size) != GAS_SUCCESS) { + return GAS_ERR_MEMORY; + } + rdm_workspace_t *workspace = plan->workspace; + memset(workspace->data, 0, local_size * sizeof(*workspace->data)); + workspace->alpha_hit_n = 0u; + for (uint32_t tt = 0; tt < plan->ntask; tt++) { + rdm12_task_accumulate(plan, bra, ket, plan->task + tt, + workspace, n1, n2); + } + const double *local = workspace->data; + memcpy(dm1a, local, n1 * sizeof(*dm1a)); + memcpy(dm1b, local + n1, n1 * sizeof(*dm1b)); + local += 2u * n1; + memcpy(dm2aa, local, n2 * sizeof(*dm2aa)); + memcpy(dm2ab, local + n2, n2 * sizeof(*dm2ab)); + memcpy(dm2bb, local + 2u * n2, n2 * sizeof(*dm2bb)); + uint32_t norb = (uint32_t)plan->gas->norb_tot; + reorder_same_spin(dm2aa, dm1a, norb); + reorder_same_spin(dm2bb, dm1b, norb); + return GAS_SUCCESS; +} + +#ifdef _OPENMP +static void rdm1_parallel(gas_rdm_plan_t *plan, + const double *bra, const double *ket, + double *dm1a, double *dm1b, size_t n1) +{ + uint32_t active_threads = 1u; + +#pragma omp parallel shared(active_threads) + { + uint32_t tid = (uint32_t)omp_get_thread_num(); + double *local = plan->workspace[tid].data; +#pragma omp single + active_threads = (uint32_t)omp_get_num_threads(); + memset(local, 0, 2u * n1 * sizeof(*local)); +#pragma omp barrier +#pragma omp for schedule(dynamic, 1) + for (int64_t tt = 0; tt < (int64_t)plan->ntask; tt++) { + rdm1_alpha_task(plan->gas, bra, ket, plan->task + tt, local); + rdm1_beta_task(plan->gas, bra, ket, plan->task + tt, + local + n1); + } +#pragma omp for schedule(static) + for (int64_t i = 0; i < (int64_t)n1; i++) { + double a = 0.0, b = 0.0; + for (uint32_t t = 0; t < active_threads; t++) { + const double *w = plan->workspace[t].data; + a += w[i]; + b += w[n1 + (size_t)i]; + } + dm1a[i] = a; + dm1b[i] = b; + } + } +} + +static void rdm12_parallel(gas_rdm_plan_t *plan, + const double *bra, const double *ket, + double *dm1a, double *dm1b, + double *dm2aa, double *dm2ab, double *dm2bb, + size_t n1, size_t n2) +{ + uint32_t active_threads = 1u; + size_t local_size = 2u * n1 + 3u * n2; + +#pragma omp parallel shared(active_threads) + { + uint32_t tid = (uint32_t)omp_get_thread_num(); + rdm_workspace_t *workspace = plan->workspace + tid; + double *local = workspace->data; +#pragma omp single + active_threads = (uint32_t)omp_get_num_threads(); + memset(local, 0, local_size * sizeof(*local)); + workspace->alpha_hit_n = 0u; +#pragma omp barrier +#pragma omp for schedule(dynamic, 1) + for (int64_t tt = 0; tt < (int64_t)plan->ntask; tt++) { + rdm12_task_accumulate(plan, bra, ket, + plan->task + tt, workspace, + n1, n2); + } +#pragma omp for schedule(static) + for (int64_t i = 0; i < (int64_t)n1; i++) { + double a = 0.0, b = 0.0; + for (uint32_t t = 0; t < active_threads; t++) { + const double *w = plan->workspace[t].data; + a += w[i]; + b += w[n1 + (size_t)i]; + } + dm1a[i] = a; + dm1b[i] = b; + } +#pragma omp for schedule(static) + for (int64_t i = 0; i < (int64_t)n2; i++) { + double aa = 0.0, ab = 0.0, bb = 0.0; + for (uint32_t t = 0; t < active_threads; t++) { + const double *w = plan->workspace[t].data + 2u * n1; + aa += w[i]; + ab += w[n2 + (size_t)i]; + bb += w[2u * n2 + (size_t)i]; + } + dm2aa[i] = aa; + dm2ab[i] = ab; + dm2bb[i] = bb; + } + } + uint32_t norb = (uint32_t)plan->gas->norb_tot; + reorder_same_spin(dm2aa, dm1a, norb); + reorder_same_spin(dm2bb, dm1b, norb); +} +#endif + +int fci_rdm_gas_plan_create(gas_rdm_plan_t **out, const gas_space_t *gas) +{ + if (out == 0 || gas == 0 || gas->link_format != GAS_LINK_RAW || + gas->norb_tot <= 0) return GAS_ERR_INVALID; + *out = 0; + gas_rdm_plan_t *plan = calloc(1, sizeof(*plan)); + if (plan == 0) return GAS_ERR_MEMORY; + plan->gas = gas; + plan->abba_t1_target_bytes = rdm_abba_t1_target_bytes; + int status = rdm_plan_build_active_ops(plan); + if (status != GAS_SUCCESS) { + fci_rdm_gas_plan_free(plan); + return status; + } + if (rdm_plan_rebuild_tasks(plan) != GAS_SUCCESS) { + fci_rdm_gas_plan_free(plan); + return GAS_ERR_MEMORY; + } + *out = plan; + return GAS_SUCCESS; +} + +int fci_rdm_gas_plan_make_rdm1s(gas_rdm_plan_t *plan, + const double *bra, const double *ket, + double *dm1a, double *dm1b) +{ + if (plan == 0 || rdm_output_validate(plan->gas, bra, ket, + dm1a, dm1b, 0, 0, 0) != GAS_SUCCESS) return GAS_ERR_INVALID; + if (rdm_plan_rebuild_tasks(plan) != GAS_SUCCESS) return GAS_ERR_MEMORY; +#ifdef _OPENMP + size_t n1 = (size_t)plan->gas->norb_tot * plan->gas->norb_tot; + uint32_t nthread = rdm_max_threads(); + if (nthread > 1u && plan->ntask > 1u) { + if (rdm_plan_resize_workspaces(plan, nthread, 2u * n1) != GAS_SUCCESS) { + return GAS_ERR_MEMORY; + } + rdm1_parallel(plan, bra, ket, dm1a, dm1b, n1); + return GAS_SUCCESS; + } +#endif + rdm1_serial(plan->gas, bra, ket, dm1a, dm1b); + return GAS_SUCCESS; +} + +int fci_rdm_gas_plan_make_rdm12s(gas_rdm_plan_t *plan, + const double *bra, const double *ket, + double *dm1a, double *dm1b, + double *dm2aa, double *dm2ab, double *dm2bb) +{ + if (plan == 0 || rdm_output_validate(plan->gas, bra, ket, + dm1a, dm1b, dm2aa, dm2ab, dm2bb) != GAS_SUCCESS) { + return GAS_ERR_INVALID; + } + if (plan->abba_t1_target_bytes != rdm_abba_t1_target_bytes) { + return GAS_ERR_INVALID; + } + if (rdm_plan_rebuild_tasks(plan) != GAS_SUCCESS) return GAS_ERR_MEMORY; + uint32_t norb = (uint32_t)plan->gas->norb_tot; + size_t n1 = (size_t)norb * norb; + size_t n2 = n1 * n1; + if (plan->ntask == 0u) { + rdm12_serial(plan->gas, bra, ket, dm1a, dm1b, + dm2aa, dm2ab, dm2bb); + return GAS_SUCCESS; + } +#ifdef _OPENMP + uint32_t nthread = rdm_max_threads(); + if (nthread > 1u && plan->ntask > 1u) { + size_t local_size = 2u * n1 + 3u * n2; + if (rdm_plan_resize_workspaces(plan, nthread, + local_size) != GAS_SUCCESS) { + return GAS_ERR_MEMORY; + } + rdm12_parallel(plan, bra, ket, dm1a, dm1b, + dm2aa, dm2ab, dm2bb, n1, n2); + return GAS_SUCCESS; + } +#endif + return rdm12_tasks_serial(plan, bra, ket, dm1a, dm1b, + dm2aa, dm2ab, dm2bb, n1, n2); +} + +void fci_rdm_gas_plan_free(gas_rdm_plan_t *plan) +{ + if (plan == 0) return; + for (uint32_t i = 0; i < plan->nworkspace; i++) { + free(plan->workspace[i].data); + free(plan->workspace[i].arena); + free(plan->workspace[i].pair); + free(plan->workspace[i].alpha_hit); + free(plan->workspace[i].op_map); + } + free(plan->workspace); + free(plan->task); + free(plan->table_ops); + free(plan->active_op); + free(plan->active_row); + free(plan); +} + +uint32_t fci_rdm_gas_plan_task_count(const gas_rdm_plan_t *plan) +{ + return plan ? plan->ntask : 0u; +} + +uint64_t fci_rdm_gas_plan_workspace_bytes(const gas_rdm_plan_t *plan) +{ + uint64_t bytes = 0; + if (plan == 0) return 0u; + for (uint32_t i = 0; i < plan->nworkspace; i++) { + const rdm_workspace_t *w = plan->workspace + i; + bytes = rdm_cost_add(bytes, rdm_cost_mul( + w->capacity, sizeof(*w->data))); + bytes = rdm_cost_add(bytes, rdm_cost_mul( + w->arena_capacity, sizeof(*w->arena))); + bytes = rdm_cost_add(bytes, rdm_cost_mul( + w->pair_capacity, sizeof(*w->pair))); + bytes = rdm_cost_add(bytes, rdm_cost_mul( + w->alpha_hit_capacity, sizeof(*w->alpha_hit))); + bytes = rdm_cost_add(bytes, rdm_cost_mul( + w->op_map_capacity, sizeof(*w->op_map))); + } + return bytes; +} diff --git a/pyscf/lib/mcscf/fci_string_gas.c b/pyscf/lib/mcscf/fci_string_gas.c new file mode 100644 index 000000000..e344a737a --- /dev/null +++ b/pyscf/lib/mcscf/fci_string_gas.c @@ -0,0 +1,1739 @@ +/* Copyright 2026 The PySCF Developers. All Rights Reserved. + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. + + * + * Author: Yi Deng + */ + +#include "fci_gas.h" + +#include +#include + +#ifdef _OPENMP +static const uint64_t LINK_BUILD_PARALLEL_MIN = 1048576u; +#endif + +/* ========================================================================== */ +/* 1. Low-level utilities */ +/* ========================================================================== */ + +static inline int popcnt64(uint64_t x) +{ + return __builtin_popcountll(x); +} + +static inline int ctz64(uint64_t x) +{ + return __builtin_ctzll(x); +} + +static inline uint64_t lowbits(int n) +{ + return (1ULL << n) - 1ULL; +} + +/* + * Branchless C(n,k) lookup; valid for 0 <= n,k <= GAS_MAX_LOCAL_ORB + */ +static const uint32_t binom_tab[(GAS_MAX_LOCAL_ORB + 1) * 32] = { + /* n = 0 */ + 1u, 0u, 0u, 0u, 0u, 0u, 0u, 0u, + 0u, 0u, 0u, 0u, 0u, 0u, 0u, 0u, + 0u, 0u, 0u, 0u, 0u, 0u, 0u, 0u, + 0u, 0u, 0u, 0u, 0u, 0u, 0u, 0u, + + /* n = 1 */ + 1u, 1u, 0u, 0u, 0u, 0u, 0u, 0u, + 0u, 0u, 0u, 0u, 0u, 0u, 0u, 0u, + 0u, 0u, 0u, 0u, 0u, 0u, 0u, 0u, + 0u, 0u, 0u, 0u, 0u, 0u, 0u, 0u, + + /* n = 2 */ + 1u, 2u, 1u, 0u, 0u, 0u, 0u, 0u, + 0u, 0u, 0u, 0u, 0u, 0u, 0u, 0u, + 0u, 0u, 0u, 0u, 0u, 0u, 0u, 0u, + 0u, 0u, 0u, 0u, 0u, 0u, 0u, 0u, + + /* n = 3 */ + 1u, 3u, 3u, 1u, 0u, 0u, 0u, 0u, + 0u, 0u, 0u, 0u, 0u, 0u, 0u, 0u, + 0u, 0u, 0u, 0u, 0u, 0u, 0u, 0u, + 0u, 0u, 0u, 0u, 0u, 0u, 0u, 0u, + + /* n = 4 */ + 1u, 4u, 6u, 4u, 1u, 0u, 0u, 0u, + 0u, 0u, 0u, 0u, 0u, 0u, 0u, 0u, + 0u, 0u, 0u, 0u, 0u, 0u, 0u, 0u, + 0u, 0u, 0u, 0u, 0u, 0u, 0u, 0u, + + /* n = 5 */ + 1u, 5u, 10u, 10u, 5u, 1u, 0u, 0u, + 0u, 0u, 0u, 0u, 0u, 0u, 0u, 0u, + 0u, 0u, 0u, 0u, 0u, 0u, 0u, 0u, + 0u, 0u, 0u, 0u, 0u, 0u, 0u, 0u, + + /* n = 6 */ + 1u, 6u, 15u, 20u, 15u, 6u, 1u, 0u, + 0u, 0u, 0u, 0u, 0u, 0u, 0u, 0u, + 0u, 0u, 0u, 0u, 0u, 0u, 0u, 0u, + 0u, 0u, 0u, 0u, 0u, 0u, 0u, 0u, + + /* n = 7 */ + 1u, 7u, 21u, 35u, 35u, 21u, 7u, 1u, + 0u, 0u, 0u, 0u, 0u, 0u, 0u, 0u, + 0u, 0u, 0u, 0u, 0u, 0u, 0u, 0u, + 0u, 0u, 0u, 0u, 0u, 0u, 0u, 0u, + + /* n = 8 */ + 1u, 8u, 28u, 56u, 70u, 56u, 28u, 8u, + 1u, 0u, 0u, 0u, 0u, 0u, 0u, 0u, + 0u, 0u, 0u, 0u, 0u, 0u, 0u, 0u, + 0u, 0u, 0u, 0u, 0u, 0u, 0u, 0u, + + /* n = 9 */ + 1u, 9u, 36u, 84u, 126u, 126u, 84u, 36u, + 9u, 1u, 0u, 0u, 0u, 0u, 0u, 0u, + 0u, 0u, 0u, 0u, 0u, 0u, 0u, 0u, + 0u, 0u, 0u, 0u, 0u, 0u, 0u, 0u, + + /* n = 10 */ + 1u, 10u, 45u, 120u, 210u, 252u, 210u, 120u, + 45u, 10u, 1u, 0u, 0u, 0u, 0u, 0u, + 0u, 0u, 0u, 0u, 0u, 0u, 0u, 0u, + 0u, 0u, 0u, 0u, 0u, 0u, 0u, 0u, + + /* n = 11 */ + 1u, 11u, 55u, 165u, 330u, 462u, 462u, 330u, + 165u, 55u, 11u, 1u, 0u, 0u, 0u, 0u, + 0u, 0u, 0u, 0u, 0u, 0u, 0u, 0u, + 0u, 0u, 0u, 0u, 0u, 0u, 0u, 0u, + + /* n = 12 */ + 1u, 12u, 66u, 220u, 495u, 792u, 924u, 792u, + 495u, 220u, 66u, 12u, 1u, 0u, 0u, 0u, + 0u, 0u, 0u, 0u, 0u, 0u, 0u, 0u, + 0u, 0u, 0u, 0u, 0u, 0u, 0u, 0u, + + /* n = 13 */ + 1u, 13u, 78u, 286u, 715u, 1287u, 1716u, 1716u, + 1287u, 715u, 286u, 78u, 13u, 1u, 0u, 0u, + 0u, 0u, 0u, 0u, 0u, 0u, 0u, 0u, + 0u, 0u, 0u, 0u, 0u, 0u, 0u, 0u, + + /* n = 14 */ + 1u, 14u, 91u, 364u, 1001u, 2002u, 3003u, 3432u, + 3003u, 2002u, 1001u, 364u, 91u, 14u, 1u, 0u, + 0u, 0u, 0u, 0u, 0u, 0u, 0u, 0u, + 0u, 0u, 0u, 0u, 0u, 0u, 0u, 0u, + + /* n = 15 */ + 1u, 15u, 105u, 455u, 1365u, 3003u, 5005u, 6435u, + 6435u, 5005u, 3003u, 1365u, 455u, 105u, 15u, 1u, + 0u, 0u, 0u, 0u, 0u, 0u, 0u, 0u, + 0u, 0u, 0u, 0u, 0u, 0u, 0u, 0u, + + /* n = 16 */ + 1u, 16u, 120u, 560u, 1820u, 4368u, 8008u, 11440u, + 12870u, 11440u, 8008u, 4368u, 1820u, 560u, 120u, 16u, + 1u, 0u, 0u, 0u, 0u, 0u, 0u, 0u, + 0u, 0u, 0u, 0u, 0u, 0u, 0u, 0u, + + /* n = 17 */ + 1u, 17u, 136u, 680u, 2380u, 6188u, 12376u, 19448u, + 24310u, 24310u, 19448u, 12376u, 6188u, 2380u, 680u, 136u, + 17u, 1u, 0u, 0u, 0u, 0u, 0u, 0u, + 0u, 0u, 0u, 0u, 0u, 0u, 0u, 0u, + + /* n = 18 */ + 1u, 18u, 153u, 816u, 3060u, 8568u, 18564u, 31824u, + 43758u, 48620u, 43758u, 31824u, 18564u, 8568u, 3060u, 816u, + 153u, 18u, 1u, 0u, 0u, 0u, 0u, 0u, + 0u, 0u, 0u, 0u, 0u, 0u, 0u, 0u, + + /* n = 19 */ + 1u, 19u, 171u, 969u, 3876u, 11628u, 27132u, 50388u, + 75582u, 92378u, 92378u, 75582u, 50388u, 27132u, 11628u, 3876u, + 969u, 171u, 19u, 1u, 0u, 0u, 0u, 0u, + 0u, 0u, 0u, 0u, 0u, 0u, 0u, 0u, + + /* n = 20 */ + 1u, 20u, 190u, 1140u, 4845u, 15504u, 38760u, 77520u, + 125970u, 167960u, 184756u, 167960u, 125970u, 77520u, 38760u, 15504u, + 4845u, 1140u, 190u, 20u, 1u, 0u, 0u, 0u, + 0u, 0u, 0u, 0u, 0u, 0u, 0u, 0u, + + /* n = 21 */ + 1u, 21u, 210u, 1330u, 5985u, 20349u, 54264u, 116280u, + 203490u, 293930u, 352716u, 352716u, 293930u, 203490u, 116280u, 54264u, + 20349u, 5985u, 1330u, 210u, 21u, 1u, 0u, 0u, + 0u, 0u, 0u, 0u, 0u, 0u, 0u, 0u, + + /* n = 22 */ + 1u, 22u, 231u, 1540u, 7315u, 26334u, 74613u, 170544u, + 319770u, 497420u, 646646u, 705432u, 646646u, 497420u, 319770u, 170544u, + 74613u, 26334u, 7315u, 1540u, 231u, 22u, 1u, 0u, + 0u, 0u, 0u, 0u, 0u, 0u, 0u, 0u, + + /* n = 23 */ + 1u, 23u, 253u, 1771u, 8855u, 33649u, 100947u, 245157u, + 490314u, 817190u, 1144066u, 1352078u, 1352078u, 1144066u, 817190u, 490314u, + 245157u, 100947u, 33649u, 8855u, 1771u, 253u, 23u, 1u, + 0u, 0u, 0u, 0u, 0u, 0u, 0u, 0u, + + /* n = 24 */ + 1u, 24u, 276u, 2024u, 10626u, 42504u, 134596u, 346104u, + 735471u, 1307504u, 1961256u, 2496144u, 2704156u, 2496144u, 1961256u, 1307504u, + 735471u, 346104u, 134596u, 42504u, 10626u, 2024u, 276u, 24u, + 1u, 0u, 0u, 0u, 0u, 0u, 0u, 0u, + + /* n = 25 */ + 1u, 25u, 300u, 2300u, 12650u, 53130u, 177100u, 480700u, + 1081575u, 2042975u, 3268760u, 4457400u, 5200300u, 5200300u, 4457400u, 3268760u, + 2042975u, 1081575u, 480700u, 177100u, 53130u, 12650u, 2300u, 300u, + 25u, 1u, 0u, 0u, 0u, 0u, 0u, 0u, + + /* n = 26 */ + 1u, 26u, 325u, 2600u, 14950u, 65780u, 230230u, 657800u, + 1562275u, 3124550u, 5311735u, 7726160u, 9657700u, 10400600u, 9657700u, 7726160u, + 5311735u, 3124550u, 1562275u, 657800u, 230230u, 65780u, 14950u, 2600u, + 325u, 26u, 1u, 0u, 0u, 0u, 0u, 0u, + + /* n = 27 */ + 1u, 27u, 351u, 2925u, 17550u, 80730u, 296010u, 888030u, + 2220075u, 4686825u, 8436285u, 13037895u, 17383860u, 20058300u, 20058300u, 17383860u, + 13037895u, 8436285u, 4686825u, 2220075u, 888030u, 296010u, 80730u, 17550u, + 2925u, 351u, 27u, 1u, 0u, 0u, 0u, 0u, + + /* n = 28 */ + 1u, 28u, 378u, 3276u, 20475u, 98280u, 376740u, 1184040u, + 3108105u, 6906900u, 13123110u, 21474180u, 30421755u, 37442160u, 40116600u, 37442160u, + 30421755u, 21474180u, 13123110u, 6906900u, 3108105u, 1184040u, 376740u, 98280u, + 20475u, 3276u, 378u, 28u, 1u, 0u, 0u, 0u, + + /* n = 29 */ + 1u, 29u, 406u, 3654u, 23751u, 118755u, 475020u, 1560780u, + 4292145u, 10015005u, 20030010u, 34597290u, 51895935u, 67863915u, 77558760u, 77558760u, + 67863915u, 51895935u, 34597290u, 20030010u, 10015005u, 4292145u, 1560780u, 475020u, + 118755u, 23751u, 3654u, 406u, 29u, 1u, 0u, 0u, + + /* n = 30 */ + 1u, 30u, 435u, 4060u, 27405u, 142506u, 593775u, 2035800u, + 5852925u, 14307150u, 30045015u, 54627300u, 86493225u, 119759850u, 145422675u, 155117520u, + 145422675u, 119759850u, 86493225u, 54627300u, 30045015u, 14307150u, 5852925u, 2035800u, + 593775u, 142506u, 27405u, 4060u, 435u, 30u, 1u, 0u, + + /* n = 31 */ + 1u, 31u, 465u, 4495u, 31465u, 169911u, 736281u, 2629575u, + 7888725u, 20160075u, 44352165u, 84672315u, 141120525u, 206253075u, 265182525u, 300540195u, + 300540195u, 265182525u, 206253075u, 141120525u, 84672315u, 44352165u, 20160075u, 7888725u, + 2629575u, 736281u, 169911u, 31465u, 4495u, 465u, 31u, 1u, +}; + +static inline uint32_t binom(int n, int k) +{ + return binom_tab[(n << 5) + k]; +} + +/* + * Fermion sign for a_cre^+ a_des |I> = sign |J> + */ +static inline int8_t cd_sign(int p, int q, uint64_t str) +{ + int hi = p > q ? p : q; + int lo = p > q ? q : p; + uint64_t mask = (1ULL << hi) - (1ULL << (lo + 1)); + return (int8_t)(1 - 2 * (popcnt64(str & mask) & 1)); +} + +/* + * Lexical address of a spin string using Cantor expansion + */ +static inline uint32_t str2addr_local(int nelec, uint64_t str) +{ + uint32_t addr = 0; + + for (int i = 1; i <= nelec; i++) { + int p = ctz64(str); + addr += binom(p, i); + str &= str - 1; + } + return addr; +} + +/* + * Spin string from lexical address using inverse Cantor expansion + */ +static inline uint64_t addr2str_local(int norb, int nelec, uint32_t addr) +{ + uint64_t str = 0; + uint32_t a = addr; + int e = nelec; + + for (int p = norb - 1; e > 0; p--) { + if (a == 0 || p < e) { + str |= lowbits(e); + break; + } + + uint32_t c = binom(p, e); + if (c <= a) { + str |= 1ULL << p; + a -= c; + e--; + } + } + return str; +} + +static inline int occ_eq_u8(const uint8_t *a, const uint8_t *b, int n) +{ + return memcmp(a, b, (size_t)n) == 0; +} + +static void gas_zero(gas_space_t *gas) +{ + memset(gas, 0, sizeof(*gas)); +} + +static inline uint8_t *sector_occ_ptr(const gas_space_t *gas, gas_sid_t s) +{ + return gas->sector_occ + (size_t)s * gas->ngas; +} + +static inline uint32_t *sector_stride_ptr(const gas_space_t *gas, gas_sid_t s) +{ + return gas->sector_stride + (size_t)s * gas->ngas; +} + +static int occ_lex_cmp_width(const uint8_t *a, const uint8_t *b, int width) +{ + for (int i = 0; i < width; i++) { + if (a[i] < b[i]) return -1; + if (a[i] > b[i]) return 1; + } + return 0; +} + +/* + * qsort comparator context; used only in serial build stages + */ +static int qsort_occ_width; + +static int cmp_occ_rows(const void *pa, const void *pb) +{ + const uint8_t *a = (const uint8_t *)pa; + const uint8_t *b = (const uint8_t *)pb; + return occ_lex_cmp_width(a, b, qsort_occ_width); +} + +static int cmp_u32(const void *pa, const void *pb) +{ + uint32_t a = *(const uint32_t *)pa; + uint32_t b = *(const uint32_t *)pb; + return (a > b) - (a < b); +} + +/* ========================================================================== */ +/* 2. Temporary row-list helpers */ +/* ========================================================================== */ + +/* + * Dynamic array for temporary storage and sort + */ +typedef struct { + uint8_t *data; + uint32_t n; + uint32_t cap; + int width; +} gas_u8_rows_t; + +static inline uint8_t *u8_row(gas_u8_rows_t *R, uint32_t i) +{ + return R->data + (size_t)i * R->width; +} + +static inline const uint8_t *u8_crow(const gas_u8_rows_t *R, uint32_t i) +{ + return R->data + (size_t)i * R->width; +} + +static inline const uint8_t *block_alpha_occ(const gas_u8_rows_t *blocks, + uint32_t b) +{ + return u8_crow(blocks, b); +} + +static inline const uint8_t *block_beta_occ(const gas_space_t *gas, + const gas_u8_rows_t *blocks, + uint32_t b) +{ + return u8_crow(blocks, b) + gas->ngas; +} + +static void free_u8_rows(gas_u8_rows_t *R) +{ + if (R != 0) { + free(R->data); + memset(R, 0, sizeof(*R)); + } +} + +static int init_u8_rows(gas_u8_rows_t *R, int width) +{ + memset(R, 0, sizeof(*R)); + R->width = width; + return 0; +} + +static int reserve_u8_rows(gas_u8_rows_t *R, uint32_t need) +{ + if (need <= R->cap) { + return 0; + } + + uint32_t cap = R->cap ? R->cap : 128u; + while (cap < need) { + if (cap > UINT32_MAX / 2u) { + cap = need; + break; + } + cap *= 2u; + } + + size_t bytes = (size_t)cap * (size_t)R->width * sizeof(*R->data); + uint8_t *p = realloc(R->data, bytes); + if (p == 0) { + return -1; + } + + R->data = p; + R->cap = cap; + return 0; +} + +static int append_u8_row(gas_u8_rows_t *R, const uint8_t *row) +{ + if (reserve_u8_rows(R, R->n + 1u) != 0) { + return -1; + } + memcpy(u8_row(R, R->n), row, (size_t)R->width); + R->n++; + return 0; +} + +static void sort_unique_occ_rows(gas_u8_rows_t *R) +{ + if (R->n == 0) { + return; + } + qsort_occ_width = R->width; + qsort(R->data, R->n, (size_t)R->width, cmp_occ_rows); + + uint32_t w = 1; + for (uint32_t r = 1; r < R->n; r++) { + uint8_t *prev = u8_row(R, w - 1u); + uint8_t *cur = u8_row(R, r); + if (!occ_eq_u8(prev, cur, R->width)) { + if (w != r) { + memcpy(u8_row(R, w), cur, (size_t)R->width); + } + w++; + } + } + R->n = w; +} + +/* ========================================================================== */ +/* 3. Public address and lookup routines */ +/* ========================================================================== */ + +uint32_t gas_str2addr_sector(const gas_space_t *gas, gas_sid_t s, uint64_t str) +{ + const uint8_t *occ = sector_occ_ptr(gas, s); + const uint32_t *stride = sector_stride_ptr(gas, s); + uint64_t addr = 0; + + for (int g = 0; g < gas->ngas; g++) { + uint64_t sub = (str >> gas->start[g]) & lowbits(gas->norb[g]); + uint32_t a = str2addr_local(occ[g], sub); + addr += (uint64_t)a * stride[g]; + } + return (uint32_t)addr; +} + +uint64_t gas_addr2str_sector(const gas_space_t *gas, gas_sid_t s, uint32_t addr) +{ + const uint8_t *occ = sector_occ_ptr(gas, s); + const uint32_t *stride = sector_stride_ptr(gas, s); + uint64_t str = 0; + uint32_t rem = addr; + + for (int g = 0; g < gas->ngas; g++) { + uint32_t q = stride[g] == 0 ? 0 : rem / stride[g]; + rem -= q * stride[g]; + str |= addr2str_local(gas->norb[g], occ[g], q) << gas->start[g]; + } + return str; +} + +uint32_t gas_block_ndet(const gas_space_t *gas, gas_bid_t b) +{ + const gas_block_t *blk = gas->block + b; + return gas->sector_nstr[blk->sa] * gas->sector_nstr[blk->sb]; +} + +uint32_t gas_det_addr(const gas_space_t *gas, gas_bid_t b, + uint32_t ia, uint32_t ib) +{ + const gas_block_t *blk = gas->block + b; + uint32_t nb = gas->sector_nstr[blk->sb]; + return blk->offset + ia * nb + ib; +} + +gas_bid_t gas_find_block(const gas_space_t *gas, gas_sid_t sa, gas_sid_t sb) +{ + if (sa >= gas->nsector) { + return GAS_INVALID_BID; + } + + const gas_row_t *row = gas->D.by_alpha_row + sa; + uint32_t off = row->off; + uint32_t n = row->n; + + if (n <= 8) { + for (uint32_t i = 0; i < n; i++) { + gas_bid_t b = off + i; + if (gas->block[b].sb == sb) { + return b; + } + } + return GAS_INVALID_BID; + } + + uint32_t lo = 0; + uint32_t hi = n; + while (lo < hi) { + uint32_t mid = lo + ((hi - lo) >> 1); + gas_sid_t x = gas->block[off + mid].sb; + if (x < sb) { + lo = mid + 1; + } else { + hi = mid; + } + } + if (lo < n && gas->block[off + lo].sb == sb) { + return off + lo; + } + return GAS_INVALID_BID; +} + +gas_tid_t gas_find_table(const gas_space_t *gas, gas_sid_t src, gas_sid_t dst) +{ + if (src >= gas->nsector) { + return GAS_INVALID_TID; + } + + const gas_row_t *row = gas->T.row + src; + const gas_sid_t *dsts = gas->T.dst + row->off; + uint32_t n = row->n; + + if (n <= 8) { + for (uint32_t i = 0; i < n; i++) { + if (dsts[i] == dst) { + return row->off + i; + } + } + return GAS_INVALID_TID; + } + + uint32_t lo = 0; + uint32_t hi = n; + while (lo < hi) { + uint32_t mid = lo + ((hi - lo) >> 1); + if (dsts[mid] < dst) { + lo = mid + 1; + } else { + hi = mid; + } + } + if (lo < n && dsts[lo] == dst) { + return row->off + lo; + } + return GAS_INVALID_TID; +} + +const gas_link_table_t *gas_get_link_table(const gas_space_t *gas, + gas_sid_t src, gas_sid_t dst) +{ + gas_tid_t tid = gas_find_table(gas, src, dst); + return tid == GAS_INVALID_TID ? 0 : gas->table + tid; +} + +/* ========================================================================== */ +/* 4. Raw block generation from D */ +/* ========================================================================== */ + +static int build_raw_blocks_from_input(gas_space_t *gas, int nblock, + const int *block_occ, + gas_u8_rows_t *blocks) +{ + if (nblock <= 0 || block_occ == 0) { + return -1; + } + init_u8_rows(blocks, 2 * gas->ngas); + if (reserve_u8_rows(blocks, (uint32_t)nblock) != 0 || + blocks->data == 0) { + return -1; + } + + for (int i = 0; i < nblock; i++) { + const int *a = block_occ + (size_t)(2 * i) * gas->ngas; + const int *b = a + gas->ngas; + uint8_t *row = u8_row(blocks, blocks->n); + int sum_a = 0; + int sum_b = 0; + + for (int g = 0; g < gas->ngas; g++) { + if (a[g] < 0 || a[g] > gas->norb[g] || + b[g] < 0 || b[g] > gas->norb[g]) { + return -1; + } + row[g] = (uint8_t)a[g]; + row[gas->ngas + g] = (uint8_t)b[g]; + sum_a += a[g]; + sum_b += b[g]; + } + if (sum_a != gas->na || sum_b != gas->nb) { + return -1; + } + if (blocks->n != 0u && + occ_lex_cmp_width(u8_crow(blocks, blocks->n - 1u), row, + 2 * gas->ngas) >= 0) { + /* D must be strictly sorted in canonical alpha-major order. */ + return -1; + } + blocks->n++; + } + return 0; +} + +/* ========================================================================== */ +/* 5. Sector collection and initialization */ +/* ========================================================================== */ + +static int append_sector_neighbors_occ(const gas_space_t *gas, + const uint8_t *occ, + gas_u8_rows_t *out, + uint8_t *work) +{ + if (append_u8_row(out, occ) != 0) { /* self-neighbor always exists */ + return -1; + } + + for (int r = 0; r < gas->ngas; r++) { + if (occ[r] == 0) { + continue; + } + for (int s = 0; s < gas->ngas; s++) { + if (r == s || occ[s] >= gas->norb[s]) { + continue; + } + memcpy(work, occ, (size_t)gas->ngas); + work[r]--; + work[s]++; + if (append_u8_row(out, work) != 0) { + return -1; + } + } + } + return 0; +} + +static int collect_sector_pool(gas_space_t *gas, const gas_u8_rows_t *blocks, + gas_u8_rows_t *sectors) +{ + init_u8_rows(sectors, gas->ngas); + uint8_t *work = malloc(sizeof(*work) * (size_t)gas->ngas); + if (work == 0) { + return -1; + } + + for (uint32_t b = 0; b < blocks->n; b++) { + const uint8_t *a = block_alpha_occ(blocks, b); + const uint8_t *bb = block_beta_occ(gas, blocks, b); + + if (append_u8_row(sectors, a) != 0 || + append_u8_row(sectors, bb) != 0 || + append_sector_neighbors_occ(gas, a, sectors, work) != 0 || + append_sector_neighbors_occ(gas, bb, sectors, work) != 0) { + free(work); + return -1; + } + } + + free(work); + sort_unique_occ_rows(sectors); + if (sectors->n > GAS_INVALID_SID) { + return -1; + } + return 0; +} + +static int init_sectors_from_pool(gas_space_t *gas, const gas_u8_rows_t *sectors) +{ + if (sectors->n == 0u || sectors->n > UINT16_MAX) { + return -1; + } + gas->nsector = (uint16_t)sectors->n; + gas->sector_nstr = malloc( + sizeof(*gas->sector_nstr) * (size_t)gas->nsector); + gas->sector_occ = malloc( + sizeof(*gas->sector_occ) * (size_t)gas->nsector * gas->ngas); + gas->sector_stride = malloc( + sizeof(*gas->sector_stride) * (size_t)gas->nsector * gas->ngas); + if (gas->sector_nstr == 0 || gas->sector_occ == 0 || + gas->sector_stride == 0) { + return -1; + } + + for (uint32_t s = 0; s < gas->nsector; s++) { + memcpy(gas->sector_occ + (size_t)s * gas->ngas, + u8_crow(sectors, s), + (size_t)gas->ngas * sizeof(uint8_t)); + } + + for (uint32_t s = 0; s < gas->nsector; s++) { + uint8_t *occ = gas->sector_occ + (size_t)s * gas->ngas; + uint32_t *stride = gas->sector_stride + (size_t)s * gas->ngas; + uint64_t st = 1; + for (int g = gas->ngas - 1; g >= 0; g--) { + stride[g] = (uint32_t)st; + st *= binom(gas->norb[g], occ[g]); + if (st > UINT32_MAX) { + return -1; + } + } + + gas->sector_nstr[s] = (uint32_t)st; + } + return 0; +} + +static int find_sector_occ(const gas_space_t *gas, const uint8_t *occ) +{ + uint32_t lo = 0; + uint32_t hi = gas->nsector; + while (lo < hi) { + uint32_t mid = lo + ((hi - lo) >> 1); + const uint8_t *m = sector_occ_ptr(gas, (gas_sid_t)mid); + int c = occ_lex_cmp_width(m, occ, gas->ngas); + if (c < 0) { + lo = mid + 1; + } else { + hi = mid; + } + } + if (lo < gas->nsector && + occ_eq_u8(sector_occ_ptr(gas, (gas_sid_t)lo), occ, gas->ngas)) { + return (int)lo; + } + return -1; +} + +/* ========================================================================== */ +/* 6. Legal block set D and two D indices */ +/* ========================================================================== */ + +static int fill_blocks_from_raw(gas_space_t *gas, const gas_u8_rows_t *blocks) +{ + gas->nblock = blocks->n; + gas->block = malloc(sizeof(*gas->block) * (size_t)gas->nblock); + if (gas->nblock && gas->block == 0) { + return -1; + } + + uint64_t off = 0; + for (uint32_t b = 0; b < gas->nblock; b++) { + const uint8_t *a = block_alpha_occ(blocks, b); + const uint8_t *bb = block_beta_occ(gas, blocks, b); + int sa = find_sector_occ(gas, a); + int sb = find_sector_occ(gas, bb); + if (sa < 0 || sb < 0) { + return -1; + } + + uint64_t nda = gas->sector_nstr[sa]; + uint64_t ndb = gas->sector_nstr[sb]; + uint64_t ndet = nda * ndb; + if (ndet > UINT32_MAX || off + ndet > UINT32_MAX) { + return -1; + } + + gas->block[b].offset = (uint32_t)off; + gas->block[b].sa = (gas_sid_t)sa; + gas->block[b].sb = (gas_sid_t)sb; + off += ndet; + } + gas->ndet = (uint32_t)off; + return 0; +} + +static int build_block_index(gas_space_t *gas) +{ + uint32_t nsec = gas->nsector; + uint32_t nb = gas->nblock; + + gas->D.by_alpha_row = calloc((size_t)nsec, sizeof(*gas->D.by_alpha_row)); + gas->D.by_beta_row = calloc((size_t)nsec, sizeof(*gas->D.by_beta_row)); + gas->D.by_beta_sid = malloc(sizeof(*gas->D.by_beta_sid) * (size_t)nb); + gas->D.by_beta_bid = malloc(sizeof(*gas->D.by_beta_bid) * (size_t)nb); + if ((nsec && (gas->D.by_alpha_row == 0 || gas->D.by_beta_row == 0)) || + (nb && (gas->D.by_beta_sid == 0 || gas->D.by_beta_bid == 0))) { + return -1; + } + + for (uint32_t b = 0; b < nb; b++) { + gas->D.by_alpha_row[gas->block[b].sa].n++; + gas->D.by_beta_row [gas->block[b].sb].n++; + } + + uint32_t off = 0; + for (uint32_t s = 0; s < nsec; s++) { + uint32_t n = gas->D.by_alpha_row[s].n; + gas->D.by_alpha_row[s].off = off; + off += n; + /* Keep n unchanged. Since block[] is canonical alpha-major, + * this row points directly into block[off : off+n). */ + } + + off = 0; + for (uint32_t s = 0; s < nsec; s++) { + uint32_t n = gas->D.by_beta_row[s].n; + gas->D.by_beta_row[s].off = off; + off += n; + gas->D.by_beta_row[s].n = 0; + } + + /* by_beta is the only materialized secondary sparse view. by_alpha is + * implicit because canonical block order makes alpha rows contiguous. */ + for (uint32_t b = 0; b < nb; b++) { + gas_sid_t sb = gas->block[b].sb; + gas_row_t *rb = gas->D.by_beta_row + sb; + uint32_t pb = rb->off + rb->n; + + gas->D.by_beta_sid[pb] = gas->block[b].sa; + gas->D.by_beta_bid[pb] = b; + rb->n++; + } + return 0; +} + +/* ========================================================================== */ +/* 7. Temporary sector-neighbor rows */ +/* ========================================================================== */ + +typedef struct { + gas_row_t *row; /* [nsector] */ + gas_sid_t *sid; /* packed generated neighbors */ + uint32_t n; + uint32_t cap; +} gas_sector_neighbors_tmp_t; + +static void free_sector_neighbors_tmp(gas_sector_neighbors_tmp_t *S) +{ + if (S != 0) { + free(S->row); + free(S->sid); + memset(S, 0, sizeof(*S)); + } +} + +static int reserve_sid_tmp(gas_sector_neighbors_tmp_t *S, uint32_t need) +{ + if (need <= S->cap) { + return 0; + } + uint32_t cap = S->cap ? S->cap : 1024u; + while (cap < need) { + if (cap > UINT32_MAX / 2u) { + cap = need; + break; + } + cap *= 2u; + } + gas_sid_t *p = realloc(S->sid, (size_t)cap * sizeof(gas_sid_t)); + if (p == 0) { + return -1; + } + S->sid = p; + S->cap = cap; + return 0; +} + +static int append_sector_neighbor_sids(const gas_space_t *gas, + gas_sid_t src, + gas_sector_neighbors_tmp_t *S, + gas_u8_rows_t *nbr_occ, + uint8_t *work) +{ + nbr_occ->n = 0; + S->row[src].off = S->n; + + if (append_sector_neighbors_occ(gas, sector_occ_ptr(gas, src), + nbr_occ, work) != 0) { + return -1; + } + sort_unique_occ_rows(nbr_occ); + + for (uint32_t k = 0; k < nbr_occ->n; k++) { + int sid = find_sector_occ(gas, u8_crow(nbr_occ, k)); + if (sid < 0) { + continue; + } + if (reserve_sid_tmp(S, S->n + 1u) != 0) { + return -1; + } + S->sid[S->n++] = (gas_sid_t)sid; + S->row[src].n++; + } + return 0; +} + +static int build_sector_neighbors_tmp(gas_space_t *gas, + gas_sector_neighbors_tmp_t *S) +{ + uint32_t nsec = gas->nsector; + uint8_t *is_legal_src = 0; + uint8_t *work = 0; + gas_u8_rows_t nbr_occ; + int ret = -1; + + memset(S, 0, sizeof(*S)); + memset(&nbr_occ, 0, sizeof(nbr_occ)); + + S->row = calloc((size_t)nsec, sizeof(*S->row)); + is_legal_src = calloc((size_t)nsec, sizeof(*is_legal_src)); + work = malloc(sizeof(*work) * (size_t)gas->ngas); + init_u8_rows(&nbr_occ, gas->ngas); + if ((nsec && (S->row == 0 || is_legal_src == 0)) || work == 0) { + goto done; + } + + for (uint32_t b = 0; b < gas->nblock; b++) { + is_legal_src[gas->block[b].sa] = 1; + is_legal_src[gas->block[b].sb] = 1; + } + + for (uint32_t s = 0; s < nsec; s++) { + S->row[s].off = S->n; + if (is_legal_src[s] && + append_sector_neighbor_sids(gas, (gas_sid_t)s, S, + &nbr_occ, work) != 0) { + goto done; + } + } + + ret = 0; + +done: + free(is_legal_src); + free(work); + free_u8_rows(&nbr_occ); + return ret; +} + +/* ========================================================================== */ +/* 8. Required table pair marker and table index */ +/* ========================================================================== */ + +typedef struct { + uint32_t *pair; /* packed (src,dst) keys */ + uint32_t n; + uint32_t cap; + + uint32_t *slot; /* open-addressing set; stores pair_key + 1 */ + uint32_t nslot; + uint32_t used; +} gas_pair_marker_t; + +static inline uint32_t pair_key(gas_sid_t src, gas_sid_t dst) +{ + return ((uint32_t)src << 16) | (uint32_t)dst; +} + +static inline gas_sid_t key_src(uint32_t key) +{ + return (gas_sid_t)(key >> 16); +} + +static inline gas_sid_t key_dst(uint32_t key) +{ + return (gas_sid_t)(key & 0xffffu); +} + +static uint32_t pair_key_hash(uint32_t x) +{ + x ^= x >> 16; + x *= UINT32_C(0x7feb352d); + x ^= x >> 15; + x *= UINT32_C(0x846ca68b); + x ^= x >> 16; + return x; +} + +static void free_pair_marker(gas_pair_marker_t *M) +{ + if (M != 0) { + free(M->pair); + free(M->slot); + memset(M, 0, sizeof(*M)); + } +} + +static int reserve_marker_pairs(gas_pair_marker_t *M, uint32_t need) +{ + if (need <= M->cap) { + return 0; + } + uint32_t cap = M->cap ? M->cap : 1024u; + while (cap < need) { + if (cap > UINT32_MAX / 2u) { + cap = need; + break; + } + cap *= 2u; + } + uint32_t *p = realloc(M->pair, (size_t)cap * sizeof(uint32_t)); + if (p == 0) { + return -1; + } + M->pair = p; + M->cap = cap; + return 0; +} + +static int marker_rehash(gas_pair_marker_t *M, uint32_t new_nslot) +{ + uint32_t *old = M->slot; + uint32_t old_n = M->nslot; + + M->slot = calloc((size_t)new_nslot, sizeof(uint32_t)); + if (M->slot == 0) { + M->slot = old; + return -1; + } + M->nslot = new_nslot; + M->used = 0; + + for (uint32_t i = 0; i < old_n; i++) { + uint32_t stored = old[i]; + if (stored == 0) { + continue; + } + uint32_t mask = new_nslot - 1u; + uint32_t h = pair_key_hash(stored) & mask; + while (M->slot[h] != 0) { + h = (h + 1u) & mask; + } + M->slot[h] = stored; + M->used++; + } + free(old); + return 0; +} + +static int init_pair_marker(gas_pair_marker_t *M) +{ + memset(M, 0, sizeof(*M)); + return marker_rehash(M, 2048u); +} + +static int mark_pair(gas_pair_marker_t *M, gas_sid_t src, gas_sid_t dst) +{ + if ((uint64_t)(M->used + 1u) * 3u >= (uint64_t)M->nslot * 2u) { + if (M->nslot > UINT32_MAX / 2u) { + return -1; + } + if (marker_rehash(M, M->nslot * 2u) != 0) { + return -1; + } + } + + uint32_t key = pair_key(src, dst); + uint32_t stored = key + 1u; + uint32_t mask = M->nslot - 1u; + uint32_t h = pair_key_hash(stored) & mask; + + while (M->slot[h] != 0) { + if (M->slot[h] == stored) { + return 0; + } + h = (h + 1u) & mask; + } + + if (reserve_marker_pairs(M, M->n + 1u) != 0) { + return -1; + } + M->slot[h] = stored; + M->used++; + M->pair[M->n++] = key; + return 0; +} + +static int prescan_required_tables(const gas_space_t *gas, + const gas_sector_neighbors_tmp_t *S, + gas_pair_marker_t *M) +{ + if (init_pair_marker(M) != 0) { + return -1; + } + + /* Required table set is the symmetric closure of legal-block one-step + * sector edges. If legal block (A,B) has a one-step neighbor X, the + * H*c contract may need both A -> X and X -> A: the forward direction + * is used for one-step/first-step contributions, and the reverse + * direction covers AA/BB second steps through illegal intermediates. */ + for (uint32_t b = 0; b < gas->nblock; b++) { + gas_sid_t A = gas->block[b].sa; + gas_sid_t B = gas->block[b].sb; + gas_row_t ar = S->row[A]; + gas_row_t br = S->row[B]; + + for (uint32_t i = 0; i < ar.n; i++) { + gas_sid_t X = S->sid[ar.off + i]; + if (mark_pair(M, A, X) != 0 || + mark_pair(M, X, A) != 0) { + return -1; + } + } + for (uint32_t i = 0; i < br.n; i++) { + gas_sid_t Y = S->sid[br.off + i]; + if (mark_pair(M, B, Y) != 0 || + mark_pair(M, Y, B) != 0) { + return -1; + } + } + } + return 0; +} + +static int build_table_index(gas_space_t *gas, gas_pair_marker_t *M) +{ + if (M->n > 1u) { + qsort(M->pair, M->n, sizeof(uint32_t), cmp_u32); + } + gas->ntable = M->n; + + uint32_t nsec = gas->nsector; + gas->T.row = calloc((size_t)nsec, sizeof(*gas->T.row)); + gas->R.row = calloc((size_t)nsec, sizeof(*gas->R.row)); + gas->T.dst = malloc(sizeof(*gas->T.dst) * (size_t)gas->ntable); + gas->R.src = malloc(sizeof(*gas->R.src) * (size_t)gas->ntable); + gas->R.tid = malloc(sizeof(*gas->R.tid) * (size_t)gas->ntable); + gas->table = calloc((size_t)gas->ntable, sizeof(*gas->table)); + if ((nsec && (gas->T.row == 0 || gas->R.row == 0)) || + (gas->ntable && (gas->T.dst == 0 || gas->R.src == 0 || + gas->R.tid == 0 || gas->table == 0))) { + return -1; + } + + for (uint32_t i = 0; i < gas->ntable; i++) { + gas_sid_t src = key_src(M->pair[i]); + gas_sid_t dst = key_dst(M->pair[i]); + gas->T.row[src].n++; + gas->R.row[dst].n++; + } + + uint32_t off = 0; + for (uint32_t s = 0; s < nsec; s++) { + uint32_t n = gas->T.row[s].n; + gas->T.row[s].off = off; + off += n; + gas->T.row[s].n = 0; + } + + off = 0; + for (uint32_t s = 0; s < nsec; s++) { + uint32_t n = gas->R.row[s].n; + gas->R.row[s].off = off; + off += n; + gas->R.row[s].n = 0; + } + + for (uint32_t tid = 0; tid < gas->ntable; tid++) { + gas_sid_t src = key_src(M->pair[tid]); + gas_sid_t dst = key_dst(M->pair[tid]); + + gas_row_t *trow = gas->T.row + src; + uint32_t tpos = trow->off + trow->n; + gas->T.dst[tpos] = dst; + trow->n++; + + gas_row_t *rrow = gas->R.row + dst; + uint32_t rpos = rrow->off + rrow->n; + gas->R.src[rpos] = src; + gas->R.tid[rpos] = tid; + rrow->n++; + } + return 0; +} + +/* ========================================================================== */ +/* 9. Sector-pair spin-string link tables */ +/* ========================================================================== */ + +static int sector_transfer(const gas_space_t *gas, gas_sid_t src, gas_sid_t dst, + int *r, int *s) +{ + const uint8_t *a = sector_occ_ptr(gas, src); + const uint8_t *b = sector_occ_ptr(gas, dst); + *r = -1; + *s = -1; + + if (src == dst) { + return 1; + } + + int nr = 0; + int ns = 0; + for (int g = 0; g < gas->ngas; g++) { + int d = (int)b[g] - (int)a[g]; + if (d == -1) { + *r = g; /* source loses one electron in r */ + nr++; + } else if (d == 1) { + *s = g; /* destination gains one electron in s */ + ns++; + } else if (d != 0) { + return 0; + } + } + return nr == 1 && ns == 1; +} + +static int nlink_pair(const gas_space_t *gas, gas_sid_t src, gas_sid_t dst, + uint32_t *out) +{ + const uint8_t *occ = sector_occ_ptr(gas, src); + int r = -1; + int s = -1; + uint64_t n = 0; + + *out = 0; + if (!sector_transfer(gas, src, dst, &r, &s)) { + return -1; + } + + if (src == dst) { + for (int g = 0; g < gas->ngas; g++) { + uint64_t nocc = occ[g]; + uint64_t nvir = (uint64_t)gas->norb[g] - occ[g]; + n += nocc * (nvir + 1u); /* p=q plus p!=q inside subspace */ + } + } else { + n = (uint64_t)occ[r] * ((uint64_t)gas->norb[s] - occ[s]); + } + if (n > UINT32_MAX) { + return -1; + } + *out = (uint32_t)n; + return 0; +} + +static inline void set_link(gas_link_entry_t *e, uint32_t addr, + int p, int q, int8_t sign) +{ + e->addr = addr; + gas_link_set_raw(e, p, q); + e->sign = sign; + e->padding = 0; +} + +static int init_link_table(gas_space_t *gas, gas_tid_t tid, + gas_sid_t src, gas_sid_t dst, + uint64_t *work) +{ + gas_link_table_t *table = gas->table + tid; + uint32_t nsrc = gas->sector_nstr[src]; + uint32_t nlink = 0; + int remove = 0; + int add = 0; + + if (nlink_pair(gas, src, dst, &nlink) != 0 || + !sector_transfer(gas, src, dst, &remove, &add)) { + return -1; + } + if (src != dst && (remove >= gas->ngas || add >= gas->ngas)) { + return -1; + } + + table->link = 0; + table->nsrc = nsrc; + table->nlink = nlink; + if (nsrc == 0 || nlink == 0) { + return 0; + } + if (nsrc > UINT32_MAX / nlink) { + return -1; + } + + size_t nentry = (size_t)nsrc * nlink; + table->link = malloc(nentry * sizeof(*table->link)); + if (table->link == 0) { + return -1; + } + if (UINT64_MAX - *work < nentry) { + return -1; + } + *work += nentry; + return 0; +} + +static int fill_link_table(const gas_space_t *gas, gas_tid_t tid, + gas_sid_t src, gas_sid_t dst) +{ + gas_link_table_t *table = gas->table + tid; + uint32_t nsrc = table->nsrc; + uint32_t nlink = table->nlink; + int r = 0; + int ss = 0; + + if (!sector_transfer(gas, src, dst, &r, &ss)) { + return -1; + } + if (nsrc == 0 || nlink == 0) { + return 0; + } + const int ngas = gas->ngas; + const uint8_t *src_occ = sector_occ_ptr(gas, src); + const uint8_t *dst_occ = sector_occ_ptr(gas, dst); + const uint32_t *src_stride = sector_stride_ptr(gas, src); + const uint32_t *dst_stride = sector_stride_ptr(gas, dst); + + for (uint32_t ia = 0; ia < nsrc; ia++) { + uint32_t subaddr[GAS_MAX_NGAS]; + uint32_t rem = ia; + uint64_t str = 0; + gas_link_entry_t *row = table->link + (size_t)ia * nlink; + uint32_t k = 0; + + for (int g = 0; g < ngas; g++) { + uint32_t q = rem / src_stride[g]; + uint64_t sub; + rem -= q * src_stride[g]; + subaddr[g] = q; + sub = addr2str_local(gas->norb[g], src_occ[g], q); + str |= sub << gas->start[g]; + } + + if (src == dst) { + for (int g = 0; g < ngas; g++) { + int beg = gas->start[g]; + uint32_t stride = src_stride[g]; + uint64_t sub = (str >> beg) & lowbits(gas->norb[g]); + uint64_t occbits = sub; + uint64_t vir0 = (~sub) & lowbits(gas->norb[g]); + + while (occbits) { + int qloc = ctz64(occbits); + int q = beg + qloc; + occbits &= occbits - 1; + + set_link(row + k, ia, q, q, 1); + k++; + + uint64_t virbits = vir0; + while (virbits) { + int ploc = ctz64(virbits); + int p = beg + ploc; + uint64_t sub1; + uint32_t olda; + uint32_t newa; + uint32_t addr1; + + virbits &= virbits - 1; + sub1 = (sub ^ (1ULL << qloc)) | + (1ULL << ploc); + olda = subaddr[g]; + newa = str2addr_local(src_occ[g], sub1); + addr1 = ia; + if (newa >= olda) { + addr1 += (newa - olda) * stride; + } else { + addr1 -= (olda - newa) * stride; + } + + set_link(row + k, addr1, p, q, + cd_sign(p, q, str)); + k++; + } + } + } + } else { + int begr = gas->start[r]; + int begs = gas->start[ss]; + uint64_t subr = (str >> begr) & lowbits(gas->norb[r]); + uint64_t subs = (str >> begs) & lowbits(gas->norb[ss]); + uint64_t occbits = subr; + uint64_t vir0 = (~subs) & lowbits(gas->norb[ss]); + uint32_t base_dst = 0; + + for (int g = 0; g < ngas; g++) { + if (g != r && g != ss) { + base_dst += subaddr[g] * dst_stride[g]; + } + } + + while (occbits) { + int qloc = ctz64(occbits); + int q = begr + qloc; + uint64_t subr1 = subr ^ (1ULL << qloc); + uint64_t virbits = vir0; + uint32_t ar = str2addr_local(dst_occ[r], subr1); + occbits &= occbits - 1; + + while (virbits) { + int ploc = ctz64(virbits); + int p = begs + ploc; + uint64_t subs1; + uint32_t as; + uint32_t addr1; + + virbits &= virbits - 1; + subs1 = subs | (1ULL << ploc); + as = str2addr_local(dst_occ[ss], subs1); + addr1 = base_dst + ar * dst_stride[r] + + as * dst_stride[ss]; + set_link(row + k, addr1, p, q, + cd_sign(p, q, str)); + k++; + } + } + } + } + return 0; +} + +static int build_all_link_tables(gas_space_t *gas) +{ + uint64_t work = 0; + + /* Allocate serially, then fill all tables in one OpenMP region. */ + + for (gas_sid_t src = 0; src < gas->nsector; src++) { + gas_row_t row = gas->T.row[src]; + for (uint32_t i = 0; i < row.n; i++) { + gas_tid_t tid = row.off + i; + gas_sid_t dst = gas->T.dst[tid]; + if (init_link_table(gas, tid, src, dst, &work) != 0) { + return -1; + } + } + } + + int failed = 0; +#ifdef _OPENMP +#pragma omp parallel for schedule(dynamic, 1) reduction(|:failed) \ + if(work >= LINK_BUILD_PARALLEL_MIN) +#endif + for (int64_t ss = 0; ss < (int64_t)gas->nsector; ss++) { + gas_sid_t src = (gas_sid_t)ss; + gas_row_t row = gas->T.row[src]; + + for (uint32_t i = 0; i < row.n; i++) { + gas_tid_t tid = row.off + i; + gas_sid_t dst = gas->T.dst[tid]; + if (fill_link_table(gas, tid, src, dst) != 0) { + failed = 1; + } + } + } + return failed ? -1 : 0; +} + +static int build_table_active_ops_raw(gas_space_t *gas) +{ + const uint32_t nnorb = (uint32_t)(gas->norb_tot * (gas->norb_tot + 1) / 2); + if (nnorb == 0u) { + return -1; + } + uint8_t *seen = malloc((size_t)nnorb * sizeof(*seen)); + if (seen == 0) { + return -1; + } + + for (uint32_t t = 0; t < gas->ntable; t++) { + gas_link_table_t *tab = gas->table + t; + size_t nentry = (size_t)tab->nsrc * tab->nlink; + uint32_t nop = 0; + + memset(seen, 0, (size_t)nnorb * sizeof(*seen)); + for (size_t i = 0; i < nentry; i++) { + const gas_link_entry_t *e = tab->link + i; + uint16_t op = gas_link_pair_index(gas_link_cre(e), gas_link_des(e)); + if (!seen[op]) { + seen[op] = 1u; + nop++; + } + } + if (nop > UINT16_MAX) { + free(seen); + return -1; + } + + free(tab->active_op); + tab->active_op = 0; + tab->nop = (uint16_t)nop; + if (nop != 0u) { + tab->active_op = malloc((size_t)nop * sizeof(*tab->active_op)); + if (tab->active_op == 0) { + free(seen); + return -1; + } + uint32_t k = 0; + for (uint32_t op = 0; op < nnorb; op++) { + if (seen[op]) { + tab->active_op[k++] = (uint16_t)op; + } + } + } + } + free(seen); + return 0; +} + +int gas_space_compress_links(gas_space_t *gas) +{ + if (gas == 0) { + return -1; + } + if (gas->link_format == GAS_LINK_COMPRESSED) { + return 0; + } + + /* Active operator sets are table properties. Build them once from + * the raw (p,q) fields before replacing those fields by pair indices. */ + if (build_table_active_ops_raw(gas) != 0) { + return -1; + } + + for (uint32_t t = 0; t < gas->ntable; t++) { + gas_link_table_t *tab = gas->table + t; + for (size_t i = 0, n = (size_t)tab->nsrc * tab->nlink; i < n; i++) { + gas_link_entry_t *e = tab->link + i; + gas_link_set_ia(e, gas_link_pair_index(gas_link_cre(e), gas_link_des(e))); + } + } + gas->link_format = GAS_LINK_COMPRESSED; + return 0; +} + +int gas_space_links_are_compressed(const gas_space_t *gas) +{ + return gas != 0 && gas->link_format == GAS_LINK_COMPRESSED; +} + +/* ========================================================================== */ +/* 10. Memory reporting */ +/* ========================================================================== */ + +void gas_memory_report(const gas_space_t *gas, gas_memory_report_t *r) +{ + memset(r, 0, sizeof(*r)); + + r->metadata = sizeof(*gas) + + (uint64_t)gas->ngas * sizeof(*gas->norb) + + (uint64_t)gas->ngas * sizeof(*gas->start); + r->sector = (uint64_t)gas->nsector * sizeof(*gas->sector_nstr) + + (uint64_t)gas->nsector * gas->ngas * sizeof(*gas->sector_occ) + + (uint64_t)gas->nsector * gas->ngas * sizeof(*gas->sector_stride); + r->block = (uint64_t)gas->nblock * sizeof(*gas->block); + r->block_index = (uint64_t)gas->nsector * sizeof(*gas->D.by_alpha_row) + + (uint64_t)gas->nsector * sizeof(*gas->D.by_beta_row) + + (uint64_t)gas->nblock * sizeof(*gas->D.by_beta_sid) + + (uint64_t)gas->nblock * sizeof(*gas->D.by_beta_bid); + r->link_table_index = (uint64_t)gas->nsector * sizeof(*gas->T.row) + + (uint64_t)gas->ntable * sizeof(*gas->T.dst) + + (uint64_t)gas->nsector * sizeof(*gas->R.row) + + (uint64_t)gas->ntable * sizeof(*gas->R.src) + + (uint64_t)gas->ntable * sizeof(*gas->R.tid); + r->link_table = (uint64_t)gas->ntable * sizeof(*gas->table); + for (uint32_t i = 0; i < gas->ntable; i++) { + r->link_table += (uint64_t)gas->table[i].nsrc * gas->table[i].nlink + * sizeof(*gas->table[i].link); + r->link_table += (uint64_t)gas->table[i].nop + * sizeof(*gas->table[i].active_op); + } + r->total = r->metadata + + r->sector + + r->block + + r->block_index + + r->link_table_index + + r->link_table; +} + +uint64_t gas_memory_bytes(const gas_space_t *gas) +{ + gas_memory_report_t r; + gas_memory_report(gas, &r); + return r.total; +} + +/* ========================================================================== */ +/* 11. Constructors and destructor */ +/* ========================================================================== */ + +static int init_common(gas_space_t *gas, int ngas, + const int *norb, int na, int nb) +{ + int off = 0; + if (gas == 0) { + return -1; + } + gas_zero(gas); + + if (ngas <= 0 || ngas > GAS_MAX_NGAS || norb == 0 || na < 0 || nb < 0) { + return -1; + } + + gas->ngas = ngas; + gas->na = na; + gas->nb = nb; + gas->norb = malloc(sizeof(*gas->norb) * (size_t)ngas); + gas->start = malloc(sizeof(*gas->start) * (size_t)ngas); + if (gas->norb == 0 || gas->start == 0) { + return -1; + } + + for (int g = 0; g < ngas; g++) { + if (norb[g] <= 0 || norb[g] > GAS_MAX_LOCAL_ORB) { + return -1; + } + if (off + norb[g] > GAS_MAX_ORB) { + return -1; + } + gas->norb[g] = norb[g]; + gas->start[g] = off; + off += norb[g]; + } + gas->norb_tot = off; + return 0; +} + +static int finish_build_from_raw(gas_space_t *gas, const gas_u8_rows_t *blocks) +{ + gas_u8_rows_t sectors; + gas_sector_neighbors_tmp_t S; + gas_pair_marker_t M; + int ret = -1; + + memset(§ors, 0, sizeof(sectors)); + memset(&S, 0, sizeof(S)); + memset(&M, 0, sizeof(M)); + + if (collect_sector_pool(gas, blocks, §ors) != 0) { + goto done; + } + if (init_sectors_from_pool(gas, §ors) != 0) { + goto done; + } + if (fill_blocks_from_raw(gas, blocks) != 0) { + goto done; + } + if (build_block_index(gas) != 0) { + goto done; + } + if (build_sector_neighbors_tmp(gas, &S) != 0) { + goto done; + } + if (prescan_required_tables(gas, &S, &M) != 0) { + goto done; + } + if (build_table_index(gas, &M) != 0) { + goto done; + } + free_pair_marker(&M); + + if (build_all_link_tables(gas) != 0) { + goto done; + } + gas->link_format = GAS_LINK_RAW; + + ret = 0; + +done: + free_u8_rows(§ors); + free_sector_neighbors_tmp(&S); + free_pair_marker(&M); + return ret; +} + +int gas_space_from_blocks(gas_space_t *gas, int ngas, const int *norb, + int na, int nb, + int nblock, const int *block_occ) +{ + gas_u8_rows_t blocks; + memset(&blocks, 0, sizeof(blocks)); + + if (init_common(gas, ngas, norb, na, nb) != 0) { + gas_space_free(gas); + return -1; + } + + if (build_raw_blocks_from_input(gas, nblock, block_occ, &blocks) != 0 || + finish_build_from_raw(gas, &blocks) != 0) { + free_u8_rows(&blocks); + gas_space_free(gas); + return -1; + } + + free_u8_rows(&blocks); + return 0; +} + +void gas_space_free(gas_space_t *gas) +{ + if (gas == 0) { + return; + } + + free(gas->norb); + free(gas->start); + free(gas->sector_nstr); + free(gas->sector_occ); + free(gas->sector_stride); + + free(gas->block); + free(gas->D.by_alpha_row); + free(gas->D.by_beta_row); + free(gas->D.by_beta_sid); + free(gas->D.by_beta_bid); + + free(gas->T.row); + free(gas->T.dst); + free(gas->R.row); + free(gas->R.src); + free(gas->R.tid); + + if (gas->table != 0) { + for (uint32_t i = 0; i < gas->ntable; i++) { + free(gas->table[i].link); + free(gas->table[i].active_op); + } + free(gas->table); + } + + gas_zero(gas); +} diff --git a/pyscf/mcscf/_gaslib.py b/pyscf/mcscf/_gaslib.py new file mode 100644 index 000000000..1911bb5e2 --- /dev/null +++ b/pyscf/mcscf/_gaslib.py @@ -0,0 +1,262 @@ +#!/usr/bin/env python +# Copyright 2026 The PySCF Developers. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# Author: Yi Deng +# + +"""Private ctypes bridge for the GAS FCI C kernels.""" + +import ctypes +import os +from pathlib import Path + + +GAS_SUCCESS = 0 +GAS_ERR_INVALID = -1 +GAS_ERR_MEMORY = -2 + +GAS_LINK_RAW = 0 +GAS_LINK_COMPRESSED = 1 + + +gas_sid_t = ctypes.c_uint16 +gas_bid_t = ctypes.c_uint32 +gas_tid_t = ctypes.c_uint32 + + +class GasBlock(ctypes.Structure): + _fields_ = [ + ("offset", ctypes.c_uint32), + ("sa", gas_sid_t), + ("sb", gas_sid_t), + ] + + +class GasRow(ctypes.Structure): + _fields_ = [ + ("off", ctypes.c_uint32), + ("n", ctypes.c_uint32), + ] + + +class GasBlockIndex(ctypes.Structure): + _fields_ = [ + ("by_alpha_row", ctypes.POINTER(GasRow)), + ("by_beta_row", ctypes.POINTER(GasRow)), + ("by_beta_sid", ctypes.POINTER(gas_sid_t)), + ("by_beta_bid", ctypes.POINTER(gas_bid_t)), + ] + + +class GasTableIndex(ctypes.Structure): + _fields_ = [ + ("row", ctypes.POINTER(GasRow)), + ("dst", ctypes.POINTER(gas_sid_t)), + ] + + +class GasTableRevIndex(ctypes.Structure): + _fields_ = [ + ("row", ctypes.POINTER(GasRow)), + ("src", ctypes.POINTER(gas_sid_t)), + ("tid", ctypes.POINTER(gas_tid_t)), + ] + + +class GasLinkEntry(ctypes.Structure): + _fields_ = [ + ("addr", ctypes.c_uint32), + ("op", ctypes.c_uint16), + ("sign", ctypes.c_int8), + ("padding", ctypes.c_uint8), + ] + + +class GasLinkTable(ctypes.Structure): + _fields_ = [ + ("link", ctypes.POINTER(GasLinkEntry)), + ("active_op", ctypes.POINTER(ctypes.c_uint16)), + ("nsrc", ctypes.c_uint32), + ("nlink", ctypes.c_uint32), + ("nop", ctypes.c_uint16), + ] + + +class GasSpaceStruct(ctypes.Structure): + _fields_ = [ + ("ngas", ctypes.c_int), + ("norb_tot", ctypes.c_int), + ("norb", ctypes.POINTER(ctypes.c_int)), + ("start", ctypes.POINTER(ctypes.c_int)), + ("na", ctypes.c_int), + ("nb", ctypes.c_int), + ("ndet", ctypes.c_uint32), + ("nsector", ctypes.c_uint16), + ("sector_padding", ctypes.c_uint16), + ("sector_nstr", ctypes.POINTER(ctypes.c_uint32)), + ("sector_occ", ctypes.POINTER(ctypes.c_uint8)), + ("sector_stride", ctypes.POINTER(ctypes.c_uint32)), + ("nblock", ctypes.c_uint32), + ("block", ctypes.POINTER(GasBlock)), + ("D", GasBlockIndex), + ("ntable", ctypes.c_uint32), + ("link_format", ctypes.c_uint8), + ("link_format_padding", ctypes.c_uint8 * 3), + ("table", ctypes.POINTER(GasLinkTable)), + ("T", GasTableIndex), + ("R", GasTableRevIndex), + ] + + +class GasMemoryReport(ctypes.Structure): + _fields_ = [ + ("metadata", ctypes.c_uint64), + ("sector", ctypes.c_uint64), + ("block", ctypes.c_uint64), + ("block_index", ctypes.c_uint64), + ("link_table", ctypes.c_uint64), + ("link_table_index", ctypes.c_uint64), + ("total", ctypes.c_uint64), + ] + + +def double_ptr(array): + return array.ctypes.data_as(ctypes.POINTER(ctypes.c_double)) + + +def _configure_library(lib): + gas_p = ctypes.POINTER(GasSpaceStruct) + gas_plan_pp = ctypes.POINTER(ctypes.c_void_p) + rdm_plan_pp = ctypes.POINTER(ctypes.c_void_p) + int_p = ctypes.POINTER(ctypes.c_int) + double_p = ctypes.POINTER(ctypes.c_double) + + lib.gas_space_from_blocks.argtypes = [ + gas_p, ctypes.c_int, int_p, ctypes.c_int, ctypes.c_int, + ctypes.c_int, int_p, + ] + lib.gas_space_from_blocks.restype = ctypes.c_int + + lib.gas_space_compress_links.argtypes = [gas_p] + lib.gas_space_compress_links.restype = ctypes.c_int + + lib.gas_space_links_are_compressed.argtypes = [gas_p] + lib.gas_space_links_are_compressed.restype = ctypes.c_int + + lib.gas_space_free.argtypes = [gas_p] + lib.gas_space_free.restype = None + + lib.gas_memory_bytes.argtypes = [gas_p] + lib.gas_memory_bytes.restype = ctypes.c_uint64 + + lib.gas_memory_report.argtypes = [gas_p, ctypes.POINTER(GasMemoryReport)] + lib.gas_memory_report.restype = None + + lib.gas_addr2str_sector.argtypes = [gas_p, gas_sid_t, ctypes.c_uint32] + lib.gas_addr2str_sector.restype = ctypes.c_uint64 + + lib.gas_str2addr_sector.argtypes = [gas_p, gas_sid_t, ctypes.c_uint64] + lib.gas_str2addr_sector.restype = ctypes.c_uint32 + + lib.gas_block_ndet.argtypes = [gas_p, gas_bid_t] + lib.gas_block_ndet.restype = ctypes.c_uint32 + + lib.fci_make_hdiag_gas.argtypes = [gas_p, double_p, double_p, double_p] + lib.fci_make_hdiag_gas.restype = ctypes.c_int + + lib.fci_contract_gas_omp_task_count.argtypes = [gas_p] + lib.fci_contract_gas_omp_task_count.restype = ctypes.c_uint32 + + lib.fci_contract_gas_parallel_units.argtypes = [gas_p] + lib.fci_contract_gas_parallel_units.restype = ctypes.c_uint32 + + lib.fci_contract_gas_plan_create.argtypes = [ + gas_plan_pp, gas_p, double_p, double_p, + ] + lib.fci_contract_gas_plan_create.restype = ctypes.c_int + + lib.fci_contract_gas_plan_execute.argtypes = [ + ctypes.c_void_p, double_p, double_p, + ] + lib.fci_contract_gas_plan_execute.restype = ctypes.c_int + + lib.fci_contract_gas_plan_free.argtypes = [ctypes.c_void_p] + lib.fci_contract_gas_plan_free.restype = None + + lib.fci_rdm_gas_plan_create.argtypes = [rdm_plan_pp, gas_p] + lib.fci_rdm_gas_plan_create.restype = ctypes.c_int + + lib.fci_rdm_gas_plan_make_rdm1s.argtypes = [ + ctypes.c_void_p, double_p, double_p, double_p, double_p, + ] + lib.fci_rdm_gas_plan_make_rdm1s.restype = ctypes.c_int + + lib.fci_rdm_gas_plan_make_rdm12s.argtypes = [ + ctypes.c_void_p, double_p, double_p, double_p, double_p, + double_p, double_p, double_p, + ] + lib.fci_rdm_gas_plan_make_rdm12s.restype = ctypes.c_int + + lib.fci_rdm_gas_plan_free.argtypes = [ctypes.c_void_p] + lib.fci_rdm_gas_plan_free.restype = None + + lib.fci_rdm_gas_plan_task_count.argtypes = [ctypes.c_void_p] + lib.fci_rdm_gas_plan_task_count.restype = ctypes.c_uint32 + + lib.fci_rdm_gas_plan_workspace_bytes.argtypes = [ctypes.c_void_p] + lib.fci_rdm_gas_plan_workspace_bytes.restype = ctypes.c_uint64 + + return lib + + +def _candidate_library_paths(): + value = os.environ.get("PYSCF_GAS_LIB") + if value: + yield Path(value) + + pyscf_dir = Path(__file__).resolve().parents[1] + yield pyscf_dir / "lib" / "libfci_gas.so" + + +_LIB = None + + +def load_library(path=None, reload=False): + """Load and configure libfci_gas. + + An explicit path or ``PYSCF_GAS_LIB`` can select a development build. + The final fallback uses PySCF's native library loader. + """ + + global _LIB + if _LIB is not None and path is None and not reload: + return _LIB + + if path is not None: + lib = ctypes.CDLL(str(path)) + else: + lib = None + for candidate in _candidate_library_paths(): + if candidate.exists(): + lib = ctypes.CDLL(str(candidate)) + break + if lib is None: + from pyscf import lib as pyscf_lib + lib = pyscf_lib.load_library("libfci_gas") + + lib = _configure_library(lib) + if path is None: + _LIB = lib + return lib diff --git a/pyscf/mcscf/addons_gas.py b/pyscf/mcscf/addons_gas.py new file mode 100644 index 000000000..b0fbf0ea9 --- /dev/null +++ b/pyscf/mcscf/addons_gas.py @@ -0,0 +1,863 @@ +#!/usr/bin/env python +# Copyright 2026 The PySCF Developers. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# Author: Yi Deng +# + +"""Small GAS helpers shared by the GASCI Python layer.""" + +import math +from collections.abc import Mapping +from functools import lru_cache + +import numpy + +from pyscf import lib +from pyscf.fci import addons as fci_addons +from pyscf.lib import logger +from pyscf.mcscf import addons + + +GAS_MAX_ORB = 63 +GAS_MAX_LOCAL_ORB = 31 +GAS_MAX_NGAS = 15 +GAS_MAX_SECTOR = 65534 +GAS_MAX_DET = 2**32 - 1 +GAS_MAX_BLOCK = 2**31 - 1 +GAS_RESTR_SPIN_SUPERGROUP = "spin-supergroup" +GAS_RESTR_SUPERGROUP = "supergroup" +GAS_RESTR_CUMULATIVE_OCC = "cumulative-occ" +GAS_RESTR_RAS = "ras" + + +def gas_restr_usage(gas_restr_type=None): + """Return the strict public input contract for GAS restrictions.""" + + header = ( + "GAS restriction input rules (only row order and duplicate rows are " + "normalized):") + rules = { + GAS_RESTR_SPIN_SUPERGROUP: ( + "gas_restr_type='spin-supergroup': gas_restr is a non-empty " + "array of shape (nrow, 2*ngas); each row is " + "(alpha occupations, beta occupations)."), + GAS_RESTR_SUPERGROUP: ( + "gas_restr_type='supergroup': gas_restr is a non-empty array " + "of shape (nrow, ngas); each row contains total occupations."), + GAS_RESTR_CUMULATIVE_OCC: ( + "gas_restr_type='cumulative-occ': gas_restr has shape " + "(ngas, 2), with non-decreasing [minimum, maximum] cumulative " + "occupations and a final bound equal to the active electrons."), + GAS_RESTR_RAS: ( + "gas_restr_type='ras': gas_orbs has exactly three entries " + "(RAS1, RAS2, RAS3), which may be zero, and gas_restr is exactly " + "{'max_holes': integer, 'max_particles': integer}."), + } + numeric = ( + "All occupations, bounds, orbital counts, and orbital indices must " + "be integers; floating-point and Boolean values are rejected.") + if gas_restr_type in rules: + return "\n".join((header, rules[gas_restr_type], numeric)) + return "\n".join((header,) + tuple(rules[key] for key in ( + GAS_RESTR_SPIN_SUPERGROUP, GAS_RESTR_SUPERGROUP, + GAS_RESTR_CUMULATIVE_OCC, GAS_RESTR_RAS)) + (numeric,)) + + +def _exact_integer_array(values, label): + """Require integer input and convert it to the C-kernel int32 dtype.""" + + raw = numpy.asarray(values, dtype=object) + if not all(lib.isinteger(value) for value in raw.flat): + raise TypeError("%s must contain integers" % label) + arr = numpy.asarray(values) + info = numpy.iinfo(numpy.int32) + if numpy.any(arr < info.min) or numpy.any(arr > info.max): + raise ValueError("%s values exceed int32 range" % label) + return numpy.asarray(arr, dtype=numpy.int32) + + +def _integer_vector(values, label, length=None): + arr = _exact_integer_array(values, label) + if arr.ndim != 1: + raise ValueError("%s must be a one-dimensional sequence" % label) + if length is not None and arr.size != int(length): + raise ValueError("%s must contain exactly %d entries" % + (label, int(length))) + if arr.size == 0: + raise ValueError("%s must not be empty" % label) + return tuple(int(value) for value in arr) + + +def _validate_state_weights(weights): + weights = numpy.asarray(weights, dtype=numpy.float64) + if weights.ndim != 1 or weights.size == 0: + raise ValueError("weights must be a non-empty one-dimensional array") + if not numpy.all(numpy.isfinite(weights)): + raise ValueError("weights must be finite") + if numpy.any(weights < 0.0): + raise ValueError("weights must be non-negative") + if abs(float(weights.sum()) - 1.0) > 1e-10: + raise ValueError("weights must sum to one") + return weights + + +class StateAverageFCISolver(addons.StateAverageFCISolver): + """PySCF state wrapper with GAS-specific multi-root RDM kernels.""" + + __name_mixin__ = "StateAverage" + + def dump_flags(self, verbose=None): + super(addons.StateAverageFCISolver, self).dump_flags(verbose) + log = logger.new_logger(self, verbose) + log.info("State-average over %d states with weights %s", + len(self.weights), self.weights) + log.info("state weights define the averaged GAS RDMs and generalized " + "Fock matrix") + return self + + def states_make_rdm1(self, ci0, norb, nelec, *args, **kwargs): + with self.make_rdm_plan(norb, nelec) as plan: + return [plan.make_rdm1(ci, ci) for ci in ci0] + + def states_make_rdm1s(self, ci0, norb, nelec, *args, **kwargs): + dm1a = [] + dm1b = [] + with self.make_rdm_plan(norb, nelec) as plan: + for ci in ci0: + values = plan.make_rdm1s(ci, ci) + dm1a.append(values[0]) + dm1b.append(values[1]) + return dm1a, dm1b + + def states_make_rdm12(self, ci0, norb, nelec, *args, **kwargs): + dm1 = [] + dm2 = [] + with self.make_rdm_plan(norb, nelec) as plan: + for ci in ci0: + values = plan.make_rdm12(ci, ci) + dm1.append(values[0]) + dm2.append(values[1]) + return dm1, dm2 + + def states_make_rdm12s(self, ci0, norb, nelec, *args, **kwargs): + dm1a = [] + dm1b = [] + dm2aa = [] + dm2ab = [] + dm2bb = [] + with self.make_rdm_plan(norb, nelec) as plan: + for ci in ci0: + dm1s, dm2s = plan.make_rdm12s(ci, ci) + dm1a.append(dm1s[0]) + dm1b.append(dm1s[1]) + dm2aa.append(dm2s[0]) + dm2ab.append(dm2s[1]) + dm2bb.append(dm2s[2]) + return (dm1a, dm1b), (dm2aa, dm2ab, dm2bb) + + def states_trans_rdm12(self, ci1, ci0, norb, nelec, *args, **kwargs): + dm1 = [] + dm2 = [] + with self.make_rdm_plan(norb, nelec) as plan: + for bra, ket in zip(ci1, ci0): + values = plan.make_rdm12(bra, ket) + dm1.append(values[0]) + dm2.append(values[1]) + return dm1, dm2 + + def states_spin_square(self, ci0, norb, nelec, *args, **kwargs): + from pyscf.mcscf import fci_gas + + nelec = fci_addons._unpack_nelec(nelec, self.spin) + values = [] + with self.make_rdm_plan(norb, nelec) as plan: + for ci in ci0: + dm1s, dm2s = plan.make_rdm12s(ci, ci) + values.append( + fci_gas.spin_square_from_rdm12s(dm1s, dm2s, nelec)) + return ([value[0] for value in values], + [value[1] for value in values]) + + def spin_square(self, ci0, norb, nelec, *args, **kwargs): + ss = numpy.asarray( + self.states_spin_square(ci0, norb, nelec, *args, **kwargs)[0]) + ss_average = float(numpy.dot(self.weights, ss)) + multiplicity = numpy.sqrt(max(0.0, 4.0 * ss_average + 1.0)) + return ss_average, multiplicity + + def undo_state_average(self): + obj = lib.view( + self, lib.drop_class(self.__class__, StateAverageFCISolver)) + del obj.weights + del obj.e_states + return obj + + +class StateAverageGASCI(addons.StateAverageMCSCFSolver): + """GAS multistate object carrying weights and root-resolved results. + + The weights define fixed-orbital RDM and generalized-Fock + postprocessing. Individual root energies remain unchanged. + """ + + __name_mixin__ = "StateAverage" + + def __init__(self, mc, fcisolver): + self.__dict__.update(mc.__dict__) + self.fcisolver = fcisolver + + @property + def weights(self): + return self.fcisolver.weights + + @weights.setter + def weights(self, value): + self.fcisolver.weights = _validate_state_weights(value) + + @property + def e_average(self): + return float(numpy.dot(self.weights, self.fcisolver.e_states)) + + @property + def e_states(self): + return self.fcisolver.e_states + + def undo_state_average(self): + obj = lib.view( + self, lib.drop_class(self.__class__, StateAverageGASCI)) + if isinstance(self.fcisolver, StateAverageFCISolver): + obj.fcisolver = self.fcisolver.undo_state_average() + return obj + + def _finalize(self): + from pyscf.mcscf import gasci + + return gasci.GASCI._finalize(self) + + +def state_average(mc, weights=(0.5, 0.5), wfnsym=None): + """Attach PySCF-style state weights to a GASCI calculation. + + GASCI orbitals remain fixed, and the individual root energies are + unchanged. The weights define averaged RDMs and generalized Fock + matrices for postprocessing. + """ + + if wfnsym is not None: + raise NotImplementedError( + "GAS wavefunction symmetry filtering is not implemented") + weights = _validate_state_weights(weights) + + if isinstance(mc, StateAverageGASCI): + mc = mc.undo_state_average() + fcisolver = mc.fcisolver + if isinstance(fcisolver, addons.StateAverageFCISolver): + fcisolver = fcisolver.undo_state_average() + fcisolver = lib.set_class( + StateAverageFCISolver(fcisolver, weights, None), + (StateAverageFCISolver, fcisolver.__class__)) + return lib.set_class( + StateAverageGASCI(mc, fcisolver), + (StateAverageGASCI, mc.__class__)) + + +def state_average_(mc, weights=(0.5, 0.5), wfnsym=None): + """In-place version of :func:`state_average`.""" + + weighted = state_average(mc, weights, wfnsym) + mc.__class__ = weighted.__class__ + mc.__dict__ = weighted.__dict__ + return mc + + +def _integer_rows_with_info(values, ncolumn, label): + """Return canonical integer rows and their normalization metadata.""" + + arr = _exact_integer_array(values, label) + if arr.ndim != 2 or arr.shape[1] != int(ncolumn): + raise ValueError( + "%s must have shape (nrow, %d)" % (label, int(ncolumn))) + if arr.shape[0] == 0: + raise ValueError("%s must contain at least one row" % label) + original = numpy.ascontiguousarray(arr, dtype=numpy.int32) + canonical = numpy.ascontiguousarray( + numpy.unique(original, axis=0), dtype=numpy.int32) + return canonical, { + "input_nrow": int(original.shape[0]), + "canonical_nrow": int(canonical.shape[0]), + "duplicates_removed": int(original.shape[0] - canonical.shape[0]), + "order_changed": not ( + original.shape == canonical.shape and + numpy.array_equal(original, canonical)), + "canonical_rows": canonical, + } + + +def _integer_rows(values, ncolumn, label): + """Return non-empty, unique, lexicographically sorted integer rows.""" + + return _integer_rows_with_info(values, ncolumn, label)[0] + + +def normalize_blocks(blocks, ngas): + """Return the canonical sorted spin-supergroup set D.""" + + return _integer_rows(blocks, 2 * int(ngas), "spin-supergroups") + + +def normalize_supergroups(supergroups, ngas): + """Return a canonical sorted spin-free supergroup set G.""" + + return _integer_rows(supergroups, int(ngas), "supergroups") + + +def normalize_cumulative_occ(bounds, ngas): + """Validate cumulative occupation bounds without reordering GAS rows.""" + + arr = _exact_integer_array(bounds, "cumulative occupation bounds") + if arr.ndim != 2 or arr.shape != (int(ngas), 2): + raise ValueError( + "cumulative occupation bounds must have shape (%d, 2)" % + int(ngas)) + return numpy.ascontiguousarray(arr, dtype=numpy.int32) + + +def _supergroup_block_count(gas_orbs, na, row): + """Count alpha partitions of one supergroup without materializing D.""" + + counts = [0] * (int(na) + 1) + counts[0] = 1 + for norb, total in zip(gas_orbs, row): + lo = max(0, int(total) - int(norb)) + hi = min(int(norb), int(total)) + updated = [0] * (int(na) + 1) + for occupied, count in enumerate(counts): + if count == 0: + continue + for alpha in range(lo, hi + 1): + if occupied + alpha <= na: + updated[occupied + alpha] += count + counts = updated + return counts[int(na)] + + +def blocks_from_supergroups(gas_orbs, nelec, supergroups): + """Expand a spin-free supergroup set G into canonical block set D.""" + + gas_orbs = _integer_vector(gas_orbs, "gas_orbs") + na, nb = fci_addons._unpack_nelec(nelec) + ngas = len(gas_orbs) + if ngas <= 0 or ngas > GAS_MAX_NGAS: + raise ValueError("number of GAS spaces exceeds C kernel limit") + if any(n <= 0 or n > GAS_MAX_LOCAL_ORB for n in gas_orbs): + raise ValueError("each GAS space must have 1..31 orbitals") + if sum(gas_orbs) > GAS_MAX_ORB: + raise ValueError("total active orbitals exceed 63") + if na < 0 or nb < 0 or na > sum(gas_orbs) or nb > sum(gas_orbs): + raise ValueError("invalid alpha/beta electron count") + + supergroups = normalize_supergroups(supergroups, ngas) + capacities = 2 * numpy.asarray(gas_orbs, dtype=numpy.int32) + if numpy.any(supergroups < 0): + raise ValueError("supergroup occupations must be non-negative") + if numpy.any(supergroups > capacities): + raise ValueError("supergroup occupation exceeds local GAS capacity") + if numpy.any(supergroups.sum(axis=1) != na + nb): + raise ValueError( + "every supergroup must sum to the requested electron count") + + nblock = sum( + _supergroup_block_count(gas_orbs, na, row) + for row in supergroups) + if nblock <= 0: + raise ValueError("supergroups generate no legal spin-supergroups") + if nblock > GAS_MAX_BLOCK: + raise ValueError("number of blocks exceeds C int32 input limit") + + blocks = numpy.empty((nblock, 2 * ngas), dtype=numpy.int32) + cursor = 0 + for row in supergroups: + lower = numpy.maximum(0, row - numpy.asarray(gas_orbs)) + upper = numpy.minimum(numpy.asarray(gas_orbs), row) + suffix_lower = numpy.zeros(ngas + 1, dtype=numpy.int32) + suffix_upper = numpy.zeros(ngas + 1, dtype=numpy.int32) + for ig in range(ngas - 1, -1, -1): + suffix_lower[ig] = suffix_lower[ig + 1] + lower[ig] + suffix_upper[ig] = suffix_upper[ig + 1] + upper[ig] + + alpha = numpy.empty(ngas, dtype=numpy.int32) + + def expand(ig, remaining): + nonlocal cursor + if ig == ngas: + if remaining == 0: + blocks[cursor, :ngas] = alpha + blocks[cursor, ngas:] = row - alpha + cursor += 1 + return + lo = max(int(lower[ig]), + int(remaining - suffix_upper[ig + 1])) + hi = min(int(upper[ig]), + int(remaining - suffix_lower[ig + 1])) + for occupied in range(lo, hi + 1): + alpha[ig] = occupied + expand(ig + 1, remaining - occupied) + + expand(0, na) + + if cursor != nblock: + raise RuntimeError("internal supergroup expansion count mismatch") + blocks = normalize_blocks(blocks, ngas) + check_kernel_limits(gas_orbs, nelec, blocks) + return blocks + + +def supergroups_from_cumulative_occ(gas_orbs, nelec, bounds, + allow_empty=False): + """Expand cumulative occupation bounds into canonical supergroups G.""" + + gas_orbs = _integer_vector(gas_orbs, "gas_orbs") + ngas = len(gas_orbs) + if ngas <= 0 or ngas > GAS_MAX_NGAS: + raise ValueError("number of GAS spaces exceeds C kernel limit") + minimum = 0 if allow_empty else 1 + if any(n < minimum or n > GAS_MAX_LOCAL_ORB for n in gas_orbs): + if allow_empty: + raise ValueError("each RAS space must have 0..31 orbitals") + raise ValueError("each GAS space must have 1..31 orbitals") + if sum(gas_orbs) <= 0: + raise ValueError("at least one active orbital is required") + if sum(gas_orbs) > GAS_MAX_ORB: + raise ValueError("total active orbitals exceed 63") + + na, nb = fci_addons._unpack_nelec(nelec) + nelectron = na + nb + if na < 0 or nb < 0 or na > sum(gas_orbs) or nb > sum(gas_orbs): + raise ValueError("invalid alpha/beta electron count") + bounds = normalize_cumulative_occ(bounds, ngas) + lower = bounds[:, 0] + upper = bounds[:, 1] + cumulative_capacity = 2 * numpy.cumsum( + numpy.asarray(gas_orbs, dtype=numpy.int32)) + + if numpy.any(lower < 0): + raise ValueError("cumulative occupations must be non-negative") + if numpy.any(lower > upper): + raise ValueError("cumulative minimum exceeds maximum") + if numpy.any(lower[1:] < lower[:-1]) or numpy.any( + upper[1:] < upper[:-1]): + raise ValueError("cumulative bounds must be non-decreasing") + if numpy.any(upper > cumulative_capacity): + raise ValueError("cumulative occupation exceeds orbital capacity") + if numpy.any(upper > nelectron): + raise ValueError("cumulative occupation exceeds active electrons") + if int(lower[-1]) != nelectron or int(upper[-1]) != nelectron: + raise ValueError( + "the final cumulative bound must equal the active electron count") + + capacities = tuple(2 * int(value) for value in gas_orbs) + suffix_capacity = [0] * (ngas + 1) + for igas in range(ngas - 1, -1, -1): + suffix_capacity[igas] = ( + suffix_capacity[igas + 1] + capacities[igas]) + + def occupation_range(igas, cumulative): + lo = max(0, int(lower[igas]) - cumulative, + nelectron - cumulative - suffix_capacity[igas + 1]) + hi = min(capacities[igas], int(upper[igas]) - cumulative, + nelectron - cumulative) + return lo, hi + + @lru_cache(maxsize=None) + def count(igas, cumulative): + if igas == ngas: + return int(cumulative == nelectron) + lo, hi = occupation_range(igas, cumulative) + if lo > hi: + return 0 + return sum(count(igas + 1, cumulative + occupation) + for occupation in range(lo, hi + 1)) + + nsupergroup = count(0, 0) + if nsupergroup <= 0: + raise ValueError("cumulative bounds generate no legal supergroups") + if nsupergroup > GAS_MAX_BLOCK: + raise ValueError("number of generated supergroups exceeds int32 limit") + + supergroups = numpy.empty((nsupergroup, ngas), dtype=numpy.int32) + row = numpy.empty(ngas, dtype=numpy.int32) + cursor = 0 + + def expand(igas, cumulative): + nonlocal cursor + if igas == ngas: + if cumulative == nelectron: + supergroups[cursor] = row + cursor += 1 + return + lo, hi = occupation_range(igas, cumulative) + for occupation in range(lo, hi + 1): + if count(igas + 1, cumulative + occupation) == 0: + continue + row[igas] = occupation + expand(igas + 1, cumulative + occupation) + + expand(0, 0) + if cursor != nsupergroup: + raise RuntimeError("internal cumulative expansion count mismatch") + return normalize_supergroups(supergroups, ngas) + + +def blocks_from_cumulative_occ(gas_orbs, nelec, bounds): + """Convert cumulative occupation bounds to canonical D.""" + + supergroups = supergroups_from_cumulative_occ( + gas_orbs, nelec, bounds) + return blocks_from_supergroups(gas_orbs, nelec, supergroups) + + +def _normalize_ras_restr(gas_restr): + if not isinstance(gas_restr, Mapping): + raise TypeError( + "ras gas_restr must be a mapping with max_holes and " + "max_particles") + required = {"max_holes", "max_particles"} + missing = required.difference(gas_restr) + extra = set(gas_restr).difference(required) + if missing: + raise ValueError("missing RAS restriction keys: %s" % + ", ".join(sorted(missing))) + if extra: + raise ValueError("unknown RAS restriction keys: %s" % + ", ".join(sorted(extra))) + max_holes = gas_restr["max_holes"] + max_particles = gas_restr["max_particles"] + if not lib.isinteger(max_holes): + raise TypeError("max_holes must be an integer") + if not lib.isinteger(max_particles): + raise TypeError("max_particles must be an integer") + max_holes = int(max_holes) + max_particles = int(max_particles) + if max_holes < 0 or max_particles < 0: + raise ValueError("RAS hole and particle limits must be non-negative") + return max_holes, max_particles + + +def _ras_spec(gas_orbs, nelec, gas_restr): + """Return non-empty kernel spaces, canonical D, and RAS metadata.""" + + gas_orbs = _integer_vector( + gas_orbs, "ras gas_orbs", length=3) + if any(value < 0 or value > GAS_MAX_LOCAL_ORB for value in gas_orbs): + raise ValueError("each RAS space must have 0..31 orbitals") + if sum(gas_orbs) <= 0: + raise ValueError("at least one RAS orbital is required") + if sum(gas_orbs) > GAS_MAX_ORB: + raise ValueError("total active orbitals exceed 63") + + max_holes, max_particles = _normalize_ras_restr(gas_restr) + nras1, nras2, nras3 = gas_orbs + if max_holes > 2 * nras1: + raise ValueError("max_holes exceeds the RAS1 electron capacity") + if max_particles > 2 * nras3: + raise ValueError("max_particles exceeds the RAS3 electron capacity") + + nelectron = sum(fci_addons._unpack_nelec(nelec)) + cumulative = numpy.array([ + [max(0, 2 * nras1 - max_holes), min(2 * nras1, nelectron)], + [max(0, nelectron - max_particles), + min(2 * (nras1 + nras2), nelectron)], + [nelectron, nelectron], + ], dtype=numpy.int32) + cumulative[1, 0] = max(cumulative[1, 0], cumulative[0, 0]) + supergroups = supergroups_from_cumulative_occ( + gas_orbs, nelec, cumulative, allow_empty=True) + + nonempty = numpy.asarray(gas_orbs, dtype=numpy.int32) > 0 + kernel_orbs = tuple(int(value) for value in + numpy.asarray(gas_orbs)[nonempty]) + kernel_supergroups = normalize_supergroups( + supergroups[:, nonempty], len(kernel_orbs)) + blocks = blocks_from_supergroups( + kernel_orbs, nelec, kernel_supergroups) + return kernel_orbs, blocks, { + "ras_orbs": gas_orbs, + "max_holes": max_holes, + "max_particles": max_particles, + "cumulative_bounds": cumulative, + "canonical_supergroups": kernel_supergroups, + "conceptual_supergroups": supergroups, + "empty_spaces_removed": tuple( + index + 1 for index, size in enumerate(gas_orbs) if size == 0), + } + + +def _spin_supergroup_spec(gas_orbs, nelec, gas_restr): + gas_orbs = tuple(int(value) for value in gas_orbs) + blocks, info = _integer_rows_with_info( + gas_restr, 2 * len(gas_orbs), "spin-supergroups") + check_kernel_limits(gas_orbs, nelec, blocks) + info["canonical_spin_supergroups"] = blocks + return gas_orbs, blocks, info + + +def _supergroup_spec(gas_orbs, nelec, gas_restr): + gas_orbs = tuple(int(value) for value in gas_orbs) + supergroups, info = _integer_rows_with_info( + gas_restr, len(gas_orbs), "supergroups") + blocks = blocks_from_supergroups(gas_orbs, nelec, supergroups) + info.update({ + "canonical_supergroups": supergroups, + "canonical_spin_supergroups": blocks, + }) + return gas_orbs, blocks, info + + +def _cumulative_occ_spec(gas_orbs, nelec, gas_restr): + gas_orbs = tuple(int(value) for value in gas_orbs) + bounds = normalize_cumulative_occ(gas_restr, len(gas_orbs)) + supergroups = supergroups_from_cumulative_occ( + gas_orbs, nelec, bounds) + blocks = blocks_from_supergroups(gas_orbs, nelec, supergroups) + return gas_orbs, blocks, { + "cumulative_bounds": bounds, + "canonical_supergroups": supergroups, + "canonical_spin_supergroups": blocks, + } + + +_GAS_RESTR_CONVERTERS = { + GAS_RESTR_SPIN_SUPERGROUP: _spin_supergroup_spec, + GAS_RESTR_SUPERGROUP: _supergroup_spec, + GAS_RESTR_CUMULATIVE_OCC: _cumulative_occ_spec, + GAS_RESTR_RAS: _ras_spec, +} + + +def normalize_gas_spec(gas_orbs, nelec, gas_restr=None, + gas_restr_type=GAS_RESTR_SPIN_SUPERGROUP, + return_info=False): + """Normalize one public GAS specification to kernel spaces and D.""" + + try: + if not isinstance(gas_restr_type, str): + raise TypeError("gas_restr_type must be a string") + if gas_restr_type not in _GAS_RESTR_CONVERTERS: + supported = ", ".join(sorted(_GAS_RESTR_CONVERTERS)) + raise NotImplementedError( + "unsupported gas_restr_type %r; supported types: %s" % + (gas_restr_type, supported)) + + expected_length = 3 if gas_restr_type == GAS_RESTR_RAS else None + user_gas_orbs = _integer_vector( + gas_orbs, "gas_orbs", length=expected_length) + info = {"gas_restr_type": gas_restr_type, + "user_gas_orbs": user_gas_orbs} + + if gas_restr is None: + if gas_restr_type != GAS_RESTR_SPIN_SUPERGROUP: + raise ValueError( + "gas_restr is required for gas_restr_type %r" % + gas_restr_type) + if len(user_gas_orbs) != 1: + raise ValueError("gas_restr is required for multi-space GAS") + gas_orbs = user_gas_orbs + blocks = numpy.asarray( + [fci_addons._unpack_nelec(nelec)], dtype=numpy.int32) + check_kernel_limits(gas_orbs, nelec, blocks) + info.update({ + "kernel_gas_orbs": gas_orbs, + "canonical_spin_supergroups": blocks, + "input_nrow": 1, + "canonical_nrow": 1, + "duplicates_removed": 0, + "order_changed": False, + }) + else: + converter = _GAS_RESTR_CONVERTERS[gas_restr_type] + gas_orbs, blocks, converter_info = converter( + user_gas_orbs, nelec, gas_restr) + info.update(converter_info) + info["kernel_gas_orbs"] = gas_orbs + except (TypeError, ValueError, NotImplementedError) as err: + usage = gas_restr_usage( + gas_restr_type if isinstance(gas_restr_type, str) else None) + raise type(err)("%s\n%s" % (err, usage)) from err + + if return_info: + return gas_orbs, blocks, info + return gas_orbs, blocks + + +def _normalize_gaslst(gaslst, gas_orbs, nmo, base): + """Validate nested GAS orbital indices and return a zero-based flat list.""" + + if not lib.isinteger(base): + raise TypeError("base must be an integer") + base = int(base) + if base not in (0, 1): + raise ValueError("base must be 0 or 1") + if not isinstance(gaslst, (tuple, list)): + raise TypeError( + "gaslst must be a sequence containing one orbital list per GAS " + "subspace") + if len(gaslst) != len(gas_orbs): + raise ValueError( + "gaslst must contain %d subspace lists matching gas_orbs %s" % + (len(gas_orbs), tuple(gas_orbs))) + + flat = [] + for igas, (indices, size) in enumerate(zip(gaslst, gas_orbs)): + values = _exact_integer_array( + indices, "gaslst[%d]" % igas) + if values.ndim != 1: + raise ValueError("gaslst[%d] must be one-dimensional" % igas) + if values.size != int(size): + raise ValueError( + "gaslst[%d] must contain %d orbital indices" % + (igas, int(size))) + flat.extend(int(value) - int(base) for value in values) + + if len(set(flat)) != len(flat): + raise ValueError("gaslst orbital indices must not contain duplicates") + if any(index < 0 or index >= int(nmo) for index in flat): + raise ValueError( + "gaslst orbital indices are outside the valid range for base=%d" % + int(base)) + return flat + + +def sort_mo(mc, mo_coeff, gaslst, base=1): + """Pick and order orbitals for each GAS subspace. + + ``gaslst`` is nested, with one orbital-index list per GAS subspace. The + order of subspaces and the order within each list are preserved. As in + :func:`pyscf.mcscf.addons.sort_mo`, ``base`` selects zero- or one-based + user indices; internal indices are always zero based. + """ + + mo_array = numpy.asarray(mo_coeff) + if mo_array.ndim != 2: + raise ValueError("mo_coeff must be a two-dimensional array") + gas_orbs = ((int(mc.ncas),) if mc.gas_orbs is None else + _integer_vector(mc.gas_orbs, "gas_orbs")) + if sum(gas_orbs) != int(mc.ncas): + raise ValueError("sum(gas_orbs) must equal ncas") + active = _normalize_gaslst( + gaslst, gas_orbs, mo_array.shape[1], base) + return addons.sort_mo(mc, mo_coeff, active, base=0) + + +def normalize_gas_restr(gas_orbs, nelec, gas_restr=None, + gas_restr_type=GAS_RESTR_SPIN_SUPERGROUP): + """Convert one public restriction type to canonical C-kernel set D. + + ``spin-supergroup`` accepts D directly. ``supergroup`` accepts spin-free + occupation rows G. ``cumulative-occ`` generates G from accumulated + electron bounds. ``ras`` maps named hole/particle limits through the + cumulative representation. Every route expands to canonical D before + any downstream GAS calculation. + """ + + return normalize_gas_spec( + gas_orbs, nelec, gas_restr, gas_restr_type)[1] + + +def is_spin_complete(gas_orbs, nelec, blocks): + """Return whether every retained spatial supergroup is spin complete. + + A spin-free supergroup fixes the total occupation of each GAS subspace. + Its explicit D representation must contain every compatible alpha/beta + partition before the determinant space is invariant under ``S^2``. + """ + + gas_orbs = _integer_vector(gas_orbs, "gas_orbs") + blocks = normalize_blocks(blocks, len(gas_orbs)) + ngas = len(gas_orbs) + supergroups = numpy.unique( + blocks[:, :ngas] + blocks[:, ngas:], axis=0) + complete = blocks_from_supergroups(gas_orbs, nelec, supergroups) + actual_rows = {tuple(int(value) for value in row) for row in blocks} + complete_rows = {tuple(int(value) for value in row) for row in complete} + return actual_rows == complete_rows + + +def _sector_string_count(norb, occ): + total = 1 + for n, k in zip(norb, occ): + total *= math.comb(int(n), int(k)) + return total + + +def check_kernel_limits(norb, nelec, blocks): + """Preflight the public C-kernel limits before allocating gas_space_t.""" + + norb = _integer_vector(norb, "gas_orbs") + na, nb = fci_addons._unpack_nelec(nelec) + ngas = len(norb) + if ngas <= 0 or ngas > GAS_MAX_NGAS: + raise ValueError("number of GAS spaces exceeds C kernel limit") + if any(n <= 0 or n > GAS_MAX_LOCAL_ORB for n in norb): + raise ValueError("each GAS space must have 1..31 orbitals") + if sum(norb) > GAS_MAX_ORB: + raise ValueError("total active orbitals exceed 63") + if na < 0 or nb < 0 or na > sum(norb) or nb > sum(norb): + raise ValueError("invalid alpha/beta electron count") + + blocks = normalize_blocks(blocks, ngas) + if blocks.shape[0] > GAS_MAX_BLOCK: + raise ValueError("number of blocks exceeds C int32 input limit") + alpha = blocks[:, :ngas] + beta = blocks[:, ngas:] + if numpy.any(blocks < 0): + raise ValueError("block occupations must be non-negative") + if numpy.any(alpha > numpy.asarray(norb)): + raise ValueError("alpha block occupation exceeds local GAS orbitals") + if numpy.any(beta > numpy.asarray(norb)): + raise ValueError("beta block occupation exceeds local GAS orbitals") + if numpy.any(alpha.sum(axis=1) != na): + raise ValueError("every alpha block must sum to the requested nelec") + if numpy.any(beta.sum(axis=1) != nb): + raise ValueError("every beta block must sum to the requested nelec") + + sectors = {tuple(row) for row in alpha.tolist()} + sectors.update(tuple(row) for row in beta.tolist()) + if len(sectors) > GAS_MAX_SECTOR: + raise ValueError("number of GAS sectors exceeds uint16_t kernel limit") + for occ in sectors: + if _sector_string_count(norb, occ) > GAS_MAX_DET: + raise ValueError("sector string count exceeds uint32_t kernel limit") + + ndet = 0 + for aocc, bocc in zip(alpha, beta): + ndet += _sector_string_count(norb, aocc) * _sector_string_count(norb, bocc) + if ndet > GAS_MAX_DET: + raise ValueError("determinant count exceeds uint32_t kernel limit") + + return { + "ngas": ngas, + "norb_total": sum(norb), + "na": na, + "nb": nb, + "nblock": int(blocks.shape[0]), + "nsector_upper": len(sectors), + "ndet_estimate": int(ndet), + } + + +def pair_index(p, q): + hi = max(int(p), int(q)) + lo = min(int(p), int(q)) + return hi * (hi + 1) // 2 + lo diff --git a/pyscf/mcscf/fci_gas.py b/pyscf/mcscf/fci_gas.py new file mode 100644 index 000000000..4412c68c1 --- /dev/null +++ b/pyscf/mcscf/fci_gas.py @@ -0,0 +1,1331 @@ +#!/usr/bin/env python +# Copyright 2026 The PySCF Developers. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# Author: Yi Deng +# + +"""GAS-aware FCISolver bindings for the frozen GAS FCI C kernels.""" + +import ctypes + +import numpy + +from pyscf import ao2mo +from pyscf import lib as pyscf_lib +from pyscf.fci import addons as fci_addons +from pyscf.fci import cistring +from pyscf.fci import direct_spin1 +from pyscf.lib import logger + +from pyscf.mcscf import _gaslib +from pyscf.mcscf import addons_gas + + +# A GAS pspace column currently costs one full C-kernel Hamiltonian product. +# Keep exact small-space diagonalization bounded until a dedicated C pspace +# builder exists. +GAS_PSPACE_MATVEC_MAX = 64 +GAS_SPIN_PSPACE_MATVEC_MAX = 400 +GAS_INIT_NOISE_NORM = 1e-3 +GAS_INIT_NOISE_CHUNK = 1 << 20 + + +def _lowest_hdiag_addresses(hdiag, count): + """Return addresses of the lowest Hamiltonian diagonal elements.""" + + hdiag = numpy.asarray(hdiag).reshape(-1) + count = int(count) + if count < 1 or count > hdiag.size: + raise ValueError("count must be between 1 and the diagonal size") + if hdiag.size <= count: + return numpy.arange(hdiag.size, dtype=numpy.intp) + try: + return numpy.argpartition(hdiag, count - 1)[:count].copy() + except AttributeError: + return numpy.argsort(hdiag)[:count].copy() + + +def _add_distributed_noise(vector, root, noise_norm=GAS_INIT_NOISE_NORM): + """Add a reproducible low-norm component over the entire GAS vector.""" + + vector = numpy.asarray(vector) + if vector.ndim != 1 or vector.size == 0: + raise ValueError("distributed noise requires a nonempty vector") + noise_norm = float(noise_norm) + if noise_norm < 0.0 or not numpy.isfinite(noise_norm): + raise ValueError("noise_norm must be finite and nonnegative") + if noise_norm == 0.0: + return vector + + mask = (1 << 64) - 1 + seed_value = ((int(root) + 1) * 0x9E3779B97F4A7C15) & mask + seed = numpy.uint64(seed_value) + scale = noise_norm / numpy.sqrt(float(vector.size)) + for start in range(0, vector.size, GAS_INIT_NOISE_CHUNK): + stop = min(vector.size, start + GAS_INIT_NOISE_CHUNK) + values = numpy.arange(start, stop, dtype=numpy.uint64) + seed + values ^= values >> numpy.uint64(30) + values *= numpy.uint64(0xBF58476D1CE4E5B9) + values ^= values >> numpy.uint64(27) + values *= numpy.uint64(0x94D049BB133111EB) + values ^= values >> numpy.uint64(31) + signs = numpy.where( + values & numpy.uint64(1), 1.0, -1.0) + vector[start:stop] += scale * signs + return vector + + +def _as_c_double(array, shape=None): + source = numpy.asarray(array) + if numpy.iscomplexobj(source): + raise TypeError("the GAS C kernels require real-valued arrays") + arr = numpy.asarray(source, dtype=numpy.float64, order="C") + if shape is not None and arr.shape != shape: + raise ValueError("expected shape %s, got %s" % (shape, arr.shape)) + return numpy.ascontiguousarray(arr, dtype=numpy.float64) + + +def _as_pair_matrix(eri, norb): + norb = int(norb) + npair = norb * (norb + 1) // 2 + source = numpy.asarray(eri) + if numpy.iscomplexobj(source): + raise TypeError("the GAS C kernels require real-valued integrals") + arr = numpy.asarray(source, dtype=numpy.float64, order="C") + if arr.shape == (npair, npair): + return numpy.ascontiguousarray(arr, dtype=numpy.float64) + if arr.shape == (norb, norb, norb, norb): + out = numpy.empty((npair, npair), dtype=numpy.float64) + for p in range(norb): + for q in range(p + 1): + pq = addons_gas.pair_index(p, q) + for r in range(norb): + for s in range(r + 1): + rs = addons_gas.pair_index(r, s) + out[pq, rs] = arr[p, q, r, s] + return out + restored = ao2mo.restore(4, arr, norb) + restored = numpy.asarray(restored, dtype=numpy.float64, order="C") + if restored.shape != (npair, npair): + raise ValueError("could not convert ERI to pair-matrix layout") + return numpy.ascontiguousarray(restored, dtype=numpy.float64) + + +def _check_status(status, name): + if status != _gaslib.GAS_SUCCESS: + raise RuntimeError("%s failed with status %d" % (name, status)) + + +def absorb_h1e(h1e, eri, norb, nelec, fac=0.5): + """Return the absorbed two-electron Hamiltonian in pair-matrix layout.""" + + h1e = _as_c_double(h1e, (int(norb), int(norb))) + eri = _as_pair_matrix(eri, norb) + h2e = direct_spin1.absorb_h1e(h1e, eri, norb, nelec, fac) + return _as_pair_matrix(h2e, norb) + + +def _fci_block_addresses(gas, block): + """Return full-FCI alpha/beta addresses for one GAS block.""" + + alpha_strings = numpy.fromiter( + (gas.addr2str(block["sa"], address) + for address in range(block["na"])), + dtype=numpy.int64, count=block["na"]) + beta_strings = numpy.fromiter( + (gas.addr2str(block["sb"], address) + for address in range(block["nb"])), + dtype=numpy.int64, count=block["nb"]) + alpha_addresses = cistring.strs2addr( + gas.norb_total, gas.nelec[0], alpha_strings) + beta_addresses = cistring.strs2addr( + gas.norb_total, gas.nelec[1], beta_strings) + return alpha_addresses, beta_addresses + + +def gas2fci(fcivec, gas): + """Embed a packed GAS CI vector in the full-FCI determinant tensor. + + The returned array uses the standard ``direct_spin1`` rectangular + ``(nstr_alpha, nstr_beta)`` ordering. Determinants outside the GAS space + are represented by zero coefficients. + + This conversion is intended for interoperability and validation when the + corresponding full-FCI tensor is small enough to construct explicitly. + """ + + ci = numpy.asarray(fcivec) + if ci.size != gas.ndet: + raise ValueError("CI vector size does not match GAS determinant count") + ci = ci.reshape(-1) + + nstr_alpha = cistring.num_strings(gas.norb_total, gas.nelec[0]) + nstr_beta = cistring.num_strings(gas.norb_total, gas.nelec[1]) + fci = numpy.zeros((nstr_alpha, nstr_beta), dtype=ci.dtype) + for block in gas.block_descriptors(): + alpha_addresses, beta_addresses = _fci_block_addresses(gas, block) + size = block["na"] * block["nb"] + gas_block = ci[block["offset"]:block["offset"] + size].reshape( + block["na"], block["nb"]) + fci[numpy.ix_(alpha_addresses, beta_addresses)] = gas_block + return fci + + +def fci2gas(fcivec, gas): + """Project a full-FCI determinant tensor into packed GAS ordering. + + Coefficients belonging to determinants outside the GAS space are + discarded. The returned array is one-dimensional and follows the + canonical GAS block ordering used by the C kernels. + """ + + nstr_alpha = cistring.num_strings(gas.norb_total, gas.nelec[0]) + nstr_beta = cistring.num_strings(gas.norb_total, gas.nelec[1]) + ci = numpy.asarray(fcivec) + if ci.size != nstr_alpha * nstr_beta: + raise ValueError("CI vector size does not match full FCI determinant count") + ci = ci.reshape(nstr_alpha, nstr_beta) + + gas_ci = numpy.empty(gas.ndet, dtype=ci.dtype) + for block in gas.block_descriptors(): + alpha_addresses, beta_addresses = _fci_block_addresses(gas, block) + size = block["na"] * block["nb"] + gas_ci[block["offset"]:block["offset"] + size] = ci[ + numpy.ix_(alpha_addresses, beta_addresses)].reshape(-1) + return gas_ci + + +def _spin_penalty_parameters(solver, norb, nelec): + """Return validated spin-penalty parameters or ``None``.""" + + if not hasattr(solver, "ss_penalty"): + return None + shift = float(solver.ss_penalty) + if not numpy.isfinite(shift) or shift < 0.0: + raise ValueError("spin penalty shift must be finite and nonnegative") + na, nb = fci_addons._unpack_nelec(nelec) + sz = 0.5 * abs(na - nb) + minimum = sz * (sz + 1.0) + target_value = getattr(solver, "ss_value", None) + target = minimum if target_value is None else float(target_value) + if not numpy.isfinite(target): + raise ValueError("target S^2 must be finite") + if target < minimum - 1e-10: + raise ValueError( + "target S^2 is below the minimum allowed by Nalpha-Nbeta") + nelectron = na + nb + max_unpaired = min(nelectron, 2 * int(norb) - nelectron) + maximum_spin = 0.5 * max_unpaired + spin_values = numpy.arange(sz, maximum_spin + 0.5, 1.0) + eigenvalues = spin_values * (spin_values + 1.0) + if not numpy.any(numpy.abs(eigenvalues - target) < 1e-10): + raise ValueError( + "target S^2 is not an allowed S(S+1) value for this space") + return shift, target, minimum, tuple(float(x) for x in eigenvalues) + + +class _GasSpinPlan: + """Reusable block-sparse ``S^2`` contraction plan. + + Raw GAS one-electron links retain excitation direction. Pairing an alpha + ``q -> p`` link with the reverse beta ``p -> q`` link applies the + opposite-spin exchange term without expanding the CI vector into the full + CAS tensor. The numerical gather/scatter is delegated to PySCF's C-backed + ``take_2d`` and ``takebak_2d`` helpers. + """ + + def __init__(self, gas): + if gas.links_are_compressed(): + raise ValueError("S^2 planning requires raw GAS link tables") + self.ndet = gas.ndet + na, nb = gas.nelec + sz = 0.5 * (na - nb) + self.diagonal = sz * (sz + 1.0) + nb + self.tasks = self._build_tasks(gas) + + @staticmethod + def _table_maps(table): + maps = {} + nsrc = int(table.nsrc) + nlink = int(table.nlink) + for source in range(nsrc): + for k in range(nlink): + entry = table.link[source * nlink + k] + op = int(entry.op) + item = maps.setdefault(op, ([], [], [])) + item[0].append(source) + item[1].append(int(entry.addr)) + item[2].append(int(entry.sign)) + return { + op: ( + numpy.asarray(values[0], dtype=numpy.int32), + numpy.asarray(values[1], dtype=numpy.int32), + numpy.asarray(values[2], dtype=numpy.float64), + ) + for op, values in maps.items() + } + + @classmethod + def _build_tasks(cls, gas): + cgas = gas._gas + table_maps = {} + outgoing = [[] for _ in range(gas.nsector)] + for source in range(gas.nsector): + row = cgas.T.row[source] + for i in range(int(row.n)): + tid = int(row.off) + i + destination = int(cgas.T.dst[tid]) + outgoing[source].append((destination, tid)) + table_maps[tid] = cls._table_maps(cgas.table[tid]) + + blocks = list(gas.block_descriptors()) + block_by_sectors = { + (block["sa"], block["sb"]): block for block in blocks + } + tasks = [] + for source in blocks: + for alpha_destination, alpha_tid in outgoing[source["sa"]]: + alpha_maps = table_maps[alpha_tid] + for beta_destination, beta_tid in outgoing[source["sb"]]: + destination = block_by_sectors.get( + (alpha_destination, beta_destination)) + if destination is None: + continue + beta_maps = table_maps[beta_tid] + for op, alpha_map in alpha_maps.items(): + p = op & 0xff + q = op >> 8 + beta_map = beta_maps.get(q | (p << 8)) + if beta_map is None: + continue + tasks.append((source, destination, + alpha_map, beta_map)) + return tuple(tasks) + + def contract(self, fcivec): + ci0 = _as_c_double(numpy.asarray(fcivec).reshape(-1)) + if ci0.size != self.ndet: + raise ValueError( + "CI vector size does not match GAS determinant count") + ci1 = self.diagonal * ci0 + for source, destination, alpha, beta in self.tasks: + source_matrix = ci0[ + source["offset"]: + source["offset"] + source["na"] * source["nb"] + ].reshape(source["na"], source["nb"]) + destination_matrix = ci1[ + destination["offset"]: + destination["offset"] + + destination["na"] * destination["nb"] + ].reshape(destination["na"], destination["nb"]) + alpha_source, alpha_destination, alpha_sign = alpha + beta_source, beta_destination, beta_sign = beta + intermediate = pyscf_lib.take_2d( + source_matrix, alpha_source, beta_source) + intermediate *= alpha_sign[:, None] + intermediate *= beta_sign + pyscf_lib.takebak_2d( + destination_matrix, -intermediate, + alpha_destination, beta_destination) + return ci1.reshape(numpy.asarray(fcivec).shape) + + def diagonal_vector(self): + """Return the exact determinant-basis diagonal of ``S^2``.""" + + diagonal = numpy.full(self.ndet, self.diagonal, dtype=numpy.float64) + for source, destination, alpha, beta in self.tasks: + if source["offset"] != destination["offset"]: + continue + alpha_source, alpha_destination, alpha_sign = alpha + beta_source, beta_destination, beta_sign = beta + alpha_mask = alpha_source == alpha_destination + beta_mask = beta_source == beta_destination + if not numpy.any(alpha_mask) or not numpy.any(beta_mask): + continue + ia = alpha_source[alpha_mask] + ib = beta_source[beta_mask] + values = -numpy.multiply.outer( + alpha_sign[alpha_mask], beta_sign[beta_mask]) + block = diagonal[ + source["offset"]: + source["offset"] + source["na"] * source["nb"] + ].reshape(source["na"], source["nb"]) + numpy.add.at( + block, + (numpy.repeat(ia, ib.size), numpy.tile(ib, ia.size)), + values.reshape(-1)) + return diagonal + + def matrix(self): + """Build the explicit ``S^2`` matrix for a bounded GAS space.""" + + matrix = numpy.empty((self.ndet, self.ndet), dtype=numpy.float64) + basis = numpy.zeros(self.ndet, dtype=numpy.float64) + for column in range(self.ndet): + basis.fill(0.0) + basis[column] = 1.0 + matrix[:, column] = self.contract(basis) + return 0.5 * (matrix + matrix.T) + + def project(self, vector, target, eigenvalues): + """Project a trial vector onto one allowed total-spin eigenspace.""" + + projected = _as_c_double(numpy.asarray(vector).reshape(-1)).copy() + for eigenvalue in eigenvalues: + if abs(eigenvalue - target) < 1e-10: + continue + projected = ( + numpy.asarray(self.contract(projected)).reshape(-1) - + eigenvalue * projected) / (target - eigenvalue) + return projected + + +class GasSpace: + """Python owner for the C ``gas_space_t`` object. + + Use this low-level object as a context manager. Normal calculations + should construct it through :meth:`FCISolver.make_space` so that the + solver's normalized restriction and loaded C library are reused. + """ + + def __init__(self, norb, nelec, blocks=None, lib=None, compress_links=False): + self.norb = tuple(int(x) for x in norb) + self.nelec = fci_addons._unpack_nelec(nelec) + if blocks is None: + if len(self.norb) != 1: + raise ValueError("blocks are required for multi-space GAS") + blocks = numpy.asarray([self.nelec], dtype=numpy.int32) + self.blocks = addons_gas.normalize_blocks(blocks, len(self.norb)) + self.limits = addons_gas.check_kernel_limits(self.norb, self.nelec, + self.blocks) + self.lib = lib or _gaslib.load_library() + self._gas = _gaslib.GasSpaceStruct() + self._norb_arr = numpy.ascontiguousarray(self.norb, dtype=numpy.int32) + self._blocks_arr = numpy.ascontiguousarray(self.blocks.reshape(-1), + dtype=numpy.int32) + status = self.lib.gas_space_from_blocks( + ctypes.byref(self._gas), + len(self.norb), + self._norb_arr.ctypes.data_as(ctypes.POINTER(ctypes.c_int)), + self.nelec[0], + self.nelec[1], + self.blocks.shape[0], + self._blocks_arr.ctypes.data_as(ctypes.POINTER(ctypes.c_int))) + _check_status(status, "gas_space_from_blocks") + if compress_links: + self.compress_links() + + @property + def c_ptr(self): + return ctypes.byref(self._gas) + + @property + def ndet(self): + """Number of determinants in the canonical GAS ordering.""" + + return int(self._gas.ndet) + + @property + def nsector(self): + """Number of distinct alpha/beta string sectors.""" + + return int(self._gas.nsector) + + @property + def nblock(self): + """Number of legal spin-supergroup blocks in D.""" + + return int(self._gas.nblock) + + @property + def ntable(self): + """Number of one-electron link tables.""" + + return int(self._gas.ntable) + + @property + def norb_total(self): + return int(self._gas.norb_tot) + + @property + def memory_bytes(self): + """Total memory owned by the C GAS-space object, in bytes.""" + + return int(self.lib.gas_memory_bytes(self.c_ptr)) + + def core_info(self): + """Return the compact core state represented by ``gas_space_t``.""" + + if self._gas is None: + raise RuntimeError("GAS space is closed") + ngas = int(self._gas.ngas) + return { + "ngas": ngas, + "norb": tuple(int(self._gas.norb[i]) for i in range(ngas)), + "start": tuple(int(self._gas.start[i]) for i in range(ngas)), + "norb_total": int(self._gas.norb_tot), + "nelec": (int(self._gas.na), int(self._gas.nb)), + "ndet": int(self._gas.ndet), + "nsector": int(self._gas.nsector), + "nblock": int(self._gas.nblock), + "ntable": int(self._gas.ntable), + } + + def memory_report(self): + """Return the C-space memory breakdown as a dictionary of bytes.""" + + report = _gaslib.GasMemoryReport() + self.lib.gas_memory_report(self.c_ptr, ctypes.byref(report)) + return {name: int(getattr(report, name)) for name, _ in report._fields_} + + def compress_links(self): + """Convert raw link entries to the frozen compressed representation.""" + + status = self.lib.gas_space_compress_links(self.c_ptr) + _check_status(status, "gas_space_compress_links") + + def links_are_compressed(self): + return bool(self.lib.gas_space_links_are_compressed(self.c_ptr)) + + def addr2str(self, sector, addr): + return int(self.lib.gas_addr2str_sector(self.c_ptr, int(sector), + int(addr))) + + def str2addr(self, sector, string): + return int(self.lib.gas_str2addr_sector(self.c_ptr, int(sector), + int(string))) + + def block_ndet(self, block_id): + return int(self.lib.gas_block_ndet(self.c_ptr, int(block_id))) + + def block_descriptors(self): + for b in range(self.nblock): + block = self._gas.block[b] + yield { + "offset": int(block.offset), + "sa": int(block.sa), + "sb": int(block.sb), + "na": int(self._gas.sector_nstr[block.sa]), + "nb": int(self._gas.sector_nstr[block.sb]), + } + + def close(self): + """Release the C object; repeated calls are safe.""" + + gas = getattr(self, "_gas", None) + if gas is not None: + self.lib.gas_space_free(self.c_ptr) + self._gas = None + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, tb): + self.close() + + def __del__(self): + try: + self.close() + except Exception: + pass + + +class GasContractPlan: + """Reusable Hamiltonian contraction plan for one compressed GAS space. + + The GAS space is retained but not closed by this object. The normalized + ERI and gos arrays are retained because the C plan stores borrowed + pointers to them for the full plan lifetime. + """ + + def __init__(self, gas, eri): + if not gas.links_are_compressed(): + raise ValueError( + "Hamiltonian contraction planning requires compressed links") + self.lib = gas.lib + self.gas = gas + self.eri = _as_pair_matrix(eri, gas.norb_total) + self.gos = numpy.ascontiguousarray( + self.eri + self.eri.T, dtype=numpy.float64) + self._plan = ctypes.c_void_p() + status = self.lib.fci_contract_gas_plan_create( + ctypes.byref(self._plan), gas.c_ptr, + _gaslib.double_ptr(self.eri), _gaslib.double_ptr(self.gos)) + _check_status(status, "fci_contract_gas_plan_create") + + @property + def ndet(self): + return self.gas.ndet + + @property + def norb(self): + return self.gas.norb_total + + @property + def nelec(self): + return self.gas.nelec + + def contract(self, fcivec): + """Contract the fixed absorbed Hamiltonian with one GAS CI vector.""" + + if self._plan is None: + raise RuntimeError("Hamiltonian contraction plan is closed") + shape = numpy.asarray(fcivec).shape + ci0 = _as_c_double(numpy.asarray(fcivec).reshape(-1)) + if ci0.size != self.ndet: + raise ValueError("CI vector size does not match GAS determinant count") + ci1 = numpy.zeros_like(ci0) + status = self.lib.fci_contract_gas_plan_execute( + self._plan, _gaslib.double_ptr(ci0), _gaslib.double_ptr(ci1)) + _check_status(status, "fci_contract_gas_plan_execute") + return ci1.reshape(shape) + + def close(self): + """Release the C plan; the borrowed GAS space remains caller-owned.""" + + plan = getattr(self, "_plan", None) + if plan is not None: + self.lib.fci_contract_gas_plan_free(plan) + self._plan = None + self.eri = None + self.gos = None + self.gas = None + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, tb): + self.close() + + def __del__(self): + try: + self.close() + except Exception: + pass + + +class GasRDMPlan: + """Reusable owner of a raw-link GAS space and its C RDM plan.""" + + def __init__(self, solver, norb, nelec): + self.lib = solver.lib + self.gas = solver.make_space(norb, nelec, compress_links=False) + self._plan = ctypes.c_void_p() + try: + status = self.lib.fci_rdm_gas_plan_create( + ctypes.byref(self._plan), self.gas.c_ptr) + _check_status(status, "fci_rdm_gas_plan_create") + except Exception: + self.gas.close() + raise + + @property + def ndet(self): + return self.gas.ndet + + @property + def norb(self): + return self.gas.norb_total + + @property + def task_count(self): + return int(self.lib.fci_rdm_gas_plan_task_count(self._plan)) + + @property + def workspace_bytes(self): + return int(self.lib.fci_rdm_gas_plan_workspace_bytes(self._plan)) + + def _vectors(self, cibra, ciket): + bra = _as_c_double(numpy.asarray(cibra).reshape(-1)) + ket = _as_c_double(numpy.asarray(ciket).reshape(-1)) + if bra.size != self.ndet or ket.size != self.ndet: + raise ValueError("CI vector size does not match GAS determinant count") + return bra, ket + + def make_rdm1s(self, cibra, ciket): + """Return alpha and beta active-space transition 1-RDMs.""" + + bra, ket = self._vectors(cibra, ciket) + dm1a = numpy.zeros((self.norb, self.norb), dtype=numpy.float64) + dm1b = numpy.zeros_like(dm1a) + status = self.lib.fci_rdm_gas_plan_make_rdm1s( + self._plan, _gaslib.double_ptr(bra), _gaslib.double_ptr(ket), + _gaslib.double_ptr(dm1a), _gaslib.double_ptr(dm1b)) + _check_status(status, "fci_rdm_gas_plan_make_rdm1s") + return dm1a, dm1b + + def make_rdm1(self, cibra, ciket): + """Return the spin-summed active-space transition 1-RDM.""" + + dm1a, dm1b = self.make_rdm1s(cibra, ciket) + return dm1a + dm1b + + def make_rdm12s(self, cibra, ciket): + """Return spin-resolved active-space transition 1- and 2-RDMs.""" + + bra, ket = self._vectors(cibra, ciket) + shape1 = (self.norb, self.norb) + shape2 = shape1 + shape1 + dm1a = numpy.zeros(shape1, dtype=numpy.float64) + dm1b = numpy.zeros(shape1, dtype=numpy.float64) + dm2aa = numpy.zeros(shape2, dtype=numpy.float64) + dm2ab = numpy.zeros(shape2, dtype=numpy.float64) + dm2bb = numpy.zeros(shape2, dtype=numpy.float64) + status = self.lib.fci_rdm_gas_plan_make_rdm12s( + self._plan, _gaslib.double_ptr(bra), _gaslib.double_ptr(ket), + _gaslib.double_ptr(dm1a), _gaslib.double_ptr(dm1b), + _gaslib.double_ptr(dm2aa), _gaslib.double_ptr(dm2ab), + _gaslib.double_ptr(dm2bb)) + _check_status(status, "fci_rdm_gas_plan_make_rdm12s") + return (dm1a, dm1b), (dm2aa, dm2ab, dm2bb) + + def make_rdm12(self, cibra, ciket): + """Return spin-summed active-space transition 1- and 2-RDMs.""" + + (dm1a, dm1b), (dm2aa, dm2ab, dm2bb) = self.make_rdm12s( + cibra, ciket) + dm1 = dm1a + dm1b + dm2 = dm2aa + dm2bb + dm2 += dm2ab + dm2 += dm2ab.transpose(2, 3, 0, 1) + return dm1, dm2 + + def close(self): + """Release the RDM plan and its GAS space; repeated calls are safe.""" + + plan = getattr(self, "_plan", None) + if plan is not None: + self.lib.fci_rdm_gas_plan_free(plan) + self._plan = None + gas = getattr(self, "gas", None) + if gas is not None: + gas.close() + self.gas = None + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, tb): + self.close() + + def __del__(self): + try: + self.close() + except Exception: + pass + + +class FCISolver(direct_spin1.FCISolver): + """Determinant-space GAS solver using the frozen C/OpenMP backend. + + The Davidson driver and its convergence controls follow + :class:`pyscf.fci.direct_spin1.FCISolver`. CI vectors are one-dimensional + arrays in canonical GAS block order rather than rectangular CAS arrays. + """ + + _keys = direct_spin1.FCISolver._keys | { + "gas_orbs", "gas_restr", "gas_restr_type", "e_spin_penalty", + "e_physical", "spin_penalty_method", + } + + def __init__(self, mol=None, gas_orbs=None, gas_restr=None, + gas_restr_type=addons_gas.GAS_RESTR_SPIN_SUPERGROUP, + lib=None): + super().__init__(mol) + self.gas_orbs = tuple(gas_orbs) if gas_orbs is not None else None + self.gas_restr = gas_restr + self.gas_restr_type = gas_restr_type + self.lib = lib or _gaslib.load_library() + self.davidson_only = False + self.e_spin_penalty = None + self.e_physical = None + self.spin_penalty_method = None + + def _space_spec(self, norb, nelec): + nelec = fci_addons._unpack_nelec(nelec, self.spin) + if self.gas_orbs is None: + gas_orbs = (int(norb),) + gas_orbs, blocks = addons_gas.normalize_gas_spec( + gas_orbs, nelec) + else: + gas_orbs = self.gas_orbs + if sum(gas_orbs) != int(norb): + raise ValueError("sum(gas_orbs) must equal norb") + gas_orbs, blocks = addons_gas.normalize_gas_spec( + gas_orbs, nelec, self.gas_restr, self.gas_restr_type) + return gas_orbs, nelec, blocks + + def dump_flags(self, verbose=None): + """Print PySCF FCI flags and GAS spin-penalty settings.""" + + direct_spin1.FCISolver.dump_flags(self, verbose) + log = logger.new_logger(self, verbose) + if hasattr(self, "ss_penalty"): + target = getattr(self, "ss_value", None) + log.info("spin penalty shift = %g", self.ss_penalty) + log.info("target S^2 = %s", "minimum" if target is None else target) + return self + + def space_info(self, norb, nelec): + """Estimate GAS dimensions and validate C-kernel limits.""" + + gas_orbs, nelec, blocks = self._space_spec(norb, nelec) + return addons_gas.check_kernel_limits(gas_orbs, nelec, blocks) + + def make_space(self, norb, nelec, compress_links=False): + """Construct an owned :class:`GasSpace` for this solver.""" + + gas_orbs, nelec, blocks = self._space_spec(norb, nelec) + return GasSpace(gas_orbs, nelec, blocks, lib=self.lib, + compress_links=compress_links) + + def make_rdm_plan(self, norb, nelec): + """Construct a reusable :class:`GasRDMPlan` context manager.""" + + return GasRDMPlan(self, norb, nelec) + + def get_init_guess(self, norb, nelec, nroots, hdiag, gas=None): + """Build PySCF-style determinant guesses in the GAS vector layout.""" + + hdiag = _as_c_double(numpy.asarray(hdiag).reshape(-1)) + ndet = (self.space_info(norb, nelec)["ndet_estimate"] + if gas is None else gas.ndet) + if hdiag.size != ndet: + raise ValueError("hdiag size does not match GAS determinant count") + nroots = int(nroots) + if nroots < 1 or nroots > hdiag.size: + raise ValueError("nroots must be between 1 and the GAS dimension") + addresses = _lowest_hdiag_addresses(hdiag, nroots) + guesses = [] + for root, address in enumerate(addresses): + guess = numpy.zeros(hdiag.size, dtype=numpy.float64) + guess[int(address)] = 1.0 + # Independent globally supported probes are required for robust + # multiroot convergence when the Hamiltonian contains invariant + # sectors that are absent from the lowest diagonal determinants. + _add_distributed_noise(guess, root) + guesses.append(guess) + return guesses + + def make_hdiag(self, h1e, eri, norb, nelec, *args, **kwargs): + """Return the Hamiltonian diagonal in canonical GAS order.""" + + compress_links = bool(kwargs.pop("compress_links", False)) + with self.make_space(norb, nelec, compress_links=compress_links) as gas: + n = gas.norb_total + h1e = _as_c_double(h1e, (n, n)) + eri = _as_pair_matrix(eri, n) + out = numpy.empty(gas.ndet, dtype=numpy.float64) + status = self.lib.fci_make_hdiag_gas( + gas.c_ptr, _gaslib.double_ptr(h1e), _gaslib.double_ptr(eri), + _gaslib.double_ptr(out)) + _check_status(status, "fci_make_hdiag_gas") + return out + + def contract_2e(self, eri, fcivec, norb, nelec, link_index=None, + *args, **kwargs): + """Contract an absorbed Hamiltonian with a GAS CI vector.""" + + plan = kwargs.pop("plan", None) + if plan is not None: + if not isinstance(plan, GasContractPlan): + raise TypeError("plan must be a GasContractPlan") + expected_nelec = fci_addons._unpack_nelec(nelec, self.spin) + if int(norb) != plan.norb or expected_nelec != plan.nelec: + raise ValueError("contraction plan does not match norb/nelec") + return plan.contract(fcivec) + + compress_links = bool(kwargs.pop("compress_links", True)) + with self.make_space(norb, nelec, compress_links=compress_links) as gas: + with GasContractPlan(gas, eri) as plan: + return plan.contract(fcivec) + + def contract_1e(self, f1e, fcivec, norb, nelec, link_index=None, + **kwargs): + raise NotImplementedError( + "GASCI uses the absorbed one-electron Hamiltonian and contract_2e") + + def energy(self, h1e, eri, fcivec, norb, nelec, link_index=None): + """Evaluate the GAS electronic energy for a normalized CI vector.""" + + h2e = absorb_h1e(h1e, eri, norb, nelec, fac=0.5) + ci = _as_c_double(numpy.asarray(fcivec).reshape(-1)) + hc = self.contract_2e(h2e, ci, norb, nelec, link_index) + return float(numpy.dot(ci, numpy.asarray(hc).reshape(-1))) + + def pspace(self, h1e, eri, norb, nelec, hdiag=None, np=400): + """Build a bounded GAS pspace Hamiltonian. + + Unlike CAS ``direct_spin1``, the frozen GAS kernel has no dedicated + pspace builder. Each column therefore uses one full Hamiltonian + product and the requested dimension is capped at + ``GAS_PSPACE_MATVEC_MAX``. + """ + + requested = int(np) + if requested < 1: + raise ValueError("np must be positive") + h1e = _as_c_double(h1e, (int(norb), int(norb))) + eri_phys = _as_pair_matrix(eri, norb) + h2e = absorb_h1e(h1e, eri_phys, norb, nelec, fac=0.5) + with self.make_space(norb, nelec, compress_links=True) as gas: + if hdiag is None: + diagonal = numpy.empty(gas.ndet, dtype=numpy.float64) + status = self.lib.fci_make_hdiag_gas( + gas.c_ptr, _gaslib.double_ptr(h1e), + _gaslib.double_ptr(eri_phys), + _gaslib.double_ptr(diagonal)) + _check_status(status, "fci_make_hdiag_gas") + else: + diagonal = _as_c_double(numpy.asarray(hdiag).reshape(-1)) + if diagonal.size != gas.ndet: + raise ValueError( + "hdiag size does not match GAS determinant count") + + size = min(requested, gas.ndet, GAS_PSPACE_MATVEC_MAX) + addresses = _lowest_hdiag_addresses(diagonal, size) + with GasContractPlan(gas, h2e) as plan: + h0 = self._pspace_with_plan(plan, addresses) + return addresses, h0 + + def gen_linkstr(self, norb, nelec, tril=True, spin=None): + raise NotImplementedError( + "GAS link tables are owned by GasSpace and are not CAS linkstr arrays") + + def contract_ss(self, fcivec, norb, nelec): + """Contract the spin-square operator with a GAS CI vector.""" + + gas_orbs, nelec, blocks = self._space_spec(norb, nelec) + if not addons_gas.is_spin_complete(gas_orbs, nelec, blocks): + raise ValueError( + "contract_ss requires a spin-complete GAS restriction") + with self.make_space(norb, nelec, compress_links=False) as gas: + return _GasSpinPlan(gas).contract(fcivec) + + def transform_ci_for_orbital_rotation(self, fcivec, norb, nelec, u): + raise NotImplementedError( + "arbitrary active-orbital rotations do not preserve a GAS space") + + def large_ci(self, fcivec, norb, nelec, tol=0.1, return_strs=True): + """Return the largest GAS CI coefficients and determinant strings.""" + + ci = _as_c_double(numpy.asarray(fcivec).reshape(-1)) + with self.make_space(norb, nelec, compress_links=False) as gas: + if ci.size != gas.ndet: + raise ValueError("CI vector size does not match GAS determinant count") + selected = numpy.flatnonzero(numpy.abs(ci) > float(tol)) + if selected.size == 0: + selected = numpy.asarray([int(numpy.argmax(numpy.abs(ci)))]) + + result = [] + blocks = list(gas.block_descriptors()) + iblock = 0 + for address in selected: + address = int(address) + while (iblock + 1 < len(blocks) and + address >= blocks[iblock + 1]["offset"]): + iblock += 1 + desc = blocks[iblock] + local = address - desc["offset"] + ia, ib = divmod(local, desc["nb"]) + alpha = gas.addr2str(desc["sa"], ia) + beta = gas.addr2str(desc["sb"], ib) + if return_strs: + alpha_out = bin(alpha) + beta_out = bin(beta) + else: + alpha_out = numpy.asarray( + [p for p in range(int(norb)) if (alpha >> p) & 1], + dtype=numpy.int32) + beta_out = numpy.asarray( + [p for p in range(int(norb)) if (beta >> p) & 1], + dtype=numpy.int32) + result.append((ci[address], alpha_out, beta_out)) + return result + + def _pspace_with_plan(self, plan, addresses): + """Construct H over selected compact GAS determinant addresses.""" + + addresses = numpy.asarray(addresses, dtype=numpy.intp).reshape(-1) + h0 = numpy.empty((addresses.size, addresses.size), + dtype=numpy.float64) + basis = numpy.zeros(plan.ndet, dtype=numpy.float64) + for column, address in enumerate(addresses): + basis.fill(0.0) + basis[int(address)] = 1.0 + product = plan.contract(basis).reshape(-1) + h0[:, column] = product[addresses] + # Roundoff in independently accumulated columns can leave tiny + # antisymmetric components. LAPACK receives an explicitly symmetric H. + return 0.5 * (h0 + h0.T) + + def kernel(self, h1e, eri, norb, nelec, ci0=None, ecore=0, + nroots=None, **kwargs): + """Solve the GASCI problem with a bounded pspace or Davidson. + + This method is intentionally low-level: h1e and eri are assumed to be + active-space integrals, and GAS restrictions are supplied through + ``gas_orbs`` and ``gas_restr`` on the solver. + """ + + nroots = int(nroots if nroots is not None + else kwargs.pop("nroots", getattr(self, "nroots", 1))) + if nroots < 1: + raise ValueError("nroots must be positive") + + h1e = _as_c_double(h1e, (int(norb), int(norb))) + eri_phys = _as_pair_matrix(eri, norb) + h2e = absorb_h1e(h1e, eri_phys, norb, nelec, fac=0.5) + + def option(name, default): + value = kwargs.pop(name, None) + return default if value is None else value + + conv_tol = option("tol", getattr(self, "conv_tol", 1e-10)) + conv_tol_residual = option( + "tol_residual", getattr(self, "conv_tol_residual", None)) + max_cycle = option("max_cycle", getattr(self, "max_cycle", 50)) + max_space = option("max_space", getattr(self, "max_space", 12)) + lindep = option("lindep", getattr(self, "lindep", 1e-14)) + verbose = option("verbose", getattr(self, "verbose", None)) + level_shift = option("level_shift", getattr(self, "level_shift", 1e-3)) + max_memory = option("max_memory", getattr(self, "max_memory", 4000)) + lessio = bool(option("lessio", getattr(self, "lessio", False))) + davidson_only = bool(option("davidson_only", self.davidson_only)) + pspace_size = int(option( + "pspace_size", getattr(self, "pspace_size", 400))) + orbsym = option("orbsym", None) + wfnsym = option("wfnsym", None) + kwargs.pop("envs", None) + if pspace_size < 0: + raise ValueError("pspace_size must be nonnegative") + if orbsym is not None or wfnsym is not None: + raise NotImplementedError( + "orbital and wavefunction symmetry filtering is not implemented") + + self.norb = int(norb) + self.nelec = fci_addons._unpack_nelec(nelec, self.spin) + spin_penalty = _spin_penalty_parameters(self, norb, self.nelec) + self.e_spin_penalty = None + self.e_physical = None + self.spin_penalty_method = None + + with self.make_space(norb, nelec, compress_links=False) as gas: + if nroots > gas.ndet: + raise ValueError("nroots exceeds GAS determinant count") + if (spin_penalty is not None and + not addons_gas.is_spin_complete( + gas.norb, self.nelec, gas.blocks)): + raise ValueError( + "fix_spin_ requires a spin-complete GAS restriction") + spin_plan = (_GasSpinPlan(gas) + if spin_penalty is not None else None) + gas.compress_links() + hdiag = numpy.empty(gas.ndet, dtype=numpy.float64) + status = self.lib.fci_make_hdiag_gas( + gas.c_ptr, _gaslib.double_ptr(h1e), _gaslib.double_ptr(eri_phys), + _gaslib.double_ptr(hdiag)) + _check_status(status, "fci_make_hdiag_gas") + + linear_spin_penalty = False + if spin_penalty is not None: + (spin_shift, spin_target, spin_minimum, + spin_eigenvalues) = spin_penalty + linear_spin_penalty = spin_target < spin_minimum + 0.1 + + plan = GasContractPlan(gas, h2e) + try: + full_pspace = ( + spin_penalty is None and ci0 is None and + not davidson_only and + gas.ndet <= pspace_size and + gas.ndet <= GAS_PSPACE_MATVEC_MAX) + full_spin_pspace = ( + spin_penalty is not None and ci0 is None and + pspace_size > 0 and gas.ndet <= pspace_size and + gas.ndet <= GAS_SPIN_PSPACE_MATVEC_MAX) + if full_pspace or full_spin_pspace: + addresses = _lowest_hdiag_addresses(hdiag, gas.ndet) + h0 = self._pspace_with_plan(plan, addresses) + if full_spin_pspace: + s2 = spin_plan.matrix() + s2 = s2[numpy.ix_(addresses, addresses)] + delta_s2 = s2 - spin_target * numpy.eye(gas.ndet) + if linear_spin_penalty: + h0 += spin_shift * delta_s2 + else: + h0 += spin_shift * numpy.dot(delta_s2, delta_s2) + self.spin_penalty_method = "exact-small-space" + eigenvalues, eigenvectors = numpy.linalg.eigh(h0) + e = eigenvalues[:nroots] + c = [] + for root in range(nroots): + vector = numpy.empty(gas.ndet, dtype=numpy.float64) + vector[addresses] = eigenvectors[:, root] + c.append(vector) + converged = numpy.ones(nroots, dtype=bool) + else: + guess_count = nroots + if spin_penalty is not None: + guess_count = min( + gas.ndet, max(4 * nroots, nroots + 4)) + default_guess = self.get_init_guess( + norb, nelec, guess_count, hdiag, gas=gas) + if ci0 is None: + guess = default_guess + else: + if isinstance(ci0, (list, tuple)): + guess = [ + _as_c_double(numpy.asarray(x).reshape(-1)) + for x in ci0] + else: + arr = numpy.asarray(ci0) + if arr.ndim == 2 and arr.shape[0] == nroots: + guess = [ + _as_c_double(arr[i].reshape(-1)) + for i in range(nroots)] + elif arr.ndim == 2 and arr.shape[1] == nroots: + guess = [ + _as_c_double(arr[:, i].reshape(-1)) + for i in range(nroots)] + else: + guess = [_as_c_double(arr.reshape(-1))] + for i, vec in enumerate(guess): + if vec.size != gas.ndet: + raise ValueError( + "initial CI vector size does not match " + "GAS determinant count") + norm = numpy.linalg.norm(vec) + if norm == 0: + raise ValueError( + "initial CI vector must be nonzero") + guess[i] = vec / norm + # direct_spin1 extends an undersized user guess so that + # multiroot Davidson can generate enough trial vectors. + if len(guess) < guess_count: + guess.extend(default_guess[len(guess):guess_count]) + if spin_penalty is not None: + projected_guess = [] + for candidate in guess: + vector = spin_plan.project( + candidate, spin_target, spin_eigenvalues) + for previous in projected_guess: + vector -= numpy.dot(previous, vector) * previous + norm = numpy.linalg.norm(vector) + if norm > max(100.0 * lindep, 1e-12): + projected_guess.append(vector / norm) + if len(projected_guess) == nroots: + break + if len(projected_guess) < nroots: + raise RuntimeError( + "could not construct enough independent " + "target-spin initial guesses") + # Retain a small unprojected globally supported probe. + # The Davidson problem must still find the lowest + # eigenvalue of the penalized Hamiltonian when the + # requested shift is too small to change spin ordering. + exploration_guess = [] + for candidate in guess[:max(1, nroots)]: + norm = numpy.linalg.norm(candidate) + if norm > 0.0: + exploration_guess.append(candidate / norm) + guess = projected_guess + exploration_guess + self.spin_penalty_method = ( + "projected-plus-global-davidson") + if nroots == 1 and len(guess) == 1: + guess = guess[0] + + def hop(vec): + result = self.contract_2e( + h2e, vec, norb, nelec, plan=plan).reshape(-1) + if spin_penalty is None: + return result + ss_vector = numpy.asarray( + spin_plan.contract(vec)).reshape(-1) + if linear_spin_penalty: + ss_vector -= spin_target * vec + result += spin_shift * ss_vector + return result + tmp = ss_vector - spin_target * vec + correction = numpy.asarray( + spin_plan.contract(tmp)).reshape(-1) + correction -= spin_target * tmp + result += spin_shift * correction + return result + + preconditioner_diagonal = hdiag + if spin_penalty is not None: + spin_diagonal = spin_plan.diagonal_vector() + delta_diagonal = spin_diagonal - spin_target + if linear_spin_penalty: + penalty_diagonal = spin_shift * delta_diagonal + else: + # This is the inexpensive diagonal approximation to + # (S^2-target)^2. The projected trial vectors + # improve access to the target sector; this term + # only improves Davidson conditioning. + penalty_diagonal = ( + spin_shift * delta_diagonal * delta_diagonal) + preconditioner_diagonal = hdiag + penalty_diagonal + + def precond(dx, e, *args): + denom = preconditioner_diagonal - e + small = numpy.abs(denom) < level_shift + denom = denom.copy() + denom[small] = level_shift + return dx / denom + + converged, e, c = pyscf_lib.davidson1( + lambda vectors: [hop(vec) for vec in vectors], + guess, precond, tol=conv_tol, max_cycle=max_cycle, + max_space=max_space, lindep=lindep, + max_memory=max_memory, nroots=nroots, lessio=lessio, + verbose=verbose, tol_residual=conv_tol_residual, + follow_state=spin_penalty is None) + finally: + plan.close() + + e = numpy.asarray(e, dtype=numpy.float64).reshape(-1) + float(ecore) + if isinstance(c, (list, tuple)): + c_list = [numpy.asarray(x, dtype=numpy.float64).reshape(-1) + for x in c] + else: + c_arr = numpy.asarray(c, dtype=numpy.float64) + if nroots == 1: + c_list = [c_arr.reshape(-1)] + elif c_arr.ndim == 2 and c_arr.shape[0] == nroots: + c_list = [c_arr[i].reshape(-1) for i in range(nroots)] + elif c_arr.ndim == 2 and c_arr.shape[1] == nroots: + c_list = [c_arr[:, i].reshape(-1) for i in range(nroots)] + else: + c_list = [c_arr.reshape(-1)] + if spin_penalty is not None: + penalty_values = [] + # The contraction GasSpace above has left its context. Rebuild a + # short-lived raw-link space for post-solver diagnostics instead + # of retaining a plan backed by released C memory. + with self.make_space( + norb, nelec, compress_links=False) as diagnostic_gas: + diagnostic_spin = _GasSpinPlan(diagnostic_gas) + for vector in c_list[:nroots]: + ss_vector = numpy.asarray( + diagnostic_spin.contract(vector)).reshape(-1) + if linear_spin_penalty: + penalty = spin_shift * ( + numpy.dot(vector, ss_vector) - spin_target) + else: + delta = ss_vector - spin_target * vector + penalty = spin_shift * numpy.dot(delta, delta) + penalty_values.append(float(penalty)) + penalty_values = numpy.asarray( + penalty_values, dtype=numpy.float64) + physical_values = e[:nroots] - penalty_values + if nroots == 1: + self.e_spin_penalty = float(penalty_values[0]) + self.e_physical = float(physical_values[0]) + else: + self.e_spin_penalty = penalty_values + self.e_physical = physical_values + if nroots == 1: + self.converged = bool(numpy.asarray(converged).reshape(-1)[0]) + self.eci = float(e[0]) + self.ci = c_list[0] + else: + self.converged = numpy.asarray(converged, dtype=bool)[:nroots] + self.eci = e[:nroots] + self.ci = c_list[:nroots] + return self.eci, self.ci + + @pyscf_lib.with_doc(direct_spin1.make_rdm1s.__doc__) + def make_rdm1s(self, ci, norb, nelec, link_index=None): + return self.trans_rdm1s(ci, ci, norb, nelec, link_index) + + @pyscf_lib.with_doc(direct_spin1.make_rdm1.__doc__) + def make_rdm1(self, ci, norb, nelec, link_index=None): + return self.trans_rdm1(ci, ci, norb, nelec, link_index) + + @pyscf_lib.with_doc(direct_spin1.make_rdm12s.__doc__) + def make_rdm12s(self, ci, norb, nelec, link_index=None, reorder=True): + if not reorder: + raise NotImplementedError("reorder=False is not supported") + with self.make_rdm_plan(norb, nelec) as plan: + return plan.make_rdm12s(ci, ci) + + @pyscf_lib.with_doc(direct_spin1.make_rdm12.__doc__) + def make_rdm12(self, ci, norb, nelec, link_index=None, reorder=True): + if not reorder: + raise NotImplementedError("reorder=False is not supported") + with self.make_rdm_plan(norb, nelec) as plan: + return plan.make_rdm12(ci, ci) + + def make_rdm2(self, ci, norb, nelec, link_index=None, reorder=True): + """Return the spin-traced GAS two-particle density matrix.""" + + return self.make_rdm12(ci, norb, nelec, link_index, reorder)[1] + + @pyscf_lib.with_doc(direct_spin1.trans_rdm1s.__doc__) + def trans_rdm1s(self, cibra, ciket, norb, nelec, link_index=None): + with self.make_rdm_plan(norb, nelec) as plan: + return plan.make_rdm1s(cibra, ciket) + + @pyscf_lib.with_doc(direct_spin1.trans_rdm1.__doc__) + def trans_rdm1(self, cibra, ciket, norb, nelec, link_index=None): + with self.make_rdm_plan(norb, nelec) as plan: + return plan.make_rdm1(cibra, ciket) + + @pyscf_lib.with_doc(direct_spin1.trans_rdm12s.__doc__) + def trans_rdm12s(self, cibra, ciket, norb, nelec, link_index=None, + reorder=True): + if not reorder: + raise NotImplementedError("reorder=False is not supported") + with self.make_rdm_plan(norb, nelec) as plan: + dm1s, (dm2aa, dm2ab, dm2bb) = plan.make_rdm12s(cibra, ciket) + _, (_, dm2ba_ji, _) = plan.make_rdm12s(ciket, cibra) + dm2ba = dm2ba_ji.transpose(3, 2, 1, 0) + return dm1s, (dm2aa, dm2ab, dm2ba, dm2bb) + + @pyscf_lib.with_doc(direct_spin1.trans_rdm12.__doc__) + def trans_rdm12(self, cibra, ciket, norb, nelec, link_index=None, + reorder=True): + if not reorder: + raise NotImplementedError("reorder=False is not supported") + with self.make_rdm_plan(norb, nelec) as plan: + return plan.make_rdm12(cibra, ciket) + + def spin_square(self, ci, norb, nelec, *args, **kwargs): + """Return ``(, 2S+1)`` from spin-resolved GAS RDMs.""" + + nelec = fci_addons._unpack_nelec(nelec, self.spin) + (dm1a, dm1b), (dm2aa, dm2ab, dm2bb) = self.make_rdm12s( + ci, norb, nelec) + return spin_square_from_rdm12s((dm1a, dm1b), (dm2aa, dm2ab, dm2bb), + nelec) + + +def spin_square_from_rdm12s(dm1s, dm2s, nelec): + """Compute from spin-resolved GAS RDMs. + + The formula uses the opposite-spin exchange trace from dm2ab and is valid + for fixed (Nalpha, Nbeta) CI vectors in a common spatial orbital basis. + """ + + na, nb = fci_addons._unpack_nelec(nelec) + dm2ab = dm2s[1] + sz = 0.5 * (na - nb) + trans = numpy.einsum("pqqp->", dm2ab) + ss = sz * (sz + 1.0) + nb - trans + ss = float(numpy.real_if_close(ss)) + mult = (max(0.0, 4.0 * ss + 1.0)) ** 0.5 + return ss, mult diff --git a/pyscf/mcscf/gasci.py b/pyscf/mcscf/gasci.py new file mode 100644 index 000000000..1cf0bd8a1 --- /dev/null +++ b/pyscf/mcscf/gasci.py @@ -0,0 +1,1155 @@ +#!/usr/bin/env python +# Copyright 2026 The PySCF Developers. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# Author: Yi Deng +# + +"""PySCF-style generalized active-space configuration interaction.""" + +import hashlib +from functools import reduce + +import numpy + +from pyscf import __config__ +from pyscf import fci +from pyscf import gto +from pyscf import lib +from pyscf.fci import addons as fci_addons +from pyscf.lib import logger +from pyscf.mcscf import addons +from pyscf.mcscf import casci + +from pyscf.mcscf import addons_gas +from pyscf.mcscf import fci_gas + + +WITH_META_LOWDIN = getattr( + __config__, "mcscf_analyze_with_meta_lowdin", True) +LARGE_CI_TOL = getattr(__config__, "mcscf_analyze_large_ci_tol", 0.1) +MO_ORTH_WARN_TOL = getattr( + __config__, "mcscf_gasci_mo_orth_warn_tol", 1e-7) +MO_ORTH_ERROR_TOL = getattr( + __config__, "mcscf_gasci_mo_orth_error_tol", 1e-5) + + +h1e_for_gas = casci.h1e_for_cas + + +def kernel(mc, mo_coeff=None, ci0=None, verbose=logger.NOTE): + """Run the fixed-orbital GASCI solver.""" + + if mo_coeff is None: + mo_coeff = mc.mo_coeff + if ci0 is None: + ci0 = mc.ci + + log = logger.new_logger(mc, verbose) + t0 = (logger.process_clock(), logger.perf_counter()) + log.debug("Start GASCI") + + eri_gas = mc.get_h2eff(mo_coeff) + t1 = log.timer("integral transformation to GAS space", *t0) + + h1eff, energy_core = mc.get_h1eff(mo_coeff) + log.debug("core energy = %.15g", energy_core) + t1 = log.timer("effective h1e in GAS space", *t1) + + if h1eff.shape[0] != mc.ncas: + raise RuntimeError( + "Active space size error. nmo=%d ncore=%d ncas=%d" % + (mo_coeff.shape[1], mc.ncore, mc.ncas)) + + max_memory = max(400, mc.max_memory - lib.current_memory()[0]) + e_tot, ci = mc.fcisolver.kernel( + h1eff, eri_gas, mc.ncas, mc.nelecas, + ci0=ci0, verbose=log, max_memory=max_memory, ecore=energy_core) + + log.timer("GASCI solver", *t1) + return e_tot, e_tot - energy_core, ci + + +def as_scanner(mc): + """Return a GASCI potential-energy-surface scanner. + + The previous CI solution is reused only when the normalized GAS problem, + determinant count, and requested root count are unchanged. + """ + + if isinstance(mc, lib.SinglePointScanner): + return mc + logger.info(mc, "Create scanner for %s", mc.__class__) + name = mc.__class__.__name__ + GASCIScanner.__name_mixin__ + return lib.set_class( + GASCIScanner(mc), (GASCIScanner, mc.__class__), name) + + +class GASCIScanner(lib.SinglePointScanner): + """Callable GASCI scanner following the PySCF CASCI scanner protocol.""" + + __name_mixin__ = "Scanner" + + def __init__(self, mc): + self.__dict__.update(mc.__dict__) + self._scf = mc._scf.as_scanner() + + def __call__(self, mol_or_geom, mo_coeff=None, ci0=None): + if isinstance(mol_or_geom, gto.MoleBase): + mol = mol_or_geom + else: + mol = self.mol.set_geom_(mol_or_geom, inplace=False) + + self.reset(mol) + for key in ("with_df", "with_x2c", "with_solvent", "with_dftd3"): + submodule = getattr(self, key, None) + if submodule: + submodule.reset(mol) + if mo_coeff is None: + self._scf(mol) + mo_coeff = self._scf.mo_coeff + self.mol = mol + if ci0 is None: + # A compact GAS space can contain disconnected spin-supergroup + # sectors. Reusing only the preceding roots may then follow an + # excited sector through a crossing. For spaces covered by the + # bounded exact pspace, solve for the lowest roots afresh; larger + # production spaces retain PySCF-style CI reuse. + try: + ndet = self._gas_problem_signature()[-1] + except (TypeError, ValueError, NotImplementedError): + pass + else: + pspace_size = int(getattr( + self.fcisolver, "pspace_size", 400)) + if (not self.fcisolver.davidson_only and + ndet <= pspace_size and + ndet <= fci_gas.GAS_PSPACE_MATVEC_MAX): + self._clear_ci_guess() + return self.kernel(mo_coeff, ci0)[0] + + +class GASCI(casci.CASCI): + """Configuration interaction in a generalized active space. + + Args: + mf : SCF object + Molecular-orbital reference used to form active-space integrals. + ncas : int + Total number of active orbitals. + nelecas : int or pair of ints + Number of active electrons, optionally resolved as alpha/beta. + gas_orbs : sequence of ints, optional + Ordered numbers of orbitals in the GAS subspaces. + gas_restr : object, optional + Restriction in the format selected by ``gas_restr_type``. + gas_restr_type : str + ``spin-supergroup``, ``supergroup``, ``cumulative-occ``, or + ``ras``. Every representation is converted to canonical D. + + Notes: + GAS CI vectors are one-dimensional arrays in canonical block order. + ``e_gas`` is the public active-space energy name; the inherited + ``e_cas`` attribute is retained only for PySCF compatibility. + """ + + _keys = casci.CASCI._keys | { + "gas_orbs", "gas_restr", "gas_restr_type", "e_spin_penalty", + "e_tot_physical", "e_gas_physical", "spin_penalty_method", + } + + def __init__(self, mf, ncas, nelecas, gas_orbs=None, gas_restr=None, + gas_restr_type=addons_gas.GAS_RESTR_SPIN_SUPERGROUP, + **kwargs): + super().__init__(mf, ncas, nelecas, **kwargs) + self.gas_orbs = tuple(gas_orbs) if gas_orbs is not None else None + self.gas_restr = gas_restr + self.gas_restr_type = gas_restr_type + self.fcisolver = fci_gas.FCISolver( + getattr(mf, "mol", None), gas_orbs=self.gas_orbs, + gas_restr=self.gas_restr, gas_restr_type=self.gas_restr_type) + self._gas_ci_signature = None + self.e_spin_penalty = None + self.e_tot_physical = None + self.e_gas_physical = None + self.spin_penalty_method = None + + @property + def ngas(self): + """Number of ordered GAS subspaces presented to the user.""" + + return 1 if self.gas_orbs is None else len(self.gas_orbs) + + @property + def e_gas(self): + """Electronic energy of the GAS problem, excluding the core energy.""" + + return self.e_cas + + def _sync_fcisolver(self): + self.fcisolver.gas_orbs = self.gas_orbs + self.fcisolver.gas_restr = self.gas_restr + self.fcisolver.gas_restr_type = self.gas_restr_type + + def _effective_nelecas(self, nelecas=None): + """Return active alpha/beta counts after applying ``fcisolver.spin``.""" + + if nelecas is None: + nelecas = self.nelecas + return fci_addons._unpack_nelec(nelecas, self.fcisolver.spin) + + def _normalized_restriction(self, return_info=False): + gas_orbs = (int(self.ncas),) if self.gas_orbs is None else self.gas_orbs + return addons_gas.normalize_gas_spec( + gas_orbs, self._effective_nelecas(), + self.gas_restr, self.gas_restr_type, + return_info=return_info) + + def gas_space_info(self): + """Return normalized GAS metadata and compact C-space information. + + ``metadata`` describes the public GAS specification after its + documented normalization. ``core`` is read from the constructed + ``gas_space_t`` object and therefore reports the space actually seen + by the C kernels. + """ + + self._sync_fcisolver() + gas_orbs, blocks, info = self._normalized_restriction( + return_info=True) + restriction_type = self.gas_restr_type + if self.gas_restr is None: + restriction = None + elif restriction_type == addons_gas.GAS_RESTR_SPIN_SUPERGROUP: + restriction = numpy.array( + info["canonical_spin_supergroups"], copy=True) + elif restriction_type == addons_gas.GAS_RESTR_SUPERGROUP: + restriction = numpy.array( + info["canonical_supergroups"], copy=True) + elif restriction_type == addons_gas.GAS_RESTR_CUMULATIVE_OCC: + restriction = numpy.array(info["cumulative_bounds"], copy=True) + elif restriction_type == addons_gas.GAS_RESTR_RAS: + restriction = { + "max_holes": int(info["max_holes"]), + "max_particles": int(info["max_particles"]), + } + else: # normalize_gas_spec rejects this before reaching this branch. + raise RuntimeError("unrecognized normalized GAS restriction type") + + user_gas_orbs = ((int(self.ncas),) if self.gas_orbs is None else + tuple(int(value) for value in self.gas_orbs)) + with fci_gas.GasSpace( + gas_orbs, self._effective_nelecas(), blocks, + lib=self.fcisolver.lib) as space: + core = space.core_info() + return { + "metadata": { + "gas_orbs": user_gas_orbs, + "gas_restr_type": restriction_type, + "gas_restr": restriction, + "kernel_gas_orbs": tuple(int(value) for value in gas_orbs), + "spin_supergroups": numpy.array(blocks, copy=True), + }, + "core": core, + } + + def _gas_problem_signature(self): + """Return the normalized GAS definition associated with a CI vector.""" + + gas_orbs, gas_restr = self._normalized_restriction() + gas_restr = numpy.ascontiguousarray(gas_restr, dtype=numpy.int32) + nelecas = self._effective_nelecas() + limits = addons_gas.check_kernel_limits( + gas_orbs, nelecas, gas_restr) + digest = hashlib.sha256(gas_restr.tobytes()).digest() + nroots = int(getattr(self.fcisolver, "nroots", 1)) + return ( + int(self.ncas), tuple(int(value) for value in nelecas), + tuple(int(value) for value in gas_orbs), + addons_gas.GAS_RESTR_SPIN_SUPERGROUP, + tuple(int(value) for value in gas_restr.shape), digest, + nroots, int(limits["ndet_estimate"]), + ) + + def _ci_matches_signature(self, ci, signature): + if ci is None or self._gas_ci_signature != signature: + return False + nroots, ndet = signature[-2:] + if isinstance(ci, (list, tuple)): + return (len(ci) == nroots and + all(numpy.asarray(root).size == ndet for root in ci)) + array = numpy.asarray(ci) + if nroots == 1: + return array.size == ndet + return array.ndim == 2 and array.shape in { + (nroots, ndet), (ndet, nroots), + } + + def _clear_ci_guess(self): + self.ci = None + self._gas_ci_signature = None + for name in ("ci", "eci"): + if hasattr(self.fcisolver, name): + setattr(self.fcisolver, name, None) + self.e_spin_penalty = None + self.e_tot_physical = None + self.e_gas_physical = None + self.spin_penalty_method = None + + def reset(self, mol=None): + """Reset molecular data while retaining only a compatible GAS CI guess.""" + + previous_signature = self._gas_ci_signature + super().reset(mol) + self._sync_fcisolver() + if previous_signature is not None: + try: + current_signature = self._gas_problem_signature() + except (TypeError, ValueError, NotImplementedError): + self._clear_ci_guess() + else: + if current_signature != previous_signature: + self._clear_ci_guess() + return self + + def dump_flags(self, verbose=None): + log = logger.new_logger(self, verbose) + log.info("") + log.info("******** %s ********", self.__class__) + ncore = self.ncore + nvir = self.mo_coeff.shape[1] - ncore - self.ncas + nelecas = self._effective_nelecas() + log.info("GAS (%de+%de, %do), ncore = %d, nvir = %d", + nelecas[0], nelecas[1], self.ncas, ncore, nvir) + if self.frozen is not None: + log.info("frozen orbitals %s", str(self.frozen)) + if self.extrasym is not None: + log.info("extra symmetry labels:\n%s", str(self.extrasym)) + log.info("gas_orbs = %s", self.gas_orbs) + log.info("gas_restr_type = %s", self.gas_restr_type) + try: + gas_orbs, gas_restr, restr_info = self._normalized_restriction( + return_info=True) + limits = addons_gas.check_kernel_limits( + gas_orbs, nelecas, gas_restr) + log.info("number of kernel GAS spaces = %d", limits["ngas"]) + if self.gas_restr_type == addons_gas.GAS_RESTR_SPIN_SUPERGROUP: + log.info("spin-supergroup ordering = alpha-sector " + "lexicographic, then beta-sector lexicographic") + log.info("input legal spin-supergroups = %d", + restr_info["input_nrow"]) + log.info("duplicate spin-supergroups removed = %d", + restr_info["duplicates_removed"]) + log.info("spin-supergroup rows reordered = %s", + "yes" if restr_info["order_changed"] else "no") + elif self.gas_restr_type == addons_gas.GAS_RESTR_SUPERGROUP: + log.info("supergroup ordering = lexicographic") + log.info("input legal supergroups = %d", + restr_info["input_nrow"]) + log.info("duplicate supergroups removed = %d", + restr_info["duplicates_removed"]) + log.info("supergroup rows reordered = %s", + "yes" if restr_info["order_changed"] else "no") + log.info("canonical legal supergroups =\n%s", + restr_info["canonical_supergroups"]) + elif (self.gas_restr_type == + addons_gas.GAS_RESTR_CUMULATIVE_OCC): + log.info("cumulative occupation bounds =\n%s", + restr_info["cumulative_bounds"]) + log.info("number of generated legal supergroups = %d", + restr_info["canonical_supergroups"].shape[0]) + elif self.gas_restr_type == addons_gas.GAS_RESTR_RAS: + log.info("RAS orbital spaces (RAS1, RAS2, RAS3) = %s", + restr_info["ras_orbs"]) + log.info("RAS maximum holes in RAS1 = %d", + restr_info["max_holes"]) + log.info("RAS maximum particles in RAS3 = %d", + restr_info["max_particles"]) + if restr_info["empty_spaces_removed"]: + log.info("empty RAS spaces omitted before C kernel = %s", + restr_info["empty_spaces_removed"]) + log.info("kernel gas_orbs = %s", gas_orbs) + log.info("equivalent cumulative occupation bounds =\n%s", + restr_info["cumulative_bounds"]) + log.info("number of generated legal supergroups = %d", + restr_info["canonical_supergroups"].shape[0]) + log.info("number of legal spin-supergroups = %d", limits["nblock"]) + log.info("estimated number of determinants = %d", + limits["ndet_estimate"]) + log.info("C kernel limit precheck = passed") + except Exception as err: + log.warn("GAS restriction precheck failed: %s", err) + log.info("natorb = %s", self.natorb) + log.info("canonicalization = %s", self.canonicalization) + log.info("sorting_mo_energy = %s", self.sorting_mo_energy) + log.info("max_memory %d MB (current use %d MB)", + self.max_memory, lib.current_memory()[0]) + self._sync_fcisolver() + self.fcisolver.dump_flags(log.verbose) + if self.mo_coeff is None: + log.error("Orbitals for GASCI are not specified. The relevant SCF " + "object may not be initialized.") + return self + + def check_sanity(self): + super().check_sanity() + if self.natorb: + raise NotImplementedError( + "automatic in-place GAS natural-orbital rotation is not " + "implemented; use get_gas_natorb or " + "get_gas_pseudo_natorb for analysis") + gas_orbs, gas_restr = self._normalized_restriction() + addons_gas.check_kernel_limits( + gas_orbs, self._effective_nelecas(), gas_restr) + return self + + def _check_mo_orthonormality(self, mo_coeff=None, verbose=None): + """Validate the MO metric before integral transformation.""" + + mo_coeff = self.mo_coeff if mo_coeff is None else mo_coeff + array = numpy.asarray(mo_coeff) + if array.ndim != 2: + raise ValueError("mo_coeff must be a two-dimensional array") + if numpy.iscomplexobj(array): + raise TypeError( + "the libfci_gas C kernels require real-valued orbitals") + if array.dtype.kind not in "fiu": + raise TypeError("mo_coeff must contain numeric values") + if array.shape[1] < self.ncore + self.ncas: + raise ValueError( + "mo_coeff has fewer columns than ncore + ncas") + if not numpy.all(numpy.isfinite(array)): + raise ValueError("mo_coeff must contain only finite values") + overlap = numpy.asarray(self._scf.get_ovlp()) + if overlap.shape != (array.shape[0], array.shape[0]): + raise ValueError( + "AO overlap shape is incompatible with mo_coeff") + metric = reduce(numpy.dot, (array.conj().T, overlap, array)) + error = float(numpy.max(numpy.abs( + metric - numpy.eye(array.shape[1], dtype=metric.dtype)))) + if not numpy.isfinite(error): + raise ValueError("MO orthonormality error is not finite") + if error > MO_ORTH_ERROR_TOL: + raise ValueError( + "mo_coeff is not orthonormal in the AO metric: " + "max|C^H S C - I| = %.6g exceeds %.6g" % + (error, MO_ORTH_ERROR_TOL)) + log = logger.new_logger(self, verbose) + if error > MO_ORTH_WARN_TOL: + log.warn("MO orthonormality max|C^H S C - I| = %.6g", error) + else: + log.debug("MO orthonormality max|C^H S C - I| = %.6g", error) + return error + + def kernel(self, mo_coeff=None, ci0=None, verbose=None): + """Run fixed-orbital GASCI and return PySCF-style results. + + Returns: + Tuple ``(e_tot, e_gas, ci, mo_coeff, mo_energy)``. For multiple + roots, energies and CI vectors follow the selected PySCF solver + convention. + """ + + self._sync_fcisolver() + if mo_coeff is None: + if self.mo_coeff is None and self._scf.mol.nelectron > 0: + self._scf.run() + self.mo_coeff = self._scf.mo_coeff + mo_coeff = self.mo_coeff + log = logger.new_logger(self, verbose) + self._check_mo_orthonormality(mo_coeff, log) + self.mo_coeff = mo_coeff + self.check_sanity() + self.dump_flags(log) + signature = self._gas_problem_signature() + if ci0 is None: + if self._ci_matches_signature(self.ci, signature): + ci0 = self.ci + else: + self._clear_ci_guess() + self.e_tot, self.e_cas, self.ci = kernel( + self, mo_coeff, ci0=ci0, verbose=log) + self._gas_ci_signature = signature + + solver_physical = getattr(self.fcisolver, "e_physical", None) + solver_penalty = getattr(self.fcisolver, "e_spin_penalty", None) + self.spin_penalty_method = getattr( + self.fcisolver, "spin_penalty_method", None) + if solver_physical is None or solver_penalty is None: + self.e_spin_penalty = None + self.e_tot_physical = self.e_tot + self.e_gas_physical = self.e_gas + else: + physical = numpy.asarray(solver_physical, dtype=numpy.float64) + penalty = numpy.asarray(solver_penalty, dtype=numpy.float64) + core = (numpy.asarray(self.e_tot, dtype=numpy.float64) - + numpy.asarray(self.e_gas, dtype=numpy.float64)) + gas_physical = physical - core + if physical.ndim == 0: + self.e_spin_penalty = float(penalty) + self.e_tot_physical = float(physical) + self.e_gas_physical = float(gas_physical) + else: + self.e_spin_penalty = penalty + self.e_tot_physical = physical + self.e_gas_physical = gas_physical + + if self.canonicalization: + gasdm1 = None + if isinstance(self.ci, (list, tuple)): + gasdm1 = self.make_gasdm1( + state=None if self._has_state_weights() else 0) + self.canonicalize_( + mo_coeff, self.ci, sort=self.sorting_mo_energy, + gas_natorb=False, gasdm1=gasdm1, verbose=log) + + if getattr(self.fcisolver, "converged", None) is not None: + self.converged = bool(numpy.all(self.fcisolver.converged)) + else: + self.converged = True + log.info("GASCI converged" if self.converged else "GASCI not converged") + self._finalize() + return self.e_tot, self.e_gas, self.ci, self.mo_coeff, self.mo_energy + + def gasci(self, mo_coeff=None, ci0=None, verbose=None): + """Alias of :meth:`kernel` using GAS terminology.""" + + return self.kernel(mo_coeff, ci0, verbose) + + def as_scanner(self): + return as_scanner(self) + + def get_h1gas(self, mo_coeff=None, ncas=None, ncore=None): + """Return the effective one-electron Hamiltonian in the GAS space.""" + + return self.get_h1eff(mo_coeff, ncas, ncore) + + def get_h2gas(self, mo_coeff=None): + """Return active-space two-electron integrals for GASCI.""" + + return self.get_h2eff(mo_coeff) + + def get_h1cas(self, *args, **kwargs): + raise NotImplementedError( + "get_h1cas is a CAS-specific name; use get_h1gas for GASCI") + + h1e_for_cas = get_h1cas + + def get_h2cas(self, *args, **kwargs): + raise NotImplementedError( + "get_h2cas is a CAS-specific name; use get_h2gas for GASCI") + + def sort_mo(self, gaslst, mo_coeff=None, base=1): + """Pick orbitals for ordered GAS subspaces. + + ``gaslst`` must contain one orbital-index list per entry of + ``gas_orbs``. ``base`` follows PySCF: 0 selects C-style indices and + 1 selects Fortran-style indices. + """ + + if mo_coeff is None: + mo_coeff = self.mo_coeff + self._check_mo_orthonormality(mo_coeff) + return addons_gas.sort_mo(self, mo_coeff, gaslst, base) + + def get_fock(self, mo_coeff=None, ci=None, eris=None, gasdm1=None, + verbose=None): + """Build the generalized Fock matrix from a GAS one-particle DM.""" + + if gasdm1 is None: + ci_source = self.ci if ci is None else ci + if isinstance(ci_source, (list, tuple)): + gasdm1 = self.make_gasdm1( + ci_source, + state=None if self._has_state_weights() else 0) + return casci.get_fock( + self, mo_coeff, ci, eris, casdm1=gasdm1, verbose=verbose) + + def canonicalize(self, mo_coeff=None, ci=None, eris=None, sort=False, + gas_natorb=False, gasdm1=None, verbose=None, + **kwargs): + """Canonicalize core/external orbitals without rotating the GAS space.""" + + if gas_natorb: + raise NotImplementedError( + "automatic in-place GAS natural-orbital rotation is not " + "implemented; use get_gas_pseudo_natorb for analysis") + return casci.canonicalize( + self, mo_coeff, ci, eris, sort=sort, cas_natorb=False, + casdm1=gasdm1, verbose=verbose, **kwargs) + + def canonicalize_(self, mo_coeff=None, ci=None, eris=None, sort=False, + gas_natorb=False, gasdm1=None, verbose=None, + **kwargs): + mo_coeff, ci, mo_energy = self.canonicalize( + mo_coeff, ci, eris, sort=sort, gas_natorb=gas_natorb, + gasdm1=gasdm1, verbose=verbose, **kwargs) + self.mo_coeff = mo_coeff + self.mo_energy = mo_energy + return mo_coeff, ci, mo_energy + + def cas_natorb(self, *args, **kwargs): + raise NotImplementedError( + "CAS natural-orbital rotation is not valid for a restricted GAS") + + cas_natorb_ = cas_natorb + + def fix_spin_(self, shift=0.2, ss=None): + """Use a PySCF-style energy penalty to target a GAS spin state. + + This optional numerical aid is only defined for a spin-complete GAS + restriction. Its success depends on the energy gaps between spin + sectors and on a suitable finite shift; it is not a spin-adapted CI + representation. Small GAS spaces are solved exactly, while larger + spaces use target-spin-projected guesses together with global probes + so that an insufficient shift does not silently change the eigenproblem. + """ + + self._sync_fcisolver() + gas_orbs, gas_restr = self._normalized_restriction() + nelecas = self._effective_nelecas() + if not addons_gas.is_spin_complete( + gas_orbs, nelecas, gas_restr): + raise ValueError( + "fix_spin_ requires a spin-complete GAS restriction") + fci.addons.fix_spin_(self.fcisolver, shift, ss) + return self + + fix_spin = fix_spin_ + + def state_average(self, weights=(0.5, 0.5), wfnsym=None): + """Return a GASCI object carrying fixed-orbital state weights. + + The weights define averaged RDMs and generalized Fock matrices for + fixed-orbital analysis. GASCI does not optimize orbitals, and every + root retains its own energy. + """ + + return addons_gas.state_average(self, weights, wfnsym) + + def state_average_(self, weights=(0.5, 0.5), wfnsym=None): + addons_gas.state_average_(self, weights, wfnsym) + return self + + def state_average_mix(self, *args, **kwargs): + raise NotImplementedError( + "state_average_mix is not implemented for GASCI") + + state_average_mix_ = state_average_mix + + def nuc_grad_method(self): + raise NotImplementedError( + "GASCI nuclear gradients are not implemented") + + def to_gpu(self, *args, **kwargs): + raise NotImplementedError( + "the libfci_gas C/OpenMP backend does not support GPU execution") + + def state_specific_(self, state=1, wfnsym=None): + if wfnsym is not None: + raise NotImplementedError( + "GAS wavefunction symmetry filtering is not implemented") + if isinstance(self, addons_gas.StateAverageGASCI): + raise ValueError( + "undo state weighting before selecting a state-specific root") + addons.state_specific_(self, state=state, wfnsym=None) + return self + + state_specific = state_specific_ + + def _has_state_weights(self): + weights = getattr(self.fcisolver, "weights", None) + return (weights is not None and + isinstance(self.fcisolver, addons.StateAverageFCISolver)) + + def _state_weights(self): + if not self._has_state_weights(): + return None + return numpy.asarray(self.fcisolver.weights, dtype=numpy.float64) + + def _base_fcisolver_method(self, name): + if self._has_state_weights(): + base = super(addons.StateAverageFCISolver, self.fcisolver) + return getattr(base, name) + return getattr(self.fcisolver, name) + + def _ci_for_rdm(self, ci, state): + ci = self.ci if ci is None else ci + if self._has_state_weights() and state is None: + if not isinstance(ci, (list, tuple)): + raise ValueError( + "weighted GAS RDMs require one CI vector per state") + if len(ci) != len(self._state_weights()): + raise ValueError("CI root count does not match state weights") + return ci, True + selected = 0 if state is None else int(state) + return self._select_ci(ci, selected), False + + def _spin_square_for_ci(self, ci, ncas, nelecas): + return self._spin_square_for_roots( + [ci], ncas, nelecas)[0] + + def _spin_square_for_roots(self, roots, ncas, nelecas): + nelecas = self._effective_nelecas(nelecas) + values = [] + with self.fcisolver.make_rdm_plan(ncas, nelecas) as plan: + for ci in roots: + dm1s, dm2s = plan.make_rdm12s(ci, ci) + values.append( + fci_gas.spin_square_from_rdm12s( + dm1s, dm2s, nelecas)) + return values + + def _finalize(self): + log = logger.Logger(self.stdout, self.verbose) + weights = self._state_weights() + if weights is not None: + totals = numpy.asarray( + self.fcisolver.e_states, dtype=numpy.float64).reshape(-1) + roots = self.ci if isinstance(self.ci, (list, tuple)) else [self.ci] + energy_core = float(self.e_tot) - float(self.e_gas) + log.note("GASCI weighted energy (fixed orbitals) = %#.15g", + self.e_tot) + log.note("GASCI energy for each state") + spin_values = self._spin_square_for_roots( + roots, self.ncas, self.nelecas) + for i, (weight, e_tot, spin) in enumerate( + zip(weights, totals, spin_values)): + log.note(" State %d weight %g E = %#.15g E(CI) = %#.15g " + "S^2 = %.7f multiplicity = %.7f", + i, weight, e_tot, e_tot - energy_core, + spin[0], spin[1]) + return self + + energies = numpy.asarray(self.e_gas).reshape(-1) + totals = numpy.asarray(self.e_tot).reshape(-1) + roots = self.ci if isinstance(self.ci, (list, tuple)) else [self.ci] + try: + spin_values = self._spin_square_for_roots( + roots, self.ncas, self.nelecas) + except NotImplementedError: + spin_values = [None] * len(roots) + for i, (e_tot, e_gas) in enumerate(zip(totals, energies)): + ss = None if spin_values[i] is None else spin_values[i][0] + state = getattr(self.fcisolver, "state", None) + if len(energies) == 1 and state is not None and ss is None: + log.note("GASCI state %3d E = %#.15g E(CI) = %#.15g", + state, e_tot, e_gas) + elif len(energies) == 1 and state is not None: + log.note("GASCI state %3d E = %#.15g E(CI) = %#.15g " + "S^2 = %.7f", state, e_tot, e_gas, ss) + elif len(energies) == 1: + if ss is None: + log.note("GASCI E = %#.15g E(CI) = %#.15g", e_tot, e_gas) + else: + log.note("GASCI E = %#.15g E(CI) = %#.15g S^2 = %.7f", + e_tot, e_gas, ss) + elif ss is None: + log.note("GASCI state %3d E = %#.15g E(CI) = %#.15g", + i, e_tot, e_gas) + else: + log.note("GASCI state %3d E = %#.15g E(CI) = %#.15g S^2 = %.7f", + i, e_tot, e_gas, ss) + if self.e_spin_penalty is not None: + penalties = numpy.asarray( + self.e_spin_penalty, dtype=numpy.float64).reshape(-1) + physical = numpy.asarray( + self.e_tot_physical, dtype=numpy.float64).reshape(-1) + target = getattr(self.fcisolver, "ss_value", None) + if target is None: + nelecas = self._effective_nelecas() + sz = 0.5 * abs(nelecas[0] - nelecas[1]) + target = sz * (sz + 1.0) + state = getattr(self.fcisolver, "state", None) + for i, (penalty, e_physical) in enumerate( + zip(penalties, physical)): + label = state if len(penalties) == 1 and state is not None else i + log.note("GASCI state %3d E(physical) = %#.15g " + "E(spin penalty) = %#.15g target S^2 = %.7f " + "method = %s", label, e_physical, penalty, target, + self.spin_penalty_method) + actual = None if spin_values[i] is None else spin_values[i][0] + if actual is not None and abs(actual - target) > 1e-6: + log.warn("GASCI state %d did not reach target S^2: " + "target = %.7f, actual = %.7f. The spin " + "penalty shift may be insufficient.", + label, target, actual) + return self + + def _select_ci(self, ci=None, state=0): + if ci is None: + ci = self.ci + if isinstance(ci, (list, tuple)): + return ci[int(state)] + arr = numpy.asarray(ci) + nroots = int(getattr(self.fcisolver, "nroots", 1)) + if arr.ndim == 2 and arr.shape[0] == nroots: + return arr[int(state)] + if arr.ndim == 2 and arr.shape[1] == nroots: + return arr[:, int(state)] + return ci + + def make_gasdm1(self, ci=None, ncas=None, nelecas=None, state=None): + """Return a spin-summed active-space GAS 1-RDM.""" + + ncas = self.ncas if ncas is None else ncas + nelecas = self.nelecas if nelecas is None else nelecas + ci_value, weighted = self._ci_for_rdm(ci, state) + method = (self.fcisolver.make_rdm1 if weighted else + self._base_fcisolver_method("make_rdm1")) + return method(ci_value, ncas, nelecas) + + def make_gasdm1s(self, ci=None, ncas=None, nelecas=None, state=None): + """Return alpha and beta active-space GAS 1-RDMs.""" + + ncas = self.ncas if ncas is None else ncas + nelecas = self.nelecas if nelecas is None else nelecas + ci_value, weighted = self._ci_for_rdm(ci, state) + method = (self.fcisolver.make_rdm1s if weighted else + self._base_fcisolver_method("make_rdm1s")) + return method(ci_value, ncas, nelecas) + + def make_gasdm12(self, ci=None, ncas=None, nelecas=None, state=None): + """Return spin-summed active-space GAS 1- and 2-RDMs.""" + + ncas = self.ncas if ncas is None else ncas + nelecas = self.nelecas if nelecas is None else nelecas + ci_value, weighted = self._ci_for_rdm(ci, state) + method = (self.fcisolver.make_rdm12 if weighted else + self._base_fcisolver_method("make_rdm12")) + return method(ci_value, ncas, nelecas) + + def make_gasdm2(self, ci=None, ncas=None, nelecas=None, state=None): + """Return the spin-summed active-space GAS 2-RDM.""" + + return self.make_gasdm12(ci, ncas, nelecas, state)[1] + + def make_gasdm12s(self, ci=None, ncas=None, nelecas=None, state=None): + """Return spin-resolved active-space GAS 1- and 2-RDMs.""" + + ncas = self.ncas if ncas is None else ncas + nelecas = self.nelecas if nelecas is None else nelecas + ci_value, weighted = self._ci_for_rdm(ci, state) + method = (self.fcisolver.make_rdm12s if weighted else + self._base_fcisolver_method("make_rdm12s")) + return method(ci_value, ncas, nelecas) + + def trans_gasdm1(self, cibra=None, ciket=None, ncas=None, nelecas=None, + bra_state=0, ket_state=1): + """Return the spin-summed GAS transition 1-RDM.""" + + ncas = self.ncas if ncas is None else ncas + nelecas = self.nelecas if nelecas is None else nelecas + bra = self._select_ci(cibra, bra_state) + ket = self._select_ci(ciket, ket_state) + return self._base_fcisolver_method("trans_rdm1")( + bra, ket, ncas, nelecas) + + def trans_gasdm1s(self, cibra=None, ciket=None, ncas=None, nelecas=None, + bra_state=0, ket_state=1): + """Return alpha and beta GAS transition 1-RDMs.""" + + ncas = self.ncas if ncas is None else ncas + nelecas = self.nelecas if nelecas is None else nelecas + bra = self._select_ci(cibra, bra_state) + ket = self._select_ci(ciket, ket_state) + return self._base_fcisolver_method("trans_rdm1s")( + bra, ket, ncas, nelecas) + + def trans_gasdm12(self, cibra=None, ciket=None, ncas=None, nelecas=None, + bra_state=0, ket_state=1): + """Return spin-summed GAS transition 1- and 2-RDMs.""" + + ncas = self.ncas if ncas is None else ncas + nelecas = self.nelecas if nelecas is None else nelecas + bra = self._select_ci(cibra, bra_state) + ket = self._select_ci(ciket, ket_state) + return self._base_fcisolver_method("trans_rdm12")( + bra, ket, ncas, nelecas) + + def trans_gasdm2(self, cibra=None, ciket=None, ncas=None, nelecas=None, + bra_state=0, ket_state=1): + """Return the spin-summed GAS transition 2-RDM.""" + + return self.trans_gasdm12( + cibra, ciket, ncas, nelecas, bra_state, ket_state)[1] + + def trans_gasdm12s(self, cibra=None, ciket=None, ncas=None, nelecas=None, + bra_state=0, ket_state=1): + """Return spin-resolved GAS transition 1- and 2-RDMs.""" + + ncas = self.ncas if ncas is None else ncas + nelecas = self.nelecas if nelecas is None else nelecas + bra = self._select_ci(cibra, bra_state) + ket = self._select_ci(ciket, ket_state) + return self._base_fcisolver_method("trans_rdm12s")( + bra, ket, ncas, nelecas) + + def make_rdm1s(self, mo_coeff=None, ci=None, ncas=None, nelecas=None, + ncore=None, state=None, **kwargs): + """Return alpha and beta AO-basis one-particle density matrices.""" + + mo_coeff = self.mo_coeff if mo_coeff is None else mo_coeff + ncas = self.ncas if ncas is None else ncas + nelecas = self.nelecas if nelecas is None else nelecas + ncore = self.ncore if ncore is None else ncore + gasdm1a, gasdm1b = self.make_gasdm1s( + ci, ncas, nelecas, state) + return self._gasdm1s_to_ao( + (gasdm1a, gasdm1b), mo_coeff, ncas, ncore) + + def _gasdm1s_to_ao(self, gasdm1s, mo_coeff, ncas, ncore): + gasdm1a, gasdm1b = gasdm1s + mocore = mo_coeff[:, :ncore] + mogas = mo_coeff[:, ncore:ncore + ncas] + dm1b = numpy.dot(mocore, mocore.conj().T) + dm1a = dm1b + reduce(numpy.dot, (mogas, gasdm1a, mogas.conj().T)) + dm1b = dm1b + reduce(numpy.dot, (mogas, gasdm1b, mogas.conj().T)) + return dm1a, dm1b + + def make_rdm1(self, mo_coeff=None, ci=None, ncas=None, nelecas=None, + ncore=None, state=None, **kwargs): + """Return the spin-summed AO-basis one-particle density matrix.""" + + dm1a, dm1b = self.make_rdm1s( + mo_coeff, ci, ncas, nelecas, ncore, state, **kwargs) + return dm1a + dm1b + + def spin_square(self, ci=None, ncas=None, nelecas=None, state=None): + """Return ``(, 2S+1)`` for one root or a weighted state set.""" + + ncas = self.ncas if ncas is None else ncas + nelecas = self.nelecas if nelecas is None else nelecas + ci_value, weighted = self._ci_for_rdm(ci, state) + if weighted: + values = self._spin_square_for_roots( + ci_value, ncas, nelecas) + ss = float(numpy.dot( + self._state_weights(), [value[0] for value in values])) + return ss, numpy.sqrt(max(0.0, 4.0 * ss + 1.0)) + return self._spin_square_for_ci(ci_value, ncas, nelecas) + + def get_gas_pseudo_natorb_occupations(self, gasdm1=None, ci=None, + state=None, sort=True): + """Return subspace-resolved pseudo-natural occupation spectra. + + Following the OpenMolcas convention for restricted active spaces, + each diagonal GAS-subspace block of the spin-summed one-particle + density matrix is diagonalized separately. These are pseudo-natural, + not true natural, occupations because no rotation across GAS + subspaces is performed. Sorting only makes the reported spectra + deterministic; it does not define an orbital transformation. + """ + + if gasdm1 is None: + gasdm1 = self.make_gasdm1(ci=ci, state=state) + gasdm1 = numpy.asarray(gasdm1) + if gasdm1.shape != (self.ncas, self.ncas): + raise ValueError("gasdm1 must have shape (ncas, ncas)") + + gas_orbs = ((self.ncas,) if self.gas_orbs is None else + tuple(int(value) for value in self.gas_orbs)) + occupations = [] + offset = 0 + for size in gas_orbs: + block = gasdm1[offset:offset + size, offset:offset + size] + values, _ = self._natural_eigensystem(block, sort) + occupations.append(values) + offset += size + return tuple(occupations) + + @staticmethod + def _natural_eigensystem(dm1, sort=True): + dm1 = numpy.asarray(dm1) + dm1 = 0.5 * (dm1 + dm1.conj().T) + occupations, rotation = numpy.linalg.eigh(dm1) + occupations = occupations.real + if sort: + order = numpy.argsort(occupations)[::-1] + occupations = occupations[order] + rotation = rotation[:, order] + return occupations, rotation + + def _rotate_gas_orbitals(self, mo_coeff, rotation): + mo_coeff = self.mo_coeff if mo_coeff is None else mo_coeff + mo_coeff = numpy.asarray(mo_coeff) + if mo_coeff.ndim != 2: + raise ValueError("mo_coeff must be a two-dimensional array") + if mo_coeff.shape[1] < self.ncore + self.ncas: + raise ValueError("mo_coeff does not contain the complete GAS space") + if rotation.shape != (self.ncas, self.ncas): + raise ValueError("GAS orbital rotation has an invalid shape") + mo_natorb = mo_coeff.copy() + gas = slice(self.ncore, self.ncore + self.ncas) + mo_natorb[:, gas] = numpy.dot(mo_coeff[:, gas], rotation) + return mo_natorb + + def get_gas_natorb(self, mo_coeff=None, gasdm1=None, ci=None, + state=0, sort=True): + """Return root-specific true natural orbitals and occupations. + + The complete spin-summed GAS one-particle density matrix is + diagonalized. For a restricted GAS, the resulting orbitals can mix + different GAS subspaces and are therefore analysis orbitals only. + This method never modifies ``mo_coeff`` or transforms the CI vector. + """ + + if gasdm1 is None and state is None: + raise ValueError( + "get_gas_natorb requires an explicit state; use " + "get_gas_average_natorb for a state-average density") + if gasdm1 is None: + gasdm1 = self.make_gasdm1(ci=ci, state=state) + gasdm1 = numpy.asarray(gasdm1) + if gasdm1.shape != (self.ncas, self.ncas): + raise ValueError("gasdm1 must have shape (ncas, ncas)") + occupations, rotation = self._natural_eigensystem(gasdm1, sort) + return self._rotate_gas_orbitals(mo_coeff, rotation), occupations + + def get_gas_average_natorb(self, mo_coeff=None, gasdm1=None, ci=None, + sort=True): + """Return natural orbitals of the explicitly state-averaged GAS DM. + + This method is available only on a state-average GASCI object. Like + :meth:`get_gas_natorb`, the returned orbitals are for analysis and do + not replace the computational GAS orbitals. + """ + + if not self._has_state_weights(): + raise ValueError( + "get_gas_average_natorb requires a state-average GASCI object") + if gasdm1 is None: + gasdm1 = self.make_gasdm1(ci=ci, state=None) + return self.get_gas_natorb( + mo_coeff=mo_coeff, gasdm1=gasdm1, sort=sort) + + def get_gas_pseudo_natorb(self, mo_coeff=None, gasdm1=None, ci=None, + state=None, sort=True): + """Return subspace-preserving pseudo-natural orbitals. + + Each diagonal GAS-subspace block of the spin-summed one-particle + density matrix is diagonalized independently. The returned orbital + coefficients preserve the GAS partition but are not installed on the + calculation object, and the CI vector is left unchanged. + """ + + if gasdm1 is None: + gasdm1 = self.make_gasdm1(ci=ci, state=state) + gasdm1 = numpy.asarray(gasdm1) + if gasdm1.shape != (self.ncas, self.ncas): + raise ValueError("gasdm1 must have shape (ncas, ncas)") + + gas_orbs = ((self.ncas,) if self.gas_orbs is None else + tuple(int(value) for value in self.gas_orbs)) + rotation = numpy.zeros_like(gasdm1) + occupations = [] + offset = 0 + for size in gas_orbs: + block = gasdm1[offset:offset + size, offset:offset + size] + values, vectors = self._natural_eigensystem(block, sort) + rotation[offset:offset + size, offset:offset + size] = vectors + occupations.append(values) + offset += size + return (self._rotate_gas_orbitals(mo_coeff, rotation), + tuple(occupations)) + + def analyze(self, mo_coeff=None, ci=None, verbose=None, + large_ci_tol=LARGE_CI_TOL, + with_meta_lowdin=WITH_META_LOWDIN, state=None, + **kwargs): + """Analyze GASCI roots without rotations across GAS subspaces.""" + + log = logger.new_logger(self, verbose) + mo_coeff = self.mo_coeff if mo_coeff is None else mo_coeff + ci_source = self.ci if ci is None else ci + if isinstance(ci_source, (list, tuple)): + roots = list(ci_source) + root_labels = list(range(len(roots))) + else: + roots = [ci_source] + root_labels = [int(getattr(self.fcisolver, "state", 0))] + if state is not None: + selected = int(state) + roots = [self._select_ci(ci_source, selected)] + root_labels = [selected] + + log.info("") + log.info("******** GASCI analysis ********") + weights = self._state_weights() + energies = numpy.asarray(self.e_states if weights is not None else + self.e_tot).reshape(-1) + spin_values = self._spin_square_for_roots( + roots, self.ncas, self.nelecas) + for label, spin in zip(root_labels, spin_values): + ss, mult = spin + energy_index = label if energies.size > 1 else 0 + if weights is None: + log.note("GASCI state %3d E = %#.15g S^2 = %.7f " + "multiplicity = %.7f", label, + energies[energy_index], ss, mult) + else: + log.note("GASCI state %3d weight = %g E = %#.15g " + "S^2 = %.7f multiplicity = %.7f", + label, weights[label], energies[energy_index], + ss, mult) + + density_state = None if weights is not None and state is None else ( + root_labels[0]) + gasdm1s = self.make_gasdm1s(ci_source, state=density_state) + gasdm1 = gasdm1s[0] + gasdm1s[1] + occupations = self.get_gas_pseudo_natorb_occupations( + gasdm1=gasdm1) + gas_orbs = ((self.ncas,) if self.gas_orbs is None else + tuple(int(value) for value in self.gas_orbs)) + offset = 0 + for igas, (size, values) in enumerate(zip(gas_orbs, occupations)): + block = gasdm1[offset:offset + size, offset:offset + size] + log.info("GAS subspace %d electron number = %.12g", igas + 1, + numpy.trace(block).real) + log.info("GAS subspace %d pseudo-natural occupations %s", + igas + 1, numpy.array2string(values, precision=10)) + offset += size + + if log.verbose >= logger.INFO and ci_source is not None: + log.info("** Largest GAS CI components **") + large_ci = self._base_fcisolver_method("large_ci") + for root, label in zip(roots, root_labels): + log.info(" [alpha occ-orbitals] [beta occ-orbitals] " + "state %-3d CI coefficient", label) + for coeff, alpha, beta in large_ci( + root, self.ncas, self.nelecas, large_ci_tol, + return_strs=False): + log.info(" %-20s %-30s % .12f", + alpha, beta, coeff) + + dm1a, dm1b = self._gasdm1s_to_ao( + gasdm1s, mo_coeff, self.ncas, self.ncore) + if log.verbose >= logger.INFO: + ovlp = self._scf.get_ovlp() + dm1 = dm1a + dm1b + spin_dm1 = dm1a - dm1b + if with_meta_lowdin: + self._scf.mulliken_meta( + self.mol, dm1, s=ovlp, verbose=log) + log.info("Mulliken spin population analysis on " + "meta-Lowdin AOs:") + self._scf.mulliken_meta( + self.mol, spin_dm1, s=ovlp, verbose=log) + else: + self._scf.mulliken_pop( + self.mol, dm1, s=ovlp, verbose=log) + log.info("Mulliken spin population analysis on AOs:") + self._scf.mulliken_pop( + self.mol, spin_dm1, s=ovlp, verbose=log) + return dm1a, dm1b diff --git a/pyscf/mcscf/test/test_gasci.py b/pyscf/mcscf/test/test_gasci.py new file mode 100644 index 000000000..8f25b2b35 --- /dev/null +++ b/pyscf/mcscf/test/test_gasci.py @@ -0,0 +1,951 @@ +#!/usr/bin/env python +# Copyright 2026 The PySCF Developers. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# Author: Yi Deng +# + +"""Tests for GAS restriction handling, the GAS FCI solver, and GASCI.""" + +import io +import unittest + +import numpy + +from pyscf import gto +from pyscf import mcscf +from pyscf import scf +from pyscf.fci import addons as fci_addons +from pyscf.fci import cistring +from pyscf.fci import direct_spin1 +from pyscf.fci import spin_op +from pyscf.lib import logger +from pyscf.mcscf import addons_gas +from pyscf.mcscf import fci_gas +from pyscf.mcscf import gasci + + +def make_integrals(norb): + """Return deterministic one- and two-electron test integrals.""" + + h1e = numpy.empty((norb, norb)) + for p in range(norb): + for q in range(norb): + h1e[p, q] = ( + 0.07 * (p + q + 2) + + 0.013 * ((p + 1) * (q + 3) % 7)) + h1e = (h1e + h1e.T) * 0.5 + + npair = norb * (norb + 1) // 2 + eri = numpy.empty((npair, npair)) + for pq in range(npair): + for rs in range(npair): + lo, hi = min(pq, rs), max(pq, rs) + eri[pq, rs] = 0.011 * (lo + 1) + 0.003 * (hi + 2) + return h1e, eri + + +class TestGASRestrictions(unittest.TestCase): + + def test_nelec_spin_inference(self): + cases = ( + (5, None, (3, 2)), + (5, 1, (3, 2)), + (5, 3, (4, 1)), + (5, -1, (2, 3)), + ((3, 2), None, (3, 2)), + ((2, 2), 2, (3, 1)), + ) + solver = fci_gas.FCISolver() + for nelec, spin, expected in cases: + with self.subTest(nelec=nelec, spin=spin): + solver.spin = spin + unpacked = fci_addons._unpack_nelec(nelec, spin) + self.assertEqual(solver._space_spec(4, nelec)[1], unpacked) + self.assertEqual(unpacked, expected) + + solver.spin = 2 + self.assertEqual(solver._space_spec(4, 4)[1], (3, 1)) + self.assertEqual(solver.space_info(4, 4)["ndet_estimate"], 16) + + def test_integer_inputs_reject_floats(self): + nelec = (2, 2) + + invalid = ( + ((2.0, 2), [[2, 2]], "supergroup"), + ((2, 2), [[2.0, 2]], "supergroup"), + ((2, 2), [[2, 2], [4, 4.0]], "cumulative-occ"), + ((1, 2, 1), {"max_holes": 1.0, "max_particles": 1}, "ras"), + ((2, True), [[2, 2]], "supergroup"), + ) + for gas_orbs, gas_restr, gas_restr_type in invalid: + with self.subTest( + gas_orbs=gas_orbs, gas_restr_type=gas_restr_type): + with self.assertRaises(TypeError): + addons_gas.normalize_gas_restr( + gas_orbs, nelec, gas_restr, gas_restr_type) + + blocks = addons_gas.normalize_gas_restr( + numpy.asarray([2, 2], dtype=numpy.int64), + nelec, + numpy.asarray([[2, 2]], dtype=numpy.int64), + "supergroup", + ) + numpy.testing.assert_array_equal( + blocks, + numpy.asarray([[0, 2, 2, 0], + [1, 1, 1, 1], + [2, 0, 0, 2]], dtype=numpy.int32), + ) + + for base in (0.0, True): + with self.subTest(base=base): + with self.assertRaises(TypeError): + addons_gas._normalize_gaslst( + [[0], [1]], (1, 1), 2, base) + + def test_equivalent_restriction_types(self): + gas_orbs = (2, 2) + nelec = (2, 2) + expected = numpy.asarray([ + [0, 2, 2, 0], + [1, 1, 1, 1], + [2, 0, 0, 2], + ], dtype=numpy.int32) + explicit = numpy.asarray([ + [2, 0, 0, 2], + [1, 1, 1, 1], + [0, 2, 2, 0], + [1, 1, 1, 1], + ], dtype=numpy.int32) + + inputs = ( + ("spin-supergroup", explicit), + ("supergroup", [[2, 2]]), + ("cumulative-occ", [[2, 2], [4, 4]]), + ) + for restriction_type, restriction in inputs: + with self.subTest(restriction_type=restriction_type): + blocks = addons_gas.normalize_gas_restr( + gas_orbs, nelec, restriction, restriction_type) + numpy.testing.assert_array_equal(blocks, expected) + + self.assertTrue( + addons_gas.is_spin_complete(gas_orbs, nelec, expected)) + + def test_ras_matches_cumulative_bounds(self): + gas_orbs = (1, 2, 1) + nelec = (2, 2) + ras = {"max_holes": 1, "max_particles": 1} + cumulative = [[1, 2], [3, 4], [4, 4]] + + kernel_orbs, ras_blocks = addons_gas.normalize_gas_spec( + gas_orbs, nelec, ras, "ras") + cumulative_blocks = addons_gas.normalize_gas_restr( + gas_orbs, nelec, cumulative, "cumulative-occ") + + self.assertEqual(kernel_orbs, gas_orbs) + numpy.testing.assert_array_equal(ras_blocks, cumulative_blocks) + + +class TestGASFCISolver(unittest.TestCase): + + @classmethod + def setUpClass(cls): + cls.norb = 4 + cls.nelec = (2, 2) + cls.h1e, cls.eri = make_integrals(cls.norb) + cls.nstr = cistring.num_strings(cls.norb, cls.nelec[0]) + + rng = numpy.random.default_rng(12) + cls.ci = rng.normal(size=cls.nstr**2) + cls.ci /= numpy.linalg.norm(cls.ci) + cls.ket = rng.normal(size=cls.nstr**2) + cls.ket /= numpy.linalg.norm(cls.ket) + + def test_gas_as_cas_hamiltonian(self): + solver = fci_gas.FCISolver() + hdiag = solver.make_hdiag( + self.h1e, self.eri, self.norb, self.nelec) + reference_hdiag = direct_spin1.make_hdiag( + self.h1e, self.eri, self.norb, self.nelec) + numpy.testing.assert_allclose( + hdiag, reference_hdiag, atol=1e-12, rtol=0) + + h2e = fci_gas.absorb_h1e( + self.h1e, self.eri, self.norb, self.nelec, fac=0.5) + hc = solver.contract_2e( + h2e, self.ci, self.norb, self.nelec) + reference_hc = direct_spin1.contract_2e( + h2e, self.ci.reshape(self.nstr, self.nstr), + self.norb, self.nelec).reshape(-1) + numpy.testing.assert_allclose( + hc, reference_hc, atol=1e-11, rtol=0) + + energy = solver.energy( + self.h1e, self.eri, self.ci, self.norb, self.nelec) + self.assertAlmostEqual(energy, numpy.dot(self.ci, hc), places=11) + + def test_init_guess_reuses_existing_gas_space(self): + solver = fci_gas.FCISolver( + gas_orbs=(1, 2, 1), + gas_restr=[[1, 2, 1], [2, 1, 1], [1, 1, 2]], + gas_restr_type="supergroup") + hdiag = solver.make_hdiag( + self.h1e, self.eri, self.norb, self.nelec) + + reference = solver.get_init_guess( + self.norb, self.nelec, 2, hdiag) + with solver.make_space( + self.norb, self.nelec, compress_links=False) as gas: + reused = solver.get_init_guess( + self.norb, self.nelec, 2, hdiag, gas=gas) + + self.assertEqual(len(reference), len(reused)) + for expected, actual in zip(reference, reused): + numpy.testing.assert_allclose(actual, expected, atol=0, rtol=0) + + def test_contract_plan_reuse(self): + solver = fci_gas.FCISolver( + gas_orbs=(1, 2, 1), + gas_restr=[[1, 2, 1], [2, 1, 1], [1, 1, 2]], + gas_restr_type="supergroup") + h2e = fci_gas.absorb_h1e( + self.h1e, self.eri, self.norb, self.nelec, fac=0.5) + + with solver.make_space( + self.norb, self.nelec, compress_links=True) as gas: + rng = numpy.random.default_rng(31) + vectors = ( + rng.normal(size=gas.ndet), + rng.normal(size=gas.ndet), + ) + with fci_gas.GasContractPlan(gas, h2e) as plan: + for ci in vectors: + planned = plan.contract(ci) + through_solver = solver.contract_2e( + h2e, ci, self.norb, self.nelec, plan=plan) + one_shot = solver.contract_2e( + h2e, ci, self.norb, self.nelec) + numpy.testing.assert_allclose( + planned, through_solver, atol=1e-12, rtol=0) + numpy.testing.assert_allclose( + planned, one_shot, atol=1e-12, rtol=0) + + def test_gas_fci_vector_converters(self): + gas_orbs = (2, 2) + nelec = (2, 1) + blocks = numpy.asarray([[1, 1, 1, 0]], dtype=numpy.int32) + solver = fci_gas.FCISolver( + gas_orbs=gas_orbs, + gas_restr=blocks, + gas_restr_type="spin-supergroup") + + with solver.make_space(self.norb, nelec) as gas: + nstr_alpha = cistring.num_strings(self.norb, nelec[0]) + nstr_beta = cistring.num_strings(self.norb, nelec[1]) + self.assertLess(gas.ndet, nstr_alpha * nstr_beta) + + rng = numpy.random.default_rng(19) + gas_ci = rng.normal(size=gas.ndet) + fci_ci = fci_gas.gas2fci(gas_ci, gas) + self.assertEqual(fci_ci.shape, (nstr_alpha, nstr_beta)) + numpy.testing.assert_array_equal( + fci_gas.fci2gas(fci_ci, gas), gas_ci) + + full_ci = numpy.arange( + 1, nstr_alpha * nstr_beta + 1, dtype=numpy.float64 + ).reshape(nstr_alpha, nstr_beta) + projected = fci_gas.gas2fci( + fci_gas.fci2gas(full_ci, gas), gas) + projected_twice = fci_gas.gas2fci( + fci_gas.fci2gas(projected, gas), gas) + numpy.testing.assert_array_equal(projected_twice, projected) + self.assertEqual(numpy.count_nonzero(projected), gas.ndet) + + def test_restricted_hamiltonian_projection(self): + cases = ( + { + "name": "spin-supergroup", + "gas_orbs": (2, 2), + "nelec": (2, 1), + "gas_restr": numpy.asarray( + [[1, 1, 1, 0]], dtype=numpy.int32), + "gas_restr_type": "spin-supergroup", + }, + { + "name": "supergroup", + "gas_orbs": (1, 2, 1), + "nelec": (2, 2), + "gas_restr": [[1, 2, 1], [2, 1, 1], [1, 1, 2]], + "gas_restr_type": "supergroup", + }, + { + "name": "ras", + "gas_orbs": (1, 2, 1), + "nelec": (2, 2), + "gas_restr": {"max_holes": 1, "max_particles": 1}, + "gas_restr_type": "ras", + }, + ) + + rng = numpy.random.default_rng(23) + for case in cases: + with self.subTest(case=case["name"]): + nelec = case["nelec"] + solver = fci_gas.FCISolver( + gas_orbs=case["gas_orbs"], + gas_restr=case["gas_restr"], + gas_restr_type=case["gas_restr_type"]) + h1e, eri = make_integrals(self.norb) + h2e = fci_gas.absorb_h1e( + h1e, eri, self.norb, nelec, fac=0.5) + + with solver.make_space(self.norb, nelec) as gas: + nstr_alpha = cistring.num_strings(self.norb, nelec[0]) + nstr_beta = cistring.num_strings(self.norb, nelec[1]) + self.assertLess(gas.ndet, nstr_alpha * nstr_beta) + + gas_ci = rng.normal(size=gas.ndet) + gas_ci /= numpy.linalg.norm(gas_ci) + fci_ci = fci_gas.gas2fci(gas_ci, gas) + + gas_hc = solver.contract_2e( + h2e, gas_ci, self.norb, nelec) + fci_hc = direct_spin1.contract_2e( + h2e, fci_ci, self.norb, nelec) + projected_hc = fci_gas.fci2gas(fci_hc, gas) + numpy.testing.assert_allclose( + gas_hc, projected_hc, atol=1e-11, rtol=0) + + def test_spin_square_projection(self): + cases = ( + { + "name": "supergroup", + "gas_orbs": (1, 2, 1), + "nelec": (2, 2), + "gas_restr": [[1, 2, 1], [2, 1, 1], [1, 1, 2]], + "gas_restr_type": "supergroup", + }, + { + "name": "ras", + "gas_orbs": (1, 2, 1), + "nelec": (2, 2), + "gas_restr": {"max_holes": 1, "max_particles": 1}, + "gas_restr_type": "ras", + }, + ) + + rng = numpy.random.default_rng(29) + for case in cases: + with self.subTest(case=case["name"]): + nelec = case["nelec"] + solver = fci_gas.FCISolver( + gas_orbs=case["gas_orbs"], + gas_restr=case["gas_restr"], + gas_restr_type=case["gas_restr_type"]) + + with solver.make_space(self.norb, nelec) as gas: + self.assertTrue(addons_gas.is_spin_complete( + gas.norb, nelec, gas.blocks)) + + gas_ci = rng.normal(size=gas.ndet) + gas_ci /= numpy.linalg.norm(gas_ci) + fci_ci = fci_gas.gas2fci(gas_ci, gas) + + gas_sc = solver.contract_ss( + gas_ci, self.norb, nelec) + fci_sc = spin_op.contract_ss( + fci_ci, self.norb, nelec) + projected_sc = fci_gas.fci2gas(fci_sc, gas) + numpy.testing.assert_allclose( + gas_sc, projected_sc, atol=1e-12, rtol=0) + + embedded_sc = fci_gas.gas2fci(gas_sc, gas) + numpy.testing.assert_allclose( + embedded_sc, fci_sc, atol=1e-12, rtol=0) + + def test_gas_as_cas_rdms(self): + solver = fci_gas.FCISolver() + ci = self.ci.reshape(self.nstr, self.nstr) + ket = self.ket.reshape(self.nstr, self.nstr) + + dm1, dm2 = solver.make_rdm12( + self.ci, self.norb, self.nelec) + reference_dm1, reference_dm2 = direct_spin1.make_rdm12( + ci, self.norb, self.nelec) + numpy.testing.assert_allclose( + dm1, reference_dm1, atol=1e-12, rtol=0) + numpy.testing.assert_allclose( + dm2, reference_dm2, atol=1e-12, rtol=0) + + tdm1, tdm2 = solver.trans_rdm12( + self.ci, self.ket, self.norb, self.nelec) + reference_tdm1, reference_tdm2 = direct_spin1.trans_rdm12( + ci, ket, self.norb, self.nelec) + numpy.testing.assert_allclose( + tdm1, reference_tdm1, atol=1e-12, rtol=0) + numpy.testing.assert_allclose( + tdm2, reference_tdm2, atol=1e-12, rtol=0) + + def test_spin_free_rdm_assembly(self): + solver = fci_gas.FCISolver() + with solver.make_rdm_plan(self.norb, self.nelec) as plan: + dm1a, dm1b = plan.make_rdm1s(self.ci, self.ket) + dm1 = plan.make_rdm1(self.ci, self.ket) + numpy.testing.assert_allclose( + dm1, dm1a + dm1b, atol=1e-12, rtol=0) + + (dm1a, dm1b), (dm2aa, dm2ab, dm2bb) = plan.make_rdm12s( + self.ci, self.ket) + dm1, dm2 = plan.make_rdm12(self.ci, self.ket) + expected_dm1 = dm1a + dm1b + expected_dm2 = dm2aa + dm2bb + expected_dm2 += dm2ab + expected_dm2 += dm2ab.transpose(2, 3, 0, 1) + numpy.testing.assert_allclose( + dm1, expected_dm1, atol=1e-12, rtol=0) + numpy.testing.assert_allclose( + dm2, expected_dm2, atol=1e-12, rtol=0) + + def test_restricted_space_and_spin_guards(self): + gas_orbs = (2, 2) + nelec = (2, 1) + blocks = numpy.asarray([[1, 1, 1, 0]], dtype=numpy.int32) + solver = fci_gas.FCISolver( + gas_orbs=gas_orbs, gas_restr=blocks) + h1e = numpy.diag((-1.0, -0.6, 0.2, 0.7)) + eri = numpy.zeros((10, 10)) + + info = solver.space_info(4, nelec) + self.assertEqual(info["ndet_estimate"], 8) + energies, roots = solver.kernel( + h1e, eri, 4, nelec, nroots=2, tol=1e-12) + self.assertEqual(numpy.asarray(energies).shape, (2,)) + self.assertEqual(len(roots), 2) + + root = numpy.asarray(roots[0]).reshape(-1) + (dm1a, dm1b), _ = solver.make_rdm12s(root, 4, nelec) + self.assertAlmostEqual(numpy.trace(dm1a), nelec[0], places=11) + self.assertAlmostEqual(numpy.trace(dm1b), nelec[1], places=11) + with self.assertRaises(ValueError): + solver.contract_ss(root, 4, nelec) + + fci_addons.fix_spin_(solver, shift=0.2, ss=0.75) + with self.assertRaises(ValueError): + solver.kernel(h1e, eri, 4, nelec) + + +# Fixed-orbital NO+ data imported from an OpenMolcas HDF5 file. + +no_plus_mol = ( + '{"atom": "[[\'N1\', [0.0, 0.0, 0.0]], [\'O2\', [0.0, 0.0, 2.1746968256946757]]]"' + ', "basis": "{\'N1\': [[0, [19730.800647, 0.00021887984991, 0.0, 0.0, 0.0, 0.0]' + ', [2957.8958745, 0.0016960708803, 0.0, 0.0, 0.0, 0.0], [673.22133595, 0.0087' + '954603538, 0.0, 0.0, 0.0, 0.0], [190.68249494, 0.035359382605, 0.0, 0.0, 0.0' + ', 0.0], [62.295441898, 0.11095789217, 0.0, 0.0, 0.0, 0.0], [22.654161182, 0.' + '24982972552, 0.0, 0.0, 0.0, 0.0], [8.9791477428, 0.0, 0.40623896148, 0.0, 0.' + '0, 0.0], [3.686300237, 0.0, 0.24338217176, 0.0, 0.0, 0.0], [0.84660076805, 0' + '.0, 0.0, 1.0, 0.0, 0.0], [0.33647133771, 0.0, 0.0, 0.0, 1.0, 0.0], [0.136476' + '53675, 0.0, 0.0, 0.0, 0.0, 1.0]], [1, [49.20038051, 0.0055552416751, 0.0, 0.' + '0], [11.346790537, 0.038052379723, 0.0, 0.0], [3.4273972411, 0.14953671029, ' + '0.0, 0.0], [1.1785525134, 0.3494930523, 0.0, 0.0], [0.41642204972, 0.0, 0.45' + '843153697, 0.0], [0.14260826011, 0.0, 0.0, 0.24428771672]], [2, [1.654, 1.0,' + " 0.0], [0.469, 0.0, 1.0]], [3, [1.093, 1.0]]], 'O2': [[0, [27032.382631, 0.0" + '0021726302465, 0.0, 0.0, 0.0, 0.0], [4052.3871392, 0.0016838662199, 0.0, 0.0' + ', 0.0, 0.0], [922.3272271, 0.0087395616265, 0.0, 0.0, 0.0, 0.0], [261.240709' + '89, 0.035239968808, 0.0, 0.0, 0.0, 0.0], [85.354641351, 0.11153519115, 0.0, ' + '0.0, 0.0, 0.0], [31.035035245, 0.25588953961, 0.0, 0.0, 0.0, 0.0], [12.26086' + '0728, 0.0, 0.39768730901, 0.0, 0.0, 0.0], [4.9987076005, 0.0, 0.2462784943, ' + '0.0, 0.0, 0.0], [1.1703108158, 0.0, 0.0, 1.0, 0.0, 0.0], [0.46474740994, 0.0' + ', 0.0, 0.0, 1.0, 0.0], [0.18504536357, 0.0, 0.0, 0.0, 0.0, 1.0]], [1, [63.27' + '4954801, 0.0060685103418, 0.0, 0.0], [14.627049379, 0.041912575824, 0.0, 0.0' + '], [4.4501223456, 0.16153841088, 0.0, 0.0], [1.5275799647, 0.35706951311, 0.' + '0, 0.0], [0.52935117943, 0.0, 0.44794207502, 0.0], [0.1747842127, 0.0, 0.0, ' + '0.24446069663]], [2, [2.314, 1.0, 0.0], [0.645, 0.0, 1.0]], [3, [1.428, 1.0]' + ']]}", "nucmod": {}, "ecp": "{}", "nucprop": {}, "magmom": [0, 0], "pseudo": ' + '"None", "_atm": [[7, 20, 1, 23, 0, 0], [8, 24, 1, 27, 0, 0]], "_bas": [[0, 0' + ', 11, 5, 0, 28, 39, 0], [0, 1, 6, 3, 0, 94, 100, 0], [0, 2, 2, 2, 0, 118, 12' + '0, 0], [0, 3, 1, 1, 0, 124, 125, 0], [1, 0, 11, 5, 0, 126, 137, 0], [1, 1, 6' + ', 3, 0, 192, 198, 0], [1, 2, 2, 2, 0, 216, 218, 0], [1, 3, 1, 1, 0, 222, 223' + ', 0]], "_env": [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, ' + '0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 2.1746' + '968256946757, 0.0, 19730.800647, 2957.8958745, 673.22133595, 190.68249494, 6' + '2.295441898, 22.654161182, 8.9791477428, 3.686300237, 0.84660076805, 0.33647' + '133771, 0.13647653675, 2.4658508374156667, 4.6034522992919324, 7.86646232650' + '279, 12.278357689028272, 16.649564930304884, 17.555273048650747, 0.0, 0.0, 0' + '.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 8.465479929725005, 2.60121028934' + '62696, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 2.229841622743' + '3675, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.1161590381537' + '6, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.5672944296303319' + ', 49.20038051, 11.346790537, 3.4273972411, 1.1785525134, 0.41642204972, 0.14' + '260826011, 4.352170395048921, 4.764479639730067, 4.192718762589188, 2.580281' + '4243677257, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.9758915092644845, 0.0, 0.0, 0.0,' + ' 0.0, 0.0, 0.0, 0.25566144791664475, 1.654, 0.469, 6.294576559932722, 0.0, 0' + '.0, 0.693556523160707, 1.093, 2.4093823583303675, 27032.382631, 4052.3871392' + ', 922.3272271, 261.24070989, 85.354641351, 31.035035245, 12.260860728, 4.998' + '7076005, 1.1703108158, 0.46474740994, 0.18504536357, 3.048118116984593, 5.69' + '1457632864213, 9.733883574443253, 15.238733819733028, 20.843228934131737, 22' + '.39104905999299, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 10.5' + '68131135849375, 3.3391469496791393, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, ' + '0.0, 0.0, 0.0, 2.8427648592056753, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0' + '.0, 0.0, 0.0, 1.422092211265869, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0' + ', 0.0, 0.0, 0.7128098301044613, 63.274954801, 14.627049379, 4.4501223456, 1.' + '5275799647, 0.52935117943, 0.1747842127, 6.257032374789428, 6.92686562359984' + '2, 6.032359926541528, 3.5035168827833356, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.31' + '7237993956345, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.3296948367394935, 2.314, 0.64' + '5, 11.32831343293501, 0.0, 0.0, 1.2113199965714336, 1.428, 4.396922678265652' + '], "_ecpbas": [], "groupname": "C1", "topgroup": "C1", "_symm_orig": null, "' + '_symm_axes": null, "_nelectron": null, "_nao": null, "_enuc": null, "_atom":' + ' [["N1", [0.0, 0.0, 0.0]], ["O2", [0.0, 0.0, 2.1746968256946757]]], "_basis"' + ': {"N1": [[0, [19730.800647, 0.00021887984991, 0.0, 0.0, 0.0, 0.0], [2957.89' + '58745, 0.0016960708803, 0.0, 0.0, 0.0, 0.0], [673.22133595, 0.0087954603538,' + ' 0.0, 0.0, 0.0, 0.0], [190.68249494, 0.035359382605, 0.0, 0.0, 0.0, 0.0], [6' + '2.295441898, 0.11095789217, 0.0, 0.0, 0.0, 0.0], [22.654161182, 0.2498297255' + '2, 0.0, 0.0, 0.0, 0.0], [8.9791477428, 0.0, 0.40623896148, 0.0, 0.0, 0.0], [' + '3.686300237, 0.0, 0.24338217176, 0.0, 0.0, 0.0], [0.84660076805, 0.0, 0.0, 1' + '.0, 0.0, 0.0], [0.33647133771, 0.0, 0.0, 0.0, 1.0, 0.0], [0.13647653675, 0.0' + ', 0.0, 0.0, 0.0, 1.0]], [1, [49.20038051, 0.0055552416751, 0.0, 0.0], [11.34' + '6790537, 0.038052379723, 0.0, 0.0], [3.4273972411, 0.14953671029, 0.0, 0.0],' + ' [1.1785525134, 0.3494930523, 0.0, 0.0], [0.41642204972, 0.0, 0.45843153697,' + ' 0.0], [0.14260826011, 0.0, 0.0, 0.24428771672]], [2, [1.654, 1.0, 0.0], [0.' + '469, 0.0, 1.0]], [3, [1.093, 1.0]]], "O2": [[0, [27032.382631, 0.00021726302' + '465, 0.0, 0.0, 0.0, 0.0], [4052.3871392, 0.0016838662199, 0.0, 0.0, 0.0, 0.0' + '], [922.3272271, 0.0087395616265, 0.0, 0.0, 0.0, 0.0], [261.24070989, 0.0352' + '39968808, 0.0, 0.0, 0.0, 0.0], [85.354641351, 0.11153519115, 0.0, 0.0, 0.0, ' + '0.0], [31.035035245, 0.25588953961, 0.0, 0.0, 0.0, 0.0], [12.260860728, 0.0,' + ' 0.39768730901, 0.0, 0.0, 0.0], [4.9987076005, 0.0, 0.2462784943, 0.0, 0.0, ' + '0.0], [1.1703108158, 0.0, 0.0, 1.0, 0.0, 0.0], [0.46474740994, 0.0, 0.0, 0.0' + ', 1.0, 0.0], [0.18504536357, 0.0, 0.0, 0.0, 0.0, 1.0]], [1, [63.274954801, 0' + '.0060685103418, 0.0, 0.0], [14.627049379, 0.041912575824, 0.0, 0.0], [4.4501' + '223456, 0.16153841088, 0.0, 0.0], [1.5275799647, 0.35706951311, 0.0, 0.0], [' + '0.52935117943, 0.0, 0.44794207502, 0.0], [0.1747842127, 0.0, 0.0, 0.24446069' + '663]], [2, [2.314, 1.0, 0.0], [0.645, 0.0, 1.0]], [3, [1.428, 1.0]]]}, "_ecp' + '": {}, "_pseudo": {}, "_built": true, "unit": "Bohr", "charge": 1, "symmetry' + '": false}' +) + +no_plus_mo_coeff = numpy.asarray([ + -0.00030993465811149129, 0.45199552674204058, -0.028904704553936636, -0.090403382320240686, + 3.3099317312990211e-23, -5.4836216379542063e-23, 2.2319201719783142e-23, 1.5234367861943547e-23, + -0.014919945575310562, -0.050585226876045802, 0.00021964788007901282, 0.63737873963625524, + -0.075317996175523913, -0.22391930706443322, 7.7663101496108472e-23, -1.2866581498817644e-22, + 6.134007397957805e-23, 4.1868757819906538e-23, -0.031952211830968344, -0.12800752151058567, + -0.0032609643711730505, 0.0068800847859985248, 0.13963789699680393, 0.3910739921457736, + -1.2545999583162285e-22, 2.0785176354004214e-22, -1.2107811154499295e-22, -8.2644017195444379e-23, + 0.051249641321393952, 0.28818667646763757, 0.0033075416811898057, -0.019087360189030803, + 0.11625833200242082, 0.52831848220213351, -1.5014963062947533e-22, 2.487555121767254e-22, + -1.4634497672650712e-22, -9.9890365142258758e-23, -0.085174620576668386, 0.41636853967767923, + -0.0031469672611320427, -0.0031292027582865865, 0.038604836398930767, 0.15078923105718686, + -4.4348773230544274e-23, 7.3473386202630555e-23, -4.3092200287570662e-23, -2.9413347268918058e-23, + -0.098222788203448541, 0.024944319955640809, 0, 0, + 0, 0, -0.067945434982550693, -0.27486984552486465, + 0.48235372755507899, -0.091000564454398664, -8.5687039038947365e-22, -2.4242021424911906e-22, + 0, 0, 0, 0, + 0.27486984552486465, -0.067945434982550706, 0.091000564454398664, 0.48235372755507905, + -8.5687039077997024e-22, -2.4242021424027878e-22, 0.0012639557204546764, 0.0080067496754596592, + -0.018257071700353976, -0.14240824605537175, 3.7000633254809653e-23, -6.1299594527524831e-23, + 3.6390182472685632e-23, 2.4838765880972595e-23, 0.35818949693892382, 0.59404179492175535, + 0, 0, 0, 0, + -0.061104899094375123, -0.24719680057269985, 0.39108520443955269, -0.073781899737669826, + -7.3177996867130668e-22, -2.1453274242608933e-22, 0, 0, + 0, 0, 0.24719680057269985, -0.06110489909437513, + 0.073781899737669812, 0.39108520443955275, -7.3177996938849402e-22, -2.145327425795134e-22, + 0.0057906197720830306, 0.0044905244275705147, -0.050359547498062866, -0.10989774503971134, + 3.8395986659881985e-23, -6.3611300856164208e-23, 3.6587167858830604e-23, 2.4973221758275792e-23, + 0.22119527201835215, 0.47949477769533205, 0, 0, + 0, 0, -0.01396644117585785, -0.056500536402334142, + 0.040432602332096335, -0.0076279904457113544, -1.2275500964325165e-22, -4.5043903833080602e-23, + 0, 0, 0, 0, + 0.056500536402334142, -0.013966441175857852, 0.0076279904457113553, 0.040432602332096348, + -1.2275500971741671e-22, -4.5043903840723138e-23, -0.00038238083713755197, 0.0013279845592333254, + -0.017755908800136749, -0.022999209207775772, 9.967481539395441e-24, -1.6513300533565073e-23, + 9.3188838241300475e-24, 6.3607697972387392e-24, 0.035002096717448826, 0.028040985470694852, + 0, 0, 0, 0, + 0, 0, 0, 0, + 0, 0, 0, 0, + 0, 0, 0.018989939867174316, -0.0046941479818857768, + 0.0014968545370685424, 0.0079341635928896221, -3.6117117263261199e-23, -1.4678334348277796e-23, + -2.3551743008472153e-05, 0.00086403660450635534, 0.0056433849698386855, -0.0030629827416551736, + -8.2746920683314983e-25, 1.3708826692233768e-24, -6.3044707156074036e-25, -4.3032285495950224e-25, + 0.027148172581965901, 0.0034952577378645397, 0, 0, + 0, 0, -0.0046941479818857759, -0.018989939867174316, + 0.0079341635928896186, -0.0014968545370685419, -3.6117117187018993e-23, -1.46783343256146e-23, + 0, 0, 0, 0, + 0, 0, 0, 0, + 0, 0, 0, 0, + 0, 0, 0, 0, + 0, 0, 0, 0, + 0, 0, 0, 0, + 0.0411345053288956, -0.010168091982423612, 0.0070672418789917996, 0.037460322234573212, + -9.6664512026001079e-23, -3.3447737295251316e-23, 0.0014997602905395015, 0.00067502139964927693, + -0.0011449273695446154, -0.0049406028118178973, 1.4457980820968987e-24, -2.3952789022211201e-24, + 1.3900830830707213e-24, 9.4882591704464335e-25, 0.01520903339622085, 0.048103483584991717, + 0, 0, 0, 0, + -0.01016809198242361, -0.0411345053288956, 0.037460322234573212, -0.0070672418789917996, + -9.6664511844836758e-23, -3.3447737231663719e-23, 0, 0, + 0, 0, 0, 0, + 0, 0, 0, 0, + 0, 0, 0, 0, + 7.2499999874677741e-10, -1.7899999969058371e-10, 4.8100000048012137e-10, 2.5480000032457044e-09, + -3.4197542448540226e-30, -7.4324802158081137e-31, 0, 0, + 0, 0, 0, 0, + 0, 0, 0, 0, + 0, 0, 0, 0, + 0.0083670497255368612, -0.002068261916424838, 0.00054552165354917993, 0.0028915689204292929, + -1.5364044981505871e-23, -6.4180798320080482e-24, -2.2886869008232982e-05, 0.00026753548184679175, + 0.0013612316299479799, 0.0016529014805612883, -7.3599720742662531e-25, 1.2193394118588906e-24, + -6.9389502040660553e-25, -4.7363038015345652e-25, 0.0072369016334001219, -0.0030670077221483846, + 0, 0, 0, 0, + -0.002068261916424838, -0.0083670497255368612, 0.0028915689204292924, -0.00054552165354917971, + -1.5364044975957611e-23, -6.4180798343126379e-24, 0, 0, + 0, 0, 0, 0, + 0, 0, 0, 0, + 0, 0, 0, 0, + 1.7899999969058368e-10, 7.2499999874677741e-10, -2.5480000032457044e-09, 4.8100000048012137e-10, + 3.4198478942988598e-30, 7.4349905418819973e-31, 0.4393247431580361, 0.0018457249590554949, + -0.082440307245391961, 0.032835096634947859, 2.0821619123975467e-23, -3.4495539612889851e-23, + 1.5598619460356965e-23, 1.0647114980004755e-23, -0.020336271781835273, 0.047334385730994001, + 0.65162213123440482, 0.0037255278528676517, -0.21656746376996816, 0.084096230873614689, + 4.8416450422560335e-23, -8.021237798474233e-23, 3.7070735687412094e-23, 2.530328957971206e-23, + -0.046972380543379266, 0.12297622487225722, 0.0079436071128575151, -0.0030155181580962718, + 0.35839646499999284, -0.14176954708538833, -6.4904984018073025e-23, 1.075292192122642e-22, + -5.126638960110416e-23, -3.4992785486762102e-23, 0.055790043610147701, -0.28073838895907421, + -0.033200168311942932, -0.0068520008268535201, 0.51531686098463858, -0.22432121813260097, + -8.9353987271549006e-23, 1.4803430938855026e-22, -6.9444510230380106e-23, -4.7400584840645831e-23, + 0.065018982623952684, -0.39655542026124913, -0.0025375225209128102, -0.00097308982350304918, + 0.15173392999789087, -0.11397918105798752, -1.5443590729720988e-23, 2.5585666163230348e-23, + -9.5534902377339056e-24, -6.5209045775134434e-24, -0.04924140514115187, -0.12183400846770261, + 0, 0, 0, 0, + -0.094587070236498413, -0.38264724233856257, -0.48671419684477685, 0.091823208129428818, + -1.3996591993569832e-22, -2.4305995029394422e-22, 0, 0, + 0, 0, 0.38264724233856257, -0.094587070236498427, + -0.091823208129428791, -0.48671419684477685, -1.3996592014462721e-22, -2.4305995020900676e-22, + -0.0088521748631843478, -0.0021438770389801847, 0.12232729499607477, -0.11144240005486751, + -8.1684090422108868e-24, 1.3532745755891694e-23, -3.3239534345426224e-24, -2.2688235022350226e-24, + -0.39294798183193108, 0.61598417392035687, 0, 0, + 0, 0, -0.081681162058807349, -0.3304370394288122, + -0.29688857455715278, 0.056010820065607377, -2.3306222063716283e-22, -2.1995631270757339e-22, + 0, 0, 0, 0, + 0.3304370394288122, -0.081681162058807363, -0.056010820065607356, -0.29688857455715278, + -2.3306222087937335e-22, -2.1995631253857797e-22, -0.00026648872209586265, 8.1205679947524303e-05, + 0.083600597499161003, -0.07641327573735239, -5.4311033514729559e-24, 8.9978036665685639e-24, + -2.1806279240103904e-24, -1.4884263516833888e-24, -0.32850953285686413, 0.38438219416587477, + 0, 0, 0, 0, + -0.025415706556066845, -0.10281796482227062, 0.01206057577249472, -0.0022753409559873156, + -1.6746212967910739e-22, -7.6954759609250789e-23, 0, 0, + 0, 0, 0.10281796482227062, -0.025415706556066849, + 0.0022753409559873182, 0.012060575772494728, -1.6746212959324084e-22, -7.6954759506770098e-23, + -0.0026062364409375283, -0.0013393631192955995, 0.03701392549892029, 0.00104693344732496, + -1.0281903630394735e-23, 1.7034209098145332e-23, -8.8625483178806873e-24, -6.0492898864790742e-24, + -0.070911916200646075, 0.015732860878165628, 0, 0, + 0, 0, 0, 0, + 0, 0, 0, 0, + 0, 0, 0, 0, + -0.012426798778519248, 0.0030717965746901456, -0.00024128447605068648, -0.0012789422564260531, + 2.0077357189300258e-23, 9.2863473077724753e-24, 0.00042642603815339612, -2.840236988567041e-06, + 0.00081366179114275724, 0.0052067273720378575, -1.3827261220261426e-24, 2.2907864857010504e-24, + -1.3634921867485244e-24, -9.3067582809012728e-25, 0.016842225986242795, 0.0090004167732186668, + 0, 0, 0, 0, + 0.0030717965746901452, 0.012426798778519248, -0.001278942256426052, 0.00024128447605068615, + 2.0077357170642162e-23, 9.2863473068702584e-24, 0, 0, + 0, 0, 0, 0, + 0, 0, 0, 0, + 0, 0, 0, 0, + 0, 0, 0, 0, + 0, 0, 0, 0, + 0, 0, -0.034208236740868231, 0.0084559785853831424, + -0.0010117741623665965, -0.005362967222734636, 5.6943337445289288e-23, 2.5713448479830125e-23, + -0.0014007224605038749, -0.0012961610192925158, 0.009343509189581517, 0.010273146103284694, + -4.8625970199805678e-24, 8.0559493035791074e-24, -4.4975543354603261e-24, -3.0698856556828762e-24, + 0.042769660961498243, 0.010560637948058595, 0, 0, + 0, 0, 0.0084559785853831407, 0.034208236740868231, + -0.0053629672227346334, 0.0010117741623665957, 5.6943337426602712e-23, 2.5713448488114281e-23, + 0, 0, 0, 0, + 0, 0, 0, 0, + 0, 0, 0, 0, + 0, 0, 1.0739999981435021e-09, -2.6499999954192564e-10, + -8.1700000142478931e-10, -4.3290000065090597e-09, 2.3018139448893698e-30, -4.4023421173914863e-31, + 0, 0, 0, 0, + 0, 0, 0, 0, + 0, 0, 0, 0, + 0, 0, 0.0067681971283006107, -0.001673039457108013, + 0.0011507847598172889, 0.0060998008355994402, -1.5846957237402312e-23, -5.4982250554805852e-24, + 0.00016707332506010045, 0.00042762110775980243, -0.0024446931399710443, -0.0015034305903884981, + 1.0047151352617039e-24, -1.6645290903649695e-24, 9.0611785430358747e-25, 6.1848684770916688e-25, + -0.0079573813825088695, -0.010525235841618416, 0, 0, + 0, 0, -0.0016730394571080128, -0.0067681971283006107, + 0.0060998008355994384, -0.0011507847598172889, -1.5846957241604296e-23, -5.4982250616607101e-24, + 0, 0, 0, 0, + 0, 0, 0, 0, + 0, 0, 0, 0, + 0, 0, 2.6499999954192559e-10, 1.0739999981435021e-09, + 4.3290000065090605e-09, -8.1700000142478942e-10, -2.2999117424193358e-30, 4.4095372155272348e-31, +], dtype=numpy.float64).reshape((62, 10)) + +no_plus_openmolcas_energies = numpy.asarray([ + -129.12863949, + -128.72314673, + -128.72314673, + -128.72083180, + -128.70307916, + -128.70307916, + -128.50342069, + -128.50342069, + -128.46044322, + -128.39715780, +]) + + +class TestGASCI(unittest.TestCase): + + @classmethod + def setUpClass(cls): + cls.mol = gto.M( + atom="N 0 0 0; N 0 0 1.10", + basis="sto-3g", + spin=0, + verbose=0, + ) + cls.mf = scf.RHF(cls.mol).run() + + @classmethod + def tearDownClass(cls): + del cls.mol, cls.mf + + def test_o2_triplet_energy(self): + mol = gto.M( + atom="O 0 0 0; O 0 0 1.21", + basis="sto-3g", + spin=2, + symmetry=False, + verbose=0, + ) + + mf = scf.ROHF(mol) + mf.conv_tol = 1e-12 + mf.kernel() + self.assertTrue(mf.converged) + + mc = gasci.GASCI( + mf, + 8, + (7, 5), + ncore=2, + gas_orbs=(2, 4, 2), + gas_restr=( + (2, 4), + (8, 10), + (12, 12), + ), + gas_restr_type="cumulative-occ", + ) + mc.verbose = 0 + mc.canonicalization = False + mc.fcisolver.spin = 2 + mc.fcisolver.nroots = 1 + mc.fcisolver.max_cycle = 300 + mc.fcisolver.conv_tol = 1e-12 + mc.kernel() + self.assertTrue(mc.converged) + + numpy.testing.assert_allclose( + mc.e_tot, + -147.5163002382864, + atol=1e-7, + rtol=0, + ) + ss, multiplicity = mc.spin_square() + self.assertAlmostEqual(ss, 2.0, places=9) + self.assertAlmostEqual(multiplicity, 3.0, places=9) + self.assertEqual(mc.gas_space_info()["core"]["ndet"], 424) + + def test_no_plus_openmolcas_multiroot(self): + mol = gto.loads(no_plus_mol) + mo_coeff = no_plus_mo_coeff.copy() + + overlap = mol.intor_symmetric("int1e_ovlp") + numpy.testing.assert_allclose( + mo_coeff.T @ overlap @ mo_coeff, + numpy.eye(mo_coeff.shape[1]), + atol=1e-10, + rtol=0, + ) + + mol.verbose = 0 + mf = scf.RHF(mol) + mf.max_cycle = 1 + mf.kernel() + mf.mo_coeff = mo_coeff + + mc = gasci.GASCI( + mf, + 8, + (5, 5), + ncore=2, + gas_orbs=(2, 4, 2), + gas_restr=( + (2, 4), + (7, 9), + (10, 10), + ), + gas_restr_type="cumulative-occ", + ) + mc.verbose = 0 + mc.canonicalization = False + mc.fcisolver.spin = 0 + mc.fcisolver.nroots = 30 + mc.fcisolver.max_space = 30 + mc.fcisolver.max_cycle = 300 + mc.fcisolver.conv_tol = 1e-10 + mc.kernel() + + self.assertTrue(numpy.all(numpy.asarray(mc.converged))) + + singlet_energies = [] + for state, energy in enumerate( + numpy.asarray(mc.e_tot).reshape(-1)): + ss, _ = mc.spin_square(state=state) + if abs(ss) < 1e-6: + singlet_energies.append(energy) + + self.assertEqual( + len(singlet_energies), + len(no_plus_openmolcas_energies), + ) + numpy.testing.assert_allclose( + numpy.asarray(singlet_energies), + no_plus_openmolcas_energies, + atol=1e-7, + rtol=0, + ) + + def test_kernel_log_label(self): + output = io.StringIO() + mc = gasci.GASCI( + self.mf, 4, (2, 2), ncore=5) + mc.stdout = output + mc.verbose = logger.DEBUG + mc.canonicalization = False + mc.kernel() + + text = output.getvalue() + self.assertIn("Start GASCI", text) + self.assertNotIn("Start CASCI", text) + + def test_gas_as_cas(self): + reference = mcscf.CASCI( + self.mf, 4, (2, 2), ncore=5) + reference.verbose = 0 + reference.canonicalization = False + reference.kernel() + + mc = gasci.GASCI( + self.mf, 4, (2, 2), ncore=5) + mc.verbose = 0 + mc.canonicalization = False + mc.kernel() + + self.assertAlmostEqual(mc.e_tot, reference.e_tot, places=9) + self.assertEqual(mc.gas_space_info()["core"]["ndet"], 36) + numpy.testing.assert_allclose( + mc.make_gasdm1(), reference.fcisolver.make_rdm1( + reference.ci, 4, (2, 2)), + atol=1e-9, rtol=0) + + def test_fcisolver_spin_propagation(self): + mc = gasci.GASCI( + self.mf, 4, 4, ncore=5) + mc.verbose = 0 + mc.canonicalization = False + mc.fcisolver.spin = 2 + mc.kernel() + + info = mc.gas_space_info()["core"] + self.assertEqual(info["nelec"], (3, 1)) + self.assertEqual(info["ndet"], 16) + self.assertEqual(mc._gas_problem_signature()[1], (3, 1)) + + dm1a, dm1b = mc.make_gasdm1s() + self.assertAlmostEqual(numpy.trace(dm1a), 3.0, places=10) + self.assertAlmostEqual(numpy.trace(dm1b), 1.0, places=10) + self.assertAlmostEqual( + mc.spin_square()[0], + mc.fcisolver.spin_square(mc.ci, 4, mc.nelecas)[0], + places=10) + + spin_signature = mc._gas_problem_signature() + mc.fcisolver.spin = 0 + self.assertNotEqual(spin_signature, mc._gas_problem_signature()) + + def test_restricted_multiroot_rdms(self): + mc = gasci.GASCI( + self.mf, 4, (2, 2), ncore=5, + gas_orbs=(2, 2), + gas_restr=[[2, 2]], + gas_restr_type="supergroup", + ) + mc.verbose = 0 + mc.canonicalization = False + mc.fcisolver.nroots = 3 + mc.kernel() + + energies = numpy.asarray(mc.e_tot).reshape(-1) + self.assertEqual(energies.shape, (3,)) + self.assertEqual(len(mc.ci), 3) + self.assertEqual(mc.gas_space_info()["core"]["ndet"], 18) + + for state in range(3): + dm1, dm2 = mc.make_gasdm12(state=state) + self.assertAlmostEqual(numpy.trace(dm1), 4.0, places=9) + self.assertAlmostEqual( + numpy.einsum("pprr", dm2), 12.0, places=8) + ss, multiplicity = mc.spin_square(state=state) + self.assertTrue(numpy.isfinite(ss)) + self.assertTrue(numpy.isfinite(multiplicity)) + + tdm1, tdm2 = mc.trans_gasdm12( + bra_state=0, ket_state=1) + reverse1, reverse2 = mc.trans_gasdm12( + bra_state=1, ket_state=0) + numpy.testing.assert_allclose( + tdm1, reverse1.T, atol=1e-10, rtol=0) + numpy.testing.assert_allclose( + tdm2, reverse2.transpose(1, 0, 3, 2), + atol=1e-10, rtol=0) + + def test_state_average(self): + weights = numpy.asarray([0.6, 0.4]) + mc = gasci.GASCI( + self.mf, 4, (2, 2), ncore=5, + gas_orbs=(2, 2), + gas_restr=[[2, 2]], + gas_restr_type="supergroup", + ).state_average_(weights) + mc.verbose = 0 + mc.canonicalization = False + mc.kernel() + + numpy.testing.assert_allclose( + mc.e_tot, numpy.dot(weights, mc.e_states), + atol=1e-10, rtol=0) + reference_dm1 = sum( + weights[state] * mc.make_gasdm1(state=state) + for state in range(2)) + numpy.testing.assert_allclose( + mc.make_gasdm1(), reference_dm1, atol=1e-10, rtol=0) + + +if __name__ == "__main__": + unittest.main()