joltax 0.1.1__tar.gz

joltax-0.1.1/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Swedish Biodiversity in Time and Space (SweBiTS)
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

joltax-0.1.1/PKG-INFO

@@ -0,0 +1,74 @@
+Metadata-Version: 2.4
+Name: joltax
+Version: 0.1.1
+Summary: A high-performance, vectorized taxonomy library for Python.
+Author-email: Daniel Svensson <daniel.svensson@umu.se>
+License: MIT
+Project-URL: Homepage, https://github.com/SweBiTS/JolTax
+Project-URL: Bug Tracker, https://github.com/SweBiTS/JolTax/issues
+Project-URL: Source Code, https://github.com/SweBiTS/JolTax
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Classifier: Intended Audience :: Science/Research
+Classifier: Topic :: Scientific/Engineering :: Bio-Informatics
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: numpy>=1.20.0
+Requires-Dist: polars>=0.20.0
+Requires-Dist: rapidfuzz>=3.0.0
+Dynamic: license-file
+
+<p align="center">
+  <img src="assets/logo.png" alt="joltax logo" width="300">
+</p>
+
+# joltax
+
+**High-performance, vectorized taxonomy library for Python.**
+
+`JolTax` is a Python library designed to handle massive taxonomies with extreme efficiency. By representing taxonomy trees as contiguous NumPy arrays and leveraging Polars for mass data handling, it achieves lightning-fast traversals, constant-time clade queries, and rapid mass annotation of large datasets.
+
+## Key Features
+- **Vectorized Performance:** Uses hardware-accelerated NumPy operations for million-scale property lookups.
+- **Memory Efficient:** Optimized string store using Polars/Arrow reduces RAM footprint.
+- **Fuzzy Name Search:** Rapid fuzzy matching using RapidFuzz to find TaxIDs from names.
+- **Instant Clade Queries:** Quickly find all descendants of any node (even millions) using optimized range indexing.
+- **Hyper-Vectorized LCA search:** Lowest Common Ancestor (LCA) search and node-to-node distance calculations at lightning speeds.
+- **Mass Annotation:** Annotate massive TaxID tables with 2,000,000+ rows in under a second using Polars.
+
+## Quick Start
+
+```python
+from joltax.joltree import JolTree
+
+# Build and process the NCBI taxonomy
+tree = JolTree(nodes_file='nodes.dmp', names_file='names.dmp')
+
+# Save for instant loading next time
+tree.save('my_taxonomy_cache')
+
+# Re-load in milliseconds (using zero-copy Arrow IPC)
+tree = JolTree.load('my_taxonomy_cache')
+
+# Batch LCA (process 10,000 pairs in <10ms)
+lcas = tree.get_lca_batch(ids1, ids2)
+
+# Fuzzy search for a name (returns a Polars DataFrame)
+results = tree.search_name('Escherchia', fuzzy=True)
+print(results)
+```
+
+## Installation
+
+```bash
+cd joltax
+pip install .
+```
+
+Requires: `numpy`, `polars`, `rapidfuzz`.
+
+## Documentation
+
+For a detailed API reference and a comprehensive "How-To" guide with example workflows, please see [USAGE.md](./USAGE.md).

joltax-0.1.1/README.md

@@ -0,0 +1,52 @@
+<p align="center">
+  <img src="assets/logo.png" alt="joltax logo" width="300">
+</p>
+
+# joltax
+
+**High-performance, vectorized taxonomy library for Python.**
+
+`JolTax` is a Python library designed to handle massive taxonomies with extreme efficiency. By representing taxonomy trees as contiguous NumPy arrays and leveraging Polars for mass data handling, it achieves lightning-fast traversals, constant-time clade queries, and rapid mass annotation of large datasets.
+
+## Key Features
+- **Vectorized Performance:** Uses hardware-accelerated NumPy operations for million-scale property lookups.
+- **Memory Efficient:** Optimized string store using Polars/Arrow reduces RAM footprint.
+- **Fuzzy Name Search:** Rapid fuzzy matching using RapidFuzz to find TaxIDs from names.
+- **Instant Clade Queries:** Quickly find all descendants of any node (even millions) using optimized range indexing.
+- **Hyper-Vectorized LCA search:** Lowest Common Ancestor (LCA) search and node-to-node distance calculations at lightning speeds.
+- **Mass Annotation:** Annotate massive TaxID tables with 2,000,000+ rows in under a second using Polars.
+
+## Quick Start
+
+```python
+from joltax.joltree import JolTree
+
+# Build and process the NCBI taxonomy
+tree = JolTree(nodes_file='nodes.dmp', names_file='names.dmp')
+
+# Save for instant loading next time
+tree.save('my_taxonomy_cache')
+
+# Re-load in milliseconds (using zero-copy Arrow IPC)
+tree = JolTree.load('my_taxonomy_cache')
+
+# Batch LCA (process 10,000 pairs in <10ms)
+lcas = tree.get_lca_batch(ids1, ids2)
+
+# Fuzzy search for a name (returns a Polars DataFrame)
+results = tree.search_name('Escherchia', fuzzy=True)
+print(results)
+```
+
+## Installation
+
+```bash
+cd joltax
+pip install .
+```
+
+Requires: `numpy`, `polars`, `rapidfuzz`.
+
+## Documentation
+
+For a detailed API reference and a comprehensive "How-To" guide with example workflows, please see [USAGE.md](./USAGE.md).

joltax-0.1.1/joltax/__init__.py

@@ -0,0 +1 @@
+from .joltree import JolTree

joltax-0.1.1/joltax/joltree.py

@@ -0,0 +1,734 @@
+#!/usr/bin/env python3
+"""
+joltax/joltree.py
+Implementation of a high-performance, vectorized taxonomy tree.
+"""
+
+__version__ = "0.1.1"
+
+# The minimum version of a saved taxonomy cache that is compatible with this software.
+# Increment this when making breaking changes to the binary layout or metadata structure.
+MINIMUM_CACHE_VERSION = "0.1.1"
+
+import logging
+import os
+import datetime
+from typing import Dict, List, Optional, Set, Union, Tuple
+from collections import namedtuple
+
+import numpy as np
+import polars as pl
+from rapidfuzz import process, fuzz, utils
+
+# Set up logging for the module
+logging.basicConfig(
+    format='%(asctime)s %(levelname)-8s %(message)s',
+    level=logging.INFO,
+    datefmt='%Y-%m-%d [%H:%M:%S]'
+)
+logger = logging.getLogger(__name__)
+
+# Standard canonical ranks in order (highest to lowest)
+# Including both superkingdom and domain for compatibility with pre/post-2025 taxonomies
+CANONICAL_RANKS = [
+    'superkingdom', 'domain', 'kingdom', 'phylum', 
+    'class', 'order', 'family', 'genus', 'species'
+]
+
+# Mapping rank names to standard Kraken-style codes
+RANK_TO_CODE = {
+    'superkingdom': 'D',
+    'domain': 'D',
+    'kingdom': 'K',
+    'phylum': 'P',
+    'class': 'C',
+    'order': 'O',
+    'family': 'F',
+    'genus': 'G',
+    'species': 'S'
+}
+
+class JolTree:
+    """
+    A high-performance taxonomy representation using vectorized arrays.
+    
+    This class replaces traditional object-oriented trees with contiguous 
+    NumPy arrays for lightning-fast lookups, traversals, and mass annotations.
+    """
+
+    def __init__(self, nodes_file: Optional[str] = None, names_file: Optional[str] = None):
+        """
+        Initialize the taxonomy tree. If files are provided, it builds from DMP files.
+        Otherwise, it can be loaded from a binary cache using `load()`.
+        
+        Args:
+            nodes_file: Path to NCBI nodes.dmp
+            names_file: Path to NCBI names.dmp
+        """
+        # Vectorized internal index mapping (sorted array of TaxIDs)
+        self._index_to_id: np.ndarray = np.array([], dtype=np.int32)
+        
+        # Primary arrays (indexed by the dense internal index)
+        self.parents: np.ndarray = np.array([], dtype=np.int32)
+        self.depths: np.ndarray = np.array([], dtype=np.int32)
+        self.ranks: np.ndarray = np.array([], dtype=np.uint8)
+        
+        # Metadata storage (Polars Series for memory efficiency)
+        self._scientific_names: pl.Series = pl.Series("scientific_name", [], dtype=pl.String)
+        self._common_names: pl.Series = pl.Series("common_name", [], dtype=pl.String)
+        self.rank_names: List[str] = []
+        self.top_rank: str = "domain"  # Default, will be detected
+        self._source_nodes: Optional[str] = None
+        self._source_names: Optional[str] = None
+        self._build_time: Optional[str] = None
+        
+        # Clade query support (Euler Tour timestamps)
+        self.entry_times: np.ndarray = np.array([], dtype=np.int32)
+        self.exit_times: np.ndarray = np.array([], dtype=np.int32)
+        
+        # Binary lifting table for LCA (initialized on demand)
+        self._up_table: Optional[np.ndarray] = None
+        
+        # Pre-calculated canonical rank maps (dense internal index -> dense internal index)
+        # Values are internal indices, not TaxIDs. -1 means no ancestor at that rank.
+        self.canonical_maps: Dict[str, np.ndarray] = {}
+        
+        # Search index (Polars DataFrame: name -> tax_id)
+        self._search_index: pl.DataFrame = pl.DataFrame(schema={"name": pl.String, "tax_id": pl.Int32})
+        
+        # Caches for vectorized lookup (prepared during build/load)
+        self._sci_names_lookup: Optional[pl.Series] = None
+        self._rank_names_series: Optional[pl.Series] = None
+        self._ranks_extended: Optional[np.ndarray] = None
+        
+        if nodes_file and names_file:
+            self.build_from_dmp(nodes_file, names_file)
+
+    def build_from_dmp(self, nodes_file: str, names_file: str) -> None:
+        """
+        Parses NCBI DMP files and builds the vectorized internal structure.
+        
+        Args:
+            nodes_file: Path to NCBI nodes.dmp
+            names_file: Path to NCBI names.dmp
+        """
+        self._source_nodes = os.path.abspath(nodes_file)
+        self._source_names = os.path.abspath(names_file)
+        self._build_time = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+        logger.info(f"Starting taxonomy build at {self._build_time}...")
+        
+        # 1. Parse Names
+        logger.info(f"Parsing names from {names_file}...")
+        scientific_names = {}
+        common_names = {}
+        search_data = [] # List of (name, tax_id)
+        
+        with open(names_file, 'r') as f:
+            for name_line in f:
+                parts = name_line.split('|')
+                name_type = parts[3].strip()
+                
+                # Only care about scientific and genbank common names for now
+                if name_type not in ['scientific name', 'genbank common name']:
+                    continue
+                
+                tax_id = int(parts[0].strip())
+                name_txt = parts[1].strip()
+                
+                if name_type == 'scientific name':
+                    scientific_names[tax_id] = name_txt
+                elif name_type == 'genbank common name':
+                    common_names[tax_id] = name_txt
+                
+                search_data.append({"name": name_txt, "tax_id": tax_id})
+        
+        # Build search index
+        self._search_index = pl.DataFrame(search_data).sort("name")
+        
+        # 2. Parse Nodes and initial parent structure
+        logger.info(f"Parsing nodes from {nodes_file}...")
+        temp_parents = {}
+        temp_ranks = {}
+        all_ranks = set()
+        
+        with open(nodes_file, 'r') as f:
+            for line in f:
+                parts = line.split('|')
+                tax_id = int(parts[0].strip())
+                parent_id = int(parts[1].strip())
+                rank = parts[2].strip()
+                
+                temp_parents[tax_id] = parent_id
+                temp_ranks[tax_id] = rank
+                all_ranks.add(rank)
+        
+        # 2.1 Detect top rank (superkingdom vs domain)
+        has_sk = 'superkingdom' in all_ranks
+        has_dm = 'domain' in all_ranks
+        if has_sk and has_dm:
+            raise ValueError("Found both 'superkingdom' and 'domain' ranks. The taxonomy must use only one as the top rank.")
+        self.top_rank = 'superkingdom' if has_sk else 'domain'
+        logger.info(f"Detected top rank: {self.top_rank}")
+
+        # 3. Create dense mapping
+        logger.info("Creating dense mapping and vectorized arrays...")
+        sorted_tax_ids = sorted(temp_parents.keys())
+        num_nodes = len(sorted_tax_ids)
+        self._index_to_id = np.array(sorted_tax_ids, dtype=np.int32)
+        
+        # Mapping rank names to indices
+        self.rank_names = sorted(list(all_ranks))
+        rank_to_idx = {r: i for i, r in enumerate(self.rank_names)}
+        
+        self.parents = np.zeros(num_nodes, dtype=np.int32)
+        self.ranks = np.zeros(num_nodes, dtype=np.uint8)
+        
+        # Temporary dict for building parent connections (will be discarded)
+        id_to_index_temp = {tid: i for i, tid in enumerate(sorted_tax_ids)}
+        
+        for tid, i in id_to_index_temp.items():
+            parent_id = temp_parents[tid]
+            # Handle root (1) which is its own parent in NCBI
+            if tid == 1:
+                self.parents[i] = i
+            else:
+                self.parents[i] = id_to_index_temp[parent_id]
+            
+            self.ranks[i] = rank_to_idx[temp_ranks[tid]]
+
+        # Populate names aligned with indices
+        logger.info("Aligning names and ranks...")
+        sci_names_list = [scientific_names.get(tid, f"Unknown_{tid}") for tid in sorted_tax_ids]
+        com_names_list = [common_names.get(tid) for tid in sorted_tax_ids]
+        self._scientific_names = pl.Series("scientific_name", sci_names_list)
+        self._common_names = pl.Series("common_name", com_names_list)
+
+        # 4. Calculate depths
+        logger.info("Calculating node depths...")
+        self.depths = np.zeros(num_nodes, dtype=np.int32)
+        for i in range(num_nodes):
+            self._calculate_depth(i)
+            
+        # 5. Build Euler Tour for clade queries
+        self._build_euler_tour()
+
+        # 6. Pre-calculate canonical rank maps
+        self._build_canonical_maps()
+        
+        # 7. Prepare caches for vectorized lookups
+        self._prepare_vectorized_caches()
+        
+        logger.info("Taxonomy build complete.")
+
+    def _prepare_vectorized_caches(self) -> None:
+        """Initializes caches used for high-performance vectorized lookups."""
+        logger.info("Preparing vectorized lookup caches...")
+        # Scientific names lookup (aligned with dense internal index + 1 for "Unknown")
+        self._sci_names_lookup = self._scientific_names.append(pl.Series([None]))
+        
+        # Rank names lookup
+        self._rank_names_series = pl.Series(self.rank_names).append(pl.Series(["unclassified"]))
+        
+        # Ranks extended with a pointer to "unclassified" for unknown nodes
+        self._ranks_extended = np.append(self.ranks, [len(self.rank_names)]).astype(np.int32)
+
+    def _build_canonical_maps(self) -> None:
+        """Pre-calculates canonical rank ancestors for all nodes."""
+        logger.info("Pre-calculating canonical rank maps...")
+        num_nodes = len(self._index_to_id)
+        
+        # Identify all canonical ranks to track
+        canonical_columns = [self.top_rank] + [r for r in CANONICAL_RANKS if r not in ['superkingdom', 'domain']]
+        
+        # Initialize maps with -1 (meaning no ancestor at that rank)
+        self.canonical_maps = {rank: np.full(num_nodes, -1, dtype=np.int32) for rank in canonical_columns}
+        
+        # Sort nodes by depth to ensure parents are processed before children
+        for i in range(num_nodes):
+            curr_idx = i
+            root_idx = 0 # TaxID 1 is always the first in sorted_tax_ids
+            while True:
+                rank_name = self.rank_names[self.ranks[curr_idx]]
+                
+                # Normalize superkingdom/domain based on detected top_rank
+                mapped_rank = rank_name
+                if rank_name in ['superkingdom', 'domain']:
+                    mapped_rank = self.top_rank
+                
+                if mapped_rank in self.canonical_maps:
+                    self.canonical_maps[mapped_rank][i] = curr_idx
+                
+                if curr_idx == root_idx:
+                    break
+                curr_idx = self.parents[curr_idx]
+
+    def _calculate_depth(self, index: int) -> int:
+        """Recursive depth calculation with memoization."""
+        if index == 0: # TaxID 1 is always index 0
+            return 0
+        if self.depths[index] != 0:
+            return self.depths[index]
+        
+        d = self._calculate_depth(self.parents[index]) + 1
+        self.depths[index] = d
+        return d
+
+    def _build_euler_tour(self) -> None:
+        """Assigns entry/exit times to enable instant clade queries."""
+        logger.info("Building Euler Tour index for clade queries...")
+        num_nodes = len(self._index_to_id)
+        self.entry_times = np.zeros(num_nodes, dtype=np.int32)
+        self.exit_times = np.zeros(num_nodes, dtype=np.int32)
+        
+        # Build adjacency list (children)
+        children = [[] for _ in range(num_nodes)]
+        root_idx = 0 # TaxID 1
+        for i, p in enumerate(self.parents):
+            if i != root_idx:
+                children[p].append(i)
+        
+        timer = 0
+        stack = [(root_idx, False)] # (index, is_processed)
+        
+        while stack:
+            idx, processed = stack.pop()
+            if not processed:
+                self.entry_times[idx] = timer
+                timer += 1
+                stack.append((idx, True))
+                for child in reversed(children[idx]):
+                    stack.append((child, False))
+            else:
+                self.exit_times[idx] = timer - 1
+
+    def _get_index(self, tax_id: int) -> int:
+        """Returns the internal index for a TaxID, or -1 if not found."""
+        idx = np.searchsorted(self._index_to_id, tax_id)
+        if idx < len(self._index_to_id) and self._index_to_id[idx] == tax_id:
+            return int(idx)
+        return -1
+    
+    def _get_indices(self, tax_ids: np.ndarray) -> np.ndarray:
+        """Returns internal indices for an array of TaxIDs, with -1 for missing."""
+        indices = np.searchsorted(self._index_to_id, tax_ids)
+        # Handle out of bounds
+        mask = indices < len(self._index_to_id)
+        # Check for actual equality
+        valid = np.zeros(len(tax_ids), dtype=bool)
+        valid[mask] = self._index_to_id[indices[mask]] == tax_ids[mask]
+        return np.where(valid, indices, -1)
+
+    def get_lineage(self, tax_id: int) -> List[int]:
+        """Returns the full lineage from root to the given TaxID."""
+        idx = self._get_index(tax_id)
+        if idx == -1:
+            logger.warning(f"TaxID {tax_id} not found in taxonomy tree.")
+            return []
+        
+        lineage = []
+        root_idx = 0
+        
+        while True:
+            lineage.append(int(self._index_to_id[idx]))
+            if idx == root_idx:
+                break
+            idx = self.parents[idx]
+            
+        return lineage[::-1]
+
+    def get_name(self, tax_id: int) -> str:
+        """Returns the scientific name of the given TaxID."""
+        idx = self._get_index(tax_id)
+        if idx != -1:
+            return self._scientific_names[idx]
+        return f"Unknown_{tax_id}"
+
+    def get_common_name(self, tax_id: int) -> Optional[str]:
+        """Returns the genbank common name of the given TaxID, if available."""
+        idx = self._get_index(tax_id)
+        if idx != -1:
+            return self._common_names[idx]
+        return None
+
+    def get_rank(self, tax_id: int) -> str:
+        """Returns the taxonomic rank of the given TaxID."""
+        idx = self._get_index(tax_id)
+        if idx == -1:
+            logger.warning(f"TaxID {tax_id} not found in taxonomy tree.")
+            return "unknown"
+        return self.rank_names[self.ranks[idx]]
+
+    def search_name(self, query: str, fuzzy: bool = False, limit: int = 10, score_cutoff: float = 60.0) -> pl.DataFrame:
+        """
+        Searches for TaxIDs by name.
+        """
+        if not fuzzy:
+            matches = self._search_index.filter(pl.col("name") == query)
+            if matches.is_empty():
+                return pl.DataFrame(schema=["tax_id", "name", "rank", "score"])
+            
+            # Vectorized rank lookup for matches
+            tids = matches["tax_id"].to_numpy()
+            indices = self._get_indices(tids)
+            ranks = [self.rank_names[self.ranks[i]] if i != -1 else "unknown" for i in indices]
+            
+            return matches.with_columns([
+                pl.Series("rank", ranks),
+                pl.lit(100.0).alias("score")
+            ])
+
+        # Fuzzy matching path
+        unique_names = self._search_index["name"].unique().to_list()
+            
+        # rapidfuzz extract
+        matches = process.extract(
+            query, 
+            unique_names, 
+            scorer=fuzz.WRatio, 
+            limit=limit, 
+            processor=utils.default_process,
+            score_cutoff=score_cutoff
+        )
+        
+        data = []
+        for match_str, score, _ in matches:
+            # Find all TaxIDs associated with this name
+            tids = self._search_index.filter(pl.col("name") == match_str)["tax_id"].to_list()
+            for tid in tids:
+                idx = self._get_index(tid)
+                rank = self.rank_names[self.ranks[idx]] if idx != -1 else "unknown"
+                
+                # Smart Ranking: Boost scores for canonical ranks
+                rank_boost = 0.0
+                if rank in CANONICAL_RANKS:
+                    rank_boost = 2.0
+                
+                data.append({
+                    "tax_id": tid,
+                    "matched_name": match_str,
+                    "scientific_name": self.get_name(tid),
+                    "rank": rank,
+                    "score": score + rank_boost
+                })
+        
+        if not data:
+            return pl.DataFrame(schema=["tax_id", "matched_name", "scientific_name", "rank", "score"])
+            
+        return pl.DataFrame(data).sort("score", descending=True)
+
+    def get_clade(self, tax_id: int) -> List[int]:
+        """Returns all TaxIDs in the clade rooted at the given TaxID."""
+        idx = self._get_index(tax_id)
+        if idx == -1:
+            logger.warning(f"TaxID {tax_id} not found in taxonomy tree.")
+            return []
+        
+        entry = self.entry_times[idx]
+        exit = self.exit_times[idx]
+        
+        mask = (self.entry_times >= entry) & (self.entry_times <= exit)
+        return self._index_to_id[mask].astype(int).tolist()
+
+    def get_clade_at_rank(self, tax_id: int, rank_name: str) -> List[int]:
+        """
+        Returns all TaxIDs of a specific rank within the clade rooted at tax_id.
+        """
+        idx = self._get_index(tax_id)
+        if idx == -1:
+            logger.warning(f"TaxID {tax_id} not found in taxonomy tree.")
+            return []
+        
+        try:
+            target_rank_idx = self.rank_names.index(rank_name)
+        except ValueError:
+            logger.warning(f"Rank '{rank_name}' not found in taxonomy. Available ranks: {self.rank_names}")
+            return []
+        
+        entry = self.entry_times[idx]
+        exit = self.exit_times[idx]
+        
+        mask = (self.entry_times >= entry) & (self.entry_times <= exit) & (self.ranks == target_rank_idx)
+        return self._index_to_id[mask].astype(int).tolist()
+
+    def get_lca(self, tax_id_1: int, tax_id_2: int) -> int:
+        """Finds the Lowest Common Ancestor using Binary Lifting."""
+        idx1 = self._get_index(tax_id_1)
+        idx2 = self._get_index(tax_id_2)
+        
+        if idx1 == -1 or idx2 == -1:
+            logger.warning(f"One or both TaxIDs ({tax_id_1}, {tax_id_2}) not found.")
+            return 1
+            
+        self._ensure_up_table()
+        
+        if self.depths[idx1] < self.depths[idx2]:
+            idx1, idx2 = idx2, idx1
+        
+        diff = self.depths[idx1] - self.depths[idx2]
+        max_log = self._up_table.shape[0]
+        
+        for i in range(max_log):
+            if (diff >> i) & 1:
+                idx1 = self._up_table[i, idx1]
+                
+        if idx1 == idx2:
+            return int(self._index_to_id[idx1])
+            
+        for i in reversed(range(max_log)):
+            up1 = self._up_table[i, idx1]
+            up2 = self._up_table[i, idx2]
+            if up1 != up2:
+                idx1 = up1
+                idx2 = up2
+                
+        return int(self._index_to_id[self.parents[idx1]])
+
+    def get_distance(self, tax_id_1: int, tax_id_2: int) -> int:
+        """Calculates distance (number of edges) between two TaxIDs."""
+        lca_id = self.get_lca(tax_id_1, tax_id_2)
+        idx1 = self._get_index(tax_id_1)
+        idx2 = self._get_index(tax_id_2)
+        idx_lca = self._get_index(lca_id)
+        return int(self.depths[idx1] + self.depths[idx2] - 2 * self.depths[idx_lca])
+
+    def get_lca_batch(self, ids1: Union[List[int], np.ndarray], ids2: Union[List[int], np.ndarray]) -> np.ndarray:
+        """
+        Calculates Lowest Common Ancestor for arrays of TaxIDs.
+        Hyper-vectorized implementation for peak performance.
+        """
+        ids1 = np.array(ids1, dtype=np.int32)
+        ids2 = np.array(ids2, dtype=np.int32)
+        
+        if ids1.shape != ids2.shape:
+            raise ValueError("Input arrays must have the same shape.")
+            
+        self._ensure_up_table()
+        
+        idx1 = self._get_indices(ids1)
+        idx2 = self._get_indices(ids2)
+        
+        # Handle missing IDs by pointing to root (index 0)
+        valid_mask = (idx1 != -1) & (idx2 != -1)
+        s_idx1 = np.where(valid_mask, idx1, 0)
+        s_idx2 = np.where(valid_mask, idx2, 0)
+        
+        # 1. Bring both nodes to the same depth
+        d1 = self.depths[s_idx1]
+        d2 = self.depths[s_idx2]
+        
+        # Ensure s_idx1 is the deeper one
+        swap = d1 < d2
+        s_idx1[swap], s_idx2[swap] = s_idx2[swap], s_idx1[swap]
+        
+        diff = np.abs(d1 - d2)
+        max_log = self._up_table.shape[0]
+        
+        for i in range(max_log):
+            mask = (diff >> i) & 1 == 1
+            if np.any(mask):
+                s_idx1[mask] = self._up_table[i, s_idx1[mask]]
+            
+        # 2. Binary search for the LCA
+        lca_indices = s_idx1.copy()
+        not_same = s_idx1 != s_idx2
+        
+        if np.any(not_same):
+            sub1 = s_idx1[not_same]
+            sub2 = s_idx2[not_same]
+            
+            for i in reversed(range(max_log)):
+                up1 = self._up_table[i, sub1]
+                up2 = self._up_table[i, sub2]
+                
+                diff_up = up1 != up2
+                sub1[diff_up] = up1[diff_up]
+                sub2[diff_up] = up2[diff_up]
+            
+            lca_indices[not_same] = self.parents[sub1]
+            
+        results = self._index_to_id[lca_indices]
+        results[~valid_mask] = 1
+        return results
+
+    def get_distance_batch(self, ids1: Union[List[int], np.ndarray], ids2: Union[List[int], np.ndarray]) -> np.ndarray:
+        """Vectorized distance calculation for arrays of TaxIDs."""
+        ids1 = np.array(ids1, dtype=np.int32)
+        ids2 = np.array(ids2, dtype=np.int32)
+        
+        lca_ids = self.get_lca_batch(ids1, ids2)
+        
+        idx1 = self._get_indices(ids1)
+        idx2 = self._get_indices(ids2)
+        idx_lca = self._get_indices(lca_ids)
+        
+        # Mask invalid lookups to avoid OOB errors
+        valid = (idx1 != -1) & (idx2 != -1) & (idx_lca != -1)
+        
+        dists = np.zeros(len(ids1), dtype=np.int32)
+        if np.any(valid):
+            v1, v2, vl = idx1[valid], idx2[valid], idx_lca[valid]
+            dists[valid] = self.depths[v1] + self.depths[v2] - 2 * self.depths[vl]
+            
+        return dists
+
+    def annotate_table(self, tax_ids: Union[List[int], np.ndarray]) -> pl.DataFrame:
+        """
+        Massively annotates a list of TaxIDs with scientific_names and canonical ranks.
+        Extremely efficient for large tables (e.g. 200k+ rows) using Polars and vectorized lookups.
+        """
+        logger.info(f"Annotating {len(tax_ids)} taxa...")
+        canonical_columns = [self.top_rank] + [r for r in CANONICAL_RANKS if r not in ['superkingdom', 'domain']]
+        
+        tax_ids_arr = np.array(tax_ids, dtype=np.int32)
+        indices = self._get_indices(tax_ids_arr)
+        valid_mask = indices != -1
+        
+        # dummy_idx points to the "Unknown/None" entry at the end of the lookup series
+        dummy_idx = len(self._index_to_id)
+        safe_indices = np.where(valid_mask, indices, dummy_idx)
+        
+        # Ensure caches are ready
+        if self._sci_names_lookup is None:
+            self._prepare_vectorized_caches()
+            
+        df_dict = {"tax_id": tax_ids_arr}
+        
+        for rank in canonical_columns:
+            # canonical_maps now store internal indices
+            ancestor_indices = np.full(len(tax_ids_arr), -1, dtype=np.int32)
+            # Map input tax_ids to their ancestor's internal index
+            ancestor_indices[valid_mask] = self.canonical_maps[rank][indices[valid_mask]]
+            
+            # Use dummy_idx for missing ancestors
+            safe_anc_indices = np.where(ancestor_indices != -1, ancestor_indices, dummy_idx)
+            
+            # Vectorized gather from Polars
+            df_dict[rank] = self._sci_names_lookup.gather(safe_anc_indices.astype(np.int32))
+
+        # Scientific name for the input TaxID
+        df_dict["scientific_name"] = self._sci_names_lookup.gather(safe_indices.astype(np.int32))
+        
+        # Rank for the input TaxID
+        target_rank_indices = self._ranks_extended[safe_indices]
+        df_dict["rank"] = self._rank_names_series.gather(target_rank_indices.astype(np.int32))
+        
+        df = pl.DataFrame(df_dict)
+        final_order = ['tax_id'] + canonical_columns + ['scientific_name', 'rank']
+        return df.select(final_order)
+
+    def _ensure_up_table(self) -> None:
+        """Lazy initialization of binary lifting table."""
+        if self._up_table is not None:
+            return
+            
+        logger.info("Initializing binary lifting table (Hyper-Vectorized)...")
+        num_nodes = len(self._index_to_id)
+        max_log = int(np.ceil(np.log2(np.max(self.depths) + 1)))
+        
+        # Shape: (max_log, num_nodes) - optimized for contiguous column access
+        self._up_table = np.zeros((max_log, num_nodes), dtype=np.int32)
+        
+        # Power 2^0 is just the parents
+        self._up_table[0, :] = self.parents
+        
+        # Power 2^j = 2^{j-1} jump from the 2^{j-1} ancestor
+        # Fully vectorized initialization
+        for j in range(1, max_log):
+            prev_ancestors = self._up_table[j-1, :]
+            self._up_table[j, :] = self._up_table[j-1, prev_ancestors]
+
+    def save(self, directory: str) -> None:
+        """Saves the vectorized tree to a directory for fast loading."""
+        if not os.path.exists(directory):
+            os.makedirs(directory)
+            
+        logger.info(f"Saving binary cache to {directory}...")
+        np.save(os.path.join(directory, "index_to_id.npy"), self._index_to_id)
+        np.save(os.path.join(directory, "parents.npy"), self.parents)
+        np.save(os.path.join(directory, "depths.npy"), self.depths)
+        np.save(os.path.join(directory, "ranks.npy"), self.ranks)
+        np.save(os.path.join(directory, "entry_times.npy"), self.entry_times)
+        np.save(os.path.join(directory, "exit_times.npy"), self.exit_times)
+        
+        # Save Polars metadata
+        self._scientific_names.to_frame().write_ipc(os.path.join(directory, "scientific_names.ipc"))
+        self._common_names.to_frame().write_ipc(os.path.join(directory, "common_names.ipc"))
+        self._search_index.write_ipc(os.path.join(directory, "search_index.ipc"))
+
+        maps_dir = os.path.join(directory, "canonical_maps")
+        if not os.path.exists(maps_dir):
+            os.makedirs(maps_dir)
+        for rank, arr in self.canonical_maps.items():
+            np.save(os.path.join(maps_dir, f"{rank}.npy"), arr)
+            
+        import pickle
+        with open(os.path.join(directory, "metadata.pkl"), 'wb') as f:
+            pickle.dump({
+                "rank_names": self.rank_names,
+                "top_rank": self.top_rank,
+                "provenance": {
+                    "build_time": self._build_time,
+                    "source_nodes": self._source_nodes,
+                    "source_names": self._source_names,
+                    "package_version": __version__,
+                    "node_count": len(self._index_to_id),
+                    "max_depth": int(np.max(self.depths))
+                }
+            }, f)
+
+    @classmethod
+    def load(cls, directory: str) -> 'JolTree':
+        """Loads the vectorized tree from a binary cache directory."""
+        logger.info(f"Loading binary cache from {directory}...")
+        
+        import pickle
+        with open(os.path.join(directory, "metadata.pkl"), 'rb') as f:
+            meta = pickle.load(f)
+            prov = meta.get("provenance", {})
+            saved_version = prov.get("package_version", "unknown")
+            def version_to_tuple(v):
+                try:
+                    return tuple(map(int, v.split('.')))
+                except (ValueError, AttributeError):
+                    return (0, 0, 0)
+            if version_to_tuple(saved_version) < version_to_tuple(MINIMUM_CACHE_VERSION):
+                raise RuntimeError(
+                    f"Incompatible taxonomy cache. Saved version: {saved_version}, "
+                    f"Minimum required: {MINIMUM_CACHE_VERSION}. Please rebuild with build_from_dmp()."
+                )
+
+            tree = cls()
+            tree.rank_names = meta["rank_names"]
+            tree.top_rank = meta.get("top_rank", "domain")
+            tree._build_time = prov.get("build_time")
+            tree._source_nodes = prov.get("source_nodes")
+            tree._source_names = prov.get("source_names")
+
+        tree._index_to_id = np.load(os.path.join(directory, "index_to_id.npy"))
+        tree.parents = np.load(os.path.join(directory, "parents.npy"))
+        tree.depths = np.load(os.path.join(directory, "depths.npy"))
+        tree.ranks = np.load(os.path.join(directory, "ranks.npy"))
+        tree.entry_times = np.load(os.path.join(directory, "entry_times.npy"))
+        tree.exit_times = np.load(os.path.join(directory, "exit_times.npy"))
+        
+        # Load Polars metadata
+        tree._scientific_names = pl.read_ipc(os.path.join(directory, "scientific_names.ipc"))["scientific_name"]
+        tree._common_names = pl.read_ipc(os.path.join(directory, "common_names.ipc"))["common_name"]
+        tree._search_index = pl.read_ipc(os.path.join(directory, "search_index.ipc"))
+
+        maps_dir = os.path.join(directory, "canonical_maps")
+        if os.path.exists(maps_dir):
+            for filename in os.listdir(maps_dir):
+                if filename.endswith(".npy"):
+                    rank = filename[:-4]
+                    tree.canonical_maps[rank] = np.load(os.path.join(maps_dir, filename))
+        
+        # Re-initialize vectorized caches
+        tree._prepare_vectorized_caches()
+        
+        logger.info("Loaded taxonomy cache:")
+        logger.info(f"  Version:       {saved_version}")
+        logger.info(f"  Build time:    {tree._build_time}")
+        logger.info(f"  Node count:    {prov.get('node_count', 'Unknown'):,}")
+        logger.info(f"  Top rank:      {tree.top_rank}")
+        return tree

joltax-0.1.1/joltax.egg-info/PKG-INFO

@@ -0,0 +1,74 @@
+Metadata-Version: 2.4
+Name: joltax
+Version: 0.1.1
+Summary: A high-performance, vectorized taxonomy library for Python.
+Author-email: Daniel Svensson <daniel.svensson@umu.se>
+License: MIT
+Project-URL: Homepage, https://github.com/SweBiTS/JolTax
+Project-URL: Bug Tracker, https://github.com/SweBiTS/JolTax/issues
+Project-URL: Source Code, https://github.com/SweBiTS/JolTax
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Classifier: Intended Audience :: Science/Research
+Classifier: Topic :: Scientific/Engineering :: Bio-Informatics
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: numpy>=1.20.0
+Requires-Dist: polars>=0.20.0
+Requires-Dist: rapidfuzz>=3.0.0
+Dynamic: license-file
+
+<p align="center">
+  <img src="assets/logo.png" alt="joltax logo" width="300">
+</p>
+
+# joltax
+
+**High-performance, vectorized taxonomy library for Python.**
+
+`JolTax` is a Python library designed to handle massive taxonomies with extreme efficiency. By representing taxonomy trees as contiguous NumPy arrays and leveraging Polars for mass data handling, it achieves lightning-fast traversals, constant-time clade queries, and rapid mass annotation of large datasets.
+
+## Key Features
+- **Vectorized Performance:** Uses hardware-accelerated NumPy operations for million-scale property lookups.
+- **Memory Efficient:** Optimized string store using Polars/Arrow reduces RAM footprint.
+- **Fuzzy Name Search:** Rapid fuzzy matching using RapidFuzz to find TaxIDs from names.
+- **Instant Clade Queries:** Quickly find all descendants of any node (even millions) using optimized range indexing.
+- **Hyper-Vectorized LCA search:** Lowest Common Ancestor (LCA) search and node-to-node distance calculations at lightning speeds.
+- **Mass Annotation:** Annotate massive TaxID tables with 2,000,000+ rows in under a second using Polars.
+
+## Quick Start
+
+```python
+from joltax.joltree import JolTree
+
+# Build and process the NCBI taxonomy
+tree = JolTree(nodes_file='nodes.dmp', names_file='names.dmp')
+
+# Save for instant loading next time
+tree.save('my_taxonomy_cache')
+
+# Re-load in milliseconds (using zero-copy Arrow IPC)
+tree = JolTree.load('my_taxonomy_cache')
+
+# Batch LCA (process 10,000 pairs in <10ms)
+lcas = tree.get_lca_batch(ids1, ids2)
+
+# Fuzzy search for a name (returns a Polars DataFrame)
+results = tree.search_name('Escherchia', fuzzy=True)
+print(results)
+```
+
+## Installation
+
+```bash
+cd joltax
+pip install .
+```
+
+Requires: `numpy`, `polars`, `rapidfuzz`.
+
+## Documentation
+
+For a detailed API reference and a comprehensive "How-To" guide with example workflows, please see [USAGE.md](./USAGE.md).

joltax-0.1.1/joltax.egg-info/SOURCES.txt

@@ -0,0 +1,11 @@
+LICENSE
+README.md
+pyproject.toml
+joltax/__init__.py
+joltax/joltree.py
+joltax.egg-info/PKG-INFO
+joltax.egg-info/SOURCES.txt
+joltax.egg-info/dependency_links.txt
+joltax.egg-info/requires.txt
+joltax.egg-info/top_level.txt
+tests/test_tree.py

joltax-0.1.1/joltax.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

joltax-0.1.1/joltax.egg-info/requires.txt

@@ -0,0 +1,3 @@
+numpy>=1.20.0
+polars>=0.20.0
+rapidfuzz>=3.0.0

joltax-0.1.1/joltax.egg-info/top_level.txt

@@ -0,0 +1 @@
+joltax

joltax-0.1.1/pyproject.toml

@@ -0,0 +1,34 @@
+[build-system]
+requires = ["setuptools>=61.0"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "joltax"
+version = "0.1.1"
+authors = [
+  { name="Daniel Svensson", email="daniel.svensson@umu.se" },
+]
+description = "A high-performance, vectorized taxonomy library for Python."
+readme = "README.md"
+requires-python = ">=3.8"
+license = {text = "MIT"}
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "License :: OSI Approved :: MIT License",
+    "Operating System :: OS Independent",
+    "Intended Audience :: Science/Research",
+    "Topic :: Scientific/Engineering :: Bio-Informatics",
+]
+dependencies = [
+    "numpy>=1.20.0",
+    "polars>=0.20.0",
+    "rapidfuzz>=3.0.0",
+]
+
+[project.urls]
+"Homepage" = "https://github.com/SweBiTS/JolTax"
+"Bug Tracker" = "https://github.com/SweBiTS/JolTax/issues"
+"Source Code" = "https://github.com/SweBiTS/JolTax"
+
+[tool.setuptools]
+packages = ["joltax"]

joltax-0.1.1/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

joltax-0.1.1/tests/test_tree.py

@@ -0,0 +1,142 @@
+import unittest
+import os
+import sys
+import numpy as np
+import polars as pl
+
+# Add the project root to sys.path
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+
+from joltax.joltree import JolTree
+
+class TestJolTree(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.names_file = 'tests/data/names.dmp'
+        cls.nodes_file = 'tests/data/nodes.dmp'
+        # Check if files exist, if not, create them (should be copied already)
+        if not os.path.exists(cls.names_file):
+            raise FileNotFoundError(f"Missing test data: {cls.names_file}")
+            
+        cls.tree = JolTree(nodes_file=cls.nodes_file, names_file=cls.names_file)
+
+    def test_lineage(self):
+        # 562 (E. coli) -> 561 (Escherichia) -> 543 -> 91347 -> 1236 -> 1224 -> 2 -> 1
+        lineage = self.tree.get_lineage(562)
+        expected = [1, 2, 1224, 1236, 91347, 543, 561, 562]
+        self.assertEqual(lineage, expected)
+
+    def test_clade(self):
+        # Clade of 561 (genus) should contain 561 and 562 (species)
+        clade = self.tree.get_clade(561)
+        self.assertIn(561, clade)
+        self.assertIn(562, clade)
+        self.assertEqual(len(clade), 2)
+
+    def test_lca(self):
+        # LCA of 562 and 561 is 561
+        lca = self.tree.get_lca(562, 561)
+        self.assertEqual(lca, 561)
+        
+        # LCA of 562 and 2 (Bacteria) is 2
+        lca = self.tree.get_lca(562, 2)
+        self.assertEqual(lca, 2)
+
+    def test_distance(self):
+        # 562 to 561 is 1 step
+        self.assertEqual(self.tree.get_distance(562, 561), 1)
+        # 562 to 2 is 6 steps
+        self.assertEqual(self.tree.get_distance(562, 2), 6)
+
+    def test_get_name_and_rank(self):
+        self.assertEqual(self.tree.get_name(562), 'Escherichia coli')
+        self.assertEqual(self.tree.get_rank(562), 'species')
+        self.assertEqual(self.tree.get_name(2), 'Bacteria')
+        self.assertEqual(self.tree.get_rank(2), 'superkingdom')
+        # Test unknown
+        self.assertEqual(self.tree.get_name(999999), 'Unknown_999999')
+        self.assertEqual(self.tree.get_rank(999999), 'unknown')
+
+    def test_annotate_table(self):
+        tax_ids = [562, 561, 2]
+        df = self.tree.annotate_table(tax_ids)
+        self.assertIsInstance(df, pl.DataFrame)
+        self.assertEqual(len(df), 3)
+        self.assertIn('species', df.columns)
+        self.assertIn('genus', df.columns)
+        
+        # Check first row (562)
+        row0 = df.row(0, named=True)
+        self.assertEqual(row0['species'], 'Escherichia coli')
+        self.assertEqual(row0['genus'], 'Escherichia')
+        self.assertEqual(row0['scientific_name'], 'Escherichia coli')
+
+    def test_name_search(self):
+        # Search by scientific name
+        df = self.tree.search_name('Escherichia coli')
+        self.assertIn(562, df['tax_id'].to_list())
+        
+        # Search by common name
+        df = self.tree.search_name('all')
+        self.assertIn(1, df['tax_id'].to_list())
+
+    def test_fuzzy_search(self):
+        # Typo: "Escherchia"
+        df = self.tree.search_name('Escherchia', fuzzy=True)
+        self.assertIsInstance(df, pl.DataFrame)
+        self.assertTrue(len(df) > 0)
+        # Top result should be Escherichia or Escherichia coli
+        top_name = df.row(0, named=True)['matched_name']
+        self.assertIn('Escherichia', top_name)
+
+    def test_save_load(self):
+        import shutil
+        cache_dir = 'tests/cache_test'
+        if os.path.exists(cache_dir):
+            shutil.rmtree(cache_dir)
+            
+        self.tree.save(cache_dir)
+        new_tree = JolTree.load(cache_dir)
+        
+        self.assertEqual(new_tree.get_lineage(562), self.tree.get_lineage(562))
+        # Check name index loaded
+        df = new_tree.search_name('Escherichia coli')
+        self.assertIn(562, df['tax_id'].to_list())
+        
+        shutil.rmtree(cache_dir)
+
+    def test_version_validation(self):
+        import shutil
+        import pickle
+        cache_dir = 'tests/version_test'
+        if os.path.exists(cache_dir):
+            shutil.rmtree(cache_dir)
+            
+        self.tree.save(cache_dir)
+        
+        # Manually corrupt metadata with old version
+        meta_path = os.path.join(cache_dir, "metadata.pkl")
+        with open(meta_path, 'rb') as f:
+            meta = pickle.load(f)
+        
+        meta["provenance"]["package_version"] = "0.0.1" # Older than 0.1.0
+        
+        with open(meta_path, 'wb') as f:
+            pickle.dump(meta, f)
+            
+        # Should raise RuntimeError
+        with self.assertRaises(RuntimeError) as cm:
+            JolTree.load(cache_dir)
+        
+        self.assertIn("Incompatible taxonomy cache", str(cm.exception))
+        shutil.rmtree(cache_dir)
+
+if __name__ == '__main__':
+    # We skip tests if dependencies aren't installed
+    try:
+        import numpy
+        import polars
+        import rapidfuzz
+        unittest.main()
+    except ImportError:
+        print("Skipping tests due to missing dependencies.")