aegis_sim.submodels.reproduction.mutation

  1import numpy as np
  2from aegis_sim import variables
  3
  4
  5class Mutator:
  6    def init(self, MUTATION_RATIO, MUTATION_METHOD, MUTATION_AGE_MULTIPLIER):
  7        self.MUTATION_RATIO = MUTATION_RATIO
  8        self.MUTATION_METHOD = MUTATION_METHOD
  9        self.MUTATION_AGE_MULTIPLIER = MUTATION_AGE_MULTIPLIER
 10        self.rate_0to1 = MUTATION_RATIO / (1 + MUTATION_RATIO)
 11        self.rate_1to0 = 1 / (1 + MUTATION_RATIO)
 12        # Set mutation method
 13        if self.MUTATION_METHOD == "by_index":
 14            self._mutate = self._mutate_by_index
 15        elif self.MUTATION_METHOD == "by_bit":
 16            self._mutate = self._mutate_by_bit
 17        else:
 18            raise ValueError("MUTATION_METHOD must be 'by_index' or 'by_bit'")
 19
 20    def _mutate_by_bit(self, genomes, muta_prob, ages, random_probabilities=None):
 21        """Induce germline mutations."""
 22
 23        if random_probabilities is None:
 24            random_probabilities = variables.rng.random(genomes.shape)
 25
 26        # Broadcast to fit [individual, chromatid, locus, bit] shape
 27        mutation_probabilities = muta_prob[:, None, None, None]
 28
 29        mutation_probabilities = self.apply_mutation_age_multiplier(
 30            mutation_probabilities=mutation_probabilities,
 31            ages=ages,
 32            MUTATION_AGE_MULTIPLIER=self.MUTATION_AGE_MULTIPLIER,
 33        )
 34        mutate_0to1 = (~genomes) & (
 35            random_probabilities < (mutation_probabilities * self.rate_0to1).astype("float32")
 36        )  # genome == 0 &
 37        mutate_1to0 = genomes & (
 38            random_probabilities < (mutation_probabilities * self.rate_1to0).astype("float32")
 39        )  # genomes == 1 &
 40
 41        genomes[mutate_0to1] = 1
 42        genomes[mutate_1to0] = 0
 43
 44        return genomes
 45
 46    def _mutate_by_index(self, genomes, muta_prob, ages):
 47        """Alternative faster method for introducing mutations.
 48
 49        Instead of generating a random probability for every bit in the array of genomes,
 50        generate random indices of bits that could be mutated."""
 51
 52        if genomes.size == 0:
 53            return genomes
 54
 55        bits_per_genome = genomes[0].size
 56
 57        muta_prob = self.apply_mutation_age_multiplier(
 58            mutation_probabilities=muta_prob,
 59            ages=ages,
 60            MUTATION_AGE_MULTIPLIER=self.MUTATION_AGE_MULTIPLIER,
 61        )
 62
 63        # Calculate number of bits to mutate
 64        n_mutations_per_individual = variables.rng.binomial(n=bits_per_genome, p=muta_prob, size=len(genomes))
 65        n_mutations_total = np.sum(n_mutations_per_individual)
 66
 67        # Generate indices to mutate
 68        mutation_indices = (
 69            np.repeat(np.arange(len(genomes)), n_mutations_per_individual),
 70            variables.rng.integers(genomes.shape[1], size=n_mutations_total),
 71            variables.rng.integers(genomes.shape[2], size=n_mutations_total),
 72            variables.rng.integers(genomes.shape[3], size=n_mutations_total),
 73        )
 74
 75        # Extract indices of 0-bits and 1-bits
 76        bits = genomes[mutation_indices]  # NOTE Use tuple for ndarray indexing
 77        bits0_indices = (~bits).nonzero()[0]
 78        bits1_indices = bits.nonzero()[0]
 79
 80        # Take into consideration the MUTATION_RATIO
 81        bits0_include = variables.rng.random(len(bits0_indices)) < self.rate_0to1
 82        bits1_include = variables.rng.random(len(bits1_indices)) < self.rate_1to0
 83        bits0_indices = bits0_indices[bits0_include]
 84        bits1_indices = bits1_indices[bits1_include]
 85
 86        # Mutate bits at mutation_indices
 87        mutation_indices = np.array(mutation_indices)
 88        genomes[tuple(mutation_indices[:, bits1_indices.T])] = False
 89        genomes[tuple(mutation_indices[:, bits0_indices.T])] = True
 90
 91        return genomes
 92
 93    @staticmethod
 94    def apply_mutation_age_multiplier(mutation_probabilities, ages, MUTATION_AGE_MULTIPLIER):
 95        """
 96        final = initial(1 + age * multiplier)
 97        """
 98        multipliers = ages * MUTATION_AGE_MULTIPLIER
 99        multipliers = multipliers.reshape(mutation_probabilities.shape)
100        return mutation_probabilities * (1 + multipliers)
101
102
103mutator = Mutator()