I need to process some very large files in the context of genomic data analysis, and I’m finding it hard to come up with an optimized solution to my task, which is as follows:

1. I have a “scores” file, that contains a collection of genomic position intervals, with a given score for each interval. They are of the form [start, end). A toy example of this file is as follows:

chromosome      start   end     score
1       12646   12647   0.173
1       12647   12648   2.52
1       12648   12649   1.75
1       12649   12650   1.34

2. I have a “bim” file, that contains a collection of genetic variants. These correspond to places in the genome (given by BP) that differ from the reference (they contain A1 instead of A0). The length of each these variants is defined as the absolute value of the difference in length between A0 and A1.
   Then, each variant defines an interval of the form [BP, BP+length].

A toy example of this file is as follows (it doesn’t have a header):

1       chr1-586001-SNV->1      0       586001  A       G
1       chr1-586045-INS->2      0       586045  ATTC       G
1       chr1-586069-SNV->3      0       586069  A       T
1       chr1-586073-DEL->4      0       586073  T       CTGAA
1       chr1-586142-SNV->5      0       586142  G       A
1       chr1-586203-COMPLEX->9  0       586203  TA      TGC

Columns are, respectively: CHR, Variant_ID (called SNP), cM (not needed), BP, A1, A0.

For example, the interval defined by the second variant would be [586045, 586045+3]= [586045, 586048].

3. The task is: for each variant, check if the interval it defines overlaps any scores interval that has a score of at least 4. If it does, annotate that variant with 1 (set GERP.RSsup4=1), otherwise set GERP.RSsup4=0.

The variant_ID contains the variant type (SNV, INS, DEL, COMPLEX), and GERP.RSsup4 also needs to be splitted by variant type, but that is easy and not essential to the task.

I have the following script that achieves the task with polars. However, it takes too long (my scores files has on the order of 188775205 lines, and my bim file on the order of 6600771 lines).

I should probably avoid using map_elements and instead use vectorized-like operations, but I’m not sure how to translate this to the native API.

Potentially useful information:

1. Variants that have a length of 0 can only overlap at most 1 scores interval.
2. Scores intervals shoud be disjoint, but variant intervals could have non-empy intersection.

Any help would be much appreciated!

import sys
import polars as pl

def read_gerp_scores(chromosome):
    gerp_file = f"scores_chr{chromosome}.txt"
    gerp_scores = pl.read_csv(gerp_file, separator='t')
    return gerp_scores

def read_bim_file(bim_file):
    bim_data = pl.read_csv(bim_file, separator='t', has_header=False, dtypes={'column_3':float}, new_columns=["CHR", "SNP", "cM", "BP", "A1", "A0"])
    return bim_data

def compute_overlap_and_score(gerp_scores, start, end):
    overlap = gerp_scores.filter((pl.col('start') < end) & (pl.col('end') > start))
    return int((overlap['score'] >= 4).any())

def process_variants(bim_file, gerp_scores):
    bim_data = read_bim_file(bim_file)

    # Calculate the end position of each variant interval
    bim_data = bim_data.with_columns([
        (pl.col('BP') + (pl.col('A0').str.len_bytes() - pl.col('A1').str.len_bytes()).abs() + 1).alias('END')
    ])

    # Function to compute GERP.RSsup4 score for a given variant type
    def get_gerp_score(snp_type):
        condition = pl.col('SNP').str.contains(snp_type)
        scores = bim_data.with_columns([
            pl.when(condition)
            .then(pl.struct(['BP', 'END']).map_elements(lambda x: compute_overlap_and_score(gerp_scores, x['BP'], x['END'])))
            .otherwise(0)
            .alias(f'GERP.RSsup4_{snp_type}')
        ])
        return scores

    # Add the computed columns to the DataFrame
    bim_data = get_gerp_score('SNV')
    bim_data = get_gerp_score('INS')
    bim_data = get_gerp_score('DEL')
    bim_data = get_gerp_score('COMPLEX')

    return bim_data.select(['CHR', 'BP', 'SNP', 'GERP.RSsup4_SNV', 'GERP.RSsup4_INS', 'GERP.RSsup4_DEL', 'GERP.RSsup4_COMPLEX'])

def main():
    if len(sys.argv) != 2:
        print("Usage: python3 script.py {num}")
        sys.exit(1)

    chromosome = sys.argv[1]
    gerp_scores = read_gerp_scores(chromosome)

    # Process BIM file
    bim_file1 = f"data_chr{chromosome}.bim"
    results1 = process_variants(bim_file1, gerp_scores)
    output_file1 = f"data_chr{chromosome}_GERP.RSsup4.txt"
    results1.write_csv(output_file1, separator='t')
    print(f"Output written to {output_file1}")

if __name__ == "__main__":
    main()

I tried solving the task with the above script.
The script seems to correctly handle the task based on a few small, examples, but it doesn’t scale to the size of my actual data. I need a much faster implementation.