lorax-arg 0.1__py3-none-any.whl

lorax/sockets/metadata.py

@@ -0,0 +1,232 @@
+"""
+Metadata event handlers for Lorax Socket.IO.
+
+Handles fetch_metadata_for_key, search_metadata, and fetch_metadata_array events.
+"""
+
+import asyncio
+
+from lorax.constants import ERROR_NO_FILE_LOADED
+from lorax.metadata.loader import (
+    get_metadata_for_key, search_samples_by_metadata, get_metadata_array_for_key
+)
+from lorax.cache import get_file_context
+from lorax.sockets.decorators import require_session
+from lorax.sockets.utils import is_csv_session_file
+
+
+def register_metadata_events(sio):
+    """Register metadata-related socket events."""
+
+    @sio.event
+    async def fetch_metadata_for_key(sid, data):
+        """Socket event to fetch metadata mapping for a specific key."""
+        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):
+                key = data.get("key")
+                if key == "sample":
+                    ctx = await get_file_context(session.file_path)
+                    if ctx is None:
+                        await sio.emit("metadata-key-result", {"error": "Failed to load CSV"}, to=sid)
+                        return
+                    # sample_names is {name: {"sample_name": name}, ...}
+                    sample_names = ctx.config.get("sample_names", {})
+                    # Return {name: name} like tskit's "sample" key
+                    result = {name: name for name in sample_names.keys()}
+                    await sio.emit("metadata-key-result", {"key": key, "data": result}, to=sid)
+                    return
+                else:
+                    await sio.emit("metadata-key-result", {
+                        "error": f"Metadata key '{key}' is not supported for CSV files."
+                    }, to=sid)
+                    return
+
+            key = data.get("key")
+            if not key:
+                await sio.emit("metadata-key-result", {
+                    "error": "Missing 'key' parameter"
+                }, to=sid)
+                return
+
+            ctx = await get_file_context(session.file_path)
+            if ctx is None:
+                await sio.emit("metadata-key-result", {
+                    "error": "Failed to load tree sequence"
+                }, to=sid)
+                return
+
+            # Pass FileContext to metadata function
+            result = await asyncio.to_thread(get_metadata_for_key, ctx, key)
+            await sio.emit("metadata-key-result", {"key": key, "data": result}, to=sid)
+        except Exception as e:
+            print(f"❌ Metadata fetch error: {e}")
+            await sio.emit("metadata-key-result", {"error": str(e)}, to=sid)
+
+    @sio.event
+    async def search_metadata(sid, data):
+        """Socket event to search for samples matching a metadata value."""
+        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):
+                key = data.get("key")
+                value = data.get("value")
+                if key == "sample":
+                    ctx = await get_file_context(session.file_path)
+                    if ctx is None:
+                        await sio.emit("search-result", {"error": "Failed to load CSV"}, to=sid)
+                        return
+                    sample_names = ctx.config.get("sample_names", {})
+                    # Return matching sample names (exact match)
+                    matching = [name for name in sample_names.keys() if name == value]
+                    await sio.emit("search-result", {"key": key, "value": value, "samples": matching}, to=sid)
+                    return
+                else:
+                    await sio.emit("search-result", {
+                        "error": f"Metadata key '{key}' is not supported for CSV files."
+                    }, to=sid)
+                    return
+
+            key = data.get("key")
+            value = data.get("value")
+
+            if not key or value is None:
+                await sio.emit("search-result", {
+                    "error": "Missing 'key' or 'value' parameter"
+                }, to=sid)
+                return
+
+            ctx = await get_file_context(session.file_path)
+            if ctx is None:
+                await sio.emit("search-result", {
+                    "error": "Failed to load tree sequence"
+                }, to=sid)
+                return
+
+            # Pass FileContext to metadata function
+            result = await asyncio.to_thread(
+                search_samples_by_metadata, ctx, key, value
+            )
+            await sio.emit("search-result", {
+                "key": key,
+                "value": value,
+                "samples": result
+            }, to=sid)
+        except Exception as e:
+            print(f"❌ Search error: {e}")
+            await sio.emit("search-result", {"error": str(e)}, to=sid)
+
+    @sio.event
+    async def fetch_metadata_array(sid, data):
+        """Socket event to fetch metadata as efficient PyArrow array format.
+
+        This is optimized for large tree sequences (1M+ samples) where JSON
+        serialization would be too slow/large. Returns binary Arrow IPC data
+        with indices that map node_id -> value index.
+        """
+        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):
+                key = data.get("key")
+                if key == "sample":
+                    ctx = await get_file_context(session.file_path)
+                    if ctx is None:
+                        await sio.emit("metadata-array-result", {"error": "Failed to load CSV"}, to=sid)
+                        return
+
+                    sample_names = ctx.config.get("sample_names", {})
+                    names_list = list(sample_names.keys())
+
+                    # Build PyArrow array where each sample maps to its own unique index
+                    import numpy as np
+                    import pyarrow as pa
+
+                    unique_values = names_list
+                    indices = np.arange(len(names_list), dtype=np.uint32)
+
+                    # Create Arrow IPC buffer
+                    table = pa.table({'idx': pa.array(indices, type=pa.uint32())})
+                    sink = pa.BufferOutputStream()
+                    writer = pa.ipc.new_stream(sink, table.schema)
+                    writer.write_table(table)
+                    writer.close()
+
+                    await sio.emit("metadata-array-result", {
+                        "key": key,
+                        "unique_values": unique_values,
+                        "sample_node_ids": list(range(len(names_list))),  # Sequential indices for CSV
+                        "buffer": sink.getvalue().to_pybytes()
+                    }, to=sid)
+                    return
+                else:
+                    await sio.emit("metadata-array-result", {
+                        "error": f"Metadata key '{key}' is not supported for CSV files."
+                    }, to=sid)
+                    return
+
+            key = data.get("key")
+            if not key:
+                await sio.emit("metadata-array-result", {
+                    "error": "Missing 'key' parameter"
+                }, to=sid)
+                return
+
+            ctx = await get_file_context(session.file_path)
+            if ctx is None:
+                await sio.emit("metadata-array-result", {
+                    "error": "Failed to load tree sequence"
+                }, to=sid)
+                return
+
+            # Pass FileContext to metadata function
+            result = await asyncio.to_thread(
+                get_metadata_array_for_key, ctx, key
+            )
+
+            # Send metadata with Arrow buffer as binary
+            await sio.emit("metadata-array-result", {
+                "key": key,
+                "unique_values": result['unique_values'],
+                "sample_node_ids": result['sample_node_ids'],
+                "buffer": result['arrow_buffer']  # Binary data
+            }, to=sid)
+
+        except Exception as e:
+            print(f"❌ Metadata array fetch error: {e}")
+            await sio.emit("metadata-array-result", {"error": str(e)}, to=sid)

lorax/sockets/mutations.py

@@ -0,0 +1,154 @@
+"""
+Mutation event handlers for Lorax Socket.IO.
+
+Handles query_mutations_window and search_mutations events.
+"""
+
+import asyncio
+
+from lorax.constants import ERROR_NO_FILE_LOADED
+from lorax.metadata.mutations import (
+    get_mutations_in_window, search_mutations_by_position
+)
+from lorax.buffer import mutations_to_arrow_buffer
+from lorax.cache import get_file_context
+from lorax.sockets.decorators import require_session
+from lorax.sockets.utils import is_csv_session_file
+
+
+def register_mutations_events(sio):
+    """Register mutation-related socket events."""
+
+    @sio.event
+    async def query_mutations_window(sid, data):
+        """Socket event to fetch mutations within a genomic window.
+
+        Returns PyArrow IPC binary data with mutations in the specified range.
+        Supports pagination via offset and limit parameters.
+        """
+        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):
+                await sio.emit("mutations-window-result", {
+                    "error": "Mutations are not supported for CSV yet."
+                }, to=sid)
+                return
+
+            start = data.get("start", 0)
+            end = data.get("end", 0)
+            offset = data.get("offset", 0)
+            limit = data.get("limit", 1000)
+
+            ctx = await get_file_context(session.file_path)
+            if ctx is None:
+                await sio.emit("mutations-window-result", {
+                    "error": "Failed to load tree sequence"
+                }, to=sid)
+                return
+
+            ts = ctx.tree_sequence
+
+            # Get mutations in the window
+            result = await asyncio.to_thread(
+                get_mutations_in_window, ts, start, end, offset, limit
+            )
+
+            # Convert to PyArrow buffer
+            buffer = await asyncio.to_thread(mutations_to_arrow_buffer, result)
+
+            await sio.emit("mutations-window-result", {
+                "buffer": buffer,
+                "total_count": result['total_count'],
+                "has_more": result['has_more'],
+                "start": start,
+                "end": end,
+                "offset": offset,
+                "limit": limit
+            }, to=sid)
+
+        except Exception as e:
+            print(f"❌ Mutations window query error: {e}")
+            await sio.emit("mutations-window-result", {"error": str(e)}, to=sid)
+
+    @sio.event
+    async def search_mutations(sid, data):
+        """Socket event to search mutations by position with configurable range.
+
+        Returns PyArrow IPC binary data with mutations sorted by distance from position.
+        Supports pagination via offset and limit parameters.
+        """
+        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):
+                await sio.emit("mutations-search-result", {
+                    "error": "Mutations search is not supported for CSV yet."
+                }, to=sid)
+                return
+
+            position = data.get("position")
+            if position is None:
+                await sio.emit("mutations-search-result", {
+                    "error": "Missing 'position' parameter"
+                }, to=sid)
+                return
+
+            range_bp = data.get("range_bp", 5000)
+            offset = data.get("offset", 0)
+            limit = data.get("limit", 1000)
+
+            ctx = await get_file_context(session.file_path)
+            if ctx is None:
+                await sio.emit("mutations-search-result", {
+                    "error": "Failed to load tree sequence"
+                }, to=sid)
+                return
+
+            ts = ctx.tree_sequence
+
+            # Search mutations around the position
+            result = await asyncio.to_thread(
+                search_mutations_by_position, ts, position, range_bp, offset, limit
+            )
+
+            # Convert to PyArrow buffer
+            buffer = await asyncio.to_thread(mutations_to_arrow_buffer, result)
+
+            await sio.emit("mutations-search-result", {
+                "buffer": buffer,
+                "total_count": result['total_count'],
+                "has_more": result['has_more'],
+                "search_start": result['search_start'],
+                "search_end": result['search_end'],
+                "position": position,
+                "range_bp": range_bp,
+                "offset": offset,
+                "limit": limit
+            }, to=sid)
+
+        except Exception as e:
+            print(f"❌ Mutations search error: {e}")
+            await sio.emit("mutations-search-result", {"error": str(e)}, to=sid)