lorax-arg 0.1__py3-none-any.whl

lorax/manager.py

@@ -0,0 +1,58 @@
+from fastapi.websockets import WebSocket, WebSocketState, WebSocketDisconnect
+
+
+class WebSocketManager:
+    def __init__(self):
+        self.connected_clients = set()
+        self.client_component = {}  # websocket -> set(component names)
+
+    async def connect(self, websocket: WebSocket):
+        client_ip = f"{websocket.client.host}:{websocket.client.port}"
+        await websocket.accept()
+        self.connected_clients.add(websocket)
+
+    async def cleanup_disconnected_clients(self):
+        dead = [ws for ws in self.connected_clients if not self.is_connected(ws)]
+        for ws in dead:
+            await self.disconnect(ws)
+
+    def get_connected_clients(self):
+        """Get list of currently connected clients"""
+        return [client for client in self.connected_clients if self.is_connected(client)]
+
+    def is_connected(self, ws: WebSocket) -> bool:
+        try:
+            return ws.client_state == WebSocketState.CONNECTED
+        except Exception:
+            return False
+
+    async def send_message(self, ws: WebSocket, message: dict):
+        if self.is_connected(ws):
+            try:
+                await ws.send_json(message)
+            except Exception as e:
+                print(f"send_message error: {e}")
+                await self.disconnect(ws)
+        else:
+            await self.disconnect(ws)
+
+    async def disconnect(self, websocket: WebSocket):
+        if websocket in self.connected_clients:
+            self.connected_clients.remove(websocket)
+        if websocket in self.client_component:
+            del self.client_component[websocket]
+
+    async def register_component(self, ws: WebSocket, component: str):
+        self.client_component.setdefault(ws, set()).add(component)
+
+    async def send_to_component(self, component: str, message: dict):
+        """Send only to sockets in *this* session that registered that component."""
+        targets = [
+            ws for ws, comps in self.client_component.items()
+            if component in comps and self.is_connected(ws)
+        ]
+        if not targets:
+            print(f"⚠️ No sockets registered for component: {component}")
+            return
+        for ws in targets:
+            await self.send_message(ws, message)

lorax/metadata/loader.py

@@ -0,0 +1,426 @@
+"""
+Metadata extraction and caching for tree sequences.
+
+Functions accept FileContext and use its nested metadata cache.
+When a FileContext is evicted, its metadata cache is evicted together.
+"""
+
+import json
+import numpy as np
+import tskit
+import pyarrow as pa
+from collections import defaultdict
+from lorax.utils import ensure_json_dict, make_json_safe, make_json_serializable
+
+
+def get_metadata_for_key(
+    ctx,
+    key,
+    sources=("individual", "node", "population"),
+    sample_name_key="name"
+):
+    """
+    Get sample-to-value mapping for a specific metadata key.
+    Results are cached in the FileContext's nested metadata cache.
+
+    Parameters
+    ----------
+    ctx : FileContext
+        The file context containing tree_sequence and metadata cache
+    key : str
+        The metadata key to extract
+    sources : tuple
+        Any of ("individual", "node", "population")
+    sample_name_key : str
+        Key in node metadata used as sample name
+
+    Returns
+    -------
+    dict
+        {sample_name: value} for the specified key
+    """
+    # Check nested cache in FileContext
+    cached = ctx.get_metadata(key)
+    if cached is not None:
+        print(f"✅ Using cached metadata for key: {key}")
+        return cached
+
+    ts = ctx.tree_sequence
+
+    # Special handling for "sample" key - each sample's value is its own name
+    if key == "sample":
+        result = {}
+        for node_id in ts.samples():
+            node = ts.node(node_id)
+            node_meta = node.metadata or {}
+            try:
+                node_meta = ensure_json_dict(node_meta)
+            except (TypeError, json.JSONDecodeError):
+                node_meta = {}
+            sample_name = str(node_meta.get(sample_name_key, f"{node_id}"))
+            result[sample_name] = sample_name
+        ctx.set_metadata(key, result)
+        return result
+
+    result = {}
+
+    for node_id in ts.samples():
+        node = ts.node(node_id)
+        node_meta = node.metadata or {}
+        node_meta = ensure_json_dict(node_meta)
+        sample_name = node_meta.get(sample_name_key, f"{node_id}")
+
+        for source in sources:
+            if source == "individual":
+                if node.individual == tskit.NULL:
+                    continue
+                meta = ts.individual(node.individual).metadata
+                meta = meta or {}
+                meta = ensure_json_dict(meta)
+
+            elif source == "node":
+                meta = node_meta
+
+            elif source == "population":
+                if node.population == tskit.NULL:
+                    continue
+                meta = ts.population(node.population).metadata
+                meta = meta or {}
+                meta = ensure_json_dict(meta)
+
+            else:
+                continue
+
+            if not meta:
+                continue
+
+            if key in meta:
+                value = meta[key]
+                if value is None:
+                    break  # Skip None values
+                if isinstance(value, (list, dict)):
+                    value = repr(value)
+                result[sample_name] = str(value)
+                break  # Found the key, move to next sample
+
+    ctx.set_metadata(key, result)
+    return result
+
+
+def search_samples_by_metadata(
+    ctx,
+    key,
+    value,
+    sources=("individual", "node", "population"),
+    sample_name_key="name"
+):
+    """
+    Search for samples that have a specific value for a metadata key.
+
+    Parameters
+    ----------
+    ctx : FileContext
+        The file context containing tree_sequence and metadata cache
+    key : str
+        The metadata key to search
+    value : str
+        The value to match
+    sources : tuple
+        Any of ("individual", "node", "population")
+    sample_name_key : str
+        Key in node metadata used as sample name
+
+    Returns
+    -------
+    list
+        List of sample names matching the criteria
+    """
+    ts = ctx.tree_sequence
+
+    # Try to use cached metadata if available
+    cached = ctx.get_metadata(key)
+
+    if cached is not None:
+        # Use cached data for fast lookup
+        return [sample for sample, val in cached.items() if str(val) == str(value)]
+
+    # If not cached, compute on the fly
+    matching_samples = []
+
+    for node_id in ts.samples():
+        node = ts.node(node_id)
+        node_meta = node.metadata or {}
+        node_meta = ensure_json_dict(node_meta)
+        sample_name = node_meta.get(sample_name_key, f"{node_id}")
+
+        for source in sources:
+            if source == "individual":
+                if node.individual == tskit.NULL:
+                    continue
+                meta = ts.individual(node.individual).metadata
+                meta = meta or {}
+                meta = ensure_json_dict(meta)
+
+            elif source == "node":
+                meta = node_meta
+
+            elif source == "population":
+                if node.population == tskit.NULL:
+                    continue
+                meta = ts.population(node.population).metadata
+                meta = meta or {}
+                meta = ensure_json_dict(meta)
+
+            else:
+                continue
+
+            if not meta:
+                continue
+
+            if key in meta:
+                meta_value = meta[key]
+                if meta_value is None:
+                    break  # Skip None values
+                if isinstance(meta_value, (list, dict)):
+                    meta_value = repr(meta_value)
+                if str(meta_value) == str(value):
+                    matching_samples.append(sample_name)
+                break
+
+    return matching_samples
+
+
+def _get_sample_metadata_value(ts, node_id, key, sources, sample_name_key="name"):
+    """
+    Helper to get a specific metadata value for a sample node.
+    Returns (sample_name, value) tuple.
+    """
+    node = ts.node(node_id)
+    node_meta = node.metadata or {}
+    node_meta = ensure_json_dict(node_meta)
+    sample_name = node_meta.get(sample_name_key, f"{node_id}")
+
+    # Special handling for "sample" key: it is not a real metadata field in tskit.
+    # Treat it as identity so "sample" searches/highlights match by sample name.
+    if key == "sample":
+        sample_name = str(sample_name)
+        return (sample_name, sample_name)
+
+    for source in sources:
+        if source == "individual":
+            if node.individual == tskit.NULL:
+                continue
+            meta = ts.individual(node.individual).metadata
+            meta = meta or {}
+            meta = ensure_json_dict(meta)
+
+        elif source == "node":
+            meta = node_meta
+
+        elif source == "population":
+            if node.population == tskit.NULL:
+                continue
+            meta = ts.population(node.population).metadata
+            meta = meta or {}
+            meta = ensure_json_dict(meta)
+
+        else:
+            continue
+
+        if not meta:
+            continue
+
+        if key in meta:
+            value = meta[key]
+            if isinstance(value, (list, dict)):
+                value = repr(value)
+            return (sample_name, value)
+
+    return (sample_name, None)
+
+
+def get_metadata_array_for_key(
+    ctx,
+    key,
+    sources=("individual", "node", "population"),
+    sample_name_key="name"
+):
+    """
+    Build efficient array-based metadata for a key using PyArrow.
+
+    Returns indices array where indices[i] is the index into unique_values
+    for the i-th sample (ordered by node_id from ts.samples()).
+
+    Parameters
+    ----------
+    ctx : FileContext
+        The file context containing tree_sequence and metadata cache
+    key : str
+        The metadata key to extract
+    sources : tuple
+        Any of ("individual", "node", "population")
+    sample_name_key : str
+        Key in node metadata used as sample name
+
+    Returns
+    -------
+    dict
+        {
+            'unique_values': [val0, val1, ...],  # Index i -> value string
+            'sample_node_ids': [node_id0, node_id1, ...],  # Sample order
+            'arrow_buffer': bytes  # PyArrow IPC serialized indices
+        }
+    """
+    cache_key = f"{key}:array"
+    cached = ctx.get_metadata(cache_key)
+    if cached is not None:
+        print(f"✅ Using cached metadata array for key: {key}")
+        return cached
+
+    ts = ctx.tree_sequence
+    sample_ids = list(ts.samples())
+    n_samples = len(sample_ids)
+
+    # Special handling for "sample" key - each sample's name is its own unique value
+    if key == "sample":
+        unique_values = []
+        value_to_idx = {}
+        indices = np.zeros(n_samples, dtype=np.uint32)
+
+        for i, node_id in enumerate(sample_ids):
+            node = ts.node(node_id)
+            node_meta = node.metadata or {}
+            try:
+                node_meta = ensure_json_dict(node_meta)
+            except (TypeError, json.JSONDecodeError):
+                node_meta = {}
+            sample_name = str(node_meta.get(sample_name_key, f"{node_id}"))
+
+            if sample_name not in value_to_idx:
+                value_to_idx[sample_name] = len(unique_values)
+                unique_values.append(sample_name)
+
+            indices[i] = value_to_idx[sample_name]
+
+        # Serialize to Arrow IPC format
+        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()
+
+        result = {
+            'unique_values': unique_values,
+            'sample_node_ids': [int(x) for x in sample_ids],
+            'arrow_buffer': sink.getvalue().to_pybytes()
+        }
+        ctx.set_metadata(cache_key, result)
+        print(f"✅ Built sample metadata array ({n_samples} samples, {len(unique_values)} unique values)")
+        return result
+
+    unique_values = []
+    value_to_idx = {}
+    indices = np.zeros(n_samples, dtype=np.uint32)
+
+    for i, node_id in enumerate(sample_ids):
+        sample_name, value = _get_sample_metadata_value(ts, node_id, key, sources, sample_name_key)
+
+        if value is None:
+            value = ""  # Handle missing values
+
+        value_str = str(value)
+
+        if value_str not in value_to_idx:
+            value_to_idx[value_str] = len(unique_values)
+            unique_values.append(value_str)
+
+        indices[i] = value_to_idx[value_str]
+
+    # Serialize to Arrow IPC format
+    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()
+
+    result = {
+        'unique_values': unique_values,
+        'sample_node_ids': [int(x) for x in sample_ids],  # Convert to Python int for JSON
+        'arrow_buffer': sink.getvalue().to_pybytes()
+    }
+
+    ctx.set_metadata(cache_key, result)
+    print(f"✅ Built metadata array for key: {key} ({n_samples} samples, {len(unique_values)} unique values)")
+    return result
+
+
+def get_metadata_schema(
+    ts,
+    sources=("individual", "node", "population"),
+    sample_name_key="name"
+):
+    """
+    Extract metadata keys only (values are fetched on-demand via get_metadata_array_for_key).
+
+    Also includes "sample" as the first key, where each sample's name/ID is its
+    own unique value (for coloring samples individually).
+
+    Parameters
+    ----------
+    ts : tskit.TreeSequence
+        The tree sequence (not FileContext - this doesn't need caching)
+    sources : tuple
+        Any of ("individual", "node", "population")
+    sample_name_key : str
+        Key in node metadata used as sample name
+
+    Returns
+    -------
+    dict
+        {
+            "metadata_keys": [key1, key2, ...]
+        }
+    """
+    keys = set()
+
+    for node_id in ts.samples():
+        node = ts.node(node_id)
+
+        # Parse node metadata
+        node_meta = node.metadata or {}
+        try:
+            node_meta = ensure_json_dict(node_meta)
+        except (TypeError, json.JSONDecodeError):
+            node_meta = {}
+
+        for source in sources:
+            if source == "individual":
+                if node.individual == tskit.NULL:
+                    continue
+                meta = ts.individual(node.individual).metadata
+                meta = meta or {}
+                meta = ensure_json_dict(meta)
+
+            elif source == "node":
+                meta = node_meta  # Reuse already parsed node metadata
+
+            elif source == "population":
+                if node.population == tskit.NULL:
+                    continue
+                meta = ts.population(node.population).metadata
+                meta = meta or {}
+                meta = ensure_json_dict(meta)
+
+            else:
+                raise ValueError(f"Unknown source: {source}")
+
+            if not meta:
+                continue
+
+            for key in meta.keys():
+                keys.add(key)
+
+    # Prepend "sample" to keys - this makes it the default colorBy option
+    return {
+        "metadata_keys": ["sample"] + sorted(list(keys))
+    }

lorax/metadata/mutations.py

@@ -0,0 +1,146 @@
+
+import numpy as np
+import tskit
+
+def get_mutations_in_window(ts, start, end, offset=0, limit=1000):
+    """
+    Get mutations within a genomic interval [start, end) with pagination.
+
+    Args:
+        ts: tskit.TreeSequence
+        start: Start genomic position (bp)
+        end: End genomic position (bp)
+        offset: Number of mutations to skip (for pagination)
+        limit: Maximum number of mutations to return
+
+    Returns:
+        dict with:
+            - 'mutations': list of mutation dicts
+            - 'total_count': total mutations in window (for pagination)
+            - 'has_more': whether there are more mutations
+    """
+    t = ts.tables
+    sites = t.sites
+    mutations = t.mutations
+
+    # Get positions for all mutations via their sites
+    positions = sites.position[mutations.site]
+
+    # Create mask for mutations in the window
+    mask = (positions >= start) & (positions < end)
+    indices = np.where(mask)[0]
+
+    total_count = len(indices)
+
+    # Apply pagination
+    paginated_indices = indices[offset:offset + limit]
+
+    # Extract mutation data
+    result_mutations = []
+    for idx in paginated_indices:
+        mut = mutations[idx]
+        site = sites[mut.site]
+        position = int(site.position)
+        site_id = int(mut.site)
+        node_id = int(mut.node)
+
+        # Get ancestral and derived states
+        ancestral_state = site.ancestral_state
+        derived_state = mut.derived_state
+
+        result_mutations.append({
+            'position': position,
+            'mutation': f"{ancestral_state}->{derived_state}",
+            'node_id': node_id,
+            'site_id': site_id,
+            'ancestral_state': ancestral_state,
+            'derived_state': derived_state,
+        })
+
+    return {
+        'mutations': result_mutations,
+        'total_count': total_count,
+        'has_more': offset + limit < total_count
+    }
+
+
+def search_mutations_by_position(ts, position, range_bp=5000, offset=0, limit=1000):
+    """
+    Search for mutations around a specific position.
+
+    Args:
+        ts: tskit.TreeSequence
+        position: Center position to search around (bp)
+        range_bp: Total range to search (searches +/- range_bp/2 around position)
+        offset: Number of mutations to skip (for pagination)
+        limit: Maximum number of mutations to return
+
+    Returns:
+        dict with:
+            - 'mutations': list of mutation dicts sorted by distance from position
+            - 'total_count': total mutations in search range
+            - 'has_more': whether there are more mutations
+            - 'search_start': actual start of search range
+            - 'search_end': actual end of search range
+    """
+    half_range = range_bp // 2
+    search_start = max(0, position - half_range)
+    search_end = min(ts.sequence_length, position + half_range)
+
+    t = ts.tables
+    sites = t.sites
+    mutations = t.mutations
+
+    # Get positions for all mutations via their sites
+    positions = sites.position[mutations.site]
+
+    # Create mask for mutations in the search range
+    mask = (positions >= search_start) & (positions < search_end)
+    indices = np.where(mask)[0]
+
+    # Calculate distances and sort by distance
+    if len(indices) > 0:
+        mutation_positions = positions[indices]
+        distances = np.abs(mutation_positions - position)
+        sorted_order = np.argsort(distances)
+        indices = indices[sorted_order]
+        sorted_distances = distances[sorted_order]
+    else:
+        sorted_distances = np.array([])
+
+    total_count = len(indices)
+
+    # Apply pagination
+    paginated_indices = indices[offset:offset + limit]
+    paginated_distances = sorted_distances[offset:offset + limit] if len(sorted_distances) > 0 else []
+
+    # Extract mutation data
+    result_mutations = []
+    for i, idx in enumerate(paginated_indices):
+        mut = mutations[idx]
+        site = sites[mut.site]
+        mut_position = int(site.position)
+        site_id = int(mut.site)
+        node_id = int(mut.node)
+
+        # Get ancestral and derived states
+        ancestral_state = site.ancestral_state
+        derived_state = mut.derived_state
+
+        result_mutations.append({
+            'position': mut_position,
+            'mutation': f"{ancestral_state}->{derived_state}",
+            'node_id': node_id,
+            'site_id': site_id,
+            'ancestral_state': ancestral_state,
+            'derived_state': derived_state,
+            'distance': int(paginated_distances[i]) if i < len(paginated_distances) else 0,
+        })
+
+    return {
+        'mutations': result_mutations,
+        'total_count': total_count,
+        'has_more': offset + limit < total_count,
+        'search_start': int(search_start),
+        'search_end': int(search_end),
+    }