Source code for ferrmion.optimize.huffman

"""Huffman-code Ternary Tree."""

import numpy as np
import numpy.typing as npt

from ferrmion.encode import TernaryTree
from ferrmion.encode.ternary_tree_node import TTNode
from ferrmion.utils import find_pauli_weight, pauli_to_symplectic, symplectic_product


def _majorana_op_frequency(
    ones: npt.NDArray[float], twos: npt.NDArray[float]
) -> npt.NDArray[float]:
    n_modes = ones.shape[0]

    # One-electron: each unique position in {i, j} gets |ones[i,j]| added once.
    # When i != j both get it; when i == j only one does.
    abs_ones = np.abs(ones)
    majorana_freq = (
        np.sum(abs_ones, axis=1) + np.sum(abs_ones, axis=0) - np.diag(abs_ones)
    )

    # Two-electron: each unique position in {i,j,k,l} gets |twos[i,j,k,l]| added once.
    # Use inclusion-exclusion to handle duplicate indices correctly.
    abs_twos = np.abs(twos)
    for p in range(n_modes):
        s1 = (
            abs_twos[p, :, :, :].sum()
            + abs_twos[:, p, :, :].sum()
            + abs_twos[:, :, p, :].sum()
            + abs_twos[:, :, :, p].sum()
        )
        s2 = (
            abs_twos[p, p, :, :].sum()
            + abs_twos[p, :, p, :].sum()
            + abs_twos[p, :, :, p].sum()
            + abs_twos[:, p, p, :].sum()
            + abs_twos[:, p, :, p].sum()
            + abs_twos[:, :, p, p].sum()
        )
        s3 = (
            abs_twos[p, p, p, :].sum()
            + abs_twos[p, p, :, p].sum()
            + abs_twos[p, :, p, p].sum()
            + abs_twos[:, p, p, p].sum()
        )
        s4 = abs_twos[p, p, p, p]
        majorana_freq[p] += s1 - s2 + s3 - s4

    return majorana_freq.repeat(2)


def _build_huffman_tree(
    n_modes: int, majorana_frequencies: npt.NDArray[float]
) -> TernaryTree:
    nodes = {i: None for i in range(len(majorana_frequencies))}
    weights = {i: j for i, j in enumerate(majorana_frequencies)}
    n_ops = len(majorana_frequencies)
    for i in range(n_ops // 2):
        parent_index = 2 * n_ops - 1 - i
        mins = sorted(weights.items(), key=lambda kv: (kv[1], kv[0]))[:3]

        parent = nodes.get(parent_index, TTNode(parent=None, qubit_label=i))

        for min, child_string in zip(mins, ["x", "y", "z"][: len(mins)]):
            possible_child = nodes[min[0]]
            if isinstance(possible_child, TTNode):
                parent.add_child(which_child=child_string, child_node=possible_child)

        new_weight = 0
        for index, weight in mins:
            new_weight += weight
            weights.pop(index)
            nodes.pop(index)

        nodes[parent_index] = parent
        weights[parent_index] = new_weight

    root_node = [*nodes.values()][0]
    huffman_tree = TernaryTree(n_modes=n_modes, root_node=root_node)
    huffman_tree.string_pairs

    return huffman_tree


def _two_e_frequency(ones, twos) -> npt.NDArray[float]:
    abs_twos = np.abs(twos)
    two_e_freq = np.sum(abs_twos, axis=(2, 3)) + np.sum(abs_twos, axis=(0, 1))
    two_e_freq = np.kron(two_e_freq, np.array([[1, 1], [1, 1]]))
    two_e_freq = np.triu(two_e_freq, k=1)
    return two_e_freq


def _mode_priority(two_e_freq):
    vaccum_frequencies = np.diag(two_e_freq, k=1)
    sorted_pairs = [
        (i, i + 1) for i in np.argsort(vaccum_frequencies)[::-1] if i % 2 == 0
    ]
    sorted_modes = [i[0] // 2 for i in sorted_pairs]
    return sorted_modes


def _operator_pair_priority(huffman_tree):
    weights = {}
    for index, pair in enumerate(huffman_tree.string_pairs.values()):
        left, right = pair
        left = huffman_tree.branch_pauli_map[left]
        right = huffman_tree.branch_pauli_map[right]

        weights[index] = {}
        left, _ = pauli_to_symplectic(left, 0)
        right, _ = pauli_to_symplectic(right, 0)
        pair_weight = find_pauli_weight(np.array([left])) + find_pauli_weight(
            np.array([right])
        )
        _, product = symplectic_product(left, right)
        product_weight = find_pauli_weight(np.array([product]))

        weights[index]["pair_weight"] = pair_weight
        weights[index]["prod_weight"] = product_weight

        operator_order = sorted(
            weights.items(), key=lambda kv: (kv[1]["prod_weight"], kv[1]["pair_weight"])
        )

        operator_order = [index for index, _ in operator_order]
    return operator_order


[docs] def huffman_ternary_tree( one_e_coeffs: npt.NDArray[float], two_e_coeffs: npt.NDArray[float] ) -> TernaryTree: """Creates a Huffman-code Ternary Tree. Li, Q. S., Liu, H. Y., Wang, Q., Wu, Y. C., & Guo, G. P. (2025). Huffman-Code-based Ternary Tree Transformation. Chinese Physics Letters. http://iopscience.iop.org/article/10.1088/0256-307X/42/10/100001 Note: Only vaccum-preserving Huffman-trees are currently supported. Args: one_e_coeffs (np.ndarray): One electron coefficients. two_e_coeffs (np.ndarray): Two electron coefficients. Return: TernaryTree: A Huffman-code ternary tree. """ n_modes = one_e_coeffs.shape[0] majorana_frequencies = _majorana_op_frequency(one_e_coeffs, two_e_coeffs) huffman_ternary_tree = _build_huffman_tree(n_modes, majorana_frequencies) huffman_ternary_tree.enumeration_scheme = ( huffman_ternary_tree.default_enumeration_scheme() ) two_e_frequencies = _two_e_frequency(one_e_coeffs, two_e_coeffs) sorted_modes = _mode_priority(two_e_frequencies) sorted_operators = _operator_pair_priority(huffman_ternary_tree) mode_op_map = [0] * len(sorted_modes) for operator_index, mode_index in enumerate(sorted_modes): mode_op_map[mode_index] = sorted_operators[operator_index] huffman_ternary_tree.default_mode_op_map = mode_op_map return huffman_ternary_tree