diff --git a/src/seqwin/cli.py b/src/seqwin/cli.py index 59b68d6..09ea4c8 100644 --- a/src/seqwin/cli.py +++ b/src/seqwin/cli.py @@ -193,6 +193,11 @@ def main( help='Number of parallel processes or threads to use.', rich_help_panel='Miscellaneous' ), + low_memory: bool = typer.Option( + False, '--low-memory', show_default=False, + help='Reduce peak memory by recomputing minimizers in a second pass.', + rich_help_panel='Miscellaneous' + ), version: bool = typer.Option( False, '--version', callback=print_version, show_default=False, expose_value=False, is_eager=True, # run this before any other options @@ -232,14 +237,15 @@ def main( run_blast=not no_blast, no_filter=no_filter, #blast_neg_only=blast_neg_only, - seed=seed, - n_cpu=n_cpu, level=level, source=source, annotated=annotated, exclude_mag=exclude_mag, gzip=not no_gzip, api_key=api_key, - download_only=download_only + download_only=download_only, + seed=seed, + n_cpu=n_cpu, + low_memory=low_memory ) _ = run(config) diff --git a/src/seqwin/config.py b/src/seqwin/config.py index 3a0fbf1..92986f4 100644 --- a/src/seqwin/config.py +++ b/src/seqwin/config.py @@ -114,6 +114,7 @@ class Config(BaseModel): seed (int): Random seed for reproducibility. [42] n_cpu (int): Number of parallel processes or threads to use. [4] + low_memory (bool): If True, reduce peak memory by recomputing minimizers in a second pass. [False] version (str): Seqwin version. """ # Inputs @@ -162,6 +163,7 @@ class Config(BaseModel): # Miscellaneous seed: int = 42 n_cpu: int = 4 + low_memory: bool = False @computed_field @cached_property diff --git a/src/seqwin/kmers.py b/src/seqwin/kmers.py index c5fb706..55e6594 100644 --- a/src/seqwin/kmers.py +++ b/src/seqwin/kmers.py @@ -2,7 +2,7 @@ K-mer Graph =========== -A core module of Seqwin. Create a k-mer graph from k-mers of all input genome assemblies. +A core module of Seqwin. Build a k-mer graph from minimizers of all input genome assemblies. Dependencies: ------------- @@ -39,14 +39,14 @@ from .graph import build, _filter_kmers from .assemblies import Assemblies from .helpers import get_subgraphs -from .utils import print_time_delta, log_and_raise +from .utils import print_time_delta from .config import Config, RunState, HAS_MASH, WORKINGDIR, EDGE_W, NODE_P class KmerGraph(object): """ - 1. Create a weighted, undirected k-mer graph, and calculate node penalty scores. - 2. Extract low-penalty subgraphs from the k-mer graph with `self.filter()`. + 1. Build a weighted, undirected k-mer graph, and calculate node penalty scores. + 2. Extract low-penalty subgraphs with `self.filter()`. Attributes: kmers (NDArray[np.void]): A 1-D NumPy structured array of k-mers from all assemblies, grouped by node range. @@ -60,7 +60,7 @@ class KmerGraph(object): Generated with `self.filter()`. """ __slots__ = ( - 'kmers', 'nodes', 'edges', 'record_offsets', 'graph', 'subgraphs', '_filtered_flag' + 'kmers', 'nodes', 'edges', 'record_offsets', 'graph', 'subgraphs', '_is_filtered' ) kmers: NDArray[np.void] nodes: NDArray[np.void] @@ -68,21 +68,22 @@ class KmerGraph(object): record_offsets: NDArray[np.uintp] graph: nx.Graph subgraphs: tuple[frozenset[np.uint64], ...] | None - _filtered_flag: bool # True if `self.filter()` is called + _is_filtered: bool # True if `self.filter()` is called - def __init__(self, assemblies: Assemblies, kmerlen: int, windowsize: int, n_cpu: int) -> None: - """ - 1. Generate minimizer sketches and weighted edges. - 2. Generate k-mer nodes and calculate node penalty scores. + def __init__(self, assemblies: Assemblies, kmerlen: int, windowsize: int, n_cpu: int, low_memory: bool) -> None: + """Build the k-mer graph and calculate node penalty scores. Args: assemblies (Assemblies): See `Assemblies` in `assemblies.py`. kmerlen (int): See `Config` in `config.py`. windowsize (int): See `Config` in `config.py`. n_cpu (int): See `Config` in `config.py`. + low_memory (bool): See `Config` in `config.py`. """ n_assemblies = len(assemblies) logger.info(f'Building minimizer graph from {n_assemblies} assemblies...') + if low_memory: + logger.warning(' - Low-memory mode is enabled; graph construction may take longer.') tik = time() kmers, nodes, edges, record_offsets, record_ids = build( @@ -91,6 +92,7 @@ def __init__(self, assemblies: Assemblies, kmerlen: int, windowsize: int, n_cpu: windowsize, assemblies.is_target, n_cpu=n_cpu, + low_memory=low_memory ) # calculate penalty for each node n_tar = sum(assemblies.is_target) @@ -113,7 +115,7 @@ def __init__(self, assemblies: Assemblies, kmerlen: int, windowsize: int, n_cpu: self.record_offsets = record_offsets self.graph = None # create graph after filtering nodes and edges self.subgraphs = None - self._filtered_flag = False + self._is_filtered = False def filter( self, penalty_th: float, edge_weight_th: float, min_nodes: int, max_nodes: int | None, rng: Random @@ -133,9 +135,9 @@ def filter( kmers = self.kmers nodes = self.nodes edges = self.edges - filtered_flag = self._filtered_flag + is_filtered = self._is_filtered - if filtered_flag: + if is_filtered: logger.error(f'K-mers are already filtered, cannot filter again.') return None @@ -164,7 +166,7 @@ def filter( self.edges = edges self.graph = graph self.subgraphs = subgraphs - self._filtered_flag = True + self._is_filtered = True @staticmethod def __filter_graph(nodes: NDArray, edges: NDArray, edge_weight_th: float) -> tuple[NDArray, NDArray, nx.Graph]: @@ -258,13 +260,14 @@ def get_kmers( max_nodes_cap = config.max_nodes_cap sketchsize = config.sketchsize n_cpu = config.n_cpu + low_memory = config.low_memory working_dir = state.working_dir rng = state.rng n_tar = state.n_tar n_neg = state.n_neg - kmers = KmerGraph(assemblies, kmerlen, windowsize, n_cpu) + kmers = KmerGraph(assemblies, kmerlen, windowsize, n_cpu, low_memory) if no_filter: return kmers, None diff --git a/tests/smoke/test_cli.py b/tests/smoke/test_cli.py index 9bc7b00..af26875 100644 --- a/tests/smoke/test_cli.py +++ b/tests/smoke/test_cli.py @@ -28,6 +28,7 @@ def test_help_shows_key_options() -> None: '--no-mash', '--no-blast', '--threads', + '--low-memory', '--prefix', ): assert opt in output @@ -63,6 +64,7 @@ def _fake_run(config): '--no-mash', '--no-blast', '--threads', '2', + '--low-memory', '--prefix', str(tmp_path), ], ) @@ -72,6 +74,7 @@ def _fake_run(config): assert cfg.run_mash is False assert cfg.run_blast is False assert cfg.n_cpu == 2 + assert cfg.low_memory is True assert cfg.prefix == tmp_path.resolve(strict=True) assert cfg.tar_paths == targets_txt.resolve(strict=True) assert cfg.neg_paths == non_targets_txt.resolve(strict=True) diff --git a/tests/smoke/test_config.py b/tests/smoke/test_config.py index b8682be..e782320 100644 --- a/tests/smoke/test_config.py +++ b/tests/smoke/test_config.py @@ -73,6 +73,7 @@ def test_json_serialization_contains_important_fields( run_mash=False, run_blast=False, n_cpu=2, + low_memory=True, ) json_text = config.model_dump_json() @@ -84,6 +85,7 @@ def test_json_serialization_contains_important_fields( assert '"run_mash":false' in json_text assert '"run_blast":false' in json_text assert '"n_cpu":2' in json_text + assert '"low_memory":true' in json_text def test_api_key_secret_and_serialization(tmp_path: Path, targets_txt: Path, non_targets_txt: Path) -> None: diff --git a/tests/smoke/test_core.py b/tests/smoke/test_core.py index d1230ed..f630e58 100644 --- a/tests/smoke/test_core.py +++ b/tests/smoke/test_core.py @@ -40,6 +40,7 @@ def test_api_run_and_overwrite_behavior( min_len=8, max_len=120, n_cpu=1, + low_memory=True, ) seq = run(Config(**run_config)) @@ -48,6 +49,7 @@ def test_api_run_and_overwrite_behavior( assert (out_dir / 'assemblies.csv').exists() assert (out_dir / 'signatures.fasta').exists() assert seq.config.run_mash is False + assert seq.config.low_memory is True with pytest.raises(FileExistsError): run(Config(**run_config)) diff --git a/tests/smoke/test_outputs.py b/tests/smoke/test_outputs.py index 5df9f61..bdd81a9 100644 --- a/tests/smoke/test_outputs.py +++ b/tests/smoke/test_outputs.py @@ -119,3 +119,18 @@ def test_graph_matches_expected(tmp_path: Path, targets_txt: Path, non_targets_t ) _assert_graph_matches_expected(out_dir / WORKINGDIR.graph, expected_graph) + + +def test_low_memory_graph_matches_expected(tmp_path: Path, targets_txt: Path, non_targets_txt: Path, expected_graph: Path) -> None: + out_dir = _run_cli( + '--tar-paths', str(targets_txt), + '--neg-paths', str(non_targets_txt), + '--prefix', str(tmp_path), + '--threads', '1', + '--title', 'no-filter-low-memory', + '--no-filter', + '--low-memory', + *_shared_config, + ) + + _assert_graph_matches_expected(out_dir / WORKINGDIR.graph, expected_graph)