add argument chunck size

Author: qclayssen

bolt/common/pcgr.py

@@ -524,12 +524,17 @@ def annotate_record(record, annotations, *, allow_missing=False):
 
     return record
 
-def split_vcf(input_vcf, output_dir):
+def split_vcf(input_vcf, output_dir, *, max_variants=None):
     """
     Splits a VCF file into multiple chunks, each containing up to max_variants variants.
     Each chunk includes the VCF header.
     Ensures no overlapping positions between chunks.
     """
+    if max_variants is None:
+        max_variants = constants.MAX_SOMATIC_VARIANTS
+    elif max_variants <= 0:
+        raise ValueError("max_variants must be a positive integer.")
+
     output_dir = pathlib.Path(output_dir / "vcf_chunks")
     output_dir.mkdir(parents=True, exist_ok=True)
     chunk_files = []
@@ -547,7 +552,7 @@ def split_vcf(input_vcf, output_dir):
     for record in vcf_in:
         current_position = record.POS
         # Check if we need to start a new chunk
-        if variant_count >= constants.MAX_SOMATIC_VARIANTS and (last_position is None or current_position != last_position):
+        if variant_count >= max_variants and (last_position is None or current_position != last_position):
             # Close the current chunk file and start a new one
             vcf_out.close()
             chunk_number += 1

bolt/workflows/smlv_somatic/annotate.py

@@ -30,6 +30,7 @@
 @click.option('--threads', required=False, default=4, type=int)
 
 @click.option('--output_dir', required=True, type=click.Path())
+@click.option('--pcgr_variant_chunk_size', required=False, type=int, help='Override maximum variants per PCGR chunk.')
 
 def entry(ctx, **kwargs):
     '''Annotate variants with information from several sources\f
@@ -99,8 +100,13 @@ def entry(ctx, **kwargs):
     print(f"Total number of variants in the input VCF: {total_variants}")
 
     # Run PCGR in chunks if exceeding the maximum allowed for somatic variants
-    if total_variants > constants.MAX_SOMATIC_VARIANTS:
-        vcf_chunks = pcgr.split_vcf(pcgr_prep_fp, output_dir)
+    chunk_size = kwargs.get('pcgr_variant_chunk_size')
+    if chunk_size is not None and chunk_size <= 0:
+        raise click.BadParameter('must be a positive integer', param_hint='--pcgr_variant_chunk_size')
+    chunk_size = chunk_size or constants.MAX_SOMATIC_VARIANTS
+
+    if total_variants > chunk_size:
+        vcf_chunks = pcgr.split_vcf(pcgr_prep_fp, output_dir, max_variants=chunk_size)
         pcgr_tsv_fp, pcgr_vcf_fp = pcgr.run_somatic_chunck(
             vcf_chunks,
             kwargs['pcgr_data_dir'],