Functional module

Functional module.

chromax.functional.cross(parents: Bool[Array, 'n 2 m d'], recombination_vec: Float[Array, 'm'], random_key: Array, mutation_probability: float = 0.0) → Bool[Array, 'n m d'][source]

Main function that computes crosses from a list of parents.

Parameters:

parents (ndarray) – parents to compute the cross. The shape of the parents is (n, 2, m, d), where n is the number of parents, m is the number of markers, and d is the ploidy.
recombination_vec – array of m probabilities. The i-th value represent the probability to recombine before the marker i.
random_key (jax.Array) – JAX random key, for reproducibility purpose.
mutation_probability (float) – The probability of having a mutation in a marker.

Returns:

offspring population of shape (n, m, d).

Return type:

ndarray

Example:

>>> from chromax import functional
>>> import numpy as np
>>> import jax
>>> n_chr, chr_len, ploidy = 10, 100, 2
>>> n_crosses = 50
>>> parents_shape = (n_crosses, 2, n_chr * chr_len, ploidy)
>>> parents = np.random.choice([False, True], size=parents_shape)
>>> rec_vec = np.full((n_chr, chr_len), 1.5 / chr_len)
>>> rec_vec[:, 0] = 0.5  # equal probability on starting haploid
>>> rec_vec = rec_vec.flatten()
>>> random_key = jax.random.key(42)
>>> f2 = functional.cross(parents, rec_vec, random_key)
>>> f2.shape
(50, 1000, 2)

chromax.functional.double_haploid(population: Bool[Array, 'n m d'], n_offspring: int, recombination_vec: Float[Array, 'm'], random_key: Array, mutation_probability: float = 0.0) → Bool[Array, 'n n_offspring m d'][source]

Computes the double haploid of the input population.

Parameters:

population (ndarray) – input population of shape (n, m, d).
n_offspring (int) – number of offspring per plant.
recombination_vec (ndarray) – array of m probabilities. The i-th value represent the probability to recombine before the marker i.
random_key (jax.Array) – JAX random key, for reproducibility purpose.
mutation_probability (float) – The probability of having a mutation in a marker.

Returns:

output population of shape (n, n_offspring, m, d). This population will be homozygote.

Return type:

ndarray

Example:

>>> from chromax import functional
>>> import numpy as np
>>> import jax
>>> n_chr, chr_len, ploidy = 10, 100, 2
>>> pop_shape = (50, n_chr * chr_len, ploidy)
>>> f1 = np.random.choice([False, True], size=pop_shape)
>>> rec_vec = np.full((n_chr, chr_len), 1.5 / chr_len)
>>> rec_vec[:, 0] = 0.5  # equal probability on starting haploid
>>> rec_vec = rec_vec.flatten()
>>> random_key = jax.random.key(42)
>>> dh = functional.double_haploid(f1, 10, rec_vec, random_key)
>>> dh.shape
(50, 10, 1000, 2)

chromax.functional.select(population: Bool[Array, 'n m d'], k: int, f_index: Callable[[Bool[Array, 'n m d']], Float[Array, 'n']]) → Tuple[Bool[Array, 'k m d'], Int[Array, 'k']][source]

Function to select individuals based on their score (index).

Parameters:

population (ndarray) – input grouped population of shape (n, m, d)
k (int) – number of individual to select.
f_index (Callable) – function that computes a score for each individual. The function accepts as input a population, i.e. and array of shape (n, m, 2) and returns an array of n float number.

Returns:

output population of shape (k, m, d), output indices of shape (k,)

Return type:

tuple of two ndarrays

Example:

>>> from chromax import functional
>>> from chromax.trait_model import TraitModel
>>> from chromax.index_functions import conventional_index
>>> import numpy as np
>>> n_chr, chr_len, ploidy = 10, 100, 2
>>> pop_shape = (50, n_chr * chr_len, ploidy)
>>> f1 = np.random.choice([False, True], size=pop_shape)
>>> marker_effects = np.random.randn(n_chr * chr_len)
>>> gebv_model = TraitModel(marker_effects[:, None])
>>> f_index = conventional_index(gebv_model)
>>> f2, selected_indices = functional.select(f1, k=10, f_index=f_index)
>>> f2.shape
(10, 1000, 2)