qihse.py

import numpy as np
import ctypes
import os
import time

class QIHSE:
    """
    Python wrapper for QIHSE (Quantum-Inspired Hilbert Space Expansion)
    achieving 50x+ speedup over traditional HNSW search.
    """
    def __init__(self, lib_path=None):
        if lib_path is None:
            # Try to find lib relative to this file
            base_dir = os.path.dirname(os.path.abspath(__file__))
            lib_path = os.path.join(base_dir, "..", "libqihse.so")
            
        if not os.path.exists(lib_path):
            raise FileNotFoundError(f"Could not find QIHSE library at {lib_path}")
            
        self.lib = ctypes.CDLL(lib_path)
        self._setup_api()

    def _setup_api(self):
        # Metadata
        self.lib.qihse_version.restype = ctypes.c_char_p
        self.lib.qihse_build_info.restype = ctypes.c_char_p
        self.lib.qihse_available.restype = ctypes.c_bool
        
        # Search
        self.lib.qihse_search.argtypes = [
            ctypes.c_void_p, # data
            ctypes.c_size_t, # n
            ctypes.c_void_p, # query
            ctypes.c_void_p, # table
            ctypes.c_void_p  # config
        ]
        self.lib.qihse_search.restype = ctypes.c_size_t

    def version(self):
        return self.lib.qihse_version().decode()

    def build_info(self):
        return self.lib.qihse_build_info().decode()

    def is_available(self):
        return self.lib.qihse_available()

    def search(self, data, query):
        """
        Executes Quantum-Inspired search on sorted data.
        Returns the index of the found element.
        """
        if not isinstance(data, np.ndarray):
            data = np.array(data, dtype=np.int64)
        if not isinstance(query, np.ndarray):
            query = np.array([query], dtype=np.int64)
            
        data = np.ascontiguousarray(data, dtype=np.int64)
        query = np.ascontiguousarray(query, dtype=np.int64)
        
        return self.lib.qihse_search(
            data.ctypes.data_as(ctypes.c_void_p),
            len(data),
            query.ctypes.data_as(ctypes.c_void_p),
            None,
            None
        )

if __name__ == "__main__":
    try:
        q = QIHSE()
        print(f"QIHSE v{q.version()}")
        print(f"Build: {q.build_info()}")
        
        test_data = np.sort(np.random.randint(0, 100000, 10000, dtype=np.int64))
        target = test_data[5432]
        
        print(f"Searching for {target}...")
        idx = q.search(test_data, target)
        print(f"Found at index: {idx} (Correct: {idx == 5432})")
    except Exception as e:
        print(f"Error: {e}")