How to model Mendelian inheritance

Generated with AI

A Rust struct for solving Punnett Square problems

The code can work with any number of genes that can independently separate during crossover. The methods provided can produce counts of children by genotype and phenotype.

I recently finished a quiz for a course in my MSc. Bioinformatics program. I spent the maximum time on questions related to Punnet Squares. Some people understand by making notes and writing, but I learn by coding. So, this post is a result of that. Also, I was tired of drawing the diagrams by hand.

Code

The code is tested on a classical genetics problem set by Khan Academy.

use itertools::Itertools;
use std::collections::{HashMap, HashSet};

#[derive(Debug, Clone)]
struct Individual {
    genotype: Vec<String>,
    phenotype: Option<Vec<String>>,
    phenotype_data: HashMap<char, String>,
}

fn count_children_by_genotype(children: Vec<Individual>) -> HashMap<String, i32> {
    let mut map: HashMap<String, i32> = HashMap::new();
    for child in &children {
        let genotype_as_string = child.genotype.join("");
        *map.entry(genotype_as_string).or_default() += 1;
    }
    map.insert("Total".to_string(), children.len() as i32);
    map
}

fn count_children_by_phenotype(children: Vec<Individual>) -> HashMap<String, i32> {
    let mut map: HashMap<String, i32> = HashMap::new();
    for child in &children {
        let phenotype_as_string = child.phenotype.clone().unwrap();
        let phenotype_as_string = phenotype_as_string.join(", ");
        *map.entry(phenotype_as_string).or_default() += 1;
    }
    map.insert("Total".to_string(), children.len() as i32);
    map
}

impl Individual {
    // Instantiate an individual and sort the genotype
    // Capital letters are at the beginning of the sorted string
    fn new(mut genotype: Vec<String>, phenotype_data: HashMap<char, String>) -> Self {
        for gene in &mut genotype {
            *gene = gene.chars().sorted().collect();
        }
        let mut instance = Self {
            genotype,
            phenotype_data,
            phenotype: None,
        };
        instance.phenotype();
        instance
    }

    fn num_genes(&self) -> usize {
        self.genotype.len()
    }

    fn can_cross(&self, other: &Self) -> bool {
        self.num_genes() == other.num_genes()
    }

    // Use the first character in the genotype to get the phenotype
    // The genotype string is sorted
    // So, if the first character is not a capital letter, the genotype does not contain any dominant alleles
    fn phenotype(&mut self) -> () {
        let phenotype_data = &self.phenotype_data;
        let mut phene: Vec<String> = vec![];
        for gene in &self.genotype {
            let allele = gene.chars().next().unwrap();
            phene.push(phenotype_data.get(&allele).unwrap().to_string());
        }
        self.phenotype = Some(phene);
    }

    fn allele_combinations(&self) -> HashSet<Vec<char>> {
        let genotype_as_strings: &Vec<String> = &self.genotype;
        let genotype_as_chars: Vec<Vec<char>> = genotype_as_strings
            .iter()
            .map(|s| s.chars().collect())
            .collect();
        let iterator = genotype_as_chars.into_iter().multi_cartesian_product();
        let combinations: HashSet<Vec<char>> = iterator.map(|v| v.into_iter().collect()).collect();
        combinations
    }

    fn cross(&self, other: &Self) -> Option<Vec<Self>> {
        let phenotype_data = &self.phenotype_data;
        if self.can_cross(other) {
            let mut children: Vec<Self> = vec![];
            let self_combinations = self.allele_combinations();
            let other_combinations = other.allele_combinations();
            for x in &self_combinations {
                for y in &other_combinations {
                    let mut child_genotype: Vec<String> = vec![];
                    for i in 0..self.num_genes() {
                        let mut gene = x[i].to_string();
                        gene.push(y[i]);
                        child_genotype.push(gene);
                    }
                    children.push(Self::new(child_genotype, phenotype_data.clone()));
                }
            }
            Some(children)
        } else {
            None
        }
    }
}

Example #1: Horses (1 gene)

Problem

In horses, black coloring (B) is dominant, and chestnut coloring (b) is recessive.

A homozygous black horse crosses with a chestnut horse.

What percentage of offspring are expected to be chestnut?

fn main() {
    let phenotype_data: HashMap<char, String> = HashMap::from([
        ('B', "black".to_string()),
        ('b', "chestnut".to_string()),
    ]);
    
    // The order in the vector should be same for both mother and father
    let father: Individual = Individual::new(
        vec!["BB".to_string()],
        phenotype_data.clone(),
    );
    let mother: Individual = Individual::new(
        vec!["bb".to_string()], 
        phenotype_data
    );
    
    // Perform a cross between mother and father and get the resulting children
    let children = mother.cross(&father).unwrap();
    
    // These methods are applied to a collection of individuals
    let count_by_genotype = count_children_by_genotype(children.clone());
    let count_by_phenotype = count_children_by_phenotype(children);
    
    println!("{:?}", count_by_genotype);
    println!("{:?}", count_by_phenotype);
}

Answer

0% (all horses are expected to be black)

Example #2: Cats (2 genes)

Problem

In cats, being awesome (A) is dominant to being average (a), and being bashful (B) is dominant to being boring (b).

A male cat that is homozygous dominant for the A trait and heterozygous for the B trait is mated with a female cat that is homozygous recessive for both. Both the A and B traits assort independently from one another.

What is the probability that they will have an offspring who is heterozygous for both traits?

fn main() {
    let phenotype_data: HashMap<char, String> = HashMap::from([
        ('A', "awesome".to_string()),
        ('a', "average".to_string()),
        ('B', "bashful".to_string()),
        ('b', "boring".to_string()),
    ]);
    
    // The order in the vector should be same for both mother and father
    let father: Individual = Individual::new(
        vec!["AA".to_string(), "Bb".to_string()],
        phenotype_data.clone(),
    );
    let mother: Individual = Individual::new(
        vec!["aa".to_string(), "bb".to_string()], 
        phenotype_data
    );
    
    // Perform a cross between mother and father and get the resulting children
    let children = mother.cross(&father).unwrap();
    
    // These methods are applied to a collection of individuals
    let count_by_genotype = count_children_by_genotype(children.clone());
    let count_by_phenotype = count_children_by_phenotype(children);
    
    println!("{:?}", count_by_genotype);
    println!("{:?}", count_by_phenotype);
}

Answer

50% (1/2)

Next Steps

• You can play with this code in the Rust Playground
• You can solve some classical genetics problems by Khan Academy