diff --git a/edrixs/_operator_builder.py b/edrixs/_operator_builder.py new file mode 100644 index 00000000..9944cd4e --- /dev/null +++ b/edrixs/_operator_builder.py @@ -0,0 +1,515 @@ +"""Backend-independent kernels for many-body operator matrix entries.""" + +from __future__ import annotations + +import numpy as np + +from .fock_basis import FockBasis, FockBinByN + + +_NUMBA_KERNELS = None + + +def _operator_terms(emat, umat, norbs, tol_e, tol_u): + if emat is not None: + emat = np.asarray(emat) + if emat.shape != (norbs, norbs): + raise ValueError( + "emat has shape {}, expected {}".format( + emat.shape, (norbs, norbs) + ) + ) + iorb, jorb = np.nonzero(np.abs(emat) > tol_e) + e_terms = np.stack((iorb, jorb), axis=-1).astype(np.int64) + e_vals = emat[iorb, jorb].astype(np.complex128) + else: + e_terms = np.empty((0, 2), dtype=np.int64) + e_vals = np.empty((0,), dtype=np.complex128) + + if umat is None: + u_terms = np.empty((0, 4), dtype=np.int64) + u_vals = np.empty((0,), dtype=np.complex128) + elif hasattr(umat, 'tocoo') and len(umat.shape) == 2: + expected = (norbs * norbs, norbs * norbs) + if umat.shape != expected: + raise ValueError( + "sparse umat has shape {}, expected {}".format(umat.shape, expected) + ) + coo = umat.tocoo() + keep = np.abs(coo.data) > tol_u + rows = np.asarray(coo.row[keep], dtype=np.int64) + cols = np.asarray(coo.col[keep], dtype=np.int64) + lorb, korb = np.divmod(rows, norbs) + jorb, iorb = np.divmod(cols, norbs) + u_terms = np.stack((lorb, korb, jorb, iorb), axis=-1) + u_vals = np.asarray(coo.data[keep], dtype=np.complex128) + else: + umat = np.asarray(umat) + expected = (norbs, norbs, norbs, norbs) + if umat.shape != expected: + raise ValueError( + "dense umat has shape {}, expected {}".format(umat.shape, expected) + ) + lorb, korb, jorb, iorb = np.nonzero(np.abs(umat) > tol_u) + u_terms = np.stack((lorb, korb, jorb, iorb), axis=-1).astype(np.int64) + u_vals = umat[lorb, korb, jorb, iorb].astype(np.complex128) + + return e_terms, e_vals, u_terms, u_vals + + +def _sign_count(state, orbital, norbs): + prefix = int(state) >> (norbs - orbital) + return 1 if prefix.bit_count() % 2 == 0 else -1 + + +def _build_entries_python(lb, rb, e_terms, e_vals, u_terms, u_vals, cstart, cend): + rows = [] + cols = [] + vals = [] + norbs = rb.norbs + + for column in range(cstart, cend): + state0 = rb.decode(column) + + for term, value in zip(e_terms, e_vals): + iorb, jorb = int(term[0]), int(term[1]) + bit_j = 1 << (norbs - 1 - jorb) + if not state0 & bit_j: + continue + sign_1 = _sign_count(state0, jorb, norbs) + state = state0 ^ bit_j + + bit_i = 1 << (norbs - 1 - iorb) + if state & bit_i: + continue + sign_2 = _sign_count(state, iorb, norbs) + state |= bit_i + + try: + row = lb.encode(state) + except KeyError: + continue + rows.append(row) + cols.append(column) + vals.append(value * sign_1 * sign_2) + + for term, value in zip(u_terms, u_vals): + lorb, korb, jorb, iorb = map(int, term) + if iorb == jorb or korb == lorb: + continue + + bit_i = 1 << (norbs - 1 - iorb) + if not state0 & bit_i: + continue + sign_1 = _sign_count(state0, iorb, norbs) + state = state0 ^ bit_i + + bit_j = 1 << (norbs - 1 - jorb) + if not state & bit_j: + continue + sign_2 = _sign_count(state, jorb, norbs) + state ^= bit_j + + bit_k = 1 << (norbs - 1 - korb) + if state & bit_k: + continue + sign_3 = _sign_count(state, korb, norbs) + state |= bit_k + + bit_l = 1 << (norbs - 1 - lorb) + if state & bit_l: + continue + sign_4 = _sign_count(state, lorb, norbs) + state |= bit_l + + try: + row = lb.encode(state) + except KeyError: + continue + rows.append(row) + cols.append(column) + vals.append(value * sign_1 * sign_2 * sign_3 * sign_4) + + return ( + np.asarray(rows, dtype=np.int64), + np.asarray(cols, dtype=np.int64), + np.asarray(vals, dtype=np.complex128), + ) + + +def _get_numba_kernels(): + global _NUMBA_KERNELS + if _NUMBA_KERNELS is not None: + return _NUMBA_KERNELS + + try: + from numba import njit + except ImportError as exc: + raise ImportError( + "use_numba=True requires numba; install numba or disable JIT construction" + ) from exc + + @njit(inline='always') + def comb_jit(n, k): + if k < 0 or k > n: + return 0 + if k == 0 or k == n: + return 1 + if k > n - k: + k = n - k + result = 1 + for i in range(k): + result = (result * (n - i)) // (i + 1) + return result + + @njit(inline='always') + def popcount_u64(value): + count = 0 + while value != np.uint64(0): + value &= value - np.uint64(1) + count += 1 + return count + + @njit(inline='always') + def hash_decoder_jit(rank, norb, nocc): + state = np.uint64(0) + j = nocc + i = norb - 1 + c = comb_jit(i, j) + while i >= 0 and j > 0: + if c <= rank: + state |= np.uint64(1) << np.uint64(i) + rank -= c + old_i, old_j = i, j + i -= 1 + j -= 1 + if j == 0 or i < 0: + break + c = (c * old_j) // old_i + else: + if i == 0: + break + c = (c * (i - j)) // i + i -= 1 + return state + + @njit(inline='always') + def hash_encoder_jit(state, norb): + if norb < 64: + state &= (np.uint64(1) << np.uint64(norb)) - np.uint64(1) + rank = 0 + k = 1 + while state != np.uint64(0): + lsb = state & (np.uint64(0) - state) + pos = 0 + tmp = lsb + while tmp > np.uint64(1): + tmp >>= np.uint64(1) + pos += 1 + rank += comb_jit(pos, k) + k += 1 + state ^= lsb + return rank + + @njit(inline='always') + def sign_count_u64(state, orbital, norbs): + bitpos = np.uint64(norbs - 1 - orbital) + if bitpos == np.uint64(63): + prefix = np.uint64(0) + else: + prefix = state >> (bitpos + np.uint64(1)) + return 1 if popcount_u64(prefix) % 2 == 0 else -1 + + @njit(inline='always') + def decode_combinadic_jit(index, norbs, noccus, offsets, sizes): + state = np.uint64(0) + for n in range(len(norbs)): + size = sizes[n] + shell_rank = index % size + index //= size + colex_rank = size - 1 - shell_rank + shell_state = hash_decoder_jit(colex_rank, norbs[n], noccus[n]) + state |= shell_state << np.uint64(offsets[n]) + return state + + @njit(inline='always') + def encode_combinadic_jit(state, norbs, noccus, offsets, sizes, strides): + index = 0 + for n in range(len(norbs)): + norb = norbs[n] + if norb == 64: + mask = np.uint64(0xffffffffffffffff) + else: + mask = (np.uint64(1) << np.uint64(norb)) - np.uint64(1) + shell_state = (state >> np.uint64(offsets[n])) & mask + if popcount_u64(shell_state) != noccus[n]: + return -1 + colex_rank = hash_encoder_jit(shell_state, norb) + shell_rank = sizes[n] - 1 - colex_rank + index += shell_rank * strides[n] + return index + + @njit + def build_combinadic( + rb_norbs, rb_noccus, rb_offsets, rb_sizes, rb_strides, + lb_norbs, lb_noccus, lb_offsets, lb_sizes, lb_strides, + e_terms, e_vals, u_terms, u_vals, cstart, cend, + ): + rows = [] + cols = [] + vals = [] + norbs = 0 + for value in rb_norbs: + norbs += value + + for column in range(cstart, cend): + state0 = decode_combinadic_jit( + column, rb_norbs, rb_noccus, rb_offsets, rb_sizes + ) + + for t in range(e_terms.shape[0]): + iorb = int(e_terms[t, 0]) + jorb = int(e_terms[t, 1]) + bit_j = np.uint64(1) << np.uint64(norbs - 1 - jorb) + if state0 & bit_j == 0: + continue + sign_1 = sign_count_u64(state0, jorb, norbs) + state = state0 ^ bit_j + + bit_i = np.uint64(1) << np.uint64(norbs - 1 - iorb) + if state & bit_i != 0: + continue + sign_2 = sign_count_u64(state, iorb, norbs) + state |= bit_i + + row = encode_combinadic_jit( + state, lb_norbs, lb_noccus, lb_offsets, lb_sizes, lb_strides + ) + if row != -1: + rows.append(row) + cols.append(column) + vals.append(e_vals[t] * sign_1 * sign_2) + + for t in range(u_terms.shape[0]): + lorb = int(u_terms[t, 0]) + korb = int(u_terms[t, 1]) + jorb = int(u_terms[t, 2]) + iorb = int(u_terms[t, 3]) + if iorb == jorb or korb == lorb: + continue + + bit_i = np.uint64(1) << np.uint64(norbs - 1 - iorb) + if state0 & bit_i == 0: + continue + sign_1 = sign_count_u64(state0, iorb, norbs) + state = state0 ^ bit_i + + bit_j = np.uint64(1) << np.uint64(norbs - 1 - jorb) + if state & bit_j == 0: + continue + sign_2 = sign_count_u64(state, jorb, norbs) + state ^= bit_j + + bit_k = np.uint64(1) << np.uint64(norbs - 1 - korb) + if state & bit_k != 0: + continue + sign_3 = sign_count_u64(state, korb, norbs) + state |= bit_k + + bit_l = np.uint64(1) << np.uint64(norbs - 1 - lorb) + if state & bit_l != 0: + continue + sign_4 = sign_count_u64(state, lorb, norbs) + state |= bit_l + + row = encode_combinadic_jit( + state, lb_norbs, lb_noccus, lb_offsets, lb_sizes, lb_strides + ) + if row != -1: + rows.append(row) + cols.append(column) + vals.append(u_vals[t] * sign_1 * sign_2 * sign_3 * sign_4) + + return ( + np.asarray(rows, dtype=np.int64), + np.asarray(cols, dtype=np.int64), + np.asarray(vals, dtype=np.complex128), + ) + + @njit + def build_explicit( + rb_states, lb_lookup, norbs, + e_terms, e_vals, u_terms, u_vals, cstart, cend, + ): + rows = [] + cols = [] + vals = [] + + for column in range(cstart, cend): + state0 = rb_states[column] + + for t in range(e_terms.shape[0]): + iorb = int(e_terms[t, 0]) + jorb = int(e_terms[t, 1]) + bit_j = np.uint64(1) << np.uint64(norbs - 1 - jorb) + if state0 & bit_j == 0: + continue + sign_1 = sign_count_u64(state0, jorb, norbs) + state = state0 ^ bit_j + + bit_i = np.uint64(1) << np.uint64(norbs - 1 - iorb) + if state & bit_i != 0: + continue + sign_2 = sign_count_u64(state, iorb, norbs) + state |= bit_i + + if state in lb_lookup: + row = lb_lookup[state] + rows.append(row) + cols.append(column) + vals.append(e_vals[t] * sign_1 * sign_2) + + for t in range(u_terms.shape[0]): + lorb = int(u_terms[t, 0]) + korb = int(u_terms[t, 1]) + jorb = int(u_terms[t, 2]) + iorb = int(u_terms[t, 3]) + if iorb == jorb or korb == lorb: + continue + + bit_i = np.uint64(1) << np.uint64(norbs - 1 - iorb) + if state0 & bit_i == 0: + continue + sign_1 = sign_count_u64(state0, iorb, norbs) + state = state0 ^ bit_i + + bit_j = np.uint64(1) << np.uint64(norbs - 1 - jorb) + if state & bit_j == 0: + continue + sign_2 = sign_count_u64(state, jorb, norbs) + state ^= bit_j + + bit_k = np.uint64(1) << np.uint64(norbs - 1 - korb) + if state & bit_k != 0: + continue + sign_3 = sign_count_u64(state, korb, norbs) + state |= bit_k + + bit_l = np.uint64(1) << np.uint64(norbs - 1 - lorb) + if state & bit_l != 0: + continue + sign_4 = sign_count_u64(state, lorb, norbs) + state |= bit_l + + if state in lb_lookup: + row = lb_lookup[state] + rows.append(row) + cols.append(column) + vals.append(u_vals[t] * sign_1 * sign_2 * sign_3 * sign_4) + + return ( + np.asarray(rows, dtype=np.int64), + np.asarray(cols, dtype=np.int64), + np.asarray(vals, dtype=np.complex128), + ) + + _NUMBA_KERNELS = build_combinadic, build_explicit + return _NUMBA_KERNELS + + +def _prepare_entries_kernel( + lb, rb, e_terms, e_vals, u_terms, u_vals, *, use_numba=False): + """Prepare a reusable entry kernel for backend-selected column ranges.""" + if not use_numba: + def build_range(cstart, cend): + return _build_entries_python( + lb, rb, e_terms, e_vals, u_terms, u_vals, cstart, cend + ) + + return build_range + + if lb.norbs > 64 or rb.norbs > 64: + raise ValueError("Numba operator construction supports at most 64 orbitals") + + build_combinadic, build_explicit = _get_numba_kernels() + + if isinstance(lb, FockBinByN) and isinstance(rb, FockBinByN): + rb_meta = rb.jit_args() + lb_meta = lb.jit_args() + + def build_range(cstart, cend): + return build_combinadic( + rb_meta[0], rb_meta[1], rb_meta[2], rb_meta[3], rb_meta[4], + lb_meta[0], lb_meta[1], lb_meta[2], lb_meta[3], lb_meta[4], + e_terms, e_vals, u_terms, u_vals, cstart, cend, + ) + + return build_range + + if isinstance(lb, FockBasis) and isinstance(rb, FockBasis): + try: + from numba import types + from numba.typed import Dict + except ImportError as exc: + raise ImportError( + "use_numba=True requires numba; install numba or disable JIT construction" + ) from exc + + rb_states = np.asarray(rb.basis_int, dtype=np.uint64) + lb_lookup = Dict.empty(key_type=types.uint64, value_type=types.int64) + for index, state in enumerate(lb.basis_int): + lb_lookup[np.uint64(state)] = np.int64(index) + + def build_range(cstart, cend): + return build_explicit( + rb_states, lb_lookup, rb.norbs, + e_terms, e_vals, u_terms, u_vals, cstart, cend, + ) + + return build_range + + raise TypeError( + "Numba construction requires matching explicit or combinadic basis representations" + ) + + +def prepare_operator_entry_kernel( + emat, umat, lb, rb=None, *, tol_e=1e-10, tol_u=1e-10, + use_numba=False): + """Prepare a reusable kernel that generates entries for requested columns. + + The caller owns the work decomposition and matrix assembly policy. Calling + the returned function as ``build_range(cstart, cend)`` produces only the + contributions for that column interval. + """ + if rb is None: + rb = lb + if lb.norbs != rb.norbs: + raise ValueError("left and right Fock bases must have the same norbs") + + e_terms, e_vals, u_terms, u_vals = _operator_terms( + emat, umat, rb.norbs, tol_e, tol_u + ) + return _prepare_entries_kernel( + lb, rb, e_terms, e_vals, u_terms, u_vals, use_numba=use_numba + ) + + +def build_operator_entries( + emat, umat, lb, rb=None, *, tol_e=1e-10, tol_u=1e-10, + cstart=0, cend=None, use_numba=False): + """Return entries for one requested column range. + + This compatibility wrapper leaves the work-range choice with the caller; + backends that need repeated bounded ranges should prepare the reusable + builder with :func:`prepare_operator_entry_kernel` instead. + """ + if rb is None: + rb = lb + if cend is None: + cend = len(rb) + + build_range = prepare_operator_entry_kernel( + emat, umat, lb, rb, tol_e=tol_e, tol_u=tol_u, use_numba=use_numba + ) + return build_range(cstart, cend) diff --git a/edrixs/fock_basis.py b/edrixs/fock_basis.py index cf471693..a4ad8bf6 100755 --- a/edrixs/fock_basis.py +++ b/edrixs/fock_basis.py @@ -3,11 +3,15 @@ from __future__ import annotations +from dataclasses import dataclass, field import itertools +from math import comb, prod + import numpy as np __all__ = [ - 'FockBasis', 'get_fock_basis_int', + 'FockBasisSpec', 'FockBasis', 'FockBinByN', + 'build_fock_basis', 'get_fock_basis_int', 'get_fock_basis_combinadic', 'fock_bin', 'get_fock_bin_by_N', 'get_fock_half_N', 'get_fock_full_N', 'get_fock_basis_by_NLz', 'get_fock_basis_by_NSz', 'get_fock_basis_by_NJz', 'get_fock_basis_by_N_abelian', @@ -15,6 +19,32 @@ ] +@dataclass(frozen=True, slots=True) +class FockBasisSpec: + """Compact description of fixed occupancies in one or more orbital shells.""" + + shapes: tuple[tuple[int, int], ...] + + def __post_init__(self): + shapes = tuple((int(norb), int(nocc)) for norb, nocc in self.shapes) + object.__setattr__(self, 'shapes', shapes) + + @classmethod + def from_args(cls, *args): + """Build a specification from ``(norbs, nocc)`` pairs.""" + if len(args) % 2 != 0: + raise ValueError("number of basis arguments must be even") + return cls(tuple(zip(args[0::2], args[1::2]))) + + @property + def norbs(self): + """Total number of spin-orbitals.""" + return sum(norb for norb, _ in self.shapes) + + def __len__(self): + return prod(comb(norb, nocc) for norb, nocc in self.shapes) + + class FockBasis: """ Integer-encoded Fock basis with constant-time state lookup. @@ -28,11 +58,15 @@ class FockBasis: Integer encodings of the basis states in matrix-index order. norbs : int Number of spin-orbitals represented by each state. + spec : FockBasisSpec or None, optional + Structured sector description when the basis was generated from fixed + shell occupancies. Arbitrary explicit bases may leave this as ``None``. """ - def __init__(self, basis_int, norbs): + def __init__(self, basis_int, norbs, spec=None): self.basis_int = [int(value) for value in basis_int] self.norbs = int(norbs) + self.spec = spec self.lookup = {value: index for index, value in enumerate(self.basis_int)} def __len__(self): @@ -47,30 +81,211 @@ def decode(self, position): return self.basis_int[position] +@dataclass(slots=True) +class FockBinByN: + """Implicit combinadic basis for a complete fixed-occupancy sector.""" + + shapes: tuple[tuple[int, int], ...] + spec: FockBasisSpec = field(init=False) + shell_norbs: tuple[int, ...] = field(init=False) + norbs: int = field(init=False) + noccus: tuple[int, ...] = field(init=False) + sizes: tuple[int, ...] = field(init=False) + num_subspaces: int = field(init=False) + num_orbitals: int = field(init=False) + dim: int = field(init=False) + offsets: tuple[int, ...] = field(init=False) + strides: tuple[int, ...] = field(init=False) + min_decode: int = field(init=False) + max_decode: int = field(init=False) + + def __post_init__(self): + self.spec = FockBasisSpec(tuple(self.shapes)) + self.shapes = self.spec.shapes + self.shell_norbs = tuple(norb for norb, _ in self.shapes) + self.norbs = self.spec.norbs + self.noccus = tuple(nocc for _, nocc in self.shapes) + self.sizes = tuple(comb(norb, nocc) for norb, nocc in self.shapes) + self.num_subspaces = len(self.shapes) + self.num_orbitals = self.norbs + self.dim = prod(self.sizes) + + running = self.norbs + offsets = [] + for norb in self.shell_norbs: + running -= norb + offsets.append(running) + self.offsets = tuple(offsets) + + # The first shell varies fastest, matching get_fock_bin_by_N. + stride = 1 + strides = [] + for size in self.sizes: + strides.append(stride) + stride *= size + self.strides = tuple(strides) + self.min_decode, self.max_decode = _min_max_decode(self.shapes) + + @classmethod + def from_spec(cls, spec): + """Build an implicit basis from a :class:`FockBasisSpec`.""" + return cls(spec.shapes) + + def __len__(self): + return self.dim + + def encode(self, state): + """Return the canonical EDRIXS index of ``state``.""" + index = _encode_combinadic( + int(state), self.shell_norbs, self.noccus, self.offsets, + self.sizes, self.strides, + ) + if index < 0: + raise KeyError(int(state)) + return index + + def decode(self, index): + """Return the integer-encoded state at the canonical EDRIXS index.""" + if index < 0: + index += self.dim + if index < 0 or index >= self.dim: + raise IndexError(index) + return _decode_combinadic( + int(index), self.shell_norbs, self.noccus, self.offsets, self.sizes + ) + + def jit_args(self): + """Return compact numeric metadata for an explicitly requested JIT path.""" + return ( + np.asarray(self.shell_norbs, dtype=np.int64), + np.asarray(self.noccus, dtype=np.int64), + np.asarray(self.offsets, dtype=np.int64), + np.asarray(self.sizes, dtype=np.int64), + np.asarray(self.strides, dtype=np.int64), + np.uint64(self.min_decode), + np.uint64(self.max_decode), + ) + + +def _hash_decoder(rank, norb, nocc): + """Decode a conventional colex combinadic rank inside one shell.""" + if nocc == 0: + return 0 + state = 0 + j = nocc + i = norb - 1 + c = comb(i, j) + while i >= 0 and j > 0: + if c <= rank: + state |= 1 << i + rank -= c + old_i, old_j = i, j + i -= 1 + j -= 1 + if j == 0 or i < 0: + break + c = (c * old_j) // old_i + else: + if i == 0: + break + c = (c * (i - j)) // i + i -= 1 + return state + + +def _hash_encoder(state, norb): + """Encode one shell into its conventional colex combinadic rank.""" + state = int(state) & ((1 << norb) - 1) + rank = 0 + k = 1 + while state: + lsb = state & -state + pos = lsb.bit_length() - 1 + rank += comb(pos, k) + k += 1 + state ^= lsb + return rank + + +def _decode_combinadic(index, norbs, noccus, offsets, sizes): + """Decode using the historical EDRIXS shell and state ordering.""" + state = 0 + for norb, nocc, offset, size in zip(norbs, noccus, offsets, sizes): + shell_rank = index % size + index //= size + colex_rank = size - 1 - shell_rank + state |= _hash_decoder(colex_rank, norb, nocc) << offset + return state + + +def _encode_combinadic(state, norbs, noccus, offsets, sizes, strides): + """Encode using the historical EDRIXS shell and state ordering.""" + index = 0 + for norb, nocc, offset, size, stride in zip( + norbs, noccus, offsets, sizes, strides): + mask = (1 << norb) - 1 + shell_state = (state >> offset) & mask + if shell_state.bit_count() != nocc: + return -1 + shell_rank = size - 1 - _hash_encoder(shell_state, norb) + index += shell_rank * stride + return index + + +def _min_max_decode(shapes): + total_norb = sum(norb for norb, _ in shapes) + state_min = 0 + state_max = 0 + running = total_norb + for norb, nocc in shapes: + running -= norb + shell_min = (1 << nocc) - 1 + shell_max = shell_min << (norb - nocc) + state_min |= shell_min << running + state_max |= shell_max << running + return state_min, state_max + + def get_fock_basis_int(*args): """ - Build an integer-encoded :class:`FockBasis` for fixed shell occupancies. + Build an explicit :class:`FockBasis` for fixed shell occupancies. - Parameters - ---------- - args : ints - ``(number_of_orbitals, occupancy)`` pairs, with the same shell ordering - convention as :func:`get_fock_bin_by_N`. - - Returns - ------- - FockBasis or None - Integer-encoded basis, or ``None`` when an odd number of arguments is - supplied. + The state ordering is the historical EDRIXS ordering produced by + :func:`get_fock_bin_by_N`. """ - basis_binary = get_fock_bin_by_N(*args) - if basis_binary is None: + if len(args) % 2 != 0: + print("Error: number of arguments is not even") return None + spec = FockBasisSpec.from_args(*args) + basis_binary = get_fock_bin_by_N(*args) basis_int = np.asarray( [int(''.join(map(str, row)), 2) for row in basis_binary], dtype=object, ) - return FockBasis(basis_int, sum(args[0::2])) + return FockBasis(basis_int, spec.norbs, spec=spec) + + +def get_fock_basis_combinadic(*args): + """Build an implicit combinadic basis for fixed shell occupancies.""" + if len(args) % 2 != 0: + print("Error: number of arguments is not even") + return None + return FockBinByN.from_spec(FockBasisSpec.from_args(*args)) + + +def build_fock_basis(basis, method='combinadic'): + """Realize a compact basis specification with the requested representation.""" + if isinstance(basis, (FockBasis, FockBinByN)): + return basis + if not isinstance(basis, FockBasisSpec): + raise TypeError("basis must be a FockBasisSpec or a realized Fock basis") + + if method == 'combinadic': + return FockBinByN.from_spec(basis) + if method == 'explicit': + args = tuple(value for shape in basis.shapes for value in shape) + return get_fock_basis_int(*args) + raise ValueError("basis method must be 'combinadic' or 'explicit'") def fock_bin(n, k): diff --git a/edrixs/models.py b/edrixs/models.py index 29d49d87..556b0ab1 100644 --- a/edrixs/models.py +++ b/edrixs/models.py @@ -1,8 +1,8 @@ """Backend-independent physical model construction for EDRIXS. The model functions in this module describe orbital-space Hamiltonians, -interactions, Fock bases, and photon-transition matrices without constructing -backend-owned many-body operators. +interactions, Fock-basis specifications, and photon-transition matrices without +constructing backend-owned many-body operators. """ from __future__ import annotations @@ -13,7 +13,7 @@ from .angular_momentum import get_wigner_dmat, rmat_to_euler from .basis_transform import cb_op, tmat_r2c from .coulomb_utensor import get_umat_slater, get_umat_slater_3shells -from .fock_basis import get_fock_basis_int +from .fock_basis import FockBasisSpec from .iostream import write_emat, write_umat from .photon_transition import get_trans_oper from .soc import atom_hsoc @@ -34,12 +34,12 @@ def model_1v1c(shell_name, *, shell_level=None, v_soc=None, c_soc=0, v_cfmat=None, v_othermat=None, loc_axis=None, verbose=0, sparse_U=False, tol=1E-10): """ - Set up orbital-space data and Fock bases for a 1v1c problem. + Set up orbital-space data and Fock-basis metadata for a 1v1c problem. This routine defines the physical one-valence-shell/one-core-shell problem independently of the numerical backend. It constructs the one-body orbital - matrices, Coulomb tensors, Fock bases, and orbital-space transition - matrices, but it does not build many-body Hamiltonians and does not + matrices, Coulomb tensors, Fock-basis specifications, and orbital-space + transition matrices, but it does not build many-body Hamiltonians and does not diagonalize anything. Parameters @@ -63,8 +63,8 @@ def model_1v1c(shell_name, *, shell_level=None, v_soc=None, c_soc=0, Returns ------- emat_i, umat_i, basis_i, emat_n, umat_n, basis_n, trans_mat - Backend-independent problem definition. trans_mat has shape - (npol, ntot, ntot). + Backend-independent problem definition. ``basis_i`` and ``basis_n`` are + compact Fock-basis specifications; trans_mat has shape (npol, ntot, ntot). """ print("edrixs >>> Setting up 1v1c problem ...") @@ -187,9 +187,9 @@ def model_1v1c(shell_name, *, shell_level=None, v_soc=None, c_soc=0, write_emat(emat_i, 'hopping_i.in') write_emat(emat_n, 'hopping_n.in') - # Fock bases. - basis_i = get_fock_basis_int(v_norb, v_noccu, c_norb, c_norb) - basis_n = get_fock_basis_int(v_norb, v_noccu + 1, c_norb, c_norb - 1) + # Fock-basis metadata. + basis_i = FockBasisSpec.from_args(v_norb, v_noccu, c_norb, c_norb) + basis_n = FockBasisSpec.from_args(v_norb, v_noccu + 1, c_norb, c_norb - 1) print("edrixs >>> Dimension of the initial Hamiltonian: ", len(basis_i)) print("edrixs >>> Dimension of the intermediate Hamiltonian: ", len(basis_n)) @@ -239,7 +239,7 @@ def model_2v1c( trans_to_which=1, loc_axis=None, verbose=0, sparse_U=False, tol=1E-10 ): """ - Set up orbital-space data and Fock bases for a 2v1c problem. + Set up orbital-space data and Fock-basis metadata for a 2v1c problem. This is the backend-neutral setup analogue of the 2-valence-shell, 1-core-shell Fortran ED/RIXS input construction. It does not build @@ -250,8 +250,7 @@ def model_2v1c( Returns ------- emat_i, umat_i, basis_i, emat_n, umat_n, basis_n, trans_mat - These can be passed directly to get_ops(..., backend='scipy') or - get_ops(..., backend='dense'). + These can be passed directly to ``get_ops`` with the SciPy or dense backend. """ print("edrixs >>> Setting up 2v1c problem ...") @@ -427,10 +426,10 @@ def model_2v1c( write_emat(emat_i, 'hopping_i.in') write_emat(emat_n, 'hopping_n.in') - basis_i = get_fock_basis_int( + basis_i = FockBasisSpec.from_args( v1v2_norb, v_tot_noccu, c_norb, c_norb ) - basis_n = get_fock_basis_int( + basis_n = FockBasisSpec.from_args( v1v2_norb, v_tot_noccu + 1, c_norb, c_norb - 1 ) @@ -463,7 +462,7 @@ def model_siam( on_which='spin', loc_axis=None, verbose=0, sparse_U=False, tol=1E-10 ): """ - Set up orbital-space data and Fock bases for a SIAM problem. + Set up orbital-space data and Fock-basis metadata for a SIAM problem. This is the backend-neutral setup analogue of ed_siam_fort. It does not search over occupancies, does not build many-body Hamiltonians, and does not @@ -474,8 +473,7 @@ def model_siam( Returns ------- emat_i, umat_i, basis_i, emat_n, umat_n, basis_n, trans_mat - These can be passed directly to get_ops(..., backend='scipy') or - get_ops(..., backend='dense'). + These can be passed directly to ``get_ops`` with the SciPy or dense backend. """ print("edrixs >>> Setting up SIAM problem ...") @@ -650,8 +648,8 @@ def model_siam( write_emat(emat_i, 'hopping_i.in') write_emat(emat_n, 'hopping_n.in') - basis_i = get_fock_basis_int(ntot_v, v_noccu, c_norb, c_norb) - basis_n = get_fock_basis_int(ntot_v, v_noccu + 1, c_norb, c_norb - 1) + basis_i = FockBasisSpec.from_args(ntot_v, v_noccu, c_norb, c_norb) + basis_n = FockBasisSpec.from_args(ntot_v, v_noccu + 1, c_norb, c_norb - 1) print("edrixs >>> Dimension of the initial Hamiltonian: ", len(basis_i)) print("edrixs >>> Dimension of the intermediate Hamiltonian: ", len(basis_n)) diff --git a/edrixs/petsc_backend/hash_basis_methods.py b/edrixs/petsc_backend/hash_basis_methods.py new file mode 100644 index 00000000..e77202d4 --- /dev/null +++ b/edrixs/petsc_backend/hash_basis_methods.py @@ -0,0 +1,28 @@ +"""Compatibility names for the unimplemented PETSc backend. + +Fock-basis representations and ranking are backend independent and live in +:mod:`edrixs.fock_basis`. ``scipy_edrixs`` intentionally keeps PETSc matrix +construction as a stub, so this module exposes only the historical basis names +without introducing a second PETSc implementation. +""" + +from __future__ import annotations + +from ..fock_basis import FockBinByN, get_fock_basis_combinadic + +__all__ = [ + 'FockBinByN', 'get_fock_basis_petsc', 'build_op_petsc_matrix', +] + + +# Backwards-compatible name used by the PETSc branch. The representation +# itself is backend independent and is shared with SciPy. +get_fock_basis_petsc = get_fock_basis_combinadic + + +def build_op_petsc_matrix(*args, **kwargs): + """Raise the standard error for unavailable PETSc matrix construction.""" + raise NotImplementedError( + "The PETSc backend contract is present, but build_op_petsc has not yet " + "been implemented" + ) diff --git a/edrixs/petsc_backend/petsc_backend.py b/edrixs/petsc_backend/petsc_backend.py index a1f6a541..f0d92042 100644 --- a/edrixs/petsc_backend/petsc_backend.py +++ b/edrixs/petsc_backend/petsc_backend.py @@ -48,7 +48,8 @@ def _not_implemented(operation): ) -def build_op_petsc(emat, umat, lb, rb=None, *, backend_kws=None): +def build_op_petsc( + emat, umat, lb, rb=None, *, use_numba=False, backend_kws=None): """Build a PETSc many-body operator (stub).""" _petsc_module() _not_implemented('build_op_petsc') diff --git a/edrixs/scipy_backend/hash_basis_methods.py b/edrixs/scipy_backend/hash_basis_methods.py new file mode 100644 index 00000000..324520e6 --- /dev/null +++ b/edrixs/scipy_backend/hash_basis_methods.py @@ -0,0 +1,46 @@ +"""SciPy-specific many-body operator construction helpers. + +Fock-basis representations and ranking live in :mod:`edrixs.fock_basis`. +This module owns SciPy sparse-matrix materialization. The shared operator code +is used only as a matrix-entry kernel over backend-selected basis ranges. +""" + +from __future__ import annotations + +import numpy as np +import scipy.sparse as sp + +from .._operator_builder import prepare_operator_entry_kernel + +__all__ = ['build_op_scipy_matrix'] + + +def _assemble_csr(rows, cols, data, shape): + """Materialize the final SciPy sparse operator using the current COO path.""" + return sp.coo_matrix( + (data, (rows, cols)), + shape=shape, + dtype=np.complex128, + ).tocsr() + + +def build_op_scipy_matrix( + emat, umat, lb, rb=None, *, tol=1e-10, use_numba=False): + """Build the final SciPy CSR matrix using the SciPy assembly policy.""" + if rb is None: + rb = lb + + build_range = prepare_operator_entry_kernel( + emat, + umat, + lb, + rb, + tol_e=tol, + tol_u=tol, + use_numba=use_numba, + ) + + # SciPy currently materializes the complete COO contribution arrays before + # converting them to CSR. This policy is intentionally backend-local. + rows, cols, data = build_range(0, len(rb)) + return _assemble_csr(rows, cols, data, (len(lb), len(rb))) diff --git a/edrixs/scipy_backend/scipy_backend.py b/edrixs/scipy_backend/scipy_backend.py index fdea5be0..d6d8e579 100644 --- a/edrixs/scipy_backend/scipy_backend.py +++ b/edrixs/scipy_backend/scipy_backend.py @@ -287,8 +287,10 @@ def four_fermion_csr_auto(umat, basis, right_basis=None, tol=1e-10): ) -def build_op_scipy(emat, umat, lb, rb=None, *, backend_kws=None): +def build_op_scipy(emat, umat, lb, rb=None, *, use_numba=False, backend_kws=None): """Build and return a SciPy CSR many-body operator.""" + from .hash_basis_methods import build_op_scipy_matrix + kws = _backend_kws(backend_kws) tol = kws.pop('tol', 1e-10) if kws: @@ -296,30 +298,15 @@ def build_op_scipy(emat, umat, lb, rb=None, *, backend_kws=None): sorted(kws) )) - if rb is None: - rb = lb - if lb.norbs != rb.norbs: - raise ValueError("left and right Fock bases must have the same norbs") - - operator = sp.csr_matrix( - (len(lb), len(rb)), - dtype=np.complex128, + return build_op_scipy_matrix( + emat, umat, lb, rb, tol=tol, use_numba=use_numba ) - if emat is not None: - operator = operator + two_fermion_csr( - emat, lb, rb, tol=tol - ) - if umat is not None: - operator = operator + four_fermion_csr_auto( - umat, lb, right_basis=rb, tol=tol - ) - return operator.tocsr() -def build_op_dense(emat, umat, lb, rb=None, *, backend_kws=None): +def build_op_dense(emat, umat, lb, rb=None, *, use_numba=False, backend_kws=None): """Compatibility dense constructor implemented through SciPy CSR.""" return build_op_scipy( - emat, umat, lb, rb, backend_kws=backend_kws + emat, umat, lb, rb, use_numba=use_numba, backend_kws=backend_kws ).toarray() diff --git a/edrixs/solvers.py b/edrixs/solvers.py index 31503803..95843c42 100644 --- a/edrixs/solvers.py +++ b/edrixs/solvers.py @@ -24,7 +24,7 @@ ) from .coulomb_utensor import get_umat_slater from .manybody_operator import two_fermion, four_fermion -from .fock_basis import get_fock_bin_by_N, write_fock_dec_by_N +from .fock_basis import build_fock_basis, get_fock_bin_by_N, write_fock_dec_by_N from .basis_transform import cb_op2, tmat_r2c, cb_op from .utils import info_atomic_shell, slater_integrals_name, boltz_dist from .rixs_utils import scattering_mat @@ -58,7 +58,8 @@ # ----------------------------------------------------------------------------- -def build_op(emat, umat, lb, rb=None, *, backend='scipy', backend_kws=None): +def build_op(emat, umat, lb, rb=None, *, backend='scipy', + basis_method='combinadic', use_numba=False, backend_kws=None): """ Build a many-body operator with the selected backend. @@ -70,13 +71,19 @@ def build_op(emat, umat, lb, rb=None, *, backend='scipy', backend_kws=None): umat : array-like, sparse matrix, or None Coefficients of the two-body part. Pass ``None`` when the operator has no two-body contribution. - lb : FockBasis - Basis for the output (left) many-body space. - rb : FockBasis or None, optional - Basis for the input (right) many-body space. When omitted, ``lb`` is + lb : FockBasisSpec or FockBasis + Basis metadata or realized basis for the output (left) many-body space. + rb : FockBasisSpec, FockBasis, or None, optional + Basis metadata or realized basis for the input (right) many-body space. + When omitted, ``lb`` is used for both sides. backend : str, optional Backend name. The default is ``'scipy'``. + basis_method : {'combinadic', 'explicit'}, optional + Representation used when ``lb``/``rb`` are compact basis specifications. + The default is the implicit combinadic representation. + use_numba : bool, optional + JIT-compile matrix-entry construction. The default is False. backend_kws : mapping, optional Backend-specific construction options. For the SciPy backend this includes ``tol``. @@ -99,18 +106,24 @@ def build_op(emat, umat, lb, rb=None, *, backend='scipy', backend_kws=None): "'petsc'".format(backend) ) + lb = build_fock_basis(lb, method=basis_method) + if rb is not None: + rb = build_fock_basis(rb, method=basis_method) + return build_op_backend( emat, umat, lb, rb, + use_numba=use_numba, backend_kws=backend_kws, ) def get_ops( emat_i, umat_i, basis_i, emat_n, umat_n, basis_n, trans_mat, *, - backend='scipy', backend_kws=None, + backend='scipy', basis_method='combinadic', use_numba=False, + backend_kws=None, ): """ Build initial/intermediate Hamiltonians and transition operators. @@ -120,8 +133,13 @@ def get_ops( emat_i, umat_i, basis_i, emat_n, umat_n, basis_n, trans_mat Backend-neutral problem definition returned by :mod:`edrixs.models` model functions. - backend : {'scipy', 'dense'}, optional + backend : {'scipy', 'dense', 'petsc'}, optional Backend used for the returned operators. The default is ``'scipy'``. + basis_method : {'combinadic', 'explicit'}, optional + Basis representation constructed from the model metadata. The default + is ``'combinadic'``. + use_numba : bool, optional + JIT-compile matrix-entry construction. The default is False. backend_kws : mapping, optional Backend-specific operator-construction options. For the SciPy and dense compatibility backends this includes ``tol``. @@ -132,13 +150,16 @@ def get_ops( Initial/final Hamiltonian, intermediate Hamiltonian, and transition operators for the selected backend. """ + basis_i = build_fock_basis(basis_i, method=basis_method) + basis_n = build_fock_basis(basis_n, method=basis_method) + hmat_i = build_op( emat_i, umat_i, basis_i, backend=backend, - backend_kws=backend_kws, + use_numba=use_numba, backend_kws=backend_kws, ) hmat_n = build_op( emat_n, umat_n, basis_n, backend=backend, - backend_kws=backend_kws, + use_numba=use_numba, backend_kws=backend_kws, ) trans_mat = np.asarray(trans_mat) @@ -151,7 +172,7 @@ def get_ops( basis_n, basis_i, backend=backend, - backend_kws=backend_kws, + use_numba=use_numba, backend_kws=backend_kws, ) for component in trans_mat ] diff --git a/tests/solver_consistency_tests/_helpers.py b/tests/solver_consistency_tests/_helpers.py index 8a54f1be..9ebdb1e2 100644 --- a/tests/solver_consistency_tests/_helpers.py +++ b/tests/solver_consistency_tests/_helpers.py @@ -4,6 +4,7 @@ import scipy.sparse as sp from numpy.testing import assert_allclose +from edrixs.fock_basis import FockBasisSpec from edrixs.solvers import get_ops @@ -40,8 +41,10 @@ def assert_problem_sparse_dense_equivalent(dense, sparse): """Compare backend-neutral setup outputs in dense and sparse-U forms.""" assert_allclose(dense[0], sparse[0]) assert_allclose(dense[3], sparse[3]) - assert dense[2].basis_int == sparse[2].basis_int - assert dense[5].basis_int == sparse[5].basis_int + assert isinstance(dense[2], FockBasisSpec) + assert isinstance(dense[5], FockBasisSpec) + assert dense[2] == sparse[2] + assert dense[5] == sparse[5] assert_allclose(dense[6], sparse[6]) for dense_u, sparse_u in ((dense[1], sparse[1]), (dense[4], sparse[4])): diff --git a/tests/solver_consistency_tests/test_setup_and_operator_consistency.py b/tests/solver_consistency_tests/test_setup_and_operator_consistency.py index 48a94ae3..12d8c6d9 100644 --- a/tests/solver_consistency_tests/test_setup_and_operator_consistency.py +++ b/tests/solver_consistency_tests/test_setup_and_operator_consistency.py @@ -1,7 +1,7 @@ """Consistency checks for model definition and many-body operator construction. These checks follow the command chain from ``model_*`` through ``get_ops`` but stop -before any eigensolver or spectral solver. They compare alternative dense and +before any eigensolver or spectral solver. They compare alternative dense and SciPy sparse representations of the same physical model. """ @@ -10,6 +10,7 @@ import scipy.sparse as sp from numpy.testing import assert_allclose +from edrixs.fock_basis import FockBasisSpec from edrixs.models import model_1v1c, model_2v1c, model_siam from edrixs.solvers import get_ops @@ -25,12 +26,7 @@ @pytest.mark.parametrize("trans_to_which", [1, 2]) def test_model_2v1c_sparse_dense_and_get_ops_equivalence(trans_to_which): - """Compare both representations through the 2v1c model-to-``get_ops`` chain. - - The check builds one model with dense and flattened sparse interactions, - verifies the same orbital data and transition target, and confirms that - dense and SciPy many-body operators act identically before spectroscopy. - """ + """Compare both representations through the 2v1c model-to-``get_ops`` chain.""" hopping = np.array( [ [0.04 + 0.02j, -0.07j], @@ -77,12 +73,7 @@ def test_model_2v1c_sparse_dense_and_get_ops_equivalence(trans_to_which): @pytest.mark.parametrize("siam_type", [0, 1]) def test_model_siam_sparse_dense_and_get_ops_equivalence(siam_type): - """Compare both representations through the SIAM model-to-``get_ops`` chain. - - The check covers both supported SIAM input forms, confirms identical dense - and sparse model data and operator actions, and verifies that the photon - transition reaches the impurity but not bath orbitals. - """ + """Compare both representations through the SIAM model-to-``get_ops`` chain.""" kwargs = siam_kwargs(siam_type) dense = model_siam(**kwargs, sparse_U=False) sparse = model_siam(**kwargs, sparse_U=True) @@ -100,12 +91,7 @@ def test_model_siam_sparse_dense_and_get_ops_equivalence(siam_type): def test_model_siam_static_core_potential_only_shifts_intermediate_impurity(): - """Check the SIAM core-hole term before operator construction and spectra. - - The static core potential should leave the initial orbital Hamiltonian - unchanged and shift only the impurity block of the intermediate model that - later enters XAS and the RIXS correction-vector solve. - """ + """The static core potential only shifts the intermediate impurity block.""" base = model_siam( ("s", "p"), 1, @@ -130,12 +116,7 @@ def test_model_siam_static_core_potential_only_shifts_intermediate_impurity(): def test_model_1v1c_sparse_u_matches_dense_u(small_1v1c_kwargs): - """Compare dense and sparse interaction outputs from ``model_1v1c``. - - This check verifies the first stage of the 1v1c SciPy path: model definition - must be unchanged when Coulomb tensors are stored sparsely for the later - ``get_ops(..., backend="scipy")`` construction. - """ + """Dense and sparse-U model outputs must retain identical basis metadata.""" dense = model_1v1c(**small_1v1c_kwargs, sparse_U=False) sparse = model_1v1c(**small_1v1c_kwargs, sparse_U=True) @@ -147,20 +128,17 @@ def test_model_1v1c_sparse_u_matches_dense_u(small_1v1c_kwargs): dense_u.reshape(norbs * norbs, norbs * norbs), ) - assert dense[2].basis_int == sparse[2].basis_int - assert dense[5].basis_int == sparse[5].basis_int + assert isinstance(dense[2], FockBasisSpec) + assert isinstance(dense[5], FockBasisSpec) + assert dense[2] == sparse[2] + assert dense[5] == sparse[5] assert_allclose(dense[0], sparse[0]) assert_allclose(dense[3], sparse[3]) assert_allclose(dense[6], sparse[6]) def test_ops_scipy_backend_matches_dense_operator_action(small_1v1c_problem): - """Compare dense and SciPy operators at the setup-to-solver boundary. - - Starting from one backend-neutral 1v1c model, this check confirms that the - two ``get_ops`` backends produce Hamiltonian and transition actions that agree - before either representation is used by ED, XAS, or RIXS. - """ + """Compare dense and SciPy operators at the setup-to-solver boundary.""" hmat_i_sp, hmat_n_sp, trans_sp = get_ops( *small_1v1c_problem, backend="scipy", diff --git a/tests/test_basis_representation_routes.py b/tests/test_basis_representation_routes.py new file mode 100644 index 00000000..261781af --- /dev/null +++ b/tests/test_basis_representation_routes.py @@ -0,0 +1,118 @@ +"""Tests for compact Fock-basis metadata and interchangeable realizations.""" + +import builtins + +import numpy as np +import pytest + +import edrixs._operator_builder as operator_builder +from edrixs.fock_basis import ( + FockBasis, + FockBasisSpec, + FockBinByN, + build_fock_basis, + get_fock_bin_by_N, +) +from edrixs.solvers import build_op + + +def _historical_integer_states(*args): + return [ + int("".join(map(str, state)), 2) + for state in get_fock_bin_by_N(*args) + ] + + +@pytest.mark.parametrize( + "args", + [ + (4, 2), + (4, 2, 2, 1), + (2, 1, 3, 2, 2, 1), + (3, 0, 2, 2), + ], +) +def test_combinadic_and_explicit_match_historical_edrixs_ordering(args): + """Both realizations must use the pre-existing EDRIXS state ordering.""" + spec = FockBasisSpec.from_args(*args) + explicit = build_fock_basis(spec, method="explicit") + combinadic = build_fock_basis(spec, method="combinadic") + expected = _historical_integer_states(*args) + + assert isinstance(explicit, FockBasis) + assert isinstance(combinadic, FockBinByN) + assert explicit.basis_int == expected + assert len(explicit) == len(combinadic) == len(expected) + + for index, state in enumerate(expected): + assert explicit.decode(index) == state + assert combinadic.decode(index) == state + assert explicit.encode(state) == index + assert combinadic.encode(state) == index + + +def test_basis_spec_is_compact_sector_metadata(): + """A basis specification stores shell structure without materializing states.""" + spec = FockBasisSpec.from_args(4, 2, 2, 1) + + assert spec.shapes == ((4, 2), (2, 1)) + assert spec.norbs == 6 + assert len(spec) == 12 + assert not hasattr(spec, "basis_int") + assert not hasattr(spec, "lookup") + + +def test_build_fock_basis_defaults_to_combinadic_and_can_select_explicit(): + spec = FockBasisSpec.from_args(4, 2) + + assert isinstance(build_fock_basis(spec), FockBinByN) + assert isinstance(build_fock_basis(spec, method="explicit"), FockBasis) + + +def test_realized_arbitrary_explicit_basis_is_preserved(): + """SciPy-compatible arbitrary explicit subspaces remain valid inputs.""" + arbitrary = FockBasis([0b1100, 0b1001, 0b0011], norbs=4) + + assert build_fock_basis(arbitrary, method="explicit") is arbitrary + assert build_fock_basis(arbitrary, method="combinadic") is arbitrary + + +def test_invalid_basis_method_is_reported(): + with pytest.raises(ValueError, match="combinadic.*explicit"): + build_fock_basis(FockBasisSpec.from_args(4, 2), method="unknown") + + +def test_numba_is_only_required_when_explicitly_requested(monkeypatch): + """The default operator route must work when importing numba is impossible.""" + spec = FockBasisSpec.from_args(2, 1) + emat = np.array([[1.0, 0.25], [0.25, -0.5]], dtype=complex) + + monkeypatch.setattr(operator_builder, "_NUMBA_KERNELS", None) + original_import = builtins.__import__ + + def guarded_import(name, *args, **kwargs): + if name == "numba" or name.startswith("numba."): + raise ImportError("numba intentionally unavailable in this test") + return original_import(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", guarded_import) + + operator = build_op( + emat, + None, + spec, + backend="scipy", + basis_method="combinadic", + use_numba=False, + ) + np.testing.assert_allclose(operator.toarray(), emat) + + with pytest.raises(ImportError, match="use_numba=True requires numba"): + build_op( + emat, + None, + spec, + backend="scipy", + basis_method="combinadic", + use_numba=True, + ) diff --git a/tests/test_model_operator_routes.py b/tests/test_model_operator_routes.py new file mode 100644 index 00000000..8f7f4ec7 --- /dev/null +++ b/tests/test_model_operator_routes.py @@ -0,0 +1,106 @@ +"""Integration tests from model metadata through all SciPy construction routes.""" + +import importlib.util +from itertools import product + +import numpy as np +import pytest +from numpy.testing import assert_allclose + +from edrixs.fock_basis import FockBasisSpec +from edrixs.models import model_1v1c, model_2v1c, model_siam +from edrixs.solvers import get_ops + + +HAS_NUMBA = importlib.util.find_spec("numba") is not None + + +# Public SciPy construction matrix: basis representation x JIT choice. +ALL_ROUTES = [ + pytest.param( + basis_method, + use_numba, + id=f"scipy-{basis_method}-{'numba' if use_numba else 'vanilla'}", + ) + for basis_method, use_numba in product( + ("combinadic", "explicit"), (False, True) + ) +] + + +def _small_1v1c_problem(): + """Return a small nontrivial model shared by all route comparisons.""" + return model_1v1c( + ("p", "s"), + shell_level=(0.2, -4.0), + v_soc=(0.05, 0.08), + v_noccu=1, + slater=([0.7], [0.9]), + v_cfmat=np.diag(np.linspace(-0.12, 0.12, 6)), + sparse_U=False, + ) + + +def _skip_unavailable(use_numba): + """Skip only an explicitly requested optional Numba route when unavailable.""" + if use_numba and not HAS_NUMBA: + pytest.skip("numba is required for the requested JIT route") + + +def test_all_model_functions_return_compact_basis_metadata(): + """No model function should materialize the many-body basis anymore.""" + problems = [ + model_1v1c(("p", "s"), v_noccu=1), + model_2v1c(("s", "s", "p"), v_tot_noccu=1, trans_to_which=1), + model_siam(("s", "p"), 1, v_noccu=1), + ] + + expected_shapes = [ + (((6, 1), (2, 2)), ((6, 2), (2, 1))), + (((4, 1), (6, 6)), ((4, 2), (6, 5))), + (((4, 1), (6, 6)), ((4, 2), (6, 5))), + ] + + for problem, (initial_shapes, intermediate_shapes) in zip(problems, expected_shapes): + basis_i, basis_n = problem[2], problem[5] + assert isinstance(basis_i, FockBasisSpec) + assert isinstance(basis_n, FockBasisSpec) + assert basis_i.shapes == initial_shapes + assert basis_n.shapes == intermediate_shapes + assert not hasattr(basis_i, "basis_int") + assert not hasattr(basis_n, "basis_int") + + +@pytest.mark.integration +@pytest.mark.parametrize("basis_method,use_numba", ALL_ROUTES) +def test_model_to_get_ops_all_scipy_routes_are_numerically_consistent( + basis_method, use_numba +): + """Compare all four public SciPy operator-construction routes numerically. + + The explicit, non-Numba SciPy route is the common reference. Every + combinadic/explicit x vanilla/Numba SciPy route must produce the same + initial Hamiltonian, intermediate Hamiltonian, and transition operators + from the same compact model metadata. + """ + _skip_unavailable(use_numba) + problem = _small_1v1c_problem() + + reference_i, reference_n, reference_t = get_ops( + *problem, + backend="scipy", + basis_method="explicit", + use_numba=False, + ) + actual_i, actual_n, actual_t = get_ops( + *problem, + backend="scipy", + basis_method=basis_method, + use_numba=use_numba, + ) + + assert_allclose(actual_i.toarray(), reference_i.toarray(), rtol=0, atol=2e-12) + assert_allclose(actual_n.toarray(), reference_n.toarray(), rtol=0, atol=2e-12) + assert len(actual_t) == len(reference_t) + for actual, reference in zip(actual_t, reference_t): + assert_allclose(actual.toarray(), reference.toarray(), rtol=0, atol=2e-12) diff --git a/tests/test_operator_route_consistency.py b/tests/test_operator_route_consistency.py new file mode 100644 index 00000000..f32f87df --- /dev/null +++ b/tests/test_operator_route_consistency.py @@ -0,0 +1,168 @@ +"""Numerical consistency tests for explicit/combinadic and vanilla/Numba builders.""" + +import importlib.util +from itertools import product + +import numpy as np +import pytest +import scipy.sparse as sp +from numpy.testing import assert_allclose + +from edrixs.fock_basis import FockBasisSpec, build_fock_basis +from edrixs.solvers import build_op + + +HAS_NUMBA = importlib.util.find_spec("numba") is not None + + +ROUTES = [ + pytest.param(method, use_numba, id=f"{method}-{'numba' if use_numba else 'vanilla'}") + for method, use_numba in product(("combinadic", "explicit"), (False, True)) +] + + +def _annihilation_operators(norbs): + """Independent full-Fock-space Jordan-Wigner annihilation matrices.""" + identity = np.eye(2, dtype=complex) + parity = np.diag([1.0, -1.0]).astype(complex) + annihilate = np.array([[0.0, 1.0], [0.0, 0.0]], dtype=complex) + + result = [] + for orbital in range(norbs): + factors = ( + [parity] * orbital + + [annihilate] + + [identity] * (norbs - orbital - 1) + ) + operator = factors[0] + for factor in factors[1:]: + operator = np.kron(operator, factor) + result.append(operator) + return result + + +def _oracle(emat, umat, left_spec, right_spec=None): + """Build an independent dense Jordan-Wigner reference matrix.""" + if right_spec is None: + right_spec = left_spec + + left = build_fock_basis(left_spec, method="explicit") + right = build_fock_basis(right_spec, method="explicit") + norbs = left.norbs + annihilators = _annihilation_operators(norbs) + full = np.zeros((2**norbs, 2**norbs), dtype=complex) + + if emat is not None: + for iorb in range(norbs): + for jorb in range(norbs): + full += ( + emat[iorb, jorb] + * annihilators[iorb].conj().T + @ annihilators[jorb] + ) + + if umat is not None: + dense_u = umat.toarray().reshape((norbs,) * 4) if sp.issparse(umat) else umat + for lorb, korb, jorb, iorb in zip(*np.nonzero(dense_u)): + full += ( + dense_u[lorb, korb, jorb, iorb] + * annihilators[lorb].conj().T + @ annihilators[korb].conj().T + @ annihilators[jorb] + @ annihilators[iorb] + ) + + return full[np.ix_(left.basis_int, right.basis_int)] + + +def _skip_missing_numba(use_numba): + if use_numba and not HAS_NUMBA: + pytest.skip("numba is required for the requested JIT route") + + +@pytest.mark.parametrize("basis_method,use_numba", ROUTES) +def test_square_one_and_two_body_operator_matches_independent_oracle( + basis_method, use_numba +): + """Every SciPy construction route must produce the same physical matrix.""" + _skip_missing_numba(use_numba) + rng = np.random.default_rng(20260817) + raw = rng.normal(size=(4, 4)) + 1j * rng.normal(size=(4, 4)) + emat = (raw + raw.conj().T) / 2 + + umat = np.zeros((4, 4, 4, 4), dtype=complex) + umat[0, 1, 1, 0] = 0.7 + umat[1, 0, 0, 1] = 0.7 + umat[3, 2, 1, 0] = -0.15 + 0.11j + umat[0, 1, 2, 3] = -0.15 - 0.11j + + spec = FockBasisSpec.from_args(4, 2) + expected = _oracle(emat, umat, spec) + actual = build_op( + emat, + umat, + spec, + backend="scipy", + basis_method=basis_method, + use_numba=use_numba, + ).toarray() + + assert_allclose(actual, expected, rtol=0, atol=2e-13) + + +@pytest.mark.parametrize("basis_method,use_numba", ROUTES) +def test_rectangular_transition_operator_matches_independent_oracle( + basis_method, use_numba +): + """Left/right sector changes must agree for all representation/JIT routes.""" + _skip_missing_numba(use_numba) + right = FockBasisSpec.from_args(2, 1, 2, 2) + left = FockBasisSpec.from_args(2, 2, 2, 1) + emat = np.zeros((4, 4), dtype=complex) + emat[0, 2] = 1.2 - 0.3j + emat[1, 3] = -0.4 + 0.2j + + expected = _oracle(emat, None, left, right) + actual = build_op( + emat, + None, + left, + right, + backend="scipy", + basis_method=basis_method, + use_numba=use_numba, + ).toarray() + + assert actual.shape == expected.shape + assert_allclose(actual, expected, rtol=0, atol=2e-13) + + +@pytest.mark.parametrize("basis_method,use_numba", ROUTES) +def test_dense_and_sparse_coulomb_inputs_are_numerically_identical( + basis_method, use_numba +): + _skip_missing_numba(use_numba) + spec = FockBasisSpec.from_args(4, 2) + umat = np.zeros((4, 4, 4, 4), dtype=complex) + umat[0, 1, 1, 0] = 0.9 + umat[3, 2, 1, 0] = 0.1 - 0.05j + sparse_u = sp.csr_matrix(umat.reshape(16, 16)) + + dense_operator = build_op( + None, + umat, + spec, + backend="scipy", + basis_method=basis_method, + use_numba=use_numba, + ) + sparse_operator = build_op( + None, + sparse_u, + spec, + backend="scipy", + basis_method=basis_method, + use_numba=use_numba, + ) + + assert_allclose(dense_operator.toarray(), sparse_operator.toarray(), atol=2e-13) diff --git a/tests/test_petsc_backend.py b/tests/test_petsc_backend.py new file mode 100644 index 00000000..dcb50657 --- /dev/null +++ b/tests/test_petsc_backend.py @@ -0,0 +1,63 @@ +"""Contract tests for the intentionally unimplemented PETSc backend stub.""" + +import numpy as np +import pytest + +import edrixs.petsc_backend.petsc_backend as backend + + +def test_owns_operator_returns_false_for_non_petsc_objects(): + """Recognition must never raise when given ordinary Python/NumPy objects.""" + assert backend.owns_operator_petsc(np.eye(2)) is False + assert backend.owns_operator_petsc(object()) is False + + +def test_public_symbols_are_exported(): + for name in ( + "owns_operator_petsc", + "build_op_petsc", + "ed_petsc", + "xas_petsc", + "rixs_petsc", + ): + assert name in backend.__all__ + assert hasattr(backend, name) + + +def test_build_op_stub_accepts_new_shared_builder_options(monkeypatch): + """The shared API may pass ``use_numba`` even while PETSc is a stub.""" + monkeypatch.setattr(backend, "_petsc_module", lambda: object()) + with pytest.raises(NotImplementedError, match="build_op_petsc"): + backend.build_op_petsc( + None, + None, + object(), + use_numba=True, + backend_kws={"unused": True}, + ) + + +@pytest.mark.parametrize( + "call,operation", + [ + (lambda: backend.ed_petsc(object()), "ed_petsc"), + ( + lambda: backend.xas_petsc( + np.array([0.0]), [object()], object(), [object()] * 3, + np.array([1.0]) + ), + "xas_petsc", + ), + ( + lambda: backend.rixs_petsc( + np.array([0.0]), [object()], object(), object(), + [object()] * 3, np.array([1.0]), np.array([0.0]) + ), + "rixs_petsc", + ), + ], +) +def test_solver_stubs_remain_unimplemented(monkeypatch, call, operation): + monkeypatch.setattr(backend, "_petsc_module", lambda: object()) + with pytest.raises(NotImplementedError, match=operation): + call()