Improve LinearOpeator performance for FermionOperator

[q-inho]

Summary

Resolve #50

Details

This PR replaces the generic FermionOperator action used by linear_operator. The change aplies when the input object is a FermionOperator that conserves particle number and spin $z$.

Previously, the implementation evaluated

$$y = \sum_t c_t \hat{g}_t x, \qquad \hat{g}_t=\hat{o}_{t,1} \hat{o}_{t,2} \cdots \hat{o}_{t, m_t}.$$

This used to store ladder-operator tuple and called the PySCF creation or annihilation for each elementary action. But this make every matvec repeated the same intermediate sector transversal, determinant admissibility checks, address computations, and ferminoic sign computation.

This implementation moves the determinant combinatorics to LinearOperator consatruction. The Python first applies the same symmetry check. Then removes terms whose orbital labels lie outside the active range $0, \dots, n_{orb}-1$, since such terms act as zero on the selected finite orbital space. The remaining operator is normal ordered and rewritten as

$$\hat{O} = \sum_t \eta_t \hat{m}_{t \alpha} \hat{m}_{t \beta}.$$

The coefficient $\eta_t$ contains the normal-ordering signs and the sign from separating alpha and beta ladder operators into fixed spin blocks.

Each spin block is compiled once into a fixed-sector transition map

$$M_{t\sigma} = \{(u_{t \sigma r}, d_{t \sigma r}, \phi_{t \sigma r})\}_{r=1}^{R_{t \sigma}}.$$

The source address $u_{t \sigma r}$, labels an input determinant in the fixed $n_\sigma$-electron spin sector. The destination address $d_{t \sigma r}$ labels the determinant reached after applying the spin block. The phase $\phi_{t \sigma r} \in \{ \pm 1 \}$ is the product of the elementtary fermionic signs obtained during the reverse traversal of that block. Empty spin blocks are represented as identities

For a mixed alpha-beta term, the packed action on the tensor $X$ is

$$Y_{d_{t \alpha r}, d_{t\beta s}} += \eta_t \phi_{t \alpha r} \phi_{t \beta s} X_{u_{t \alpha r}, u_{t \beta s}}.$$

Alpha-only, beta-only, scalar terms are the corresponding reductions of this update. The Python layer stores scalar coefficients, compressed one-spin transitions, mixed-term coefficients, mixed transition arrays, and per-term index pointers. One-spin transitions with equal source and destination addresses are summed before entering the accumulation kernel. Mixed terms remain factored into alpha and beta maps, so the full tenrsor product list of flattened addresses is not materialized.

Rust is implement for hot accumulation loop. So druing matvec, Python converts the input vector to a contiguous complex array, reshape it to $(D_\alpha, D_\beta)$, allocates an output tensor, and passes both tensors with the packed arrays to Rust. Rust then contracts scalar terms, alpha-only transitions, beta-only transitions, and mixed transitions into the output buffer. The mixed contraction chooses the output spin loop by comparing the number of alpha and beta transitions for the term.

The rmatvec implementation reuses the same packed data and update

$$Y_d += s X_u,$$

and update adjoint

$$Y_u += \bar{s}X_d.$$

The Rust kernel implements this rule by conjugating the stored scalar and swapping source and destination addresses when the reverse flag is set.

Algorithm 1: Staged fixed-sector action for a fermionic operator

Require: $\hat{O} = \sum_{t} c_{t} \hat{g_t}$, active orbital count $n_{orb}$, sector $(n_\alpha, n_\beta)$, amplitues $X$, direction $d \in \{forward,adjoint\}$

Notation: $|u \rangle, |d \rangle$ denote fixed-spin determinants, $\phi \in \{ \pm 1 \}$ is the fermionic sign, and $M_{t \sigma}$ is the transition map for spin block $m_{t\sigma}$. $\mathcal{S}, \mathcal{A}, \mathcal{B}$, and $\mathcal{M}$ denote scalar, alpha-only, beta-only, and mixed alpha-beta terms. The tensor indices are $a$ for alpha addresses and $b$ for beta addresses.

Benchamrk

The benchmark has been implemented compared to fqe.
Note that the benchmark was made in Macbook m4 air, single core.

ffsim

====== ============= =============
--           filling_fraction     
------ ---------------------------
 norb       0.25          0.5     
====== ============= =============
  4     3.27±0.03μs   3.78±0.01μs 
  8      13.0±0.3μs    51.2±0.2μs 
  12      264±2μs     7.61±0.07ms 
  16     19.9±0.7ms    4.31±0.01s 
====== ============= =============

fqe

====== ============= =============
--           filling_fraction     
------ ---------------------------
 norb       0.25          0.5     
====== ============= =============
  4     14.5±0.04ms   15.1±0.06ms 
  8     35.7±0.1ms    43.2±0.3ms 
 12     27.2±0.2ms    55.0±0.2ms 
 16     74.4±0.4ms    6.09±0.9s  
====== ============= =============

[kevinsung]

Thanks Inho, here are a few initial comments.

[kevinsung]

We shouldn't import anything from _slow since that's subject to deletion at any time. Instead, copy and paste the old implementation as a function in this test file.

[kevinsung]

In the main branch you can now do bound_operator.normal_ordered(group_by_spin=True) (see #668). Then, you can get rid of the inversion counting logic.

[kevinsung]

Nit: ffsim style is to unpack n_alpha, n_beta = nelec near the top of the function.

[kevinsung]

Instead of hard-coding operators, we should generate them randomly. In the main branch, you can now do

op = ffsim.random.random_fermion_operator(norb, num_and_spin_conserving=True)

[kevinsung]

Let's move all the code specific to FermionOperator to python/ffsim/operators/fermion_operator.py and just import the _fermion_operator_to_linear_operator in this file.