lorax-arg 0.1__py3-none-any.whl

lorax/csv/config.py

@@ -0,0 +1,250 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Dict, List, Tuple
+
+import pandas as pd
+
+from lorax.utils import extract_sample_names, max_branch_length_from_newick
+
+
+@dataclass(frozen=True)
+class CsvConfigOptions:
+    window_size: int = 50_000
+
+
+REQUIRED_COLUMNS = ("genomic_positions", "newick")
+MAX_BRANCH_COL = "max_branch_length"
+
+
+def _dedupe_preserve_order(values: List[Any]) -> List[str]:
+    seen = set()
+    out: List[str] = []
+    for v in values:
+        s = str(v)
+        if s in seen:
+            continue
+        seen.add(s)
+        out.append(s)
+    return out
+
+
+def _validate_csv_df(df: pd.DataFrame) -> None:
+    missing = [c for c in REQUIRED_COLUMNS if c not in df.columns]
+    if missing:
+        raise ValueError(
+            f"CSV missing required column(s): {missing}. "
+            f"Expected columns: {list(REQUIRED_COLUMNS)}"
+        )
+
+    if len(df) == 0:
+        raise ValueError("CSV contains no rows.")
+
+
+def _sorted_reset(df: pd.DataFrame) -> pd.DataFrame:
+    # Deterministic ordering: global_index == row index after this.
+    out = df.copy()
+    out["genomic_positions"] = pd.to_numeric(out["genomic_positions"], errors="raise")
+    out = out.sort_values("genomic_positions", kind="mergesort")
+    out = out.reset_index(drop=True)
+    return out
+
+
+def _compute_intervals(genomic_positions: List[int], window_size: int) -> List[int]:
+    # Frontend expects N+1 breakpoints for N trees (see localBackendWorker logic).
+    # Each tree i spans [intervals[i], intervals[i+1]).
+    if not genomic_positions:
+        return []
+
+    positions = [int(p) for p in genomic_positions]
+    intervals = positions[:]  # N items
+    last_end = int(positions[-1]) + int(window_size)
+    intervals.append(last_end)  # N+1
+    return intervals
+
+
+def _is_empty_value(value: Any) -> bool:
+    if value is None:
+        return True
+    if isinstance(value, float) and pd.isna(value):
+        return True
+    if isinstance(value, str) and value.strip() == "":
+        return True
+    return False
+
+
+def extract_csv_metadata(df: pd.DataFrame) -> Dict[str, List[str]]:
+    """Extract file-level metadata from a CSV DataFrame.
+
+    Metadata columns are all columns after max_branch_length. File-level metadata
+    values come only from metadata-only rows (rows with empty tree columns).
+    """
+    if MAX_BRANCH_COL not in df.columns:
+        return {}
+
+    max_idx = list(df.columns).index(MAX_BRANCH_COL)
+    tree_cols = list(df.columns[: max_idx + 1])
+    metadata_cols = list(df.columns[max_idx + 1 :])
+
+    if not metadata_cols:
+        return {}
+
+    file_level: Dict[str, List[str]] = {c: [] for c in metadata_cols}
+
+    for _, row in df.iterrows():
+        if any(not _is_empty_value(row[col]) for col in tree_cols):
+            continue
+        for col in metadata_cols:
+            value = row[col]
+            if _is_empty_value(value):
+                continue
+            file_level[col].append(str(value))
+
+    return file_level
+
+
+def build_csv_config(
+    df: pd.DataFrame,
+    file_path: str,
+    *,
+    options: CsvConfigOptions | None = None,
+) -> Dict[str, Any]:
+    """Build a Lorax-compatible config for Newick-per-row CSV.
+
+    CSV schema:
+    - genomic_positions: int (tree start position)
+    - newick: str (Newick tree)
+
+    Contract notes:
+    - `intervals` must be N+1 breakpoints (required by frontend binning logic).
+    - `times` uses branch length for CSV: {type: "branch length", values: [0, max]}.
+    """
+    options = options or CsvConfigOptions()
+
+    _validate_csv_df(df)
+    file_metadata = extract_csv_metadata(df)
+    df2 = _sorted_reset(df)
+
+    # Optional per-tree metadata: "tree_info" / "tree info" column.
+    # This is intended for frontend-side per-tree UI (e.g., default per-tree colors)
+    # and is sent eagerly with config on load_file.
+    tree_info_map: Dict[str, str] = {}
+    tree_info_col = None
+    for cand in ("tree_info", "tree info"):
+        if cand in df2.columns:
+            tree_info_col = cand
+            break
+
+    if tree_info_col is not None:
+        # Use row index as tree_idx (this matches backend CSV layout access: df.iloc[tree_idx])
+        for i, row in df2.iterrows():
+            nwk = row.get("newick")
+            if isinstance(nwk, float) and pd.isna(nwk):
+                continue
+            if _is_empty_value(nwk):
+                continue
+            v = row.get(tree_info_col)
+            if _is_empty_value(v):
+                continue
+            tree_info_map[str(int(i))] = str(v)
+
+    # Compute max tree height (branch length) and sample names (best-effort, lightweight).
+    #
+    # Prefer the CSV-provided per-tree `max_branch_length` column when present, since
+    # regex parsing only captures max *edge* length, not the full root→tip height.
+    max_branch_length_all = 0.0
+    samples_set = set()
+    has_max_col = MAX_BRANCH_COL in df2.columns
+    saw_valid_max_col_value = False
+
+    for _, row in df2.iterrows():
+        nwk = row["newick"]
+        if isinstance(nwk, float) and pd.isna(nwk):
+            continue
+        if _is_empty_value(nwk):
+            continue
+        nwk = str(nwk)
+
+        if has_max_col:
+            v = row.get(MAX_BRANCH_COL)
+            if not _is_empty_value(v):
+                try:
+                    max_br = float(v)
+                    saw_valid_max_col_value = True
+                    if max_br > max_branch_length_all:
+                        max_branch_length_all = max_br
+                except Exception:
+                    # Ignore invalid per-row values; may fallback to regex below.
+                    pass
+
+        try:
+            sample_names = extract_sample_names(nwk)
+            samples_set.update(sample_names)
+        except Exception:
+            pass
+
+    # Fallback: if the CSV doesn't provide usable per-tree max heights, derive a
+    # best-effort global max from Newick text (max edge length).
+    if not saw_valid_max_col_value:
+        for _, row in df2.iterrows():
+            nwk = row["newick"]
+            if isinstance(nwk, float) and pd.isna(nwk):
+                continue
+            if _is_empty_value(nwk):
+                continue
+            nwk = str(nwk)
+
+            try:
+                max_br = float(max_branch_length_from_newick(nwk))
+                if max_br > max_branch_length_all:
+                    max_branch_length_all = max_br
+            except Exception:
+                # Keep config resilient; downstream can still load.
+                pass
+
+    genomic_positions = df2["genomic_positions"].astype(int).tolist()
+    intervals = _compute_intervals(genomic_positions, options.window_size)
+    genome_length = int(intervals[-1]) if intervals else int(genomic_positions[-1])
+
+    # Compute centered initial position (10% of genome, minimum 1kb) like tskit loader.
+    window_size = max(genome_length * 0.1, 1000)
+    midpoint = genome_length / 2.0
+    start = max(0, midpoint - window_size / 2.0)
+    end = min(genome_length, midpoint + window_size / 2.0)
+
+    # Build a deterministic, file-level sample order for stable tip node IDs across trees.
+    # Prefer explicit file metadata (if provided), otherwise derive from all Newicks.
+    samples_list_meta = file_metadata.get("samples", [])
+    if samples_list_meta:
+        samples_order = _dedupe_preserve_order(samples_list_meta)
+    else:
+        samples_order = sorted(str(s) for s in samples_set)
+
+    # Temporary workaround: CSV inputs may include the outgroup sample "etal" (any case).
+    # The layout/parsing pipeline prunes it from the Newick tree, so also remove it
+    # from the file-level samples list to keep UI/search options consistent.
+    samples_order = [s for s in samples_order if str(s).lower() != "etal"]
+
+    sample_names_map = {str(s): {"sample_name": s} for s in samples_order}
+
+    config: Dict[str, Any] = {
+        "genome_length": genome_length,
+        "initial_position": [int(start), int(end)],
+        "times": {"type": "branch length", "values": [0.0, float(max_branch_length_all)]},
+        "intervals": intervals,
+        "filename": str(file_path).split("/")[-1],
+        "sample_names": sample_names_map,
+        "samples": samples_order,
+        # Present but empty for compatibility; CSV doesn’t have these yet.
+        "metadata_schema": {
+            # Always expose "sample" as the only supported metadata key for CSV.
+            # (CSV metadata Socket.IO handlers currently support only key == "sample".)
+            "metadata_keys": ["sample"]
+        },
+    }
+
+    if tree_info_map:
+        # Kept as "tree_info" regardless of whether CSV column is "tree info" or "tree_info".
+        config["tree_info"] = tree_info_map
+
+    return config

lorax/csv/layout.py

@@ -0,0 +1,182 @@
+from __future__ import annotations
+
+import struct
+from typing import Any, Dict, List
+
+import numpy as np
+import pandas as pd
+import pyarrow as pa
+
+
+def build_empty_tree_layout_arrow_ipc() -> bytes:
+    """Return a valid (possibly empty) PyArrow IPC stream for tree layout.
+
+    Frontend expects a table with these columns:
+    - node_id:int32
+    - parent_id:int32
+    - is_tip:bool
+    - tree_idx:int32
+    - x:float32
+    - y:float32
+    """
+    empty_table = pa.table(
+        {
+            "node_id": pa.array([], type=pa.int32()),
+            "parent_id": pa.array([], type=pa.int32()),
+            "is_tip": pa.array([], type=pa.bool_()),
+            "tree_idx": pa.array([], type=pa.int32()),
+            "x": pa.array([], type=pa.float32()),
+            "y": pa.array([], type=pa.float32()),
+            "name": pa.array([], type=pa.string()),
+        }
+    )
+    sink = pa.BufferOutputStream()
+    writer = pa.ipc.new_stream(sink, empty_table.schema)
+    writer.write_table(empty_table)
+    writer.close()
+    return sink.getvalue().to_pybytes()
+
+
+def build_empty_layout_response(tree_indices: List[int] | None = None) -> Dict[str, Any]:
+    """Build the response shape returned by `handle_tree_graph_query` for CSV."""
+    return {
+        "buffer": build_empty_tree_layout_arrow_ipc(),
+        "global_min_time": 0.0,
+        "global_max_time": 0.0,
+        "tree_indices": tree_indices or [],
+    }
+
+
+def build_csv_layout_response(
+    df: pd.DataFrame,
+    tree_indices: List[int],
+    max_branch_length: float,
+    samples_order: List[str] | None = None,
+    pre_parsed_graphs: Dict[int, Any] | None = None,
+    shift_tips_to_one: bool = False,
+) -> Dict[str, Any]:
+    """Build PyArrow IPC buffer for CSV trees.
+
+    Parses Newick strings from the DataFrame and generates the same buffer format
+    as construct_trees_batch() for tskit files.
+
+    Args:
+        df: DataFrame with 'newick' column containing Newick tree strings
+        tree_indices: List of row indices (global_index) to process
+        max_branch_length: From config times.values[1], for y normalization
+
+    Returns:
+        Same format as construct_trees_batch():
+        {buffer, global_min_time, global_max_time, tree_indices}
+    """
+    from lorax.csv.newick_tree import parse_newick_to_tree
+
+    # Collect all nodes from all trees
+    all_node_ids: List[np.ndarray] = []
+    all_parent_ids: List[np.ndarray] = []
+    all_is_tip: List[np.ndarray] = []
+    all_tree_idx: List[np.ndarray] = []
+    all_x: List[np.ndarray] = []
+    all_y: List[np.ndarray] = []
+    all_names: List[np.ndarray] = []
+
+    processed_indices: List[int] = []
+
+    for tree_idx in tree_indices:
+        tree_idx = int(tree_idx)
+        if tree_idx < 0 or tree_idx >= len(df):
+            continue
+
+        newick_str = df.iloc[tree_idx].get("newick")
+        if pd.isna(newick_str):
+            continue
+
+        tree_max_branch_length = None
+        if "max_branch_length" in df.columns:
+            try:
+                v = df.iloc[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
+
+        graph = pre_parsed_graphs.get(tree_idx) if pre_parsed_graphs else None
+        if graph is None:
+            try:
+                graph = 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 as e:
+                # Log error but continue with other trees
+                print(f"Failed to parse Newick for tree {tree_idx}: {e}")
+                continue
+
+        n = len(graph.node_id)
+        if n == 0:
+            continue
+
+        # Collect nodes for this tree
+        # SWAP coordinates to match tskit convention: time -> x, layout -> y
+        all_node_ids.append(graph.node_id)
+        all_parent_ids.append(graph.parent_id)
+        all_is_tip.append(graph.is_tip)
+        all_tree_idx.append(np.full(n, tree_idx, dtype=np.int32))
+        all_x.append(graph.y.astype(np.float32))  # SWAP: time -> x
+        all_y.append(graph.x.astype(np.float32))  # SWAP: layout -> y
+        all_names.append(np.array(graph.name, dtype=object))
+
+        processed_indices.append(tree_idx)
+
+    # Build PyArrow table
+    if not all_node_ids:
+        # Return empty buffer
+        return build_empty_layout_response(list(tree_indices))
+
+    # Concatenate all arrays
+    node_table = pa.table(
+        {
+            "node_id": pa.array(np.concatenate(all_node_ids), type=pa.int32()),
+            "parent_id": pa.array(np.concatenate(all_parent_ids), type=pa.int32()),
+            "is_tip": pa.array(np.concatenate(all_is_tip), type=pa.bool_()),
+            "tree_idx": pa.array(np.concatenate(all_tree_idx), type=pa.int32()),
+            "x": pa.array(np.concatenate(all_x), type=pa.float32()),
+            "y": pa.array(np.concatenate(all_y), type=pa.float32()),
+            "name": pa.array(np.concatenate(all_names), type=pa.string()),
+        }
+    )
+
+    # Empty mutation table (CSV has no mutations)
+    mut_table = pa.table(
+        {
+            "mut_x": pa.array([], type=pa.float32()),
+            "mut_y": pa.array([], type=pa.float32()),
+            "mut_tree_idx": pa.array([], type=pa.int32()),
+        }
+    )
+
+    # Serialize to IPC format (same as tree_graph.py)
+    node_sink = pa.BufferOutputStream()
+    node_writer = pa.ipc.new_stream(node_sink, node_table.schema)
+    node_writer.write_table(node_table)
+    node_writer.close()
+    node_bytes = node_sink.getvalue().to_pybytes()
+
+    mut_sink = pa.BufferOutputStream()
+    mut_writer = pa.ipc.new_stream(mut_sink, mut_table.schema)
+    mut_writer.write_table(mut_table)
+    mut_writer.close()
+    mut_bytes = mut_sink.getvalue().to_pybytes()
+
+    # Combine with length prefix (same format as tree_graph.py:559)
+    combined = struct.pack("<I", len(node_bytes)) + node_bytes + mut_bytes
+
+    return {
+        "buffer": combined,
+        "global_min_time": 0.0,
+        "global_max_time": float(max_branch_length),
+        "tree_indices": processed_indices,
+    }

lorax/csv/newick_tree.py

@@ -0,0 +1,234 @@
+"""Parse Newick strings and compute x,y layout coordinates.
+
+This module provides tree layout computation for CSV files containing Newick strings.
+It mirrors the approach in tree_graph/tree_graph.py but sources data from Newick parsing
+instead of tskit tables.
+
+Coordinate system:
+- y (time): anchored time in [0,1] where tips are always 1.0 and the root is
+  1 - (tree_height / global_max_height)
+- x (layout): tips get sequential x, internal nodes get (min+max)/2 of children
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import List, Optional
+
+import numpy as np
+
+# ete3 is used for Newick parsing
+from ete3 import Tree
+
+
+def prune_outgroup_sample(tree: Tree, outgroup: str = "etal") -> None:
+    """Prune an outgroup leaf from an ete3 Tree (best-effort).
+
+    This is a temporary compatibility workaround for CSV Newick inputs that include
+    a known outgroup sample (e.g. "etal") that should not be displayed.
+
+    Behavior:
+    - If outgroup is missing, do nothing.
+    - Otherwise, reroot on outgroup, delete it, and clear the root branch lengths
+    """
+    try:
+        leaf_names = tree.get_leaf_names()
+        if not leaf_names:
+            return
+
+        # Match case-insensitively (e.g. "etal" vs "Etal")
+        target = None
+        out_l = str(outgroup).lower()
+        for name in leaf_names:
+            if str(name).lower() == out_l:
+                target = str(name)
+                break
+        if target is None:
+            return
+
+        tree.set_outgroup(target)
+        (tree & target).delete()
+
+        # Remove root branch distortion after rerooting/pruning
+        tree.dist = 0.0
+        for child in tree.get_children():
+            child.dist = 0.0
+    except Exception:
+        # Best-effort: CSV loading should remain resilient.
+        return
+
+
+@dataclass
+class NewickTreeGraph:
+    """Tree structure from parsed Newick with layout coordinates.
+
+    Attributes:
+        node_id: Sequential node IDs assigned via post-order traversal
+        parent_id: Parent node ID (-1 for root)
+        is_tip: Boolean array indicating leaf nodes
+        name: Leaf names from the Newick string
+        branch_length: Distance to parent
+        x: Layout position [0,1] - tips spread, internal = (min+max)/2 of children
+        y: Normalized time [0,1] - cumulative distance from root
+    """
+
+    node_id: np.ndarray  # int32
+    parent_id: np.ndarray  # int32
+    is_tip: np.ndarray  # bool
+    name: List[str]  # leaf names
+    branch_length: np.ndarray  # float32
+    x: np.ndarray  # float32, layout position
+    y: np.ndarray  # float32, normalized time
+
+
+def shift_tree_tips_to_one(y: np.ndarray, is_tip: np.ndarray) -> np.ndarray:
+    """Shift y so the maximum tip height is exactly 1.0 (clipped to [0,1])."""
+    if y.size == 0:
+        return y
+    if is_tip is None or is_tip.size == 0:
+        return y
+    tip_max = float(y[is_tip].max()) if np.any(is_tip) else 1.0
+    if not np.isfinite(tip_max) or tip_max == 1.0:
+        return y
+    return np.clip(y + (1.0 - tip_max), 0.0, 1.0)
+
+
+def parse_newick_to_tree(
+    newick_str: str,
+    max_branch_length: float,
+    samples_order: Optional[List[str]] = None,
+    *,
+    tree_max_branch_length: float | None = None,
+    shift_tips_to_one: bool = False,
+) -> NewickTreeGraph:
+    """Parse Newick string and compute x,y coordinates.
+
+    Uses ete3 to parse the Newick string, then computes layout coordinates
+    using the same algorithm as tree_graph.py:
+    - y: cumulative distance from root, normalized by max_branch_length
+    - x: post-order layout where tips get sequential x, internals get (min+max)/2
+
+    Args:
+        newick_str: Newick format tree string
+        max_branch_length: Global max height for y anchoring (from config times.values[1])
+        tree_max_branch_length: Optional per-tree max height. If provided, this is
+            used to anchor the root at 1 - (tree_height / global_max). If omitted,
+            tree height is derived from the parsed tree (max cumulative root distance).
+
+    Returns:
+        NewickTreeGraph with layout coordinates normalized to [0,1]
+    """
+    # Parse with ete3 - format=1 includes branch lengths and internal node names
+    tree = Tree(newick_str, format=1)
+    # Prune the outgroup before any layout or time computation.
+    # TODO: remove in feature.
+    prune_outgroup_sample(tree, outgroup="etal")
+
+    # Assign traversal order via post-order traversal (after pruning)
+    nodes = list(tree.traverse("postorder"))
+    if not nodes:
+        empty_i32 = np.array([], dtype=np.int32)
+        empty_bool = np.array([], dtype=np.bool_)
+        empty_f32 = np.array([], dtype=np.float32)
+        return NewickTreeGraph(
+            node_id=empty_i32,
+            parent_id=empty_i32,
+            is_tip=empty_bool,
+            name=[],
+            branch_length=empty_f32,
+            x=empty_f32,
+            y=empty_f32,
+        )
+    node_index = {node: idx for idx, node in enumerate(nodes)}
+
+    num_nodes = len(nodes)
+
+    # Build arrays
+    node_id = np.arange(num_nodes, dtype=np.int32)
+    parent_id = np.full(num_nodes, -1, dtype=np.int32)
+    is_tip = np.zeros(num_nodes, dtype=np.bool_)
+    branch_length = np.zeros(num_nodes, dtype=np.float32)
+    name = [""] * num_nodes
+
+    assigned_ids = {}
+    next_id = 0
+    if samples_order:
+        sample_id_map = {str(name): idx for idx, name in enumerate(samples_order)}
+        for node in nodes:
+            if node.is_leaf():
+                node_name = node.name if node.name else ""
+                if node_name not in sample_id_map:
+                    raise ValueError(
+                        f"Leaf sample name '{node_name}' not found in samples_order. "
+                        "Ensure CSV config provides a complete, file-level samples list."
+                    )
+                assigned_ids[node] = sample_id_map[node_name]
+        next_id = len(samples_order)
+        for node in nodes:
+            if not node.is_leaf():
+                assigned_ids[node] = next_id
+                next_id += 1
+    else:
+        for node in nodes:
+            assigned_ids[node] = next_id
+            next_id += 1
+
+    for node in nodes:
+        idx = node_index[node]
+        if node.up is not None:
+            parent_id[idx] = assigned_ids[node.up]
+        is_tip[idx] = node.is_leaf()
+        branch_length[idx] = node.dist if node.dist else 0.0
+        name[idx] = node.name if node.name else ""
+        node_id[idx] = assigned_ids[node]
+
+    # Compute y_raw (time): cumulative distance from root
+    # Root at 0, tips at tree_height
+    y_raw = np.zeros(num_nodes, dtype=np.float32)
+    for node in tree.traverse("preorder"):  # Root first
+        idx = node_index[node]
+        if node.up is not None:
+            parent_idx = node_index[node.up]
+            y_raw[idx] = y_raw[parent_idx] + branch_length[idx]
+
+    # Anchor time so tips are always 1.0 and the root is 1 - tree_height/global_max.
+    #
+    # If the per-tree max height is known (e.g. CSV 'max_branch_length' column), use
+    # it; otherwise derive it from the parsed tree.
+    tree_height = float(tree_max_branch_length) if (tree_max_branch_length or 0.0) > 0 else float(y_raw.max())
+    if max_branch_length > 0:
+        y = 1.0 - (tree_height - y_raw) / float(max_branch_length)
+        y = np.clip(y, 0.0, 1.0)
+    else:
+        # Degenerate case: no global height information. Keep tips at 1.
+        y = np.ones(num_nodes, dtype=np.float32)
+
+    if shift_tips_to_one:
+        y = shift_tree_tips_to_one(y, is_tip)
+
+    # Compute x (layout): tips get sequential x, internals = (min+max)/2
+    x = np.zeros(num_nodes, dtype=np.float32)
+    tip_counter = 0
+
+    for node in tree.traverse("postorder"):
+        idx = node_index[node]
+        if node.is_leaf():
+            x[idx] = tip_counter
+            tip_counter += 1
+        else:
+            child_xs = [x[node_index[child]] for child in node.children]
+            x[idx] = (min(child_xs) + max(child_xs)) / 2.0
+
+    # Normalize x to [0,1]
+    if tip_counter > 1:
+        x = x / (tip_counter - 1)
+
+    return NewickTreeGraph(
+        node_id=node_id,
+        parent_id=parent_id,
+        is_tip=is_tip,
+        name=name,
+        branch_length=branch_length,
+        x=x,
+        y=y,
+    )