agglovar 0.0.1.dev8__tar.gz → 0.0.1.dev10__tar.gz

{agglovar-0.0.1.dev8/src/agglovar.egg-info → agglovar-0.0.1.dev10}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: agglovar
-Version: 0.0.1.dev8
+Version: 0.0.1.dev10
 Summary: Toolkit for fast genomic variant transformations and intersects
 License-Expression: MIT
 Requires-Python: >=3.12

{agglovar-0.0.1.dev8 → agglovar-0.0.1.dev10}/src/agglovar/__init__.py

@@ -1,6 +1,6 @@
 """Agglovar: A toolkit for fast genomic variant transformations and intersects."""
 
-__version__ = '0.0.1.dev8'
+__version__ = '0.0.1.dev10'
 
 __all__ = [
     'align',
@@ -13,6 +13,7 @@ __all__ = [
     'meta',
     'schema',
     'seqmatch',
+    'util',
 ]
 
 from . import align
@@ -25,3 +26,4 @@ from . import merge
 from . import meta
 from . import schema
 from . import seqmatch
+from . import util

{agglovar-0.0.1.dev8 → agglovar-0.0.1.dev10}/src/agglovar/bed/intersect.py

@@ -86,6 +86,17 @@ def as_proportion(
     if name is None or not (name := name.strip()):
         raise ValueError('Name must be a non-empty string')
 
+    if isinstance(df_a, pl.DataFrame):
+        df_a = df_a.lazy()
+
+    if isinstance(df_b, pl.DataFrame):
+        df_b = df_b.lazy()
+
+    col_set_a = set(df_a.collect_schema().names())
+
+    if '_index' not in col_set_a:
+        df_a = df_a.with_row_index('_index')
+
     # Collapse b, index a
     df_b_nr = (
         merge_depth(df_b, 0, col_names_b)
@@ -96,8 +107,6 @@ def as_proportion(
 
     df_a = (
         df_a
-        .drop('_index', strict=False)
-        .with_row_index('_index')
         .filter(
             pl.col(col_names_a.pos).is_not_null(),
             pl.col(col_names_a.end).is_not_null(),
@@ -128,11 +137,12 @@ def as_proportion(
         df_a
         .select(
             '_index',
+            '_index_a_bed_as_prop',
             (col_expr_a.end - col_expr_a.pos).alias('len')
         )
         .join(
             df_join.lazy(),
-            on='_index',
+            on='_index_a_bed_as_prop',
             how='left',
         )
         .select(

{agglovar-0.0.1.dev8 → agglovar-0.0.1.dev10}/src/agglovar/bed/join.py

@@ -46,6 +46,15 @@ class _JoinResources:
         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')
+
+        if '_index' not in cols_b:
+            df_b = df_b.with_row_index('_index')
+
         # Set column names
         ref_cols = get_coord_cols('ref')
 
@@ -65,16 +74,18 @@ class _JoinResources:
         # Prepare tables
         df_a = (
             df_a
-            .select(*col_select_a)
-            .drop('_index_a', strict=False)
-            .with_row_index('_index_a')
+            .select(
+                pl.col('_index').alias('_index_a'),
+                *col_select_a
+            )
         )
 
         df_b = (
             df_b
-            .select(*col_select_b)
-            .drop('_index_b', strict=False)
-            .with_row_index('_index_b')
+            .select(
+                pl.col('_index').alias('_index_b'),
+                *col_select_b
+            )
         )
 
         self.df_a = df_a
@@ -322,13 +333,10 @@ def pairwise_join_tree(
 
     for df_batch_b in (
         df_b
-        .drop('_index_b', strict=False)
-        .with_row_index('_index_b')
         .select(pl.col('_index_b'), *col_b.exprs())
         .collect_batches(chunk_size=chunk_size)
     ):
         for index_b, chrom_b, pos_b, end_b in df_batch_b.iter_rows():
-            # print(f'Adding: {chrom_b} - {pos_b} - {end_b}')
             itree[chrom_b].addi(pos_b - distance, end_b + distance, index_b)
 
     # Intersect
@@ -336,8 +344,6 @@ def pairwise_join_tree(
 
     for df_chunk in (
         df_a
-        .drop('_index_a', strict=False)
-        .with_row_index('_index_a')
         .select(pl.col('_index_a'), *col_a.exprs())
         .collect_batches(chunk_size=chunk_size)
     ):

{agglovar-0.0.1.dev8 → agglovar-0.0.1.dev10}/src/agglovar/dev/imports.py

@@ -6,13 +6,17 @@ __all__ = [
     'find_imports',
     'get_defined_names',
     'get_module_definitions',
+    'librl',
 ]
 
 import ast
+from collections.abc import Iterable
 from dataclasses import dataclass, field
+from graphlib import TopologicalSorter
 import importlib
 import inspect
 import re
+import sys
 from types import ModuleType
 from typing import Any, Generator, Optional
 
@@ -439,3 +443,40 @@ def get_module_definitions(
                     if include_dunder or not target.id.startswith('__'):
                         yield target.id
 
+def librl(
+        packages: Iterable[str] | str = 'agglovar',
+) -> None:
+    """Reload libraries.
+
+    Reloads all modules in the given packages and their dependencies. Attempts to resolve import
+    order by inspecting members of each package.
+
+    :param packages: Names of packages to reload. Must be top-level package names (name
+        before the first dot).
+    """
+    if isinstance(packages, str):
+        packages = {packages,}
+    else:
+        packages = set(packages)
+
+    deps: dict[str, set[str]] = {}
+
+    for key in sorted(sys.modules.keys()):
+        if key.split('.', 1)[0] not in packages:
+            continue
+
+        deps.setdefault(key, set())
+
+        for attr_name in dir(sys.modules[key]):
+            attr = getattr(sys.modules[key], attr_name)
+
+            if '__module__' in dir(attr):
+                module = attr.__module__
+
+                if not attr.__module__.split('.', 1)[0] in packages or module == key:
+                    continue
+
+                deps[key].add(attr.__module__)
+
+    for key in TopologicalSorter(deps).static_order():
+        importlib.reload(sys.modules[key])

{agglovar-0.0.1.dev8 → agglovar-0.0.1.dev10}/src/agglovar/kmer/plot.py

@@ -32,6 +32,19 @@ def dotplot(
 ) -> t_mpl_Figure:
     """Generate a dotplot for two sequences.
 
+    Configuration items:
+
+    * label_x, label_y: Axis labels (default "Sequence", and "Reference")
+    * start_x, start_y: Start position along each axis. Should be the first base position in the extracted region (Default 0).
+    * linewidth: Width of the lines (default 1)
+    * plot_width, plot_height, plot_dpi: Plot dimensions (default 7, 7, 300)
+    * invert_y: Invert the y axis (default False)
+    * color: A tuple of matplotlib colors for forward and reverse k-mer dots (default ("black", "red"))
+      * May be "reverse" or "rev" to reverse the order of the default colors.
+    * color_reverse: Reverse the order of the default colors (defaults to invert_y)
+    * kmer_util: A kmer_util.KmerUtil object (default None)
+    * k: K-mer size if kmer_util is missing (default 32)
+
     :params seq_x: Sequence 1 (horizontal axis).
     :params seq_y: Sequence 2 (vertical axis).
     :params config: Configuration dictionary.
@@ -56,10 +69,8 @@ def dotplot(
     # Get k-mer util.
     if 'kutil' in config:
         kutil = config['kutil']
-    elif 'k' in config:
-        kutil = kmer_util.KmerUtil(int(config['k']))
     else:
-        kutil = kmer_util.KmerUtil(32)
+        kutil = kmer_util.KmerUtil(int(config.get('k', 32)))
 
     if anno_list is None:
         anno_list = list()
@@ -72,8 +83,10 @@ def dotplot(
     elif isinstance(color, str):
         if color.lower() in {'reverse', 'rev'}:
             color = ('red', 'black')
+        elif color.lower() in {'forward', 'fwd'}:
+            color = ('black', 'red')
         else:
-            raise ValueError(f'Argument color is a string: Must be "reverse" or "rev": {color}')
+            raise ValueError(f'Argument color is a string: Must be "reverse" or "rev" or a tuple of two colors: {color}')
 
     elif isinstance(color, (list, tuple)):
         if len(color) != 2:
@@ -99,9 +112,12 @@ def dotplot(
 
     plot_props = dict()  # Plot properties, such as x and y limits
 
+    if config.get('color_reverse', invert_y):
+        color = color[::-1]
+
     # Process annotation dict
-    anno_list_background = list()
-    anno_list_foreground = list()
+    anno_list_background = []
+    anno_list_foreground = []
 
     if anno_list is None:
         anno_list = list()

{agglovar-0.0.1.dev8 → agglovar-0.0.1.dev10}/src/agglovar/merge/base.py

@@ -21,8 +21,8 @@ from ..util.str import collision_rename
 CallsetDefType: TypeAlias = (
     pl.DataFrame
     | pl.LazyFrame
-    | tuple[
-        pl.DataFrame | pl.LazyFrame, Optional[Any]
+    | Iterable[
+        pl.DataFrame | pl.LazyFrame | str
     ]
 )
 """Alias for acceptable types."""
@@ -38,12 +38,17 @@ class MergeBase(ABC):
     @abstractmethod
     def __call__(
             self,
-            callsets: Iterable[CallsetDefType]
+            callsets: Iterable[CallsetDefType],
+            retain_index: bool = False,
+            pre_filter: Optional[Iterable[pl.Expr]] = None,
     ) -> pl.LazyFrame:
         """
         Intersect callsets.
 
         :param callsets: Callsets to intersect.
+        :param retain_index: If `True`, do not drop an existing "_index" column if it exists.
+        :param pre_filter: If set, filter each table with these expressions. Filter is applied
+            last (after "_index" is set).
 
         :return: A merged callset table.
         """
@@ -52,7 +57,8 @@ class MergeBase(ABC):
     @staticmethod
     def get_intersect_tuples(
             callsets: Iterable[CallsetDefType],
-            retain_index: bool = False
+            retain_index: bool = False,
+            pre_filter: Optional[Iterable[pl.Expr]] = None,
     ) -> list[tuple[pl.LazyFrame, str, int]]:
         """
         Transform input arguments to a list of tuples with set fields.
@@ -72,6 +78,8 @@ class MergeBase(ABC):
         :param callsets: Callsets parameter. May be an iterable of DataFrames, LazyFrames, or
             tuples of (DataFrame, name).
         :param retain_index: If `True`, do not drop an existing "_index" column if it exists.
+        :param pre_filter: If set, filter each table with these expressions. Filter is applied
+            last (after "_index" is set).
 
         :return: A list of tuples where each tuple element represents one input source.
         """
@@ -84,30 +92,43 @@ class MergeBase(ABC):
 
         i = 0
 
+        pre_filter = list(pre_filter if pre_filter is not None else [])
+
         if callsets is None:
             raise ValueError('Missing callsets')
 
         for callset in callsets:
-            if isinstance(callset, tuple):
-                if not len(callset) == 2:
+            # Get callset table and name
+
+            if isinstance(callset, (pl.LazyFrame, pl.DataFrame)):
+                callset_table = callset
+                callset_name = None
+            else:
+                try:
+                    callset = tuple(callset)
+                except TypeError:
                     raise ValueError(
-                        f'Callset at index {i} tuple must have exactly 2 elements: {callset}'
+                        f'Callset at index {i}: Callset must be a Polars table or an iterable '
+                        f'of two elements (table, name): Received {type(callset)}'
                     )
 
-                callset_table = callset[0]
-                callset_name = _get_name(callset[1], i, name_set)
+                if len(callset) == 1:
+                    callset_table = callset[0]
+                    callset_name = None
 
-            else:
-                callset_table = callset
-                callset_name = _get_name(None, i, name_set)
+                elif len(callset) == 2:
+                    callset_table = callset[0]
+                    callset_name = str(callset[1])
 
-            if isinstance(callset_table, pl.DataFrame):
-                callset_table = callset_table.lazy()
-            elif not isinstance(callset_table, pl.LazyFrame):
-                raise TypeError(
-                    f'Callset at index {i} must be a DataFrame or LazyFrame, got {type(callset_table)}'
-                )
+                else:
+                    raise ValueError(f'Callset at index {i}: Input callsets must be a sequence of one or two elements: Received {len(callset)}')
 
+            if not isinstance(callset_table, (pl.LazyFrame, pl.DataFrame)):
+                raise ValueError(f'Callset at index {i}: First element Polars table: Received {type(callset_table)}')
+
+            callset_name = _get_name(callset_name, i, name_set)
+
+            # Set index if missing or forced
             if not (retain_index and '_index' in callset_table.collect_schema().names()):
                 callset_table = (
                     callset_table
@@ -118,22 +139,27 @@ class MergeBase(ABC):
             # Add missing IDs
             callset_table = (
                 callset_table
+                .rename({'_index': '_mg_src_index'})
                 .with_columns(
                     pl.coalesce(
                         pl.col(r'^id$'),  # Existing ID if present
-                        pl.concat_str(pl.lit('var'), pl.col('_index'))  # Fill missing/null
+                        pl.concat_str(pl.lit('var'), pl.col('_mg_src_index'))  # Fill missing/null
                     ).cast(pl.String).alias('id')
                 )
                 .with_columns(
                     (
                         pl.when(pl.col('id').is_null())
-                        .then(pl.concat_str(pl.lit('var'),  pl.col('_index')))
+                        .then(pl.concat_str(pl.lit('var'),  pl.col('_mg_src_index')))
                         .otherwise(pl.col('id'))
                     ).cast(pl.String).alias('id')
                 )
+                .filter(
+                    *pre_filter
+                )
             )
 
-            callset_tuple_list.append((callset_table, callset_name, i))
+            # Finalize
+            callset_tuple_list.append(tuple((callset_table, callset_name, i)))
             name_set.add(callset_name)
 
             i += 1

{agglovar-0.0.1.dev8 → agglovar-0.0.1.dev10}/src/agglovar/merge/cumulative.py

@@ -15,11 +15,13 @@ where multiple join choices are possible.
 """
 
 __all__ = [
+    'LeadStrategy',
     'MergeCumulative',
 ]
 
 from collections.abc import Iterable
 from enum import Enum
+from typing import Optional
 
 import polars as pl
 import polars.selectors as cs
@@ -72,12 +74,11 @@ def _init_cumulative(
     :return: A table ready to be used as a cumulative table or appended to an existing one.
     """
     mg_cols = ('_mg_src', '_mg_src_pos', '_mg_stat')
-    required_cols = set(required_cols) - set(mg_cols)
+    required_cols = (set(required_cols) - set(mg_cols)) | {'_mg_src_index'}
 
     return (
         df_next
         .select(
-            pl.col('_index').alias('_mg_index'),
             *required_cols
         )
         .with_columns(  # Source columns
@@ -85,7 +86,7 @@ def _init_cumulative(
                 pl.struct(
                     pl.lit(src_index).cast(pl.Int32).alias('index'),
                     pl.lit(src_name).cast(pl.String).alias('id'),
-                    pl.col('_mg_index').alias('var_index'),
+                    pl.col('_mg_src_index').alias('var_index'),
                     pl.col('id').cast(pl.String).alias('var_id'),
                 )
             ).alias('_mg_src'),
@@ -99,11 +100,11 @@ def _init_cumulative(
             # May be a Polars bug.
             pl.col('_mg_src_pos').list.concat(pl.col('pos').cast(pl.Int64))
         )
+        .drop('_mg_src_index')
     )
 
 def _get_match(
         df_cumulative: pl.LazyFrame,
-        df_next: pl.LazyFrame,
         df_join: pl.DataFrame,
         src_index: int,
         src_name: str,
@@ -111,9 +112,10 @@ def _get_match(
 ) -> pl.LazyFrame:
     return (  # Update matching variants
         df_cumulative
+        .with_row_index('_index')
         .join(
             df_join.lazy().select(pl.all().name.prefix('_mg_join_')),
-            left_on='_mg_index', right_on='_mg_join_index_a', how='inner'
+            left_on='_index', right_on='_mg_join_index_a', how='inner'
         )
         .with_columns(
             pl.col('_mg_src').list.concat(
@@ -132,7 +134,7 @@ def _get_match(
                 })
             )
         )
-        .drop('_mg_index', cs.starts_with('_mg_join_'))
+        .drop('_index', cs.starts_with('_mg_join_'))
     )
 
 
@@ -161,7 +163,10 @@ class MergeCumulative(MergeBase):
     def __call__(
             self,
             callsets: Iterable[CallsetDefType],
-            retain_index: bool = False
+            retain_index: bool = False,
+            pre_filter: Optional[Iterable[pl.Expr]] = None,
+            sort: bool = True,
+            add_id: bool = True,
     ) -> pl.LazyFrame:
         """
         Intersect callsets.
@@ -169,10 +174,12 @@ class MergeCumulative(MergeBase):
         :param callsets: Callsets to intersect.
         :param retain_index: If `True`, do not drop an existing "_index" column in callset tables
             if they exist.
+        :param pre_filter: If set, filter each table with these expressions. Filter is applied
+            last (after "_index" is set).
 
         :return: A merged callset table.
         """
-        callsets = self.get_intersect_tuples(callsets, retain_index)
+        callsets = self.get_intersect_tuples(callsets, retain_index, pre_filter)
 
         if len(callsets) == 0:
             raise ValueError('No callsets to intersect.')
@@ -195,7 +202,7 @@ class MergeCumulative(MergeBase):
         df_next, src_name, src_index = callsets[0]
 
         merge_stat_cols = {  # Empty join, gets names of columns this join will return
-            col: dtype for col, dtype in self.pairwise_join.join(df_next.head(0).drop('_index'), df_next.head(0).drop('_index')).collect_schema().items()
+            col: dtype for col, dtype in self.pairwise_join.join(df_next.head(0), df_next.head(0)).collect_schema().items()
                 if col not in {'index_a', 'index_b', 'id_a', 'id_b'}
         }
 
@@ -216,16 +223,23 @@ class MergeCumulative(MergeBase):
             # Intersect
             df_join = (
                 (
-                    self.pairwise_join.join(df_cumulative, df_next.drop('_index'))
+                    self.pairwise_join.join(df_cumulative, df_next)
                     .unique('index_a', keep='first')
                     .unique('index_b', keep='first')
                     .sort('index_a', 'index_b')
                 )
                 .join(
-                    df_next.select('_index', 'pos', 'id'),
+                    (
+                        df_next
+                        .with_row_index('_index')
+                        .select('_index', '_mg_src_index', 'pos', 'id')
+                    ),
                     left_on='index_b', right_on='_index',
                     how='left'
                 )
+                .with_columns(
+                    pl.col('_mg_src_index').alias('index_b')
+                )
                 .collect()
             )
 
@@ -235,36 +249,36 @@ class MergeCumulative(MergeBase):
             # 3) New: New records in this callset
 
             df_match = _get_match(
-                df_cumulative, df_next, df_join, src_index, src_name, merge_stat_cols
+                df_cumulative, df_join, src_index, src_name, merge_stat_cols
             )
 
             df_nomatch = (  # Gather records that were not updated
                 df_cumulative
+                .with_row_index('_index')
                 .join(
                     df_join.lazy(),
-                    left_on='_mg_index', right_on='index_a', how='anti'
+                    left_on='_index', right_on='index_a', how='anti'
                 )
-                .drop('_mg_index')
+                .drop('_index')
             )
 
-            df_new = (
+            df_new = (  # New records
                 _init_cumulative(
                     (
                         df_next
                         .join(
                             df_join.lazy(),
-                            left_on='_index', right_on='index_b', how='anti'
+                            left_on='_mg_src_index', right_on='index_b', how='anti'
                         )
                     ),
                     src_index, src_name, required_cols, merge_stat_cols
                 )
-                .drop('_mg_index')
             )
 
             df_cumulative = (
                 pl.concat([df_match, df_nomatch, df_new])
                 .sort('chrom', 'pos', 'end')
-                .with_row_index('_mg_index')
+                .with_row_index('_mg_src_index')
                 .collect()
                 .lazy()
             )
@@ -272,7 +286,7 @@ class MergeCumulative(MergeBase):
         # Subset to merge columns (drop columns needed for join)
         df_cumulative = (
             df_cumulative
-            .select(cs.starts_with('_mg_').exclude('_mg_index'))
+            .select(cs.starts_with('_mg_').exclude('_mg_src_index'))
             .rename(lambda col: col.removeprefix('_'))
         )
 
@@ -320,31 +334,43 @@ class MergeCumulative(MergeBase):
                         .filter(pl.col('mg_src_lead').struct.field('index') == src_index)
                         .select(
                             '_mg_index',
-                            pl.col('mg_src_lead').struct.field('var_index').alias('_mg_src_var_index'),
+                            pl.col('mg_src_lead').struct.field('var_index').alias('_mg_src_index'),
                         )
                     ),
-                    left_on='_index', right_on='_mg_src_var_index', how='inner'
+                    on='_mg_src_index', how='inner'
                 )
-                .drop('_index')
+                .drop('_mg_src_index')
             )
 
             for col in set(table_cols) - df_next_cols:
                 df_next = df_next.with_columns(pl.lit(None).cast(all_col_dict[col]).alias(col))
 
-            df_next = df_next.with_columns(
-                pl.col('filter').fill_null([])
+            df_next = (
+                df_next.with_columns(
+                    pl.col('filter').fill_null([])
+                )
+                .select('_mg_index', *table_cols)
             )
 
             lead_list.append(df_next)
 
-        return (
+        df_merge = (
             pl.concat(lead_list)
             .join(
                 df_cumulative.lazy(),
                 on='_mg_index', how='inner'
             )
-            .with_columns(id_version_expr())
+        )
+
+        if add_id:
+            df_merge = df_merge.with_columns(id_version_expr())
+
+        if sort:
+            df_merge = df_merge.sort('chrom', 'pos', 'end', 'id')
+
+
+        return (
+            df_merge
             .drop(*(['_mg_index'] + ['filter'] if drop_filter else []))
             .select(col_order)
-            .sort('chrom', 'pos', 'end', 'id')
         )

{agglovar-0.0.1.dev8 → agglovar-0.0.1.dev10}/src/agglovar/meta/descriptors.py

@@ -14,6 +14,7 @@ Classes:
 __all__ = [
     'AutoInitBase',
     'CheckedBool',
+    'CheckedObject',
     'CheckedString',
     'OneWayBool',
     'BoundedInt',

{agglovar-0.0.1.dev8 → agglovar-0.0.1.dev10}/src/agglovar/pairwise/base.py

@@ -9,7 +9,6 @@ __all__ = [
 
 from abc import ABC, abstractmethod
 from collections.abc import Iterator, Iterable
-from typing import Optional
 
 import polars as pl
 
@@ -40,12 +39,15 @@ class PairwiseJoin(ABC):
     def join_iter(
             self,
             df_a: pl.DataFrame | pl.LazyFrame,
-            df_b: pl.DataFrame | pl.LazyFrame
+            df_b: pl.DataFrame | pl.LazyFrame,
+            retain_index: bool = False,
     ) -> Iterator[pl.LazyFrame]:
         """Find all pairs of variants in two sources that meet a set of criteria.
 
         :param df_a: Source dataframe.
         :param df_b: Target dataframe.
+        :param retain_index: If True, do not drop an existing "_index" column in callset tables
+            if they exist.
 
         :yields: A LazyFrame for each chunk.
         """
@@ -54,7 +56,8 @@ class PairwiseJoin(ABC):
     def join(
             self,
             df_a: pl.DataFrame | pl.LazyFrame,
-            df_b: pl.DataFrame | pl.LazyFrame
+            df_b: pl.DataFrame | pl.LazyFrame,
+            retain_index: bool = False,
     ) -> pl.LazyFrame:
         """Find all pairs of variants in two sources that meet a set of criteria.
 
@@ -62,11 +65,13 @@ class PairwiseJoin(ABC):
 
         :param df_a: Table A.
         :param df_b: Table B.
+        :param retain_index: If True, do not drop an existing "_index" column in callset tables
+            if they exist.
 
         :returns: A join table.
         """
         return pl.concat(
-            self.join_iter(df_a, df_b)
+            self.join_iter(df_a, df_b, retain_index=retain_index)
         )
 
     @property

{agglovar-0.0.1.dev8 → agglovar-0.0.1.dev10}/src/agglovar/pairwise/overlap/_overlap.py

@@ -8,10 +8,14 @@ from collections.abc import (
     Iterable,
     Iterator,
     Mapping,
+    MutableMapping,
 )
 import functools
 import operator
-from typing import Optional
+from typing import (
+    Any,
+    Optional,
+)
 from warnings import warn
 
 import polars as pl
@@ -189,21 +193,26 @@ class PairwiseOverlap(PairwiseJoin):
     def weight_expr(self):
         return self.weight_strategy.expr
 
-
     def join_iter(
             self,
             df_a: pl.DataFrame | pl.LazyFrame,
             df_b: pl.DataFrame | pl.LazyFrame,
+            retain_index: bool = False,
     ) -> Iterator[pl.LazyFrame]:
         """Find all pairs of variants in two sources that meet a set of criteria.
 
         :param df_a: Source dataframe.
         :param df_b: Target dataframe.
+        :param retain_index: If True, do not drop an existing "_index" column in callset tables
+            if they exist.
 
         :yields: A LazyFrame for each chunk.
         """
 
-        if self.is_equi_offset or self.chunk_size == 0:
+        # Prepare tables
+        df_a, df_b = self._prepare_tables(df_a, df_b, warn_on_reserved=True, retain_index=retain_index)
+
+        if self.chunk_size == 0 or (self.is_equi_offset and self.chunk_size is None):
             return self._join_iter_notchunked(df_a, df_b)
 
         return self._join_iter_chunked(df_a, df_b)
@@ -215,22 +224,16 @@ class PairwiseOverlap(PairwiseJoin):
     ) -> Iterator[pl.LazyFrame]:
         """Find all pairs of variants in two sources that meet a set of criteria.
 
-        :param df_a: Source dataframe.
-        :param df_b: Target dataframe.
+        :param df_a: Source dataframe after `_prepare_tables()`.
+        :param df_b: Target dataframe after `_prepare_tables()`.
 
         :yields: A LazyFrame for each chunk.
         """
-
-        # print('_join_iter_chunked(): Starting')
-
         join_empty = True  # Detects if no joins were written
 
         chunk_range = self._get_chunk_range()
         chunk_size = self.chunk_size if self.chunk_size is not None else DEFAULT_CHUNK_SIZE
 
-        # Prepare tables
-        df_a, df_b = self._prepare_tables(df_a, df_b, warn_on_reserved=True)
-
         for chrom, last_index_a in (
             df_a
             .group_by('chrom_a')
@@ -308,18 +311,13 @@ class PairwiseOverlap(PairwiseJoin):
     ) -> Iterator[pl.LazyFrame]:
         """Find all pairs of variants in two sources that meet a set of criteria.
 
-        :param df_a: Source dataframe.
-        :param df_b: Target dataframe.
+        :param df_a: Source dataframe after `_prepare_tables()`.
+        :param df_b: Target dataframe after `_prepare_tables()`.
 
         :yields: A LazyFrame for each chunk.
         """
-        # print('_join_iter_notchunked(): Starting')
-
         join_empty = True  # Detects if no joins were written
 
-        # Prepare tables
-        df_a, df_b = self._prepare_tables(df_a, df_b, warn_on_reserved=True)
-
         chrom_list = sorted(
             set(df_a.select('chrom_a').unique().collect().to_series().to_list())
             | set(df_b.select('chrom_b').unique().collect().to_series().to_list())
@@ -415,7 +413,8 @@ class PairwiseOverlap(PairwiseJoin):
             self,
             df_a: pl.DataFrame | pl.LazyFrame,
             df_b: pl.DataFrame | pl.LazyFrame,
-            warn_on_reserved: bool = False
+            retain_index: bool = False,
+            warn_on_reserved: bool = False,
     ) -> tuple[pl.LazyFrame, pl.LazyFrame]:
         """Prepares tables for join.
 
@@ -424,6 +423,8 @@ class PairwiseOverlap(PairwiseJoin):
 
         :param df_a: Table A.
         :param df_b: Table B.
+        :param retain_index: If True, do not drop an existing "_index" column in callset tables
+            if they exist.
         :param warn_on_reserved: If True, generate a warning if reserved columns are found and
             drop them. If false, raise an error.
 
@@ -460,25 +461,37 @@ class PairwiseOverlap(PairwiseJoin):
             )
 
         # Drop reserved columns
-        if reserved_cols := sorted(self.check_reserved_cols(columns_a)):
-            err_str = f'Reserved columns in table "A": {", ".join(reserved_cols)}'
+        reserved_col_a = self.check_reserved_cols(columns_a)
+        reserved_col_b = self.check_reserved_cols(columns_b)
+
+        retain_index_a = '_index' in reserved_col_a and retain_index
+        retain_index_b = '_index' in reserved_col_b and retain_index
+
+        if retain_index:
+            reserved_col_a -= {'_index', }
+            reserved_col_b -= {'_index', }
+
+        if reserved_col_a:
+            reserved_col_a = sorted(reserved_col_a)
+            err_str = f'Reserved columns in table "A": {", ".join(reserved_col_a)}'
 
             if warn_on_reserved:
                 warn(f'{err_str}: Dropping column(s)')
             else:
                 raise ValueError(err_str)
 
-            df_a = df_a.drop(reserved_cols, strict=False)
+            df_a = df_a.drop(reserved_col_a, strict=False)
 
-        if reserved_cols := sorted(self.check_reserved_cols(columns_b)):
-            err_str = f'Reserved columns in table "B": {", ".join(reserved_cols)}'
+        if reserved_col_b:
+            reserved_col_b = sorted(reserved_col_b)
+            err_str = f'Reserved columns in table "B": {", ".join(reserved_col_b)}'
 
             if warn_on_reserved:
                 warn(f'{err_str}: Dropping column(s)')
             else:
                 raise ValueError(err_str)
 
-            df_b = df_b.drop(reserved_cols, strict=False)
+            df_b = df_b.drop(reserved_col_b, strict=False)
 
         # Cast columns
         try:
@@ -526,15 +539,19 @@ class PairwiseOverlap(PairwiseJoin):
         )
 
         # Set index
-        df_a = (
-            df_a
-            .with_row_index('_index')
-        )
+        if not retain_index_a:
+            df_a = (
+                df_a
+                .drop('_index', strict=False)
+                .with_row_index('_index')
+            )
 
-        df_b = (
-            df_b
-            .with_row_index('_index')
-        )
+        if not retain_index_b:
+            df_b = (
+                df_b
+                .drop('_index', strict=False)
+                .with_row_index('_index')
+            )
 
         # Set ID
         if 'id' not in columns_a:
@@ -644,7 +661,7 @@ class PairwiseOverlap(PairwiseJoin):
     def _append_join_cols(
             self,
             exprs: Iterable[pl.Expr | str] | pl.Expr | str,
-            join_expr_map: Mapping[str, Optional[pl.Expr]],
+            join_expr_map: MutableMapping[str, Optional[pl.Expr]],
             no_replace: bool = False,
     ) -> None:
         """Append expressions to the list of columns included in the join table.
@@ -660,7 +677,7 @@ class PairwiseOverlap(PairwiseJoin):
 
         :param exprs: A join column or a list of join columns where each may be defined as a string
             (known column name or a field already in the join table) or a Polars expression.
-        :param join_col_exprs: A dictionary to which the expressions will be added.
+        :param join_expr_map: A dictionary to which the expressions will be added.
         :param no_replace: Do not replace existing join column expressions if True.
         """
         if isinstance(exprs, (pl.Expr, str)):
@@ -937,7 +954,7 @@ class PairwiseOverlap(PairwiseJoin):
     @classmethod
     def from_definiton(
             cls,
-            overlap_def,
+            overlap_def: Mapping[str, Any] | Iterable[Mapping[str, Any]],
     ):
         keys: Optional[set[str]] = None
 
@@ -959,12 +976,12 @@ class PairwiseOverlap(PairwiseJoin):
                     f'of stage definitions: {type(overlap_def)}'
                 )
 
-            overlap_def = {
+            overlap_def: Mapping[str, Any] | Iterable[Mapping[str, Any]] = {
                 'stages': stage_list,
             }
 
         # Parse arguments
-        parsed_def = {}
+        parsed_def: dict[str, Any] = {}
 
         for key, value in overlap_def.items():
             if key == 'stages':

{agglovar-0.0.1.dev8 → agglovar-0.0.1.dev10}/src/agglovar/pairwise/overlap/_stage.py

@@ -38,6 +38,10 @@ class PairwiseOverlapStage():
     :ivar match_ref: "REF" column must match between two variants.
     :ivar match_alt: "ALT" column must match between two variants.
     :ivar match_vartype: "VARTYPE" column must match between two variants.
+    :ivar match_filter_pass: If true, then matched variants must either pass filters ("filter"
+        column emtpy for both variants) or must fail filters ("filter" column non-empty for both
+        variants). This does not compare the contents of the "filter" column, just the presence or
+        absence of a filter.
     :ivar match_prop_min: Minimum matched base proportion in alignment or None to not match.
     """
     ro_min: Optional[float] = BoundedFloat(min_val=0.0, max_val=1.0)
@@ -82,6 +86,7 @@ class PairwiseOverlapStage():
         self.match_ref = match_ref
         self.match_alt = match_alt
         self.match_vartype = match_vartype
+        self.match_filter_pass = match_filter_pass
         self.seg_ro_min = seg_ro_min
         self.match_prop_min = match_prop_min
 
@@ -185,6 +190,11 @@ class PairwiseOverlapStage():
                 pl.col('seg_ro') >= self.seg_ro_min
             )
 
+        if self.match_filter_pass:
+            join_predicates_list.append(
+                (pl.col('filter_a').list.len() == 0) ^ (pl.col('filter_b').list.len() == 0)
+            )
+
         if self.match_prop_min is not None and self.match_prop_min > 0.0:
             join_filters_list.append(
                 pl.col('match_prop') >= self.match_prop_min

{agglovar-0.0.1.dev8 → agglovar-0.0.1.dev10}/src/agglovar/util/var.py

@@ -1,5 +1,9 @@
 """Variant utilites."""
 
+__all__ = [
+    'id_version_expr',
+]
+
 from typing import Any
 
 import polars as pl

{agglovar-0.0.1.dev8 → agglovar-0.0.1.dev10/src/agglovar.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: agglovar
-Version: 0.0.1.dev8
+Version: 0.0.1.dev10
 Summary: Toolkit for fast genomic variant transformations and intersects
 License-Expression: MIT
 Requires-Python: >=3.12