lorax-arg 0.1__py3-none-any.whl

lorax/sockets/node_search.py

@@ -0,0 +1,535 @@
+"""
+Node search event handlers for Lorax Socket.IO.
+
+Handles search_nodes and get_highlight_positions_event events.
+"""
+
+import asyncio
+
+from lorax.context import tree_graph_cache, csv_tree_graph_cache
+from lorax.constants import ERROR_NO_FILE_LOADED
+from lorax.handlers import (
+    should_shift_csv_tips,
+    search_nodes_in_trees,
+    get_highlight_positions,
+    get_multi_value_highlight_positions,
+)
+from lorax.cache import get_file_context
+from lorax.sockets.decorators import require_session
+from lorax.sockets.utils import is_csv_session_file
+
+
+async def _get_or_parse_csv_tree_graph(ctx, session_id: str, tree_idx: int):
+    """
+    Return a parsed CSV Newick tree graph, using CsvTreeGraphCache.
+
+    The cache is populated by the layout pipeline, but node search/highlight
+    should also be able to populate it on demand.
+    """
+    cached = await csv_tree_graph_cache.get(session_id, int(tree_idx))
+    if cached is not None:
+        return cached
+
+    # Import lazily to avoid CSV dependencies on tskit paths.
+    import pandas as pd
+    from lorax.csv.newick_tree import parse_newick_to_tree
+
+    df = ctx.tree_sequence
+    if not isinstance(df, pd.DataFrame):
+        return None
+
+    try:
+        newick_str = df.iloc[int(tree_idx)].get("newick")
+    except Exception:
+        newick_str = None
+    if newick_str is None or pd.isna(newick_str):
+        return None
+
+    # Parse with same normalization settings as the layout pipeline.
+    times_values = ctx.config.get("times", {}).get("values", [0.0, 1.0])
+    max_branch_length = float(times_values[1]) if len(times_values) > 1 else 1.0
+    samples_order = ctx.config.get("samples") or []
+    shift_tips_to_one = should_shift_csv_tips(ctx.file_path)
+    tree_max_branch_length = None
+    if "max_branch_length" in df.columns:
+        try:
+            v = df.iloc[int(tree_idx)].get("max_branch_length")
+            if v is not None and not (isinstance(v, float) and pd.isna(v)) and str(v).strip() != "":
+                tree_max_branch_length = float(v)
+        except Exception:
+            tree_max_branch_length = None
+
+    try:
+        graph = await asyncio.to_thread(
+            parse_newick_to_tree,
+            str(newick_str),
+            max_branch_length,
+            samples_order=samples_order,
+            tree_max_branch_length=tree_max_branch_length,
+            shift_tips_to_one=shift_tips_to_one,
+        )
+    except Exception:
+        return None
+
+    await csv_tree_graph_cache.set(session_id, int(tree_idx), graph)
+    return graph
+
+
+def _find_node_index(graph, node_id: int):
+    """Find the array index in a NewickTreeGraph for a given node_id."""
+    # node_id is stored as a numpy array; use vectorized equality.
+    try:
+        idxs = (graph.node_id == int(node_id)).nonzero()[0]
+        if idxs.size == 0:
+            return None
+        return int(idxs[0])
+    except Exception:
+        return None
+
+
+def _build_node_id_to_index(graph):
+    """Build a node_id -> array index dict for a NewickTreeGraph."""
+    try:
+        return {int(nid): i for i, nid in enumerate(graph.node_id)}
+    except Exception:
+        return {}
+
+
+def _compute_csv_lineage_path(graph, seed_node_id: int, node_id_to_index: dict):
+    """
+    Compute ancestry path (tip -> root) for a CSV NewickTreeGraph.
+
+    Returns a list of node_ids starting at seed_node_id, following parent_id until -1.
+    Includes a cycle guard to avoid infinite loops on malformed graphs.
+    """
+    path = []
+    current = int(seed_node_id)
+    visited = set()
+
+    while current != -1:
+        if current in visited:
+            break
+        visited.add(current)
+        path.append(current)
+
+        idx = node_id_to_index.get(current)
+        if idx is None:
+            break
+        parent = int(graph.parent_id[idx])
+        current = parent
+
+    return path
+
+
+def register_node_search_events(sio):
+    """Register node search socket events."""
+
+    @sio.event
+    async def search_nodes(sid, data):
+        """Socket event to search for nodes matching metadata values in trees.
+
+        This is used for highlighting nodes when searching/filtering by metadata.
+        Returns node_ids for matching samples in each tree, and optionally lineage paths.
+        Frontend computes positions using the post-order layout data.
+
+        data: {
+            lorax_sid: str,
+            sample_names: [str],    # Sample names to search for
+            tree_indices: [int],    # Tree indices to search in
+            show_lineages: bool,    # Whether to compute lineage paths
+            sample_colors: dict     # Optional {sample_name: [r,g,b,a]}
+        }
+
+        Returns: {
+            highlights: {tree_idx: [{node_id, name}]},
+            lineage: {tree_idx: [{path_node_ids: [int], color}]}
+        }
+        """
+        try:
+            lorax_sid = data.get("lorax_sid")
+            session = await require_session(lorax_sid, sid, sio)
+            if not session:
+                return
+
+            if not session.file_path:
+                print(f"⚠️ No file loaded for session {lorax_sid}")
+                await sio.emit("error", {
+                    "code": ERROR_NO_FILE_LOADED,
+                    "message": "No file loaded. Please load a file first."
+                }, to=sid)
+                return
+
+            if is_csv_session_file(session.file_path):
+                sample_names = data.get("sample_names", [])
+                tree_indices = data.get("tree_indices", [])
+                show_lineages = data.get("show_lineages", False)
+
+                # CSV mode currently supports only tip highlights (no lineage paths).
+                if show_lineages:
+                    show_lineages = False
+
+                if not sample_names or not tree_indices:
+                    await sio.emit("search-nodes-result", {"highlights": {}, "lineage": {}}, to=sid)
+                    return
+
+                ctx = await get_file_context(session.file_path)
+                if ctx is None:
+                    await sio.emit("search-nodes-result", {"error": "Failed to load CSV"}, to=sid)
+                    return
+
+                samples_order = ctx.config.get("samples") or []
+                sample_id_map = {str(name): idx for idx, name in enumerate(samples_order)}
+
+                highlights = {}
+
+                for tree_idx in tree_indices:
+                    tree_idx = int(tree_idx)
+                    graph = await _get_or_parse_csv_tree_graph(ctx, lorax_sid, tree_idx)
+                    if graph is None:
+                        continue
+
+                    hits = []
+                    for name in sample_names:
+                        node_id = sample_id_map.get(str(name))
+                        if node_id is None:
+                            continue
+                        if _find_node_index(graph, node_id) is None:
+                            continue
+                        hits.append({"node_id": int(node_id), "name": str(name)})
+
+                    if hits:
+                        highlights[tree_idx] = hits
+
+                await sio.emit("search-nodes-result", {"highlights": highlights, "lineage": {}}, to=sid)
+                return
+
+            sample_names = data.get("sample_names", [])
+            tree_indices = data.get("tree_indices", [])
+            show_lineages = data.get("show_lineages", False)
+            sample_colors = data.get("sample_colors", {})
+
+            # Check cache and warn if trees not found (they should be cached from layout)
+            if tree_indices:
+                uncached = []
+                for tree_idx in tree_indices:
+                    cached = await tree_graph_cache.get(lorax_sid, int(tree_idx))
+                    if cached is None:
+                        uncached.append(tree_idx)
+                if uncached:
+                    print(f"⚠️ WARNING: Trees {uncached} not in cache for session {lorax_sid[:8]}... "
+                          f"(expected from layout render)")
+
+            if not sample_names or not tree_indices:
+                await sio.emit("search-nodes-result", {
+                    "highlights": {},
+                    "lineage": {}
+                }, to=sid)
+                return
+
+            ctx = await get_file_context(session.file_path)
+            if ctx is None:
+                await sio.emit("search-nodes-result", {
+                    "error": "Failed to load tree sequence"
+                }, to=sid)
+                return
+
+            ts = ctx.tree_sequence
+
+            result = await asyncio.to_thread(
+                search_nodes_in_trees,
+                ts,
+                sample_names,
+                tree_indices,
+                show_lineages,
+                sample_colors
+            )
+
+            await sio.emit("search-nodes-result", result, to=sid)
+        except Exception as e:
+            print(f"❌ Search nodes error: {e}")
+            await sio.emit("search-nodes-result", {"error": str(e)}, to=sid)
+
+    @sio.event
+    async def get_highlight_positions_event(sid, data):
+        """Socket event to get positions for all tip nodes matching a metadata value.
+
+        Returns positions for ALL matching nodes, ignoring sparsification.
+        Used for highlighting nodes that may not be currently rendered.
+
+        data: {
+            lorax_sid: str,
+            metadata_key: str,      # Metadata key to filter by
+            metadata_value: str,    # Metadata value to match
+            tree_indices: [int]     # Tree indices to compute positions for
+        }
+
+        Returns: {
+            positions: [{node_id, tree_idx, x, y}, ...]
+        }
+        """
+        try:
+            lorax_sid = data.get("lorax_sid")
+            session = await require_session(lorax_sid, sid, sio)
+            if not session:
+                return
+
+            if not session.file_path:
+                print(f"⚠️ No file loaded for session {lorax_sid}")
+                await sio.emit("error", {
+                    "code": ERROR_NO_FILE_LOADED,
+                    "message": "No file loaded. Please load a file first."
+                }, to=sid)
+                return
+
+            if is_csv_session_file(session.file_path):
+                metadata_key = data.get("metadata_key")
+                metadata_value = data.get("metadata_value")
+                tree_indices = data.get("tree_indices", [])
+
+                if metadata_key != "sample":
+                    await sio.emit("highlight-positions-result", {"positions": []}, to=sid)
+                    return
+
+                if metadata_value is None or not tree_indices:
+                    await sio.emit("highlight-positions-result", {"positions": []}, to=sid)
+                    return
+
+                ctx = await get_file_context(session.file_path)
+                if ctx is None:
+                    await sio.emit("highlight-positions-result", {"error": "Failed to load CSV"}, to=sid)
+                    return
+
+                samples_order = ctx.config.get("samples") or []
+                sample_id_map = {str(name): idx for idx, name in enumerate(samples_order)}
+                target_node_id = sample_id_map.get(str(metadata_value))
+                if target_node_id is None:
+                    await sio.emit("highlight-positions-result", {"positions": []}, to=sid)
+                    return
+
+                positions = []
+                for tree_idx in tree_indices:
+                    tree_idx = int(tree_idx)
+                    graph = await _get_or_parse_csv_tree_graph(ctx, lorax_sid, tree_idx)
+                    if graph is None:
+                        continue
+
+                    arr_idx = _find_node_index(graph, target_node_id)
+                    if arr_idx is None:
+                        continue
+
+                    # Match the coordinate convention used by CSV layout buffer:
+                    # build_csv_layout_response swaps (time->x, layout->y).
+                    positions.append(
+                        {
+                            "node_id": int(target_node_id),
+                            "tree_idx": int(tree_idx),
+                            # Match tskit highlight convention: x=layout, y=time.
+                            # (tskit highlight uses TreeGraph internal coords, not the swapped PyArrow coords)
+                            "x": float(graph.x[arr_idx]),
+                            "y": float(graph.y[arr_idx]),
+                        }
+                    )
+
+                await sio.emit("highlight-positions-result", {"positions": positions}, to=sid)
+                return
+
+            metadata_key = data.get("metadata_key")
+            metadata_value = data.get("metadata_value")
+            tree_indices = data.get("tree_indices", [])
+
+            if not metadata_key or metadata_value is None:
+                await sio.emit("highlight-positions-result", {
+                    "error": "Missing metadata_key or metadata_value"
+                }, to=sid)
+                return
+
+            if not tree_indices:
+                await sio.emit("highlight-positions-result", {"positions": []}, to=sid)
+                return
+
+            ctx = await get_file_context(session.file_path)
+            if ctx is None:
+                await sio.emit("highlight-positions-result", {
+                    "error": "Failed to load tree sequence"
+                }, to=sid)
+                return
+
+            ts = ctx.tree_sequence
+
+            result = await get_highlight_positions(
+                ts,
+                session.file_path,
+                metadata_key,
+                metadata_value,
+                tree_indices,
+                lorax_sid,
+                tree_graph_cache
+            )
+
+            await sio.emit("highlight-positions-result", result, to=sid)
+        except Exception as e:
+            print(f"❌ Get highlight positions error: {e}")
+            await sio.emit("highlight-positions-result", {"error": str(e)}, to=sid)
+
+    @sio.event
+    async def search_metadata_multi_event(sid, data):
+        """Socket event for multi-value metadata search.
+
+        Returns positions for tip nodes matching ANY of the metadata values,
+        grouped by value for per-value coloring with OR logic.
+
+        data: {
+            lorax_sid: str,
+            metadata_key: str,          # Metadata key to filter by
+            metadata_values: [str],     # Array of values (OR logic)
+            tree_indices: [int],        # Tree indices to compute positions for
+            show_lineages: bool         # Whether to compute lineage paths
+        }
+
+        Emits: "search-metadata-multi-result" with:
+        {
+            positions_by_value: {"Africa": [{node_id, tree_idx, x, y}, ...], ...},
+            lineages: {"Africa": {tree_idx: [{path_node_ids, color}]}} if show_lineages,
+            total_count: int
+        }
+        """
+        try:
+            lorax_sid = data.get("lorax_sid")
+            session = await require_session(lorax_sid, sid, sio)
+            if not session:
+                return
+
+            if not session.file_path:
+                print(f"⚠️ No file loaded for session {lorax_sid}")
+                await sio.emit("error", {
+                    "code": ERROR_NO_FILE_LOADED,
+                    "message": "No file loaded. Please load a file first."
+                }, to=sid)
+                return
+
+            if is_csv_session_file(session.file_path):
+                metadata_key = data.get("metadata_key")
+                metadata_values = data.get("metadata_values", [])
+                tree_indices = data.get("tree_indices", [])
+                show_lineages = bool(data.get("show_lineages", False))
+
+                if metadata_key != "sample":
+                    await sio.emit(
+                        "search-metadata-multi-result",
+                        {"positions_by_value": {}, "lineages": {}, "total_count": 0},
+                        to=sid,
+                    )
+                    return
+
+                if not metadata_values or not tree_indices:
+                    await sio.emit(
+                        "search-metadata-multi-result",
+                        {"positions_by_value": {}, "lineages": {}, "total_count": 0},
+                        to=sid,
+                    )
+                    return
+
+                ctx = await get_file_context(session.file_path)
+                if ctx is None:
+                    await sio.emit("search-metadata-multi-result", {"error": "Failed to load CSV"}, to=sid)
+                    return
+
+                samples_order = ctx.config.get("samples") or []
+                sample_id_map = {str(name): idx for idx, name in enumerate(samples_order)}
+
+                # Deduplicate and stringify values
+                unique_values = list({str(v) for v in metadata_values})
+                positions_by_value = {v: [] for v in unique_values}
+                lineages = {} if show_lineages else {}
+                total_count = 0
+
+                for value in unique_values:
+                    node_id = sample_id_map.get(value)
+                    if node_id is None:
+                        continue
+
+                    value_lineages = {} if show_lineages else None
+
+                    for tree_idx in tree_indices:
+                        tree_idx = int(tree_idx)
+                        graph = await _get_or_parse_csv_tree_graph(ctx, lorax_sid, tree_idx)
+                        if graph is None:
+                            continue
+                        node_id_to_index = _build_node_id_to_index(graph) if show_lineages else None
+                        arr_idx = _find_node_index(graph, node_id)
+                        if arr_idx is None:
+                            continue
+
+                        positions_by_value[value].append(
+                            {
+                                "node_id": int(node_id),
+                                "tree_idx": int(tree_idx),
+                                # Match tskit highlight convention: x=layout, y=time.
+                                "x": float(graph.x[arr_idx]),
+                                "y": float(graph.y[arr_idx]),
+                            }
+                        )
+                        total_count += 1
+
+                        if show_lineages:
+                            path_node_ids = _compute_csv_lineage_path(graph, int(node_id), node_id_to_index)
+                            if len(path_node_ids) > 1:
+                                # Emit root -> tip to match frontend L-shape construction.
+                                path_node_ids = list(reversed(path_node_ids))
+                                value_lineages.setdefault(tree_idx, []).append(
+                                    {"path_node_ids": path_node_ids, "color": None}
+                                )
+
+                    if show_lineages and value_lineages:
+                        lineages[value] = value_lineages
+
+                await sio.emit(
+                    "search-metadata-multi-result",
+                    {"positions_by_value": positions_by_value, "lineages": lineages, "total_count": total_count},
+                    to=sid,
+                )
+                return
+
+            metadata_key = data.get("metadata_key")
+            metadata_values = data.get("metadata_values", [])
+            tree_indices = data.get("tree_indices", [])
+            show_lineages = data.get("show_lineages", False)
+
+            if not metadata_key:
+                await sio.emit("search-metadata-multi-result", {
+                    "error": "Missing metadata_key"
+                }, to=sid)
+                return
+
+            if not metadata_values or not tree_indices:
+                await sio.emit("search-metadata-multi-result", {
+                    "positions_by_value": {},
+                    "lineages": {},
+                    "total_count": 0
+                }, to=sid)
+                return
+
+            ctx = await get_file_context(session.file_path)
+            if ctx is None:
+                await sio.emit("search-metadata-multi-result", {
+                    "error": "Failed to load tree sequence"
+                }, to=sid)
+                return
+
+            ts = ctx.tree_sequence
+
+            result = await get_multi_value_highlight_positions(
+                ts,
+                session.file_path,
+                metadata_key,
+                metadata_values,
+                tree_indices,
+                lorax_sid,
+                tree_graph_cache,
+                show_lineages
+            )
+
+            await sio.emit("search-metadata-multi-result", result, to=sid)
+        except Exception as e:
+            print(f"❌ Search metadata multi error: {e}")
+            await sio.emit("search-metadata-multi-result", {"error": str(e)}, to=sid)

lorax/sockets/tree_layout.py

@@ -0,0 +1,117 @@
+"""
+Tree layout event handlers for Lorax Socket.IO.
+
+Handles process_postorder_layout and cache_trees events.
+"""
+
+from lorax.context import tree_graph_cache, csv_tree_graph_cache
+from lorax.handlers import handle_tree_graph_query, ensure_trees_cached
+from lorax.sockets.decorators import require_session
+from lorax.sockets.utils import is_csv_session_file
+
+
+def register_tree_layout_events(sio):
+    """Register tree layout socket events."""
+
+    @sio.event
+    async def process_postorder_layout(sid, data):
+        """Socket event to get post-order tree traversal for efficient rendering.
+
+        Returns PyArrow IPC binary data with post-order node arrays.
+        Frontend computes layout using stack-based reconstruction.
+
+        Uses Socket.IO acknowledgement callback pattern - returns result directly
+        instead of emitting to ensure request-response correlation.
+        """
+        try:
+            lorax_sid = data.get("lorax_sid")
+            session = await require_session(lorax_sid, sid, sio)
+            if not session:
+                return {"error": "Session not found", "request_id": data.get("request_id")}
+
+            if not session.file_path:
+                print(f"⚠️ No file loaded for session {lorax_sid}")
+                return {"error": "No file loaded for session", "request_id": data.get("request_id")}
+
+            display_array = data.get("displayArray", [])
+            actual_display_array = data.get("actualDisplayArray", display_array)
+            sparsification = data.get("sparsification", False)
+            request_id = data.get("request_id")
+
+            # handle_tree_graph_query returns dict with PyArrow buffer (Numba-optimized)
+            # Pass session_id and tree_graph_cache for caching TreeGraph objects
+            # actual_display_array contains all visible trees for cache eviction
+            result = await handle_tree_graph_query(
+                session.file_path,
+                display_array,
+                sparsification=sparsification,
+                session_id=lorax_sid,
+                tree_graph_cache=tree_graph_cache,
+                csv_tree_graph_cache=csv_tree_graph_cache,
+                actual_display_array=actual_display_array
+            )
+
+            if "error" in result:
+                return {"error": result["error"], "request_id": request_id}
+            else:
+                # Return result directly - Socket.IO sends as acknowledgement callback
+                return {
+                    "buffer": result["buffer"],  # Binary PyArrow IPC data
+                    "global_min_time": result["global_min_time"],
+                    "global_max_time": result["global_max_time"],
+                    "tree_indices": result["tree_indices"],
+                    "request_id": request_id
+                }
+        except Exception as e:
+            print(f"❌ Postorder layout query error: {e}")
+            return {"error": str(e), "request_id": data.get("request_id")}
+
+    @sio.event
+    async def cache_trees(sid, data):
+        """Socket event to pre-cache TreeGraph objects for lineage operations.
+
+        Call this after process_postorder_layout to enable subsequent lineage queries.
+
+        data: {
+            lorax_sid: str,
+            tree_indices: [int]  # Tree indices to cache
+        }
+
+        Returns: {
+            cached_count: int,  # Number of trees newly cached
+            total_cached: int   # Total trees now in cache for session
+        }
+        """
+        try:
+            lorax_sid = data.get("lorax_sid")
+            session = await require_session(lorax_sid, sid, sio)
+            if not session:
+                return {"error": "Session not found", "cached_count": 0}
+
+            if not session.file_path:
+                return {"error": "No file loaded", "cached_count": 0}
+
+            if is_csv_session_file(session.file_path):
+                return {"error": "Lineage not supported for CSV", "cached_count": 0}
+
+            tree_indices = data.get("tree_indices", [])
+            if not tree_indices:
+                return {"cached_count": 0, "total_cached": 0}
+
+            newly_cached = await ensure_trees_cached(
+                session.file_path,
+                tree_indices,
+                lorax_sid,
+                tree_graph_cache
+            )
+
+            # Get total cached
+            all_cached = await tree_graph_cache.get_all_for_session(lorax_sid)
+
+            return {
+                "cached_count": newly_cached,
+                "total_cached": len(all_cached)
+            }
+        except Exception as e:
+            print(f"❌ Cache trees error: {e}")
+            return {"error": str(e), "cached_count": 0}

lorax/sockets/utils.py

@@ -0,0 +1,10 @@
+"""
+Socket utility functions for Lorax.
+
+Common helpers shared across socket event handlers.
+"""
+
+
+def is_csv_session_file(file_path: str | None) -> bool:
+    """Check if the session file is a CSV file."""
+    return bool(file_path) and str(file_path).lower().endswith(".csv")

lorax/tree_graph/__init__.py

@@ -0,0 +1,12 @@
+"""
+tree_graph - Optimized tree construction with Numba JIT compilation.
+
+This module provides fast tree construction from tskit tables with:
+- Numba-compiled post-order traversal (50-100x faster than Python)
+- CSR format for efficient children access
+- PyArrow serialization for frontend rendering
+"""
+
+from .tree_graph import TreeGraph, construct_tree, construct_trees_batch
+
+__all__ = ['TreeGraph', 'construct_tree', 'construct_trees_batch']