Convert GVL object to MSA

[bschilder]

Thinking through the best way to convert the GVL dataset to a Multiple Sequence Alignment so that I can compute a consensus sequence.

If this ends up being useful for you, I can make a PR to GVL directly.

Best,
Brian

[bschilder]

Set up GVL db

import genvarloader as gvl
import numpy as np

Download data

# GRCh38 chromosome 22 sequence
reference = pooch.retrieve(
    url="https://ftp.ensembl.org/pub/release-112/fasta/homo_sapiens/dna/Homo_sapiens.GRCh38.dna.chromosome.22.fa.gz",
    known_hash="sha256:974f97ac8ef7ffae971b63b47608feda327403be40c27e391ee4a1a78b800df5",
    progressbar=True,
)
if not Path(f"{reference[:-3]}.bgz").exists():
    !gzip -dc {reference} | bgzip > {reference[:-3]}.bgz
reference = reference[:-3] + ".bgz"

# PLINK 2 files
variants = pooch.retrieve(
    url="doi:10.5281/zenodo.13656224/1kGP.chr22.pgen",
    known_hash="md5:31aba970e35f816701b2b99118dfc2aa",
    progressbar=True,
    fname="1kGP.chr22.pgen",
)
pooch.retrieve(
    url="doi:10.5281/zenodo.13656224/1kGP.chr22.psam",
    known_hash="md5:eefa7aad5acffe62bf41df0a4600129c",
    progressbar=True,
    fname="1kGP.chr22.psam",
)
pooch.retrieve(
    url="doi:10.5281/zenodo.13656224/1kGP.chr22.pvar",
    known_hash="md5:5f922af91c1a2f6822e2f1bb4469d12b",
    progressbar=True,
    fname="1kGP.chr22.pvar",
)

# BED
bed_path = pooch.retrieve(
    url="doi:10.5281/zenodo.13656224/chr22_egenes.bed",
    known_hash="md5:ccb55548e4ddd416d50dbe6638459421",
    progressbar=True,
)

Create GVL db

ds_path = TemporaryDirectory(suffix=".gvl").name

bed_chrom = bed.filter(pl.col('chrom').str.replace("chr","")=='22')
force = False
if not os.path.exists(ds_path) or force:
    gvl.write(
        path=ds_path,
        bed=gvl.with_length(bed_chrom[:20], 2**18),  # change region length to 131,072 bp
        variants=variants,
        overwrite=True,
    )

Open GVL db

ds = (
    gvl.Dataset.open(ds_path, reference=reference)
    .with_seqs("haplotypes") 
    .with_len(2**18)
)

Create consensus sequence

First, I can do this the dumb way and stack all the haplotypes into a 2D array ((sample x plod) x seq length) and find the most frequent nucleotides at each sequence position.

# Get an MSA

# reshape the 
seq_arr = ds[0,:10,]

# Stack the ploid dimension (dimension 2 in seq_arr)

# Stack the ploid dimension to get all haplotypes in one array
seq_arr_stacked = np.stack([seq_arr[:,0,:], seq_arr[:,1,:]], axis=1).reshape(-1, seq_arr.shape[-1])
seq_arr_stacked

# Find most common value in each column
# Convert string array to integer array using ord() to get ASCII values
consensus_seq = np.apply_along_axis(
    lambda x: np.bincount([ord(c) for c in x]).argmax(),
    axis=0,
    arr=seq_arr_stacked
)

# Convert ASCII values back to bytes
consensus_seq = np.array([bytes([b]) for b in consensus_seq], dtype='|S1')
consensus_seq

Obviously, this doesn't account for haplotypes that contain indels, which can shift the entire alignment. Next, I'll work on a way to take this into account.

[bschilder]

I think what I need are these, which tell me how much to shift the sequences:

ds._seqs.variants.ilens

For each sequence,

• when there's a deletion (negative value) I should add "-"*len_del to that sequence
• when there's an insertion (positive value) I should add "-"*len_ins to every other sequence.
• when there's no indel (0) I should do nothing

[d-laub]

IMO to call this an MSA it must be an optimal MSA given some score function. This would be quite the undertaking. I think with SNPs only, the approach you've outlined is probably very close to an optimal sequence alignment, although I'm not sure if it's provably so in the general case. With indels, I think it's much more likely to create a suboptimal MSA since gap placements that do not correspond to the observed indels wrt the reference genome could be a more optimal way to align the sequences. Another way to intuit this is that for arbitrary sets of haplotypes, the reference genome can have a greater edit distance from the haps than they do to each other.

Maybe this too high a standard tho and it would be useful even to have suboptimal MSAs, in which case should this be called something else?

I also am concerned that such a feature could scale very poorly with hap length and # of samples as full MSA is $O(L^S)$ where $L$ is seq len and $S$ is # of sequences.

[bschilder]

This seems to work ok, though it is probably very inefficient atm:

def stack_ploidy(arr):
    return np.stack([arr[:,0,:], arr[:,1,:]], axis=1).reshape(-1, arr.shape[-1])

def create_msa(ref_coords, seq_arr):
    """
    Create a gappy Multiple Sequence Alignment from a ragged 2D array using reference coordinates.
    
    This will take a 2D array of the reference coordinates:
        ref_coords = [
        [100, 101, 103],  # First sequence has bases at positions 100, 101, 103
        [100, 102, 103]   # Second sequence has bases at positions 100, 102, 103
        ]
    And a 2D array of the sequences (left-aligned):
        seq_arr = [
            ['A', 'T', 'G'],  # First sequence
            ['A', 'C', 'G']   # Second sequence
        ] 
    And return a gappy Multiple Sequence Alignment that is aligned at all positions.  
        A T - G  # First sequence
        A - C G  # Second sequence
    Args:
        ref_coords: 2D array of reference coordinates (sequence x position) indicating where each position maps in the reference genome
        seq_arr: 2D array of sequences (sequence x position) with ragged right end
        
    Returns:
        2D array containing the gappy MSA
    
    Example:
        # Test the function with sample data
        test_ref_coords = np.array([[100, 101, 103], [100, 102, 103]])
        test_seq_arr = np.array([[b'A', b'T', b'G'], 
                                [b'A', b'C', b'G']], dtype='|S1')
        msa = create_msa(test_ref_coords, test_seq_arr)
    """
    assert ref_coords.shape == seq_arr.shape, f"Shapes do not match: ref_coords.shape ({ref_coords.shape}) != seq_arr.shape ({seq_arr.shape})"

    # Find the min and max reference coordinates across all sequences
    min_coord = min(min(coords) for coords in ref_coords)
    max_coord = max(max(coords) for coords in ref_coords)
    alignment_length = max_coord - min_coord + 1
    
    # Initialize output array with gaps
    msa = np.full((len(seq_arr), alignment_length), b'-', dtype='|S1')
    
    # Fill in the sequences using reference coordinates
    for i, (coords, seq) in enumerate(zip(ref_coords, seq_arr)):
        # Fill in the sequence at the correct positions based on reference coordinates
        for j, (coord, base) in enumerate(zip(coords, seq)):
            # Convert base to bytes if it isn't already
            if isinstance(base, str):
                base = base.encode('ascii')
            msa[i, coord - min_coord] = base
            
    return msa

region_idx = 0
sample_idx = range(10)

ref_coords = stack_ploidy(site_ds[region_idx, ][0].ref_coords[sample_idx,])
seq_arr = stack_ploidy(ds.subset_to(samples=sample_idx)[region_idx,])
msa = create_msa(ref_coords, seq_arr)

Wrote a function to find the portion of the sequence that begins to differ substantially:

# Convert byte arrays to strings and find columns with differences
aligned_seqs_str = [''.join(seq.astype(str)) for seq in msa]

# Find columns where sequences differ
diff_cols = []
for col in range(len(aligned_seqs_str[0])):
    col_chars = [seq[col] for seq in aligned_seqs_str]
    if len(set(col_chars)) > 1:  # If more than one unique character in column
        diff_cols.append(col)

# Print sequences showing only columns with differences
for seq in aligned_seqs_str:
    print(''.join(seq[i] for i in diff_cols))

Here's the very end of the alignment. Gappy (probably bc we're not filling in adjacent seqs), but still seems to be aligned properly:

[bschilder]

Takes 92 seconds for 901 sequences (451 samples x 2 haplotypes)

[bschilder]

Got it down to 52 seconds with some Cursor-assisted numba implementation:

@nb.njit(parallel=True)
def _create_msa_fast(ref_coords, seq_arr, min_coord, max_coord):
    """Numba-accelerated core MSA creation function"""
    n_seqs = ref_coords.shape[0]
    alignment_length = max_coord - min_coord + 1
    msa = np.full((n_seqs, alignment_length), ord('-'), dtype=np.int8)
    
    for i in nb.prange(n_seqs):
        coords = ref_coords[i]
        seq = seq_arr[i]
        for j in range(len(coords)):
            coord = coords[j]
            base = seq[j]
            msa[i, coord - min_coord] = base
            
    return msa

def create_msa(ref_coords, seq_arr):
    """
    Create a gappy Multiple Sequence Alignment from a ragged 2D array using reference coordinates.
    
    This will take a 2D array of the reference coordinates:
        ref_coords = [
        [100, 101, 103],  # First sequence has bases at positions 100, 101, 103
        [100, 102, 103]   # Second sequence has bases at positions 100, 102, 103
        ]
    And a 2D array of the sequences (left-aligned):
        seq_arr = [
            ['A', 'T', 'G'],  # First sequence
            ['A', 'C', 'G']   # Second sequence
        ] 
    And return a gappy Multiple Sequence Alignment that is aligned at all positions.  
        A T - G  # First sequence
        A - C G  # Second sequence
    Args:
        ref_coords: 2D array of reference coordinates (sequence x position) indicating where each position maps in the reference genome
        seq_arr: 2D array of sequences (sequence x position) with ragged right end
        
    Returns:
        2D array containing the gappy MSA
    """
    import time
    start_time = time.time()
    
    assert ref_coords.shape == seq_arr.shape, f"Shapes do not match: ref_coords.shape ({ref_coords.shape}) != seq_arr.shape ({seq_arr.shape})"

    # Find the min and max reference coordinates across all sequences
    min_coord = np.min(ref_coords)
    max_coord = np.max(ref_coords)
    
    # Convert input arrays to numpy arrays if they aren't already
    ref_coords = np.asarray(ref_coords)
    seq_arr = np.asarray(seq_arr)
    
    # Convert bytes to integers before passing to numba function
    seq_arr_int = np.zeros_like(seq_arr, dtype=np.int8)
    for i in range(seq_arr.shape[0]):
        for j in range(seq_arr.shape[1]):
            if isinstance(seq_arr[i,j], bytes):
                seq_arr_int[i,j] = seq_arr[i,j][0]
            else:
                seq_arr_int[i,j] = ord(seq_arr[i,j])
    
    # Create MSA using numba-accelerated function
    msa = _create_msa_fast(ref_coords, seq_arr_int, min_coord, max_coord)
    
    # Convert back to bytes array
    msa = msa.view('|S1')
    
    end_time = time.time()
    print(f"MSA for {seq_arr.shape[0]} sequences done {end_time - start_time:.2f} seconds")
            
    return msa

[bschilder]

Yelling at Cursor apparently helps a lot.

Now down to 0.19 seconds for 901 sequences!!

@nb.njit(cache=True)
def stack_ploidy(arr):
    """Optimized ploidy stacking using numba"""
    n_samples = arr.shape[0]
    seq_len = arr.shape[2]
    result = np.empty((n_samples * 2, seq_len), dtype=arr.dtype)
    for i in range(n_samples):
        result[i*2] = arr[i,0]
        result[i*2+1] = arr[i,1]
    return result

@nb.njit(parallel=True, cache=True)
def _create_msa_fast(ref_coords, seq_arr, min_coord, max_coord):
    """Numba-accelerated core MSA creation function"""
    n_seqs = ref_coords.shape[0]
    alignment_length = max_coord - min_coord + 1
    msa = np.full((n_seqs, alignment_length), ord('-'), dtype=np.int8)
    
    # Pre-compute coordinate offsets
    coord_offsets = ref_coords - min_coord
    
    for i in nb.prange(n_seqs):
        coords = coord_offsets[i]
        seq = seq_arr[i]
        seq_len = len(coords)
        for j in range(seq_len):
            msa[i, coords[j]] = seq[j]
            
    return msa

def create_msa(ref_coords, seq_arr):
    """
    Create a gappy Multiple Sequence Alignment from a ragged 2D array using reference coordinates.
    
    This will take a 2D array of the reference coordinates:
        ref_coords = [
        [100, 101, 103],  # First sequence has bases at positions 100, 101, 103
        [100, 102, 103]   # Second sequence has bases at positions 100, 102, 103
        ]
    And a 2D array of the sequences (left-aligned):
        seq_arr = [
            ['A', 'T', 'G'],  # First sequence
            ['A', 'C', 'G']   # Second sequence
        ] 
    And return a gappy Multiple Sequence Alignment that is aligned at all positions.  
        A T - G  # First sequence
        A - C G  # Second sequence
    Args:
        ref_coords: 2D array of reference coordinates (sequence x position) indicating where each position maps in the reference genome
        seq_arr: 2D array of sequences (sequence x position) with ragged right end
        
    Returns:
        2D array containing the gappy MSA
    """
    import time
    start_time = time.time()
    
    assert ref_coords.shape == seq_arr.shape, f"Shapes do not match: ref_coords.shape ({ref_coords.shape}) != seq_arr.shape ({seq_arr.shape})"

    # Convert input arrays to numpy arrays if they aren't already
    ref_coords = np.asarray(ref_coords)
    seq_arr = np.asarray(seq_arr)
    
    # Find the min and max reference coordinates across all sequences
    min_coord = np.min(ref_coords)
    max_coord = np.max(ref_coords)
    
    # Optimize byte conversion using vectorized operations
    if seq_arr.dtype == np.dtype('|S1'):
        seq_arr_int = np.frombuffer(seq_arr.tobytes(), dtype=np.int8).reshape(seq_arr.shape)
    else:
        seq_arr_int = np.array([[ord(c) for c in row] for row in seq_arr], dtype=np.int8)
    
    # Create MSA using numba-accelerated function
    msa = _create_msa_fast(ref_coords, seq_arr_int, min_coord, max_coord)
    
    # Convert back to bytes array using view
    msa = msa.view('|S1')
    
    end_time = time.time()
    print(f"MSA for {seq_arr.shape[0]} sequences done {end_time - start_time:.2f} seconds")
            
    return msa

[bschilder]

I don't mind just keeping these functions in my other code bases for now. We can always revisit this later if it's something we want to work on further. Feel free to close for now!

[d-laub]

I think this could be useful depending on how things pan out. And then for incorporating it into GVL it can just have a caution sign saying this isn't necessarily an optimal MSA. Can you put this in a new branch? e.g. bschilder/msa That way it'll be easy-ish for us to keep it up to date with changes in main.

[bschilder]

Np, will do!