NJIT / 64bit整数用の混合ハッシュ

jitst_cpx.py

#        _.-=-._     ＿＿＿＿＿＿＿＿＿＿＿＿＿＿＿＿＿
#      .'  \\\\ //  '.  /  !/ Mmmmmm!!
#     /     .-.-.    \/￣￣￣￣￣￣￣￣￣￣￣￣￣￣
#    |    .'  _  '.    |
#    |   /   (_)   \   |
#     \  |  .___ .  |  /
#      '._\  \_/  /_.'
#         /`-.___.-`\\
#        /  /  |  \  \\
#       /__/   |   \__\\
#
#_______________________________________________

from numba import njit
import numpy as np

# 制御バイト
_EMPTY = np.uint8(0x80)    # 空スロット
_DELETED = np.uint8(0xFE)  # 削除済みスロット（墓石）
_GROUP = 8                 # 1回の探索で確認する要素数

@njit(cache=True, inline="always")
def _hash64_int(x):
    # 64bit整数用の混合ハッシュ
    h = np.uint64(x)
    h ^= h >> np.uint64(33)
    h *= np.uint64(0xff51afd7ed558ccd)
    h ^= h >> np.uint64(33)
    h *= np.uint64(0xc4ceb9fe1a85ec53)
    h ^= h >> np.uint64(33)
    return h

@njit(cache=True, inline="always")
def _fp(h):
    # ハッシュ値から7bit指紋を生成
    fp = np.uint8((h >> np.uint64(57)) & np.uint64(0x7F))
    return np.uint8(0x7E) if fp == np.uint8(0x7F) else fp

@njit(cache=True)
def _find_slot(keys, ctrl, mask, key):
    # 挿入位置と検索結果を同時に求める
    h = _hash64_int(key); fp = _fp(h)
    start = np.int64(h & np.uint64(mask))
    first_deleted = -1; group = 0;cap = mask + 1
    # テーブル全体をグループ単位で走査
    while group < cap:
        base = (start + group) & mask
        for off in range(_GROUP):
            i = (base + off) & mask; c = ctrl[i]
            # 空スロット発見時は探索終了
            if c == _EMPTY:
                if first_deleted != -1:
                    return (first_deleted, False)
                return (i, False)
            # 墓石は再利用候補として記録
            if c == _DELETED:
                if first_deleted == -1: first_deleted = i
                continue
            # 指紋一致後にキー比較
            if c == fp and keys[i] == key:
                return (i, True)
        group += _GROUP
    if first_deleted != -1:
        return (first_deleted, False)
    return (0, False)

@njit(cache=True)
def _rehash(old_keys, old_values, old_ctrl, old_cap, new_cap):
    # 墓石を除去しながら再配置する
    keys = np.zeros(new_cap, dtype=np.int64)
    values = np.zeros(new_cap, dtype=np.complex128)
    ctrl = np.full(new_cap, _EMPTY, dtype=np.uint8)
    mask = new_cap - 1
    for i in range(old_cap):
        c = old_ctrl[i]
        if c != _EMPTY and c != _DELETED:
            key = old_keys[i]; val = old_values[i]
            h = _hash64_int(key); fp = _fp(h)
            start = np.int64(h & np.uint64(mask))
            group = 0
            while True:
                base = (start + group) & mask; placed = False
                for off in range(_GROUP):
                    j = (base + off) & mask
                    if ctrl[j] == _EMPTY:
                        ctrl[j] = fp; keys[j] = key
                        values[j] = val; placed = True
                        break
                if placed:
                    break
                group += _GROUP
    return keys, values, ctrl


class JITSwissTableComplex:
    __slots__ = ("cap", "mask", "size", "tombstones", "ctrl", "keys", "values")
    MAX_LOAD = 0.875; MAX_TOMBSTONES = 0.125

    def __init__(self, capacity=16):
        # 容量は2のべき乗へ切り上げる
        cap = 1
        while cap < capacity: cap <<= 1
        if cap < 8: cap = 8
        self.cap = cap; self.mask = cap - 1
        self.size = 0; self.tombstones = 0
        self.ctrl = np.full(cap, _EMPTY, dtype=np.uint8)
        self.keys = np.zeros(cap, dtype=np.int64)
        self.values = np.zeros(cap, dtype=np.complex128)

    def _maybe_resize(self):
        # 実使用率または墓石率が閾値を超えたら再配置する
        used = self.size + self.tombstones
        if used / self.cap > self.MAX_LOAD or self.tombstones / self.cap > self.MAX_TOMBSTONES:
            new_cap = self.cap * 2
            self.keys, self.values, self.ctrl = _rehash(self.keys, self.values, self.ctrl, self.cap, new_cap)
            self.cap = new_cap; self.mask = new_cap - 1; self.tombstones = 0

    def __setitem__(self, key, value):
        # Python型をNumPy互換型へ変換する
        key = np.int64(key); value = np.complex128(value)
        self._maybe_resize()
        i, found = _find_slot(self.keys, self.ctrl, self.mask, key)
        fp = _fp(_hash64_int(key))
        if found:
            self.values[i] = value
            return
        if self.ctrl[i] == _DELETED: self.tombstones -= 1
        self.keys[i] = key; self.values[i] = value
        self.ctrl[i] = fp; self.size += 1

    def __getitem__(self, key):
        key = np.int64(key)
        i, found = _find_slot(self.keys, self.ctrl, self.mask, key)
        if not found:
            raise KeyError(key)
        return self.values[i]

    def __contains__(self, key):
        key = np.int64(key)
        _, found = _find_slot(self.keys, self.ctrl, self.mask, key)
        return found

    def get(self, key, default=None):
        key = np.int64(key)
        i, found = _find_slot(self.keys, self.ctrl, self.mask, key)
        return self.values[i] if found else default

    def pop(self, key, default=...):
        key = np.int64(key)
        i, found = _find_slot(self.keys, self.ctrl, self.mask, key)
        if not found:
            if default is ...:
                raise KeyError(key)
            return default
        value = self.values[i]; self.ctrl[i] = _DELETED
        self.size -= 1; self.tombstones += 1
        if self.tombstones / self.cap > self.MAX_TOMBSTONES:
            self.keys, self.values, self.ctrl = _rehash(self.keys, self.values, self.ctrl, self.cap, self.cap)
            self.tombstones = 0
        return value

    def __delitem__(self, key):
        self.pop(key)

    def __len__(self):
        return self.size

# •••••••••••••••••••••••••••••••••••••••••••••••••••••
# 例1: 基本的な格納・取得
m = JITSwissTableComplex()
m[1] = 1.0 + 2.0j
m[2] = -3.0 + 4.0j
m[3] = 5.0 - 6.0j
print(m[1])  # (1+2j)
print(m[2])  # (-3+4j)
print(m[3])  # (5-6j)

# •••••••••••••••••••••••••••••••••••••••••••••••••••••
# 例2: 実部の符号による分割
pos = JITSwissTableComplex()
neg = JITSwissTableComplex()
for k in range(-5, 6):
    if k == 0:
        continue
    z = complex(k, -k)
    if z.real >= 0: pos[k] = z
    else: neg[k] = z
print(len(pos), len(neg))
print(pos[1])   # (1-1j)
print(neg[-1])  # (-1+1j)

# •••••••••••••••••••••••••••••••••••••••••••••••••••••
# 例3: 振幅による分割
short_len = JITSwissTableComplex()
mid_len = JITSwissTableComplex()
long_len = JITSwissTableComplex()
for k in range(1, 40):
    z = complex(k, k / 2); r = abs(z)
    if r < 10: short_len[k] = z
    elif r < 25: mid_len[k] = z
    else: long_len[k] = z
print(len(short_len), len(mid_len), len(long_len))

# •••••••••••••••••••••••••••••••••••••••••••••••••••••
# 例4: 条件付き更新
for k in range(1, 10):
    z = complex(k, -k)
    if abs(z) > 5: m[k] = -z
    else: m[k] = z
print(m[7])

# •••••••••••••••••••••••••••••••••••••••••••••••••••••
# 例5: 複素数ベクトルを用いた簡易SVD風投影
# 目的: 複素係数を高速参照しながら、信号成分を基底へ写像する
basis = JITSwissTableComplex()
coeff = JITSwissTableComplex()
projection = JITSwissTableComplex()
for k in range(1, 33):
    basis[k] = complex(np.cos(k * 0.1), np.sin(k * 0.1))
    coeff[k] = complex(1.0 / (k + 1), -1.0 / (k + 2))
# SVDの厳密計算ではなく、基底係数の加重和を高速に保持する例
for k in range(1, 33):
    projection[k] = basis[k] * coeff[k]
print(projection[1])
print(projection[16])

# •••••••••••••••••••••••••••••••••••••••••••••••••••••
# 例6: 誘導・抵抗を含む複素応答の高速格納
# 目的: 周波数ごとのインピーダンス応答 Z = R + jX を保持する
impedance = JITSwissTableComplex()
response = JITSwissTableComplex()
for k in range(1, 50):
    f = np.float64(k)
    resistance = 0.05 * f
    reactance = 0.2 * f
    z = complex(resistance, reactance)
    impedance[k] = z
    # 伝達応答の簡易モデル
    response[k] = z / complex(1.0 + 0.01 * f, 0.02 * f)
print(impedance[10])
print(response[10])

# •••••••••••••••••••••••••••••••••••••••••••••••••••••
# 例7: 墓石生成とコンパクション挙動の確認
for k in range(1, 20): m[k] = complex(k, k)
for k in range(1, 10): del m[k]
print(len(m))
m[100] = 100 + 100j
print(m[100])

# •••••••••••••••••••••••••••••••••••••••••••••••••••••
# 例8: 高頻度更新パターン
fast = JITSwissTableComplex()
for round_id in range(3):
    for k in range(1, 100):
        fast[k] = complex(k * 0.5, -k * 0.25)
    for k in range(1, 50):
        fast[k] = complex(-k * 0.75, k * 0.125)
print(len(fast))
print(fast[75])