| """ |
| MutationEncoder for DNA mutation encoding |
| """ |
| import torch |
|
|
| class MutationEncoder: |
| """Encode DNA mutations for neural network input""" |
| |
| MUTATION_TYPES = { |
| ('A','T'): 0, ('A','C'): 1, ('A','G'): 2, |
| ('T','A'): 3, ('T','C'): 4, ('T','G'): 5, |
| ('C','A'): 6, ('C','T'): 7, ('C','G'): 8, |
| ('G','A'): 9, ('G','T'): 10, ('G','C'): 11 |
| } |
|
|
| def __init__(self): |
| self.nucl_to_idx = {'A': 0, 'T': 1, 'G': 2, 'C': 3, 'N': 4} |
|
|
| def encode_sequence(self, sequence, max_len=99): |
| """One-hot encode DNA sequence""" |
| sequence = sequence.upper()[:max_len] |
| sequence = sequence + 'N' * (max_len - len(sequence)) |
| encoded = torch.zeros(max_len, 5) |
| for i, nucl in enumerate(sequence): |
| idx = self.nucl_to_idx.get(nucl, 4) |
| encoded[i, idx] = 1.0 |
| return encoded |
|
|
| def encode_mutation(self, ref_seq, mut_seq): |
| """ |
| Encode mutation as 1101-dimensional tensor |
| |
| Args: |
| ref_seq: Reference DNA sequence (99bp) |
| mut_seq: Mutated DNA sequence (99bp) |
| |
| Returns: |
| torch.Tensor of shape (1101,) |
| """ |
| max_len = 99 |
| ref_encoded = self.encode_sequence(ref_seq, max_len) |
| mut_encoded = self.encode_sequence(mut_seq, max_len) |
| |
| |
| diff_mask = torch.zeros(max_len, 1) |
| mut_pos = None |
| ref_base = None |
| mut_base = None |
| |
| for i in range(min(len(ref_seq), len(mut_seq), max_len)): |
| if ref_seq[i] != mut_seq[i]: |
| diff_mask[i, 0] = 1.0 |
| if mut_pos is None: |
| mut_pos = i |
| ref_base = ref_seq[i].upper() |
| mut_base = mut_seq[i].upper() |
| |
| |
| mutation_onehot = torch.zeros(12) |
| if ref_base is not None and mut_base is not None: |
| key = (ref_base, mut_base) |
| idx = self.MUTATION_TYPES.get(key, None) |
| if idx is not None: |
| mutation_onehot[idx] = 1.0 |
| |
| |
| seq_part = torch.cat([ref_encoded, mut_encoded, diff_mask], dim=1) |
| seq_flat = seq_part.flatten() |
| return torch.cat([seq_flat, mutation_onehot]) |
|
|
