agglovar 0.0.1.dev13__tar.gz → 0.0.1.dev14__tar.gz

{agglovar-0.0.1.dev13/src/agglovar.egg-info → agglovar-0.0.1.dev14}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: agglovar
-Version: 0.0.1.dev13
+Version: 0.0.1.dev14
 Summary: Toolkit for fast genomic variant transformations and intersects
 License-Expression: MIT
 Requires-Python: >=3.12
@@ -8,7 +8,6 @@ Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: biopython>=1.85
 Requires-Dist: edlib>=1.3.9.post1
-Requires-Dist: intervaltree>=3.1.0
 Requires-Dist: numpy>=2.3.0
 Requires-Dist: ply>=3.11
 Requires-Dist: polars[numpy,pandas,pyarrow]>=1.35.2

{agglovar-0.0.1.dev13 → agglovar-0.0.1.dev14}/pyproject.toml

@@ -13,7 +13,6 @@ requires-python = ">=3.12"
 dependencies = [
     "biopython>=1.85",
     "edlib>=1.3.9.post1",
-    "intervaltree>=3.1.0",
     "numpy>=2.3.0",
     "ply>=3.11",
     "polars[numpy,pandas,pyarrow]>=1.35.2",
@@ -108,6 +107,8 @@ max-line-length = 120
 max-line-length = 120
 docstring-convention = "pep257"
 per-file-ignores = [
-    "src/agglovar/join/config/parser.py: D200, D205, D212, D300, D403, D415"
+    "src/agglovar/merge/config/parser.py: D200, D205, D212, D300, D400, D401, D403, D415",
+    "src/agglovar/vcf/__init__.py: DALL002, F401, F403",
+    "src/agglovar/pairwise/overlap/__init__.py: DALL002, F401, F403",
 ]
-rst-roles = ["class", "meth", "const", "ref", "func", "mod"]
+rst-roles = ["class", "meth", "const", "ref", "func", "mod", "attr", "data"]

{agglovar-0.0.1.dev13 → agglovar-0.0.1.dev14}/src/agglovar/__init__.py

@@ -1,6 +1,6 @@
 """Agglovar: A toolkit for fast genomic variant transformations and intersects."""
 
-__version__ = '0.0.1.dev13'
+__version__ = '0.0.1.dev14'
 
 __all__ = [
     'align',

{agglovar-0.0.1.dev13 → agglovar-0.0.1.dev14}/src/agglovar/bed/col.py

@@ -14,6 +14,7 @@ from typing import Optional, Self, Any
 
 import polars as pl
 
+
 @dataclass(frozen=True)
 class CoordColExpr(Iterable[pl.Expr]):
     """Represents coordinate column expressions.
@@ -22,6 +23,7 @@ class CoordColExpr(Iterable[pl.Expr]):
     :ivar pos: Position column expression.
     :ivar end: End column expression.
     """
+
     chrom: pl.Expr
     pos: pl.Expr
     end: pl.Expr
@@ -49,11 +51,14 @@ class CoordColExpr(Iterable[pl.Expr]):
         return self.end.meta.output_name()
 
     def __iter__(self):
+        """Iterate over chrom, pos, and end expressions."""
         return iter((self.chrom, self.pos, self.end))
 
     def __repr__(self):
+        """Return a string representation."""
         return f"CoordColExpr({self.chrom!r}, {self.pos!r}, {self.end!r})"
 
+
 @dataclass(frozen=True, order=True)
 class CoordCol(Iterable[str], Container[str]):
     """Represents coordinate columns to use.
@@ -62,6 +67,7 @@ class CoordCol(Iterable[str], Container[str]):
     :ivar pos: Position column name.
     :ivar end: End column name.
     """
+
     chrom: str
     pos: str
     end: str
@@ -89,25 +95,30 @@ class CoordCol(Iterable[str], Container[str]):
         return CoordColExpr(
             *(
                 pl.col(col).alias(alias + suffix)
-                    for col, alias in zip(self, alias)
+                for col, alias in zip(self, alias)
             )
         )
 
     def __iter__(self):
+        """Iterate over chrom, pos, and end column names."""
         return iter((self.chrom, self.pos, self.end))
 
     def __contains__(self, o: Optional[Any]):
+        """Test whether ``o`` is one of the coordinate column names."""
         return o in (self.chrom, self.pos, self.end)
 
     def __repr__(self):
+        """Return a string representation."""
         return f"CoordCol({self.chrom!r}, {self.pos!r}, {self.end!r})"
 
+
 COL_CHROM: CoordCol = CoordCol('chrom', 'pos', 'end')
 """Standard chromosome columns."""
 
 COL_QRY: CoordCol = CoordCol('qry_id', 'qry_pos', 'qry_end')
 """Standard query columns."""
 
+
 def get_coord_cols(
         col_names: Optional[CoordCol | str | Iterable[str]] = None,
 ) -> CoordCol:
@@ -140,10 +151,12 @@ def get_coord_cols(
         ))
 
         if len(col_names_tuple) != 3:
-            raise ValueError(f'Invalid cols: Expected 3 elements, found {len(col_names_tuple)}: {col_names:r}')
+            raise ValueError(
+                f'Invalid cols: Expected 3 elements, found {len(col_names_tuple)}: {col_names!r}'
+            )
 
         if any(col is None or col == '' for col in col_names_tuple):
-            raise ValueError(f'Columns is missing values: {col_names:r}')
+            raise ValueError(f'Columns is missing values: {col_names!r}')
 
         return CoordCol(*col_names_tuple)
 

agglovar-0.0.1.dev14/src/agglovar/bed/intersect.py

@@ -0,0 +1,186 @@
+"""Table intersects."""
+
+from pathlib import Path
+from typing import Iterable, Optional
+
+import polars as pl
+
+from .join import pairwise_join, pairwise_join_iter
+from .merge import merge_depth
+from .col import CoordCol, get_coord_cols
+
+__all__ = [
+    'as_bool',
+    'as_proportion',
+]
+
+
+def as_bool(
+        df_a: pl.LazyFrame | pl.DataFrame,
+        df_b: pl.LazyFrame | pl.DataFrame,
+        name: str,
+        distance: int = 0,
+        negate: bool = False,
+        col_names_a: Optional[CoordCol | Iterable[str] | str] = None,
+        col_names_b: Optional[CoordCol | Iterable[str] | str] = None,
+        temp_dir: bool | str | Path = False,
+) -> pl.LazyFrame:
+    """Add a boolean column to df_a indicating whether each record intersects with df_b.
+
+    :param df_a: Table a.
+    :param df_b: Table b.
+    :param name: Name of the column to add.
+    :param distance: Maximum distance between two records. May be negative to require overlap.
+    :param negate: If True, negate the boolean column to annotate misses instead of hits.
+    :param col_names_a: Columns in a (chromosome or query ID, pos, end).
+    :param col_names_b: Columns in b (chromosome or query ID, pos, end).
+    :param temp_dir: How to materialise the prepared tables before iterating. See
+        :func:`agglovar.bed.join.pairwise_join`.
+
+    :return: A LazyFrame with two columns: ``_index`` and ``name``.
+    """
+    if name is None or not (name := name.strip()):
+        raise ValueError('Name must be a non-empty string')
+
+    col_names_a = get_coord_cols(col_names_a)
+    col_names_b = get_coord_cols(col_names_b)
+
+    if isinstance(df_a, pl.DataFrame):
+        df_a = df_a.lazy()
+
+    if isinstance(df_b, pl.DataFrame):
+        df_b = df_b.lazy()
+
+    hit_val = not negate
+
+    join_list = []
+
+    for df_join in pairwise_join_iter(
+            df_a=df_a,
+            df_b=df_b,
+            distance=distance,
+            col_names_a=col_names_a,
+            col_names_b=col_names_b,
+            temp_dir=temp_dir,
+    ):
+        join_list.append(
+            df_join
+            .select(pl.col('index_a').alias('_index'))
+            .collect()
+            .lazy()
+        )
+
+    if '_index' in df_a.collect_schema().names():
+        df_a_indexed = df_a
+    else:
+        df_a_indexed = df_a.with_row_index('_index').with_columns(pl.col('_index').cast(pl.UInt64))
+
+    if not join_list:
+        return (
+            df_a_indexed
+            .select(
+                '_index',
+                pl.lit(not hit_val).alias(name),
+            )
+        )
+
+    hits = pl.concat(join_list).unique('_index')
+
+    return (
+        df_a_indexed
+        .select('_index')
+        .join(
+            hits.with_columns(pl.lit(hit_val).alias(name)),
+            on='_index', how='left',
+        )
+        .select(
+            '_index',
+            pl.col(name).fill_null(not hit_val),
+        )
+    )
+
+
+def as_proportion(
+        df_a: pl.LazyFrame | pl.DataFrame,
+        df_b: pl.LazyFrame | pl.DataFrame,
+        name: str,
+        col_names_a: Optional[CoordCol | Iterable[str] | str] = None,
+        col_names_b: Optional[CoordCol | Iterable[str] | str] = None,
+        temp_dir: bool | str | Path = False,
+) -> pl.LazyFrame:
+    """Compute the proportion of each interval in ``df_a`` covered by intervals in ``df_b``.
+
+    Rows in ``df_a`` with null ``pos`` or ``end`` are preserved in the output with a null
+    proportion. Zero-length intervals (``pos == end``) produce ``NaN`` (0 / 0).
+
+    :param df_a: Table a.
+    :param df_b: Table b.
+    :param name: Name of the column to add.
+    :param col_names_a: Columns in a (chromosome or query ID, pos, end).
+    :param col_names_b: Columns in b (chromosome or query ID, pos, end).
+    :param temp_dir: How to materialise the prepared tables before iterating. See
+        :func:`agglovar.bed.join.pairwise_join`.
+
+    :return: A LazyFrame with two columns: ``_index`` and ``name``.
+    """
+    if name is None or not (name := name.strip()):
+        raise ValueError('Name must be a non-empty string')
+
+    col_names_a = get_coord_cols(col_names_a)
+    col_names_b = get_coord_cols(col_names_b)
+
+    col_expr_a = col_names_a.exprs()
+
+    if isinstance(df_a, pl.DataFrame):
+        df_a = df_a.lazy()
+
+    if isinstance(df_b, pl.DataFrame):
+        df_b = df_b.lazy()
+
+    if '_index' not in df_a.collect_schema().names():
+        df_a = df_a.with_row_index('_index').with_columns(pl.col('_index').cast(pl.UInt64))
+
+    # Records with null pos/end are kept and assigned a null proportion at the end.
+    df_a_clean = df_a.filter(
+        pl.col(col_names_a.pos).is_not_null(),
+        pl.col(col_names_a.end).is_not_null(),
+    )
+
+    # Collapse overlapping b intervals so coverage is not double-counted.
+    df_b_nr = (
+        merge_depth(df_b, 0, col_names_b)
+        .select(*col_names_b.exprs())
+        .collect()
+    )
+
+    df_join = pairwise_join(
+        df_a=df_a_clean,
+        df_b=df_b_nr,
+        col_names_a=col_names_a,
+        col_names_b=col_names_b,
+        temp_dir=temp_dir,
+    )
+
+    df_overlap = (
+        df_join
+        .filter(pl.col('end') > pl.col('pos'))
+        .select(
+            pl.col('index_a').alias('_index'),
+            (pl.col('end') - pl.col('pos')).alias('_overlap'),
+        )
+        .group_by('_index')
+        .agg(pl.col('_overlap').sum())
+    )
+
+    return (
+        df_a
+        .select(
+            '_index',
+            (col_expr_a.end - col_expr_a.pos).alias('_len'),
+        )
+        .join(df_overlap.lazy(), on='_index', how='left')
+        .select(
+            '_index',
+            (pl.col('_overlap').fill_null(0.0) / pl.col('_len')).alias(name),
+        )
+    )

agglovar-0.0.1.dev14/src/agglovar/bed/join.py

@@ -0,0 +1,323 @@
+"""Base join operations for intersects."""
+
+__all__ = [
+    'pairwise_join',
+    'pairwise_join_iter',
+]
+
+from pathlib import Path
+from typing import Iterable, Iterator, Optional
+
+import polars as pl
+
+from .col import CoordCol, get_coord_cols
+from ..util.lazy import materialize_pair
+
+CHUNK_SIZE: int = 2_500
+"""Default size of join chunks. Breaks up tables into batches of this size or less."""
+
+
+class _JoinResources:
+    """Resources for joining tables."""
+
+    df_a: pl.LazyFrame
+    df_b: pl.LazyFrame
+    distance: int
+    chunk_size: int
+    col_a: CoordCol
+    col_b: CoordCol
+
+    def __init__(
+            self,
+            df_a: pl.LazyFrame,
+            df_b: pl.LazyFrame,
+            distance: int = 0,
+            chunk_size: int = CHUNK_SIZE,
+            col_names_a: Optional[CoordCol | Iterable[str] | str] = None,
+            col_names_b: Optional[CoordCol | Iterable[str] | str] = None,
+    ):
+        if chunk_size < 1:
+            raise ValueError('chunk_size must be greater than 0')
+
+        if isinstance(df_a, pl.DataFrame):
+            df_a = df_a.lazy()
+
+        if isinstance(df_b, pl.DataFrame):
+            df_b = df_b.lazy()
+
+        cols_a = set(df_a.collect_schema().keys())
+        cols_b = set(df_b.collect_schema().keys())
+
+        if '_index' not in cols_a:
+            df_a = df_a.with_row_index('_index').with_columns(pl.col('_index').cast(pl.UInt64))
+
+        if '_index' not in cols_b:
+            df_b = df_b.with_row_index('_index').with_columns(pl.col('_index').cast(pl.UInt64))
+
+        # Set column names
+        ref_cols = get_coord_cols('ref')
+
+        try:
+            col_select_a = get_coord_cols(col_names_a).exprs(alias=ref_cols, suffix='_a')
+        except (ValueError, TypeError) as e:
+            raise ValueError(f'col_names_a: {e}')
+
+        try:
+            col_select_b = get_coord_cols(col_names_b).exprs(alias=ref_cols, suffix='_b')
+        except (ValueError, TypeError) as e:
+            raise ValueError(f'col_names_b: {e}')
+
+        col_a = col_select_a.col_names()
+        col_b = col_select_b.col_names()
+
+        # Prepare tables
+        df_a = (
+            df_a
+            .select(
+                pl.col('_index').alias('_index_a'),
+                *col_select_a
+            )
+        )
+
+        df_b = (
+            df_b
+            .select(
+                pl.col('_index').alias('_index_b'),
+                *col_select_b
+            )
+        )
+
+        self.df_a = df_a
+        self.df_b = df_b
+        self.distance = distance
+        self.chunk_size = chunk_size
+        self.col_a = col_a
+        self.col_b = col_b
+
+
+def pairwise_join(
+        df_a: pl.LazyFrame | pl.DataFrame,
+        df_b: pl.LazyFrame | pl.DataFrame,
+        distance: int = 0,
+        chunk_size: int = CHUNK_SIZE,
+        col_names_a: Optional[CoordCol | Iterable[str] | str] = None,
+        col_names_b: Optional[CoordCol | Iterable[str] | str] = None,
+        temp_dir: bool | str | Path = False,
+) -> pl.LazyFrame:
+    """Join two tables.
+
+    Thin wrapper around :func:`pairwise_join_iter` that concatenates all yielded
+    chunks into a single table.
+
+    Returns a table with columns:
+
+        * index_a: Index in table a.
+        * index_b: Index in table b.
+        * chrom: Chromosome matched.
+        * pos: Start position of intersection.
+        * end: End position of intersection.
+        * distance: Distance between the two intervals with negative values representing overlapping intervals.
+
+    Note that if padding is greater than 0, the "pos" and "end" will have been modified to include padding.
+
+    :param df_a: Table a.
+    :param df_b: Table b.
+    :param distance: Maximum distance between two records. May be negative to force overlapping.
+    :param chunk_size: Chunk A by this size per chromosome to bound the IEJoin working set.
+    :param col_names_a: Columns to select from `df_a` if not None, otherwise, use object defaults.
+    :param col_names_b: Columns to select from `df_b` if not None, otherwise, use object defaults.
+    :param temp_dir: How to materialise the prepared tables before the chunked loop.
+        ``False`` (default) collects both into memory; ``True`` writes them to the
+        system temp directory as parquet files; a ``str``/``Path`` writes them to
+        that directory. Temp files are always removed on exit.
+
+    :return: A LazyFrame with the joined tables.
+    """
+    return pl.concat(list(pairwise_join_iter(
+        df_a=df_a,
+        df_b=df_b,
+        distance=distance,
+        chunk_size=chunk_size,
+        col_names_a=col_names_a,
+        col_names_b=col_names_b,
+        temp_dir=temp_dir,
+    )))
+
+
+def pairwise_join_iter(
+        df_a: pl.LazyFrame | pl.DataFrame,
+        df_b: pl.LazyFrame | pl.DataFrame,
+        distance: int = 0,
+        chunk_size: int = CHUNK_SIZE,
+        col_names_a: Optional[CoordCol | Iterable[str] | str] = None,
+        col_names_b: Optional[CoordCol | Iterable[str] | str] = None,
+        temp_dir: bool | str | Path = False,
+) -> Iterator[pl.LazyFrame]:
+    """Join two tables, yielding one LazyFrame per chunk.
+
+    Returns chunks with the same columns as :func:`pairwise_join`. At least one chunk
+    is always yielded; an empty schema-only frame is yielded when no chunk would otherwise
+    have been produced (so callers can safely call ``pl.concat`` on the result).
+
+    Each chunk is the result of a per-chromosome ``join_where`` (Polars IEJoin) over a
+    bounded slice of A and the chrom-matched B records pre-filtered to that slice's
+    range.
+
+    :param df_a: Table a.
+    :param df_b: Table b.
+    :param distance: Maximum distance between two records. May be negative to force overlapping.
+    :param chunk_size: Chunk A by this size per chromosome to bound the IEJoin working set.
+    :param col_names_a: Columns to select from `df_a` if not None, otherwise, use object defaults.
+    :param col_names_b: Columns to select from `df_b` if not None, otherwise, use object defaults.
+    :param temp_dir: How to materialise the prepared tables before the chunked loop.
+        See :func:`pairwise_join`.
+
+    :return: An iterator of LazyFrames.
+    """
+    resources = _JoinResources(
+        df_a=df_a,
+        df_b=df_b,
+        distance=distance,
+        chunk_size=chunk_size,
+        col_names_a=col_names_a,
+        col_names_b=col_names_b,
+    )
+
+    with materialize_pair(
+        resources.df_a, resources.df_b, temp_dir, prefix='bed_join_prep_',
+    ) as (df_a_mat, df_b_mat):
+        resources.df_a = df_a_mat
+        resources.df_b = df_b_mat
+
+        yielded = False
+
+        for chunk in _join_chunks(resources):
+            yielded = True
+            yield chunk
+
+        if not yielded:
+            yield _empty_join(resources)
+
+
+def _build_join_pair(
+        df_a: pl.LazyFrame,
+        df_b: pl.LazyFrame,
+        distance: int,
+        col_a: CoordCol,
+        col_b: CoordCol,
+) -> pl.LazyFrame:
+    """Apply ``join_where`` (IEJoin) and the standard output projection.
+
+    Both ``df_a`` and ``df_b`` must already have ``_index_a``/``_index_b`` columns
+    and the coordinate columns named per ``col_a``/``col_b``. The chrom equality is
+    enforced by the caller filtering both inputs to the same chromosome upstream;
+    only the inequality predicates are passed to ``join_where`` (the shape Polars's
+    IEJoin path is optimised for).
+    """
+    return (
+        df_a
+        .join_where(
+            df_b,
+            pl.col(col_b.pos) - distance <= pl.col(col_a.end),
+            pl.col(col_b.end) + distance >= pl.col(col_a.pos),
+        )
+        .select(
+            pl.col('_index_a').alias('index_a'),
+            pl.col('_index_b').alias('index_b'),
+            pl.col(col_a.chrom).alias('chrom'),
+            pl.max_horizontal(col_a.pos, col_b.pos).alias('pos'),
+            pl.min_horizontal(col_a.end, col_b.end).alias('end'),
+        )
+        .with_columns(
+            pl.min_horizontal('pos', 'end').alias('pos'),
+            pl.max_horizontal('pos', 'end').alias('end'),
+            (pl.col('pos') - pl.col('end')).alias('distance'),
+        )
+    )
+
+
+def _empty_join(join_resources: _JoinResources) -> pl.LazyFrame:
+    """Schema-only LazyFrame for the empty-input case."""
+    return _build_join_pair(
+        join_resources.df_a.head(0),
+        join_resources.df_b.head(0),
+        join_resources.distance,
+        join_resources.col_a,
+        join_resources.col_b,
+    )
+
+
+def _join_chunks(
+        join_resources: _JoinResources
+) -> Iterator[pl.LazyFrame]:
+    """Iterate over join results by chrom + A chunk via IEJoin.
+
+    Yields nothing if no chunks have rows; callers (``pairwise_join_iter``) handle
+    the empty-case schema frame.
+    """
+    df_a = join_resources.df_a
+    df_b = join_resources.df_b
+    distance = join_resources.distance
+    chunk_size = join_resources.chunk_size
+    col_a = join_resources.col_a
+    col_b = join_resources.col_b
+
+    # Restrict the chrom loop to chroms present in both tables. Chroms unique to A
+    # would otherwise drive a wasted A chunk loop with no possible matches.
+    for chrom, last_index_a in (
+        df_a
+        .group_by(col_a.chrom)
+        .agg(pl.len().alias('last_index'))
+        .join(
+            df_b.select(pl.col(col_b.chrom)).unique(),
+            left_on=col_a.chrom, right_on=col_b.chrom,
+            how='inner',
+        )
+        .sort(col_a.chrom)
+    ).collect().rows():
+        df_a_chrom = (
+            df_a.filter(pl.col(col_a.chrom) == chrom)
+            .with_row_index('_index_chrom_a')
+        )
+        df_b_chrom = df_b.filter(pl.col(col_b.chrom) == chrom)
+
+        start_index_a = 0
+        while start_index_a < last_index_a:
+            end_index_a = start_index_a + chunk_size
+
+            df_a_chunk = df_a_chrom.filter(
+                pl.col('_index_chrom_a') >= start_index_a,
+                pl.col('_index_chrom_a') < end_index_a,
+            )
+
+            end_max, pos_min = (
+                df_a_chunk
+                .select(
+                    pl.col(col_a.end).max().alias('end_max'),
+                    pl.col(col_a.pos).min().alias('pos_min'),
+                )
+                .collect()
+                .row(0)
+            )
+
+            if end_max is None or pos_min is None:
+                start_index_a = end_index_a
+                continue
+
+            # Pre-filter B to records that could possibly match any A record in this
+            # chunk. Bounds use <= and >= to match the inner predicate (touching cases
+            # at distance == 0 must survive the prefilter).
+            df_b_chunk = df_b_chrom.filter(
+                pl.col(col_b.pos) - distance <= end_max,
+                pl.col(col_b.end) + distance >= pos_min,
+            )
+
+            yield _build_join_pair(
+                df_a_chunk,
+                df_b_chunk,
+                distance,
+                col_a,
+                col_b,
+            ).collect().lazy()
+
+            start_index_a = end_index_a

{agglovar-0.0.1.dev13 → agglovar-0.0.1.dev14}/src/agglovar/dev/imports.py

@@ -209,7 +209,7 @@ def alias_to_obj_dict(
                             )
                         ) + (
                             ('.' + module_path)
-                                if module_path is not None else ''
+                            if module_path is not None else ''
                         )
                     )
 
@@ -350,7 +350,6 @@ def get_defined_names(
         imports (None for others) and is the level of the import with 0 being the top-level and 1 or more being a
         relative import.
     """
-
     tree = ast.parse(inspect.getsource(mod))
 
     for node in tree.body:
@@ -371,8 +370,10 @@ def get_defined_names(
         elif isinstance(node, ast.ImportFrom):
 
             module_name = (
-                '.'.join(mod.__name__.split('.')[:-node.level]
-                + (node.module.split('.') if node.module else []))
+                '.'.join(
+                    mod.__name__.split('.')[:-node.level]
+                    + (node.module.split('.') if node.module else [])
+                )
             )
 
             for node_name in node.names:
@@ -443,6 +444,7 @@ def get_module_definitions(
                     if include_dunder or not target.id.startswith('__'):
                         yield target.id
 
+
 def librl(
         packages: Iterable[str] | str = 'agglovar',
 ) -> None:
@@ -455,7 +457,7 @@ def librl(
         before the first dot).
     """
     if isinstance(packages, str):
-        packages = {packages,}
+        packages = {packages, }
     else:
         packages = set(packages)